feat: 实现 Gin 限流中间件

- Middleware 函数返回 Gin 中间件
- keyFunc 参数支持灵活提取限流 key(IP/用户 ID 等)
- 限流触发时返回 HTTP 429 + Retry-After header
- 支持 nil limiter(跳过限流)和空 key(跳过限流)
- 完整单元测试(6 个测试用例,全部通过)
- 测试覆盖:允许、拒绝、nil limiter、空 key、keyFunc、Retry-After 舍入
This commit is contained in:
hhs
2026-06-20 23:56:44 +08:00
parent b74fb3564d
commit ea00939c13
2 changed files with 238 additions and 0 deletions

View File

@@ -0,0 +1,42 @@
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()
}
}