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()
|
|||
|
|
}
|
|||
|
|
}
|