hubproxy/src/ratelimiter.go

package main

import (
	"fmt"
	"net"
	"strings"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
	"golang.org/x/time/rate"
)

const (
	// 清理间隔
	CleanupInterval = 10 * time.Minute
	MaxIPCacheSize = 10000
)

// IPRateLimiter IP限流器结构体
type IPRateLimiter struct {
	ips       map[string]*rateLimiterEntry // IP到限流器的映射
	mu        *sync.RWMutex                // 读写锁，保证并发安全
	r         rate.Limit                   // 速率限制（每秒允许的请求数）
	b         int                          // 令牌桶容量（突发请求数）
	whitelist []*net.IPNet                 // 白名单IP段
	blacklist []*net.IPNet                 // 黑名单IP段
}

// rateLimiterEntry 限流器条目
type rateLimiterEntry struct {
	limiter    *rate.Limiter
	lastAccess time.Time
}

// initGlobalLimiter 初始化全局限流器
func initGlobalLimiter() *IPRateLimiter {
	cfg := GetConfig()
	
	whitelist := make([]*net.IPNet, 0, len(cfg.Security.WhiteList))
	for _, item := range cfg.Security.WhiteList {
		if item = strings.TrimSpace(item); item != "" {
			if !strings.Contains(item, "/") {
				item = item + "/32" // 单个IP转为CIDR格式
			}
			_, ipnet, err := net.ParseCIDR(item)
			if err == nil {
				whitelist = append(whitelist, ipnet)
			} else {
				fmt.Printf("警告: 无效的白名单IP格式: %s\n", item)
			}
		}
	}
	
	// 解析黑名单IP段
	blacklist := make([]*net.IPNet, 0, len(cfg.Security.BlackList))
	for _, item := range cfg.Security.BlackList {
		if item = strings.TrimSpace(item); item != "" {
			if !strings.Contains(item, "/") {
				item = item + "/32" // 单个IP转为CIDR格式
			}
			_, ipnet, err := net.ParseCIDR(item)
			if err == nil {
				blacklist = append(blacklist, ipnet)
			} else {
				fmt.Printf("警告: 无效的黑名单IP格式: %s\n", item)
			}
		}
	}
	
	// 计算速率：将 "每N小时X个请求" 转换为 "每秒Y个请求"
	ratePerSecond := rate.Limit(float64(cfg.RateLimit.RequestLimit) / (cfg.RateLimit.PeriodHours * 3600))
	
	burstSize := cfg.RateLimit.RequestLimit
	if burstSize < 1 {
		burstSize = 1
	}
	
	limiter := &IPRateLimiter{
		ips:       make(map[string]*rateLimiterEntry),
		mu:        &sync.RWMutex{},
		r:         ratePerSecond,
		b:         burstSize,
		whitelist: whitelist,
		blacklist: blacklist,
	}
	
	// 启动定期清理goroutine
	go limiter.cleanupRoutine()
	
	return limiter
}

// initLimiter 初始化限流器
func initLimiter() {
	globalLimiter = initGlobalLimiter()
}

// cleanupRoutine 定期清理过期的限流器
func (i *IPRateLimiter) cleanupRoutine() {
	ticker := time.NewTicker(CleanupInterval)
	defer ticker.Stop()
	
	for range ticker.C {
		now := time.Now()
		expired := make([]string, 0)
		
		// 查找过期的条目
		i.mu.RLock()
		for ip, entry := range i.ips {
			// 如果最后访问时间超过1小时，认为过期
			if now.Sub(entry.lastAccess) > 1*time.Hour {
				expired = append(expired, ip)
			}
		}
		i.mu.RUnlock()
		
		// 如果有过期条目或者缓存过大，进行清理
		if len(expired) > 0 || len(i.ips) > MaxIPCacheSize {
			i.mu.Lock()
			// 删除过期条目
			for _, ip := range expired {
				delete(i.ips, ip)
			}
			
			// 如果缓存仍然过大，全部清理
			if len(i.ips) > MaxIPCacheSize {
				i.ips = make(map[string]*rateLimiterEntry)
			}
			i.mu.Unlock()
		}
	}
}

// extractIPFromAddress 从地址中提取纯IP，去除端口号
func extractIPFromAddress(address string) string {
	// 处理IPv6地址 [::1]:8080 格式
	if strings.HasPrefix(address, "[") {
		if endIndex := strings.Index(address, "]"); endIndex != -1 {
			return address[1:endIndex]
		}
	}
	
	// 处理IPv4地址 192.168.1.1:8080 格式
	if lastColon := strings.LastIndex(address, ":"); lastColon != -1 {
		return address[:lastColon]
	}
	
	return address
}

// isIPInCIDRList 检查IP是否在CIDR列表中
func isIPInCIDRList(ip string, cidrList []*net.IPNet) bool {
	// 先提取纯IP地址
	cleanIP := extractIPFromAddress(ip)
	parsedIP := net.ParseIP(cleanIP)
	if parsedIP == nil {
		return false
	}
	
	for _, cidr := range cidrList {
		if cidr.Contains(parsedIP) {
			return true
		}
	}
	return false
}

// GetLimiter 获取指定IP的限流器，同时返回是否允许访问
func (i *IPRateLimiter) GetLimiter(ip string) (*rate.Limiter, bool) {
	// 提取纯IP地址
	cleanIP := extractIPFromAddress(ip)
	
	// 检查是否在黑名单中
	if isIPInCIDRList(cleanIP, i.blacklist) {
		return nil, false
	}
	
	// 检查是否在白名单中
	if isIPInCIDRList(cleanIP, i.whitelist) {
		return rate.NewLimiter(rate.Inf, i.b), true
	}
	
	now := time.Now()
	
	i.mu.RLock()
	entry, exists := i.ips[cleanIP]
	i.mu.RUnlock()
	
	if exists {
		i.mu.Lock()
		if entry, stillExists := i.ips[cleanIP]; stillExists {
			entry.lastAccess = now
			i.mu.Unlock()
			return entry.limiter, true
		}
		i.mu.Unlock()
	}
	
	i.mu.Lock()
	if entry, exists := i.ips[cleanIP]; exists {
		entry.lastAccess = now
		i.mu.Unlock()
		return entry.limiter, true
	}
	
	entry = &rateLimiterEntry{
		limiter:    rate.NewLimiter(i.r, i.b),
		lastAccess: now,
	}
	i.ips[cleanIP] = entry
	i.mu.Unlock()
	
	return entry.limiter, true
}

// RateLimitMiddleware 速率限制中间件
func RateLimitMiddleware(limiter *IPRateLimiter) gin.HandlerFunc {
	return func(c *gin.Context) {
		// 获取客户端真实IP
		var ip string
		
		// 优先尝试从请求头获取真实IP
		if forwarded := c.GetHeader("X-Forwarded-For"); forwarded != "" {
			// X-Forwarded-For可能包含多个IP，取第一个
			ips := strings.Split(forwarded, ",")
			ip = strings.TrimSpace(ips[0])
		} else if realIP := c.GetHeader("X-Real-IP"); realIP != "" {
			// 如果有X-Real-IP头
			ip = realIP
		} else if remoteIP := c.GetHeader("X-Original-Forwarded-For"); remoteIP != "" {
			// 某些代理可能使用此头
			ips := strings.Split(remoteIP, ",")
			ip = strings.TrimSpace(ips[0])
		} else {
			// 回退到ClientIP方法
			ip = c.ClientIP()
		}
		
		// 提取纯IP地址（去除端口号）
		cleanIP := extractIPFromAddress(ip)
		
		// 日志记录请求IP和头信息
		fmt.Printf("请求IP: %s (去除端口后: %s), X-Forwarded-For: %s, X-Real-IP: %s\n", 
			ip, 
			cleanIP,
			c.GetHeader("X-Forwarded-For"), 
			c.GetHeader("X-Real-IP"))
		
		// 获取限流器并检查是否允许访问
		ipLimiter, allowed := limiter.GetLimiter(cleanIP)
		
		// 如果IP在黑名单中
		if !allowed {
			c.JSON(403, gin.H{
				"error": "您已被限制访问",
			})
			c.Abort()
			return
		}
		
		// 智能限流判断：检查是否应该跳过限流计数
		shouldSkip := smartLimiter.ShouldSkipRateLimit(cleanIP, c.Request.URL.Path)
		
		// 只有在不跳过的情况下才检查限流
		if !shouldSkip && !ipLimiter.Allow() {
			c.JSON(429, gin.H{
				"error": "请求频率过快，暂时限制访问",
			})
			c.Abort()
			return
		}
		
		c.Next()
	}
}

// ApplyRateLimit 应用限流到特定路由
func ApplyRateLimit(router *gin.Engine, path string, method string, handler gin.HandlerFunc) {
	// 使用全局限流器
	limiter := globalLimiter
	if limiter == nil {
		limiter = initGlobalLimiter()
	}
	
	// 根据HTTP方法应用限流
	switch method {
	case "GET":
		router.GET(path, RateLimitMiddleware(limiter), handler)
	case "POST":
		router.POST(path, RateLimitMiddleware(limiter), handler)
	case "PUT":
		router.PUT(path, RateLimitMiddleware(limiter), handler)
	case "DELETE":
		router.DELETE(path, RateLimitMiddleware(limiter), handler)
	default:
		router.Any(path, RateLimitMiddleware(limiter), handler)
	}
}