tamabox/internal/route/ratelimit.go

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)
		}
	}
}