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:
42
backend/internal/ratelimit/middleware.go
Normal file
42
backend/internal/ratelimit/middleware.go
Normal 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()
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user