[feature] migrate to monorepo

backend/internal/middleware/ratelimit.go

@@ -0,0 +1,107 @@
+package middleware
+
+import (
+	"net/http"
+	"sync"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"golang.org/x/time/rate"
+	"tss-rocks-be/internal/types"
+)
+
+// ipLimiter IP限流器
+type ipLimiter struct {
+	limiter   *rate.Limiter
+	lastSeen  time.Time
+}
+
+// rateLimiter 限流器管理器
+type rateLimiter struct {
+	ips    map[string]*ipLimiter
+	mu     sync.RWMutex
+	config *types.RateLimitConfig
+	routes map[string]*rate.Limiter
+}
+
+// newRateLimiter 创建新的限流器
+func newRateLimiter(config *types.RateLimitConfig) *rateLimiter {
+	// 初始化路由限流器
+	routes := make(map[string]*rate.Limiter)
+	for path, cfg := range config.RouteRates {
+		routes[path] = rate.NewLimiter(rate.Limit(cfg.Rate), cfg.Burst)
+	}
+
+	rl := &rateLimiter{
+		ips:    make(map[string]*ipLimiter),
+		config: config,
+		routes: routes,
+	}
+
+	// 启动清理过期IP限流器的goroutine
+	go rl.cleanupIPLimiters()
+
+	return rl
+}
+
+// cleanupIPLimiters 清理过期的IP限流器
+func (rl *rateLimiter) cleanupIPLimiters() {
+	for {
+		time.Sleep(time.Hour) // 每小时清理一次
+
+		rl.mu.Lock()
+		for ip, limiter := range rl.ips {
+			if time.Since(limiter.lastSeen) > time.Hour {
+				delete(rl.ips, ip)
+			}
+		}
+		rl.mu.Unlock()
+	}
+}
+
+// getLimiter 获取IP限流器
+func (rl *rateLimiter) getLimiter(ip string) *rate.Limiter {
+	rl.mu.Lock()
+	defer rl.mu.Unlock()
+
+	v, exists := rl.ips[ip]
+	if !exists {
+		limiter := rate.NewLimiter(rate.Limit(rl.config.IPRate), rl.config.IPBurst)
+		rl.ips[ip] = &ipLimiter{limiter: limiter, lastSeen: time.Now()}
+		return limiter
+	}
+
+	v.lastSeen = time.Now()
+	return v.limiter
+}
+
+// RateLimit 创建限流中间件
+func RateLimit(config *types.RateLimitConfig) gin.HandlerFunc {
+	rl := newRateLimiter(config)
+
+	return func(c *gin.Context) {
+		// 检查路由限流
+		path := c.Request.URL.Path
+		if limiter, ok := rl.routes[path]; ok {
+			if !limiter.Allow() {
+				c.JSON(http.StatusTooManyRequests, gin.H{
+					"error": "too many requests for this route",
+				})
+				c.Abort()
+				return
+			}
+		}
+
+		// 检查IP限流
+		limiter := rl.getLimiter(c.ClientIP())
+		if !limiter.Allow() {
+			c.JSON(http.StatusTooManyRequests, gin.H{
+				"error": "too many requests from this IP",
+			})
+			c.Abort()
+			return
+		}
+
+		c.Next()
+	}
+}