98 lines
2.4 KiB
Go
98 lines
2.4 KiB
Go
package middleware
|
||
|
||
import (
|
||
"net/http"
|
||
"sync"
|
||
"time"
|
||
|
||
"fashionapi/internal/config"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
)
|
||
|
||
// rateLimiter 基于令牌桶的按 IP 限流（进程内，单实例足够；多实例需换集中式如 redis）。
|
||
//
|
||
// 设计取舍：公开接口「首页公开、翻页需登录」已天然抬高批量枚举门槛；本中间件作为
|
||
// 兜底，按客户端 IP 限制总量，兜住 casual 爬虫与突发流量，避免单 IP 打垮服务。
|
||
type rateLimiter struct {
|
||
mu sync.Mutex
|
||
visitors map[string]*visitor
|
||
rate float64 // 每秒补充令牌数
|
||
burst float64 // 桶容量
|
||
lastSeen map[string]time.Time
|
||
}
|
||
|
||
type visitor struct {
|
||
tokens float64
|
||
last time.Time
|
||
}
|
||
|
||
func newRateLimiter(rps, burst int) *rateLimiter {
|
||
rl := &rateLimiter{
|
||
visitors: make(map[string]*visitor),
|
||
lastSeen: make(map[string]time.Time),
|
||
rate: float64(rps),
|
||
burst: float64(burst),
|
||
}
|
||
go rl.cleanup(10 * time.Minute)
|
||
return rl
|
||
}
|
||
|
||
// cleanup 周期清理长时间无活动的客户端，避免 map 无限增长。
|
||
func (rl *rateLimiter) cleanup(interval time.Duration) {
|
||
ticker := time.NewTicker(interval)
|
||
defer ticker.Stop()
|
||
for range ticker.C {
|
||
rl.mu.Lock()
|
||
for ip, t := range rl.lastSeen {
|
||
if time.Since(t) > 3*interval {
|
||
delete(rl.visitors, ip)
|
||
delete(rl.lastSeen, ip)
|
||
}
|
||
}
|
||
rl.mu.Unlock()
|
||
}
|
||
}
|
||
|
||
func (rl *rateLimiter) allow(ip string) bool {
|
||
rl.mu.Lock()
|
||
defer rl.mu.Unlock()
|
||
now := time.Now()
|
||
v, ok := rl.visitors[ip]
|
||
if !ok {
|
||
// 首次访问直接发放 burst-1 个令牌（允许一个突发请求），后续按速率补充
|
||
rl.visitors[ip] = &visitor{tokens: rl.burst - 1, last: now}
|
||
rl.lastSeen[ip] = now
|
||
return true
|
||
}
|
||
elapsed := now.Sub(v.last).Seconds()
|
||
v.tokens += elapsed * rl.rate
|
||
if v.tokens > rl.burst {
|
||
v.tokens = rl.burst
|
||
}
|
||
v.last = now
|
||
rl.lastSeen[ip] = now
|
||
if v.tokens >= 1 {
|
||
v.tokens--
|
||
return true
|
||
}
|
||
return false
|
||
}
|
||
|
||
// RateLimit 按 IP 的速率限制中间件。Enabled=false 时为 noop（便于灰度与本地开发）。
|
||
func RateLimit(cfg config.RateLimitConfig) gin.HandlerFunc {
|
||
if !cfg.Enabled || cfg.RPS <= 0 {
|
||
return func(c *gin.Context) { c.Next() }
|
||
}
|
||
rl := newRateLimiter(cfg.RPS, cfg.Burst)
|
||
return func(c *gin.Context) {
|
||
if !rl.allow(c.ClientIP()) {
|
||
c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
|
||
"error": "too many requests",
|
||
})
|
||
return
|
||
}
|
||
c.Next()
|
||
}
|
||
}
|