dual check v0.1

server/internal/middleware/ratelimit.go

@@ -0,0 +1,64 @@
+package middleware
+
+import (
+	"net/http"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"github.com/ulule/limiter/v3"
+	"github.com/ulule/limiter/v3/drivers/store/memory"
+)
+
+// RateLimitConfig 限流配置
+type RateLimitConfig struct {
+	GlobalRPS  float64
+	PerIPRPS   float64
+	PerIPBurst int
+}
+
+// RateLimiter 双维度限流：全局令牌桶 + 单 IP 令牌桶
+// 开发/测试环境可用内存存储，生产环境建议换 Redis 存储
+func RateLimiter(cfg RateLimitConfig) gin.HandlerFunc {
+	store := memory.NewStore()
+
+	// 全局限流
+	globalRate := limiter.Rate{
+		Period: time.Second,
+		Limit:  int64(cfg.GlobalRPS),
+	}
+	globalLimiter := limiter.New(store, globalRate)
+
+	// 单 IP 限流
+	perIPRate := limiter.Rate{
+		Period: time.Second,
+		Limit:  int64(cfg.PerIPRPS),
+	}
+	perIPLimiter := limiter.New(store, perIPRate)
+
+	return func(c *gin.Context) {
+		// 1. 检查全局限流
+		globalCtx, err := globalLimiter.Get(c.Request.Context(), "global")
+		if err != nil || globalCtx.Reached {
+			c.JSON(http.StatusTooManyRequests, gin.H{
+				"code":    4290,
+				"message": "server busy, please retry later",
+			})
+			c.Abort()
+			return
+		}
+
+		// 2. 检查单 IP 限流
+		ip := c.ClientIP()
+		ipCtx, err := perIPLimiter.Get(c.Request.Context(), "ip:"+ip)
+		if err != nil || ipCtx.Reached {
+			c.JSON(http.StatusTooManyRequests, gin.H{
+				"code":    4291,
+				"message": "too many requests from this IP",
+			})
+			c.Abort()
+			return
+		}
+
+		c.Next()
+	}
+}