package route import ( "fmt" "net/http" "strconv" "sync" "time" "github.com/flamego/flamego" "github.com/flamego/session" "github.com/tamakyi/TamaBox/internal/i18n" "github.com/tamakyi/TamaBox/internal/shared" "github.com/tamakyi/TamaBox/route/vote" ) // tooManyRequests 是 429 的唯一出口。 // // 原来两处中间件各写一遍,只有一句英文硬编码 —— 既没走 i18n(中文站给访客看英文), // 也没有 Retry-After,客户端只能靠猜什么时候可以重试(多数会立刻重试,把限流放大)。 // Retry-After 取窗口长度:滑动窗口下真实可重试时间可能更早,这里是安全上界。 // // 语言固定用站点默认:限流中间件跑在 Contexter 之前,此时拿不到已解析的界面语言, // 而 429 的 body 通常由前端 JS 展示、很少直接被阅读。 func tooManyRequests(ctx flamego.Context, window time.Duration) { w := ctx.ResponseWriter() retry := int(window.Seconds()) if retry < 1 { retry = 1 } w.Header().Set("Retry-After", strconv.Itoa(retry)) w.Header().Set("Content-Type", "text/plain; charset=utf-8") w.WriteHeader(http.StatusTooManyRequests) _, _ = w.Write([]byte(i18n.T(i18n.Default, "common.tooManyRequests"))) } // rateLimiter provides a simple in-memory rate limiter using a sliding window. type rateLimiter struct { mu sync.Mutex visitors map[string][]time.Time limit int window time.Duration } // ---------- 全局清理 ---------- // // 所有 rateLimiter 实例注册到全局注册表,由单个常驻 goroutine 统一清理, // 避免每个中间件实例各自启动一个永久 goroutine(旧实现会随路由数量线性泄漏 goroutine)。 var ( rateLimitersMu sync.Mutex rateLimiters = make(map[*rateLimiter]struct{}) rateLimiterClean sync.Once ) func registerRateLimiter(rl *rateLimiter) { rateLimitersMu.Lock() rateLimiters[rl] = struct{}{} rateLimitersMu.Unlock() // 仅启动一次全局清理循环 rateLimiterClean.Do(func() { go func() { for { time.Sleep(5 * time.Minute) rateLimitersMu.Lock() for rl := range rateLimiters { rl.cleanup() } rateLimitersMu.Unlock() } }() }) } func newRateLimiter(limit int, window time.Duration) *rateLimiter { rl := &rateLimiter{ visitors: make(map[string][]time.Time), limit: limit, window: window, } registerRateLimiter(rl) return rl } func (rl *rateLimiter) allow(ip string) bool { rl.mu.Lock() defer rl.mu.Unlock() // 恶意伪造 X-Forwarded-For 可撑爆 visitors map,超过阈值时快速清理 // 窗口内无活动的 key(其最新时间戳已过期) if len(rl.visitors) > 10000 { cutoff2 := time.Now().Add(-rl.window) for ip, ts := range rl.visitors { if len(ts) == 0 || ts[len(ts)-1].Before(cutoff2) { delete(rl.visitors, ip) } } } now := time.Now() cutoff := now.Add(-rl.window) times := rl.visitors[ip] // Remove expired entries. valid := times[:0] for _, t := range times { if t.After(cutoff) { valid = append(valid, t) } } if len(valid) >= rl.limit { rl.visitors[ip] = valid return false } valid = append(valid, now) rl.visitors[ip] = valid return true } func (rl *rateLimiter) cleanup() { rl.mu.Lock() defer rl.mu.Unlock() cutoff := time.Now().Add(-rl.window) for ip, times := range rl.visitors { valid := times[:0] for _, t := range times { if t.After(cutoff) { valid = append(valid, t) } } if len(valid) == 0 { delete(rl.visitors, ip) } else { rl.visitors[ip] = valid } } } // rateLimitMiddleware returns a Flamego handler that applies rate limiting. // limit: max requests per window; window: time window duration. // // 匿名 / 登录分级限流:登录用户(session 中存在 uid)使用 3 倍额度, // 未登录访客使用基础额度。分级使用独立限流器统计,互不污染。 func rateLimitMiddleware(limit int, window time.Duration) flamego.Handler { anonymousRL := newRateLimiter(limit, window) loggedRL := newRateLimiter(limit*3, window) return func(ctx flamego.Context, s session.Session) { // 统一通过 shared.ResolveClientIP 解析,trust_proxy=false 时忽略伪造的代理头。 // 再经 shared.LimitKey 归一化:IPv4 原样,IPv6 聚合到 /64, // 防止在单个 /64 内轮换海量地址绕过阈值(详见 LimitKey 注释)。 ip := shared.LimitKey(shared.ResolveClientIP(ctx.Request().Request)) rl := anonymousRL if _, ok := s.Get("uid").(uint); ok { rl = loggedRL } if !rl.allow(ip) { tooManyRequests(ctx, window) return } } } // voteSubmitRateLimitMiddleware 对投票提交同时执行两层限流: // 匿名身份层按 token 计数,IP 层使用更高阈值作为批量换 token 的风控兜底。 // 这样共享公网出口的多个匿名浏览器不会共享低阈值,同时仍能限制异常突发流量。 func voteSubmitRateLimitMiddleware(limit int, window time.Duration) flamego.Handler { anonymousTokenRL := newRateLimiter(limit, window) loggedUserRL := newRateLimiter(limit*3, window) anonymousIPRL := newRateLimiter(limit*10, window) tooMany := func(ctx flamego.Context) { tooManyRequests(ctx, window) } return func(ctx flamego.Context, s session.Session) { r := ctx.Request().Request if uid, ok := s.Get("uid").(uint); ok && uid > 0 { if !loggedUserRL.allow(fmt.Sprintf("user:%d", uid)) { tooMany(ctx) } return } // 匿名请求始终经过高阈值 IP 兜底;共享出口下每个有效 token 仍独立计数。 // IP 键同样经 LimitKey 聚合,避免 IPv6 轮换绕过兜底阈值。 if !anonymousIPRL.allow("ip:" + shared.LimitKey(shared.ResolveClientIP(r))) { tooMany(ctx) return } key := vote.VoterTokenRateLimitKey(r) if key != "" && !anonymousTokenRL.allow("token:"+key) { tooMany(ctx) } } }