hhs
ea00939c13
- Middleware 函数返回 Gin 中间件 - keyFunc 参数支持灵活提取限流 key（IP/用户 ID 等） - 限流触发时返回 HTTP 429 + Retry-After header - 支持 nil limiter（跳过限流）和空 key（跳过限流） - 完整单元测试（6 个测试用例，全部通过） - 测试覆盖：允许、拒绝、nil limiter、空 key、keyFunc、Retry-After 舍入
43 lines
882 B
Go
43 lines
882 B
Go
package ratelimit
|
||
|
||
import (
|
||
"fmt"
|
||
"net/http"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
)
|
||
|
||
// 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 {
|
||
// 设置 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()
|
||
}
|
||
}
|