package middleware import ( "math" "net/http" "strconv" "github.com/gin-gonic/gin" "golang.org/x/time/rate" "github.com/wangjia/jiu/backend/config" "github.com/wangjia/jiu/backend/internal/ratelimit" ) // 限流状态存储在 internal/ratelimit(默认内存;REDIS_ADDR 配置后外置 redis, // 跨重启保状态、支持多实例。2026-07 外置改造,原 keyedLimiter 迁入该包)。 // PerMinute 把「每分钟 n 次」转成 rate.Limit(令牌/秒)。 func PerMinute(n int) rate.Limit { return rate.Limit(float64(n) / 60.0) } // PerSecond 把「每秒 n 次」转成 rate.Limit。 func PerSecond(n int) rate.Limit { return rate.Limit(n) } // rateLimit 通用工厂:keyFn 抽取限流维度的 key(返回空串表示无法判定 → 放行,不误伤)。 // config.C.RateLimit.Enabled=false 时整体放行(应急/测试开关)。 func rateLimit(r rate.Limit, burst int, keyFn func(*gin.Context) string) gin.HandlerFunc { lim := ratelimit.Default().NewLimiter(r, burst) return func(c *gin.Context) { if !config.C.RateLimit.Enabled { c.Next() return } key := keyFn(c) if key == "" { c.Next() return } if res := lim.Allow(key); !res.Allowed { retry := int(math.Ceil(res.RetryAfter.Seconds())) if retry <= 0 { retry = 60 } c.Header("Retry-After", strconv.Itoa(retry)) c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{"error": "请求过于频繁,请稍后再试"}) return } c.Next() } } // RateLimitByIP 按真实客户端 IP 限流(依赖 main.go 的 SetTrustedProxies + RemoteIPHeaders, // 使 c.ClientIP() 返回不可伪造的真实 IP)。 func RateLimitByIP(r rate.Limit, burst int) gin.HandlerFunc { return rateLimit(r, burst, func(c *gin.Context) string { return "ip:" + c.ClientIP() }) } // RateLimitByShop 按门店(shop_id,取自 JWT)限流,须挂在 JWT 中间件之后。 // 未取到 shop_id 时放行(交由 JWT 中间件拦截非法 token)。 func RateLimitByShop(r rate.Limit, burst int) gin.HandlerFunc { return rateLimit(r, burst, func(c *gin.Context) string { id := GetShopID(c) if id == 0 { return "" } return "shop:" + strconv.FormatUint(id, 10) }) }