xlgo-core/middleware/ratelimit.go

package middleware

import (
	"context"
	"fmt"
	"sync"
	"time"

	"github.com/EthanCodeCraft/xlgo-core/database"
	"github.com/EthanCodeCraft/xlgo-core/response"
	"github.com/gin-gonic/gin"
)

// RateLimiter 速率限制器（内存版，单实例使用）
type RateLimiter struct {
	visitors map[string]*visitor
	mu       sync.RWMutex
	rate     int           // 每分钟允许的请求数
	window   time.Duration // 时间窗口
	ctx      context.Context
	cancel   context.CancelFunc
	wg       sync.WaitGroup
}

type visitor struct {
	lastSeen time.Time
	count    int
}

// NewRateLimiter 创建速率限制器（内存版）
func NewRateLimiter(rate int, window time.Duration) *RateLimiter {
	ctx, cancel := context.WithCancel(context.Background())
	limiter := &RateLimiter{
		visitors: make(map[string]*visitor),
		rate:     rate,
		window:   window,
		ctx:      ctx,
		cancel:   cancel,
	}

	limiter.wg.Add(1)
	go limiter.cleanupVisitors()

	return limiter
}

// Allow 检查是否允许请求
func (rl *RateLimiter) Allow(ip string) bool {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	v, exists := rl.visitors[ip]
	if !exists {
		rl.visitors[ip] = &visitor{
			lastSeen: time.Now(),
			count:    1,
		}
		return true
	}

	if time.Since(v.lastSeen) > rl.window {
		v.count = 1
		v.lastSeen = time.Now()
		return true
	}

	if v.count >= rl.rate {
		return false
	}

	v.count++
	v.lastSeen = time.Now()
	return true
}

// cleanupVisitors 清理过期的访问者记录
func (rl *RateLimiter) cleanupVisitors() {
	defer rl.wg.Done()

	ticker := time.NewTicker(time.Minute)
	defer ticker.Stop()

	for {
		select {
		case <-rl.ctx.Done():
			return
		case <-ticker.C:
			rl.mu.Lock()
			for ip, v := range rl.visitors {
				if time.Since(v.lastSeen) > rl.window {
					delete(rl.visitors, ip)
				}
			}
			rl.mu.Unlock()
		}
	}
}

// Stop 停止限流器（释放资源）
func (rl *RateLimiter) Stop() {
	rl.cancel()
	rl.wg.Wait()
}

// ===== Redis 分布式限流器 =====

// RedisRateLimiter Redis 分布式限流器
type RedisRateLimiter struct {
	keyPrefix string        // 键名前缀
	rate      int           // 每分钟允许的请求数
	window    time.Duration // 时间窗口
}

// slidingWindowLua 滑动窗口限流 Lua 脚本
// 返回: 当前窗口内的请求数
const slidingWindowLua = `
local key = KEYS[1]
local now = tonumber(ARGV[1])
local window = tonumber(ARGV[2])
local rate = tonumber(ARGV[3])

-- 移除窗口外的旧记录
redis.call('ZREMRANGEBYSCORE', key, 0, now - window)

-- 获取当前窗口内的请求数
local count = redis.call('ZCARD', key)

if count < rate then
    -- 添加当前请求
    redis.call('ZADD', key, now, now .. '-' .. math.random())
    redis.call('PEXPIRE', key, window)
    return 0
else
    return count
end
`

// NewRedisRateLimiter 创建 Redis 分布式限流器
func NewRedisRateLimiter(keyPrefix string, rate int, window time.Duration) *RedisRateLimiter {
	return &RedisRateLimiter{
		keyPrefix: keyPrefix,
		rate:      rate,
		window:    window,
	}
}

// Allow 检查是否允许请求
func (rl *RedisRateLimiter) Allow(ctx context.Context, identifier string) (bool, error) {
	if database.RedisClient == nil {
		// Redis 未启用，默认允许
		return true, nil
	}

	key := rl.keyPrefix + ":" + identifier
	now := float64(time.Now().UnixMilli())
	windowMs := float64(rl.window.Milliseconds())

	result, err := database.RedisClient.Eval(ctx, slidingWindowLua, []string{key}, now, windowMs, rl.rate).Result()
	if err != nil {
		return true, err // 出错时允许请求，避免影响业务
	}

	count := result.(int64)
	return count == 0, nil
}

// GetCount 获取当前窗口内的请求数
func (rl *RedisRateLimiter) GetCount(ctx context.Context, identifier string) (int64, error) {
	if database.RedisClient == nil {
		return 0, nil
	}

	key := rl.keyPrefix + ":" + identifier
	now := time.Now().UnixMilli()
	windowStart := now - rl.window.Milliseconds()

	// 移除旧记录并获取当前计数
	database.RedisClient.ZRemRangeByScore(ctx, key, "0", fmt.Sprintf("%d", windowStart))
	return database.RedisClient.ZCard(ctx, key).Result()
}

// Reset 重置限流计数
func (rl *RedisRateLimiter) Reset(ctx context.Context, identifier string) error {
	if database.RedisClient == nil {
		return nil
	}

	key := rl.keyPrefix + ":" + identifier
	return database.RedisClient.Del(ctx, key).Err()
}

// ===== 全局限速器 =====

var (
	loginLimiter  *RateLimiter
	apiLimiter    *RateLimiter
	uploadLimiter *RateLimiter
	redisLimiters map[string]*RedisRateLimiter
	limitersMu    sync.Mutex
)

func init() {
	redisLimiters = make(map[string]*RedisRateLimiter)
}

// InitRateLimiters 初始化限速器
func InitRateLimiters() {
	limitersMu.Lock()
	defer limitersMu.Unlock()

	// 先停止旧的限流器
	if loginLimiter != nil {
		loginLimiter.Stop()
	}
	if apiLimiter != nil {
		apiLimiter.Stop()
	}
	if uploadLimiter != nil {
		uploadLimiter.Stop()
	}

	// 内存限流器（单实例）
	loginLimiter = NewRateLimiter(10, time.Minute)
	apiLimiter = NewRateLimiter(100, time.Minute)
	uploadLimiter = NewRateLimiter(20, time.Minute)
}

// StopRateLimiters 停止所有限速器（应用关闭时调用）
func StopRateLimiters() {
	limitersMu.Lock()
	defer limitersMu.Unlock()

	if loginLimiter != nil {
		loginLimiter.Stop()
		loginLimiter = nil
	}
	if apiLimiter != nil {
		apiLimiter.Stop()
		apiLimiter = nil
	}
	if uploadLimiter != nil {
		uploadLimiter.Stop()
		uploadLimiter = nil
	}
}

// RateLimit 通用速率限制中间件（内存版）
func RateLimit(limiter *RateLimiter) gin.HandlerFunc {
	return func(c *gin.Context) {
		ip := c.ClientIP()
		if !limiter.Allow(ip) {
			response.RateLimit(c)
			c.Abort()
			return
		}
		c.Next()
	}
}

// LoginRateLimit 登录接口速率限制
func LoginRateLimit() gin.HandlerFunc {
	limitersMu.Lock()
	if loginLimiter == nil {
		loginLimiter = NewRateLimiter(10, time.Minute)
	}
	limiter := loginLimiter
	limitersMu.Unlock()

	return RateLimit(limiter)
}

// APIRateLimit 普通 API 速率限制
func APIRateLimit() gin.HandlerFunc {
	limitersMu.Lock()
	if apiLimiter == nil {
		apiLimiter = NewRateLimiter(100, time.Minute)
	}
	limiter := apiLimiter
	limitersMu.Unlock()

	return RateLimit(limiter)
}

// UploadRateLimit 上传接口速率限制
func UploadRateLimit() gin.HandlerFunc {
	limitersMu.Lock()
	if uploadLimiter == nil {
		uploadLimiter = NewRateLimiter(20, time.Minute)
	}
	limiter := uploadLimiter
	limitersMu.Unlock()

	return RateLimit(limiter)
}

// CustomRateLimit 自定义速率限制（内存版）
func CustomRateLimit(rate int, window time.Duration) gin.HandlerFunc {
	limiter := NewRateLimiter(rate, window)
	return RateLimit(limiter)
}

// ===== Redis 分布式限流中间件 =====

// RedisRateLimit Redis 分布式限流中间件
// 参数: keyPrefix 键名前缀（如 "login_limit"），rate 每分钟请求数
func RedisRateLimit(keyPrefix string, rate int) gin.HandlerFunc {
	limiter := NewRedisRateLimiter(keyPrefix, rate, time.Minute)

	return func(c *gin.Context) {
		identifier := c.ClientIP()

		// 可选：使用用户ID作为标识（登录后）
		// userID := GetUserID(c)
		// if userID > 0 {
		//     identifier = fmt.Sprintf("user:%d", userID)
		// }

		allowed, err := limiter.Allow(c.Request.Context(), identifier)
		if err != nil {
			// Redis 错误时允许请求，避免影响业务
			c.Next()
			return
		}

		if !allowed {
			response.RateLimit(c)
			c.Abort()
			return
		}

		c.Next()
	}
}

// RedisRateLimitWithIdentifier 自定义标识的 Redis 分布式限流
// 参数: keyPrefix 键名前缀，rate 每分钟请求数，identifierFunc 标识获取函数
func RedisRateLimitWithIdentifier(keyPrefix string, rate int, identifierFunc func(c *gin.Context) string) gin.HandlerFunc {
	limiter := NewRedisRateLimiter(keyPrefix, rate, time.Minute)

	return func(c *gin.Context) {
		identifier := identifierFunc(c)
		if identifier == "" {
			identifier = c.ClientIP()
		}

		allowed, err := limiter.Allow(c.Request.Context(), identifier)
		if err != nil {
			c.Next()
			return
		}

		if !allowed {
			response.RateLimit(c)
			c.Abort()
			return
		}

		c.Next()
	}
}

// LoginRedisRateLimit 登录接口 Redis 分布式限流
func LoginRedisRateLimit() gin.HandlerFunc {
	return RedisRateLimit("login_limit", 10)
}

// APIRedisRateLimit API Redis 分布式限流
func APIRedisRateLimit() gin.HandlerFunc {
	return RedisRateLimit("api_limit", 100)
}

// UploadRedisRateLimit 上传接口 Redis 分布式限流
func UploadRedisRateLimit() gin.HandlerFunc {
	return RedisRateLimit("upload_limit", 20)
}

// CustomRedisRateLimit 自定义 Redis 分布式限流
func CustomRedisRateLimit(keyPrefix string, rate int, window time.Duration) gin.HandlerFunc {
	limiter := NewRedisRateLimiter(keyPrefix, rate, window)

	return func(c *gin.Context) {
		identifier := c.ClientIP()

		allowed, err := limiter.Allow(c.Request.Context(), identifier)
		if err != nil {
			c.Next()
			return
		}

		if !allowed {
			response.RateLimit(c)
			c.Abort()
			return
		}

		c.Next()
	}
}