wangjia
3c687e5e0b
新包 internal/ratelimit：内存/redis(GCRA+Lua) 双实现 + 出错逐调用降级内存 （fail-open 到内存不 fail-closed）。REDIS_ADDR 空=内存模式，行为与既往一致； 配置后跨重启保状态、支持多实例。miniredis 全覆盖测试，零真实外部依赖。 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
74 lines
2.2 KiB
Go
74 lines
2.2 KiB
Go
package middleware
|
||
|
||
import (
|
||
"math"
|
||
"net/http"
|
||
"strconv"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
"golang.org/x/time/rate"
|
||
|
||
"github.com/wangjia/jiu/backend/config"
|
||
"github.com/wangjia/jiu/backend/internal/ratelimit"
|
||
)
|
||
|
||
// 限流状态存储在 internal/ratelimit（默认内存；REDIS_ADDR 配置后外置 redis，
|
||
// 跨重启保状态、支持多实例。2026-07 外置改造，原 keyedLimiter 迁入该包）。
|
||
|
||
// PerMinute 把「每分钟 n 次」转成 rate.Limit（令牌/秒）。
|
||
func PerMinute(n int) rate.Limit {
|
||
return rate.Limit(float64(n) / 60.0)
|
||
}
|
||
|
||
// PerSecond 把「每秒 n 次」转成 rate.Limit。
|
||
func PerSecond(n int) rate.Limit {
|
||
return rate.Limit(n)
|
||
}
|
||
|
||
// rateLimit 通用工厂：keyFn 抽取限流维度的 key（返回空串表示无法判定 → 放行，不误伤）。
|
||
// config.C.RateLimit.Enabled=false 时整体放行（应急/测试开关）。
|
||
func rateLimit(r rate.Limit, burst int, keyFn func(*gin.Context) string) gin.HandlerFunc {
|
||
lim := ratelimit.Default().NewLimiter(r, burst)
|
||
return func(c *gin.Context) {
|
||
if !config.C.RateLimit.Enabled {
|
||
c.Next()
|
||
return
|
||
}
|
||
key := keyFn(c)
|
||
if key == "" {
|
||
c.Next()
|
||
return
|
||
}
|
||
if res := lim.Allow(key); !res.Allowed {
|
||
retry := int(math.Ceil(res.RetryAfter.Seconds()))
|
||
if retry <= 0 {
|
||
retry = 60
|
||
}
|
||
c.Header("Retry-After", strconv.Itoa(retry))
|
||
c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{"error": "请求过于频繁，请稍后再试"})
|
||
return
|
||
}
|
||
c.Next()
|
||
}
|
||
}
|
||
|
||
// RateLimitByIP 按真实客户端 IP 限流（依赖 main.go 的 SetTrustedProxies + RemoteIPHeaders，
|
||
// 使 c.ClientIP() 返回不可伪造的真实 IP）。
|
||
func RateLimitByIP(r rate.Limit, burst int) gin.HandlerFunc {
|
||
return rateLimit(r, burst, func(c *gin.Context) string {
|
||
return "ip:" + c.ClientIP()
|
||
})
|
||
}
|
||
|
||
// RateLimitByShop 按门店（shop_id，取自 JWT）限流，须挂在 JWT 中间件之后。
|
||
// 未取到 shop_id 时放行（交由 JWT 中间件拦截非法 token）。
|
||
func RateLimitByShop(r rate.Limit, burst int) gin.HandlerFunc {
|
||
return rateLimit(r, burst, func(c *gin.Context) string {
|
||
id := GetShopID(c)
|
||
if id == 0 {
|
||
return ""
|
||
}
|
||
return "shop:" + strconv.FormatUint(id, 10)
|
||
})
|
||
}
|