package ratelimit import ( "fmt" "net/http" "github.com/gin-gonic/gin" "github.com/hhs/camtalk/internal/trace" ) // Middleware 返回 Gin 中间件,按 key 维度限流。 // keyFunc 从请求中提取限流 key(如 IP、用户 ID)。 func Middleware(limiter Limiter, keyFunc func(*gin.Context) string) gin.HandlerFunc { return func(c *gin.Context) { if limiter == nil { c.Next() return } key := keyFunc(c) if key == "" { // key 为空时跳过限流 c.Next() return } allowed, retryAfter := limiter.Allow(c.Request.Context(), key) if !allowed { log := trace.FromContext(c.Request.Context()) log.Warnw("rate limited", "client_ip", c.ClientIP(), "path", c.Request.URL.Path, "limit_key", key, "retry_after_sec", int(retryAfter.Seconds()+0.5)) // 设置 Retry-After header(秒) c.Header("Retry-After", fmt.Sprintf("%d", int(retryAfter.Seconds()+0.5))) c.JSON(http.StatusTooManyRequests, gin.H{ "code": "RATE_LIMITED", "message": fmt.Sprintf("too many requests, retry after %s", retryAfter.Round(1)), }) c.Abort() return } c.Next() } }