wangjia
118 lines
3.6 KiB
Go
118 lines
3.6 KiB
Go
// Package middleware holds gin-agnostic-ish HTTP cross-cutting concerns
|
|
// (currently: per-IP rate limiting) shared by the v2 gateway router.
|
|
package middleware
|
|
|
|
import (
|
|
"net/http"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
// bucket 单个 IP 的令牌桶状态。
|
|
type bucket struct {
|
|
mu sync.Mutex
|
|
tokens float64
|
|
updated time.Time
|
|
}
|
|
|
|
// maxBuckets 是 buckets map 的硬上限——防御性 backstop:哪怕 SetTrustedProxies(nil)
|
|
// 已经堵死"伪造 X-Forwarded-For 绕过限流"这条路,socket 对端 IP 本身也可能被
|
|
// 大量真实/伪装源地址(如僵尸网络)刷出海量不同 IP,惰性 sweep(见 sweepLocked)
|
|
// 只回收超过 10min 不活跃的桶,短时间集中攻击仍可能在 sweep 生效前把 map 撑爆。
|
|
// 100k 条 * 单桶几十字节量级,内存占用可控(几 MB),对当前单机部署(阿里云,
|
|
// ~1-2GB 内存假设,与 pangolin 单机 Runbook 同量级)足够安全。超限时宁可误限
|
|
// (安全侧),也不无界建桶 OOM。
|
|
const maxBuckets = 100_000
|
|
|
|
// IPRateLimiter 简单内存令牌桶,按客户端 IP 分桶限速——单实例部署足够(jiu 反馈波,
|
|
// P3 v2 限流纵深)。多实例横向扩展需换外置存储(Redis INCR+EXPIRE 等)统一计数,
|
|
// 当前 pay 只跑单进程,内存方案够用且零额外依赖。
|
|
type IPRateLimiter struct {
|
|
mu sync.Mutex
|
|
buckets map[string]*bucket
|
|
rate float64 // 每秒补充的令牌数
|
|
burst float64 // 桶容量(=突发上限,取每分钟额度)
|
|
sweepCounter uint64
|
|
}
|
|
|
|
// NewIPRateLimiter 按每分钟额度构造限流器;perMinute<=0 时按 30 兜底(与 config 里
|
|
// "零值=默认开、默认 30/min" 的约定一致,调用方无需重复判空)。
|
|
func NewIPRateLimiter(perMinute int) *IPRateLimiter {
|
|
if perMinute <= 0 {
|
|
perMinute = 30
|
|
}
|
|
return &IPRateLimiter{
|
|
buckets: make(map[string]*bucket),
|
|
rate: float64(perMinute) / 60.0,
|
|
burst: float64(perMinute),
|
|
}
|
|
}
|
|
|
|
// allow 令牌桶判定:按经过时间匀速补充令牌(封顶 burst),够 1 个则放行并扣 1。
|
|
func (l *IPRateLimiter) allow(ip string) bool {
|
|
now := time.Now()
|
|
l.mu.Lock()
|
|
b, ok := l.buckets[ip]
|
|
if !ok {
|
|
if len(l.buckets) >= maxBuckets {
|
|
// map 已到硬上限:先尝试一次 sweep 腾陈旧桶的位置。
|
|
l.sweepLocked(now)
|
|
}
|
|
if len(l.buckets) >= maxBuckets {
|
|
// sweep 后仍超限(短时间内大量活跃桶):拒绝新建桶,按限流命中处理
|
|
// (429)而不是无界增长——安全侧,宁可误限未知新 IP 也不 OOM。
|
|
l.mu.Unlock()
|
|
return false
|
|
}
|
|
b = &bucket{tokens: l.burst, updated: now}
|
|
l.buckets[ip] = b
|
|
}
|
|
l.sweepCounter++
|
|
if l.sweepCounter%1000 == 0 {
|
|
l.sweepLocked(now)
|
|
}
|
|
l.mu.Unlock()
|
|
|
|
b.mu.Lock()
|
|
defer b.mu.Unlock()
|
|
elapsed := now.Sub(b.updated).Seconds()
|
|
b.tokens += elapsed * l.rate
|
|
if b.tokens > l.burst {
|
|
b.tokens = l.burst
|
|
}
|
|
b.updated = now
|
|
if b.tokens < 1 {
|
|
return false
|
|
}
|
|
b.tokens--
|
|
return true
|
|
}
|
|
|
|
// sweepLocked 惰性回收长期不活跃(>10min 无请求)的 IP 桶,防长期运行的单实例
|
|
// 在"短连接、海量不同 IP"场景下内存无界增长。调用方须已持 l.mu。
|
|
func (l *IPRateLimiter) sweepLocked(now time.Time) {
|
|
for ip, b := range l.buckets {
|
|
b.mu.Lock()
|
|
stale := now.Sub(b.updated) > 10*time.Minute
|
|
b.mu.Unlock()
|
|
if stale {
|
|
delete(l.buckets, ip)
|
|
}
|
|
}
|
|
}
|
|
|
|
// Gin 返回 gin 中间件:超速回 429,不中断其它路由。
|
|
func (l *IPRateLimiter) Gin() gin.HandlerFunc {
|
|
return func(c *gin.Context) {
|
|
if !l.allow(c.ClientIP()) {
|
|
c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
|
|
"code": "rate_limited", "message": "请求过于频繁,请稍后重试",
|
|
})
|
|
return
|
|
}
|
|
c.Next()
|
|
}
|
|
}
|