0904

internal/route/ratelimit.go

@@ -0,0 +1,183 @@
+package route
+
+import (
+	"fmt"
+	"net/http"
+	"sync"
+	"time"
+
+	"github.com/flamego/flamego"
+	"github.com/flamego/session"
+
+	"github.com/tamakyi/TamaBox/internal/shared"
+	"github.com/tamakyi/TamaBox/route/vote"
+)
+
+// 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) {
+			ctx.ResponseWriter().WriteHeader(http.StatusTooManyRequests)
+			_, _ = ctx.ResponseWriter().Write([]byte("Too many requests. Please try again later."))
+			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) {
+		ctx.ResponseWriter().WriteHeader(http.StatusTooManyRequests)
+		_, _ = ctx.ResponseWriter().Write([]byte("Too many requests. Please try again later."))
+	}
+
+	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)
+		}
+	}
+}