wangjia
7bbc944ae2
Deploy Server / release-deploy-server (push) Successful in 51s
服务端安全加固：多维限流（按 IP/按门店）+ 敏感接口独立速率上限抵御 DDoS/刷接口； 登录暴力破解新增按来源 IP 锁定；反代后正确识别真实客户端 IP； 门店 custom_fields 轻量配置（录入默认值）透传保存。 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Y2Wdwo7SmgBJU37cBrkhPK
128 lines
3.5 KiB
Go
128 lines
3.5 KiB
Go
package middleware
|
||
|
||
import (
|
||
"net/http"
|
||
"strconv"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
"golang.org/x/time/rate"
|
||
|
||
"github.com/wangjia/jiu/backend/config"
|
||
)
|
||
|
||
// 限流器内存条目的清理参数：每 5 分钟扫一次，淘汰超过 10 分钟未活动的 key。
|
||
// 保证 map 不随攻击者构造的随机 key（IP/shop）无限增长。
|
||
const (
|
||
rateLimitSweep = 5 * time.Minute
|
||
rateLimitIdleTTL = 10 * time.Minute
|
||
rateLimitRetryHdr = "60" // Retry-After 秒数（提示客户端退避）
|
||
)
|
||
|
||
// keyedLimiter 按任意字符串 key（IP 或 shop_id）维护独立令牌桶，内存有界（带 janitor）。
|
||
// 单实例进程内状态，重启即清零；多实例水平扩展时需改为 Redis（见方案「暂不做」）。
|
||
type keyedLimiter struct {
|
||
mu sync.Mutex
|
||
entries map[string]*limiterBucket
|
||
r rate.Limit
|
||
burst int
|
||
}
|
||
|
||
type limiterBucket struct {
|
||
lim *rate.Limiter
|
||
lastSeen time.Time
|
||
}
|
||
|
||
func newKeyedLimiter(r rate.Limit, burst int) *keyedLimiter {
|
||
kl := &keyedLimiter{entries: map[string]*limiterBucket{}, r: r, burst: burst}
|
||
go kl.janitor()
|
||
return kl
|
||
}
|
||
|
||
// get 取（或惰性创建）该 key 的令牌桶并刷新活动时间。
|
||
func (kl *keyedLimiter) get(key string) *rate.Limiter {
|
||
kl.mu.Lock()
|
||
defer kl.mu.Unlock()
|
||
b := kl.entries[key]
|
||
if b == nil {
|
||
b = &limiterBucket{lim: rate.NewLimiter(kl.r, kl.burst)}
|
||
kl.entries[key] = b
|
||
}
|
||
b.lastSeen = time.Now()
|
||
return b.lim
|
||
}
|
||
|
||
func (kl *keyedLimiter) janitor() {
|
||
t := time.NewTicker(rateLimitSweep)
|
||
defer t.Stop()
|
||
for range t.C {
|
||
kl.sweep(rateLimitIdleTTL)
|
||
}
|
||
}
|
||
|
||
// sweep 淘汰超过 ttl 未活动的 key。拆出便于测试。
|
||
func (kl *keyedLimiter) sweep(ttl time.Duration) {
|
||
now := time.Now()
|
||
kl.mu.Lock()
|
||
defer kl.mu.Unlock()
|
||
for k, b := range kl.entries {
|
||
if now.Sub(b.lastSeen) > ttl {
|
||
delete(kl.entries, k)
|
||
}
|
||
}
|
||
}
|
||
|
||
// 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 {
|
||
kl := newKeyedLimiter(r, burst)
|
||
return func(c *gin.Context) {
|
||
if !config.C.RateLimit.Enabled {
|
||
c.Next()
|
||
return
|
||
}
|
||
key := keyFn(c)
|
||
if key == "" {
|
||
c.Next()
|
||
return
|
||
}
|
||
if !kl.get(key).Allow() {
|
||
c.Header("Retry-After", rateLimitRetryHdr)
|
||
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)
|
||
})
|
||
}
|