dev 分支暂存

backend/gateway/internal/middleware/ratelimit.go

@@ -0,0 +1,102 @@
+package middleware
+
+import (
+	"net/http"
+	"sync"
+	"time"
+
+	"github.com/gin-gonic/gin"
+)
+
+// RateLimiter 基于内存令牌桶的限流中间件
+// MVP阶段使用内存实现，后续可迁移到Redis
+type RateLimiter struct {
+	mu      sync.Mutex
+	buckets map[string]*tokenBucket
+	rate    float64 // 每秒生成的令牌数
+	burst   int     // 桶容量
+}
+
+type tokenBucket struct {
+	tokens   float64
+	lastTime time.Time
+}
+
+// NewRateLimiter 创建限流器
+func NewRateLimiter(rate float64, burst int) *RateLimiter {
+	rl := &RateLimiter{
+		buckets: make(map[string]*tokenBucket),
+		rate:    rate,
+		burst:   burst,
+	}
+
+	// 定期清理过期桶
+	go rl.cleanup()
+
+	return rl
+}
+
+// Handler 返回Gin中间件
+func (rl *RateLimiter) Handler() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		key := c.ClientIP() // 按IP限流
+
+		if !rl.allow(key) {
+			c.JSON(http.StatusTooManyRequests, gin.H{
+				"error": "请求过于频繁，请稍后再试",
+			})
+			c.Abort()
+			return
+		}
+
+		c.Next()
+	}
+}
+
+func (rl *RateLimiter) allow(key string) bool {
+	rl.mu.Lock()
+	defer rl.mu.Unlock()
+
+	bucket, ok := rl.buckets[key]
+	now := time.Now()
+
+	if !ok {
+		rl.buckets[key] = &tokenBucket{
+			tokens:   float64(rl.burst) - 1,
+			lastTime: now,
+		}
+		return true
+	}
+
+	// 补充令牌
+	elapsed := now.Sub(bucket.lastTime).Seconds()
+	bucket.tokens += elapsed * rl.rate
+	if bucket.tokens > float64(rl.burst) {
+		bucket.tokens = float64(rl.burst)
+	}
+	bucket.lastTime = now
+
+	// 消耗令牌
+	if bucket.tokens >= 1 {
+		bucket.tokens--
+		return true
+	}
+
+	return false
+}
+
+// cleanup 定期清理长时间未使用的桶
+func (rl *RateLimiter) cleanup() {
+	for {
+		time.Sleep(5 * time.Minute)
+
+		rl.mu.Lock()
+		cutoff := time.Now().Add(-10 * time.Minute)
+		for key, bucket := range rl.buckets {
+			if bucket.lastTime.Before(cutoff) {
+				delete(rl.buckets, key)
+			}
+		}
+		rl.mu.Unlock()
+	}
+}