FileShare/server/internal/middleware/ratelimit.go

package middleware

import (
	"context"
	"errors"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/gin-gonic/gin"

	"fileshare/internal/cache"
	"fileshare/internal/response"
)

// 限流类别（对齐参考 apps/base/utils.py 的 ip_limit）。
const (
	LimitError  = "error"    // 取件错误（密码错误、取件失败）
	LimitUpload = "upload"   // 上传次数
	LimitLogin  = "login"    // 管理员登录失败
	LimitMeta   = "metadata" // 分享元信息查询
)

// LimitRule 限流规则：window 内最多 count 次。
type LimitRule struct {
	Count  int           // 允许次数
	Window time.Duration // 时间窗口
}

// clientIP 解析客户端真实 IP：仅当直连地址属于可信代理时才采信 X-Forwarded-For / X-Real-IP。
// 语义对齐参考 apps/base/dependencies.py 的 get_client_ip。
func clientIP(c *gin.Context, trustedProxies []*net.IPNet) string {
	remote := net.ParseIP(c.RemoteIP())
	parse := func(s string) net.IP {
		ip := net.ParseIP(strings.TrimSpace(s))
		return ip
	}
	isTrusted := func(ip net.IP) bool {
		if ip == nil {
			return false
		}
		for _, n := range trustedProxies {
			if n.Contains(ip) {
				return true
			}
		}
		return false
	}

	if !isTrusted(remote) {
		return remote.String()
	}
	// X-Forwarded-For：从右往左找第一个非可信代理地址
	if xff := c.GetHeader("X-Forwarded-For"); xff != "" {
		parts := strings.Split(xff, ",")
		for i := len(parts) - 1; i >= 0; i-- {
			candidate := parse(parts[i])
			if candidate == nil {
				return remote.String()
			}
			if !isTrusted(candidate) {
				return candidate.String()
			}
		}
		return strings.TrimSpace(parts[0])
	}
	if xr := c.GetHeader("X-Real-IP"); xr != "" {
		if ip := parse(xr); ip != nil {
			return ip.String()
		}
	}
	return remote.String()
}

// ParseTrustedProxies 把 CIDR/单 IP 字符串解析为网络列表。
func ParseTrustedProxies(items []string) []*net.IPNet {
	var out []*net.IPNet
	for _, item := range items {
		item = strings.TrimSpace(item)
		if item == "" {
			continue
		}
		if !strings.Contains(item, "/") {
			item += "/32"
			if strings.Contains(item, ":") { // IPv6
				item = item[:len(item)-3] + "/128"
			}
		}
		_, network, err := net.ParseCIDR(item)
		if err != nil {
			continue
		}
		out = append(out, network)
	}
	return out
}

// ClientIP 中间件：解析真实 IP 并写入上下文（ctxClientIP）。
func ClientIP(trustedProxies []*net.IPNet) gin.HandlerFunc {
	return func(c *gin.Context) {
		c.Set("ctxClientIP", clientIP(c, trustedProxies))
		c.Next()
	}
}

// GetClientIP 从 gin 上下文取解析后的客户端 IP。
func GetClientIP(c *gin.Context) string {
	if v, ok := c.Get("ctxClientIP"); ok {
		if s, ok := v.(string); ok {
			return s
		}
	}
	if ip := net.ParseIP(c.RemoteIP()); ip != nil {
		return ip.String()
	}
	return c.RemoteIP()
}

// RateLimiter 基于 cache.Cache 的固定窗口 IP 限流器。
// 计数语义对齐参考实现：check 通过时放行，业务方在发生"计数事件"（如失败/成功上传）后调用 Add。
//
// L9：缓存故障降级——此前 cache.Get/Incr 失败（如 Redis 宕机）时一律放行，
// 登录爆破防护随之失效。现降级为进程内固定窗口计数（单实例语义），
// 缓存恢复后自动回到共享缓存计数。降级期间计数独立于缓存，不叠加。
type RateLimiter struct {
	cache  cache.Cache
	limits map[string]LimitRule
	prefix string

	fbMu      sync.Mutex
	fallback  map[string]*fallbackEntry // 进程内降级计数
	lastPrune time.Time
}

type fallbackEntry struct {
	count   int64
	expires time.Time
}

// fallbackMaxEntries 降级计数表上限（超出即整体重置，防内存增长）。
const fallbackMaxEntries = 8192

// NewRateLimiter 构造限流器；limits 为各类别规则（来自 settings 的 errorCount/errorMinute 等）。
func NewRateLimiter(cache cache.Cache, limits map[string]LimitRule) *RateLimiter {
	if limits == nil {
		limits = map[string]LimitRule{}
	}
	return &RateLimiter{cache: cache, limits: limits, prefix: "fcb:rl", fallback: map[string]*fallbackEntry{}}
}

// SetRule 运行时更新规则（settings KV 变更后调用）。
func (r *RateLimiter) SetRule(kind string, rule LimitRule) {
	r.limits[kind] = rule
}

func (r *RateLimiter) windowKey(kind, ip string, now time.Time) string {
	// 固定窗口：按窗口起点分桶
	bucket := now.Unix() / int64(r.limits[kind].Window/time.Second)
	return r.prefix + ":" + kind + ":" + ip + ":" + itoa64(bucket)
}

// Check 只读检查该 IP 在当前窗口内是否仍被允许（不计数）。
// 对齐参考 check_ip：已用次数 >= 上限即拒绝。
// 缓存键不存在（ErrNotFound）视为 0 次；缓存故障时降级为进程内计数。
func (r *RateLimiter) Check(c *gin.Context, kind string) (bool, int64) {
	rule, ok := r.limits[kind]
	if !ok || rule.Count <= 0 || rule.Window <= 0 {
		return true, 0
	}
	ip := GetClientIP(c)
	ctx, cancel := context.WithTimeout(c.Request.Context(), 2*time.Second)
	defer cancel()
	now := time.Now()
	raw, err := r.cache.Get(ctx, r.windowKey(kind, ip, now))
	if err != nil {
		if errors.Is(err, cache.ErrNotFound) {
			return true, 0 // 键不存在：窗口内尚无计数
		}
		// 缓存故障：降级进程内计数判定
		return r.fallbackCount(kind, ip, rule, now) < int64(rule.Count), 0
	}
	n := parseInt64(raw)
	return n < int64(rule.Count), n
}

// Add 记录一次计数事件（对齐参考 add_ip：调用即计数，如上传成功/登录失败/取件错误）。
// 缓存故障时降级为进程内计数。
func (r *RateLimiter) Add(c *gin.Context, kind string) {
	rule, ok := r.limits[kind]
	if !ok || rule.Count <= 0 || rule.Window <= 0 {
		return
	}
	ip := GetClientIP(c)
	ctx, cancel := context.WithTimeout(c.Request.Context(), 2*time.Second)
	defer cancel()
	if _, err := r.cache.Incr(ctx, r.windowKey(kind, ip, time.Now()), rule.Window); err != nil &&
		!errors.Is(err, cache.ErrNotFound) {
		// Incr 正常情况下不会因键不存在失败（缺键即从 0 起）；
		// 其余错误视为缓存故障 → 进程内计数
		r.fallbackIncr(kind, ip, rule, time.Now())
	}
}

// —— 进程内降级计数（L9）——

func (r *RateLimiter) fallbackIncr(kind, ip string, rule LimitRule, now time.Time) {
	key := r.windowKey(kind, ip, now)
	expires := now.Add(rule.Window)
	r.fbMu.Lock()
	defer r.fbMu.Unlock()
	r.pruneFallbackLocked(now)
	if len(r.fallback) >= fallbackMaxEntries {
		r.fallback = map[string]*fallbackEntry{} // 极端情况整体重置，防内存无限增长
	}
	e, ok := r.fallback[key]
	if !ok || now.After(e.expires) {
		r.fallback[key] = &fallbackEntry{count: 1, expires: expires}
		return
	}
	e.count++
}

func (r *RateLimiter) fallbackCount(kind, ip string, rule LimitRule, now time.Time) int64 {
	key := r.windowKey(kind, ip, now)
	r.fbMu.Lock()
	defer r.fbMu.Unlock()
	e, ok := r.fallback[key]
	if !ok || now.After(e.expires) {
		return 0
	}
	return e.count
}

// pruneFallbackLocked 清理已过窗口的降级计数（低频触发：每 1000 条或 10 分钟一次）。
func (r *RateLimiter) pruneFallbackLocked(now time.Time) {
	if r.lastPrune.IsZero() || len(r.fallback) >= 1024 || now.Sub(r.lastPrune) >= 10*time.Minute {
		for k, e := range r.fallback {
			if now.After(e.expires) {
				delete(r.fallback, k)
			}
		}
		r.lastPrune = now
	}
}

// RequireRateLimit 中间件：请求进入即检查，请求完成即计数。
// 适用于"每次访问都计数"的类别（如 metadata 查询）；
// 上传/登录等"仅成功/失败才计数"的场景由 handler 显式调用 Check/Add。
func (r *RateLimiter) RequireRateLimit(kind string) gin.HandlerFunc {
	return func(c *gin.Context) {
		allowed, _ := r.Check(c, kind)
		if !allowed {
			response.Fail(c, http.StatusLocked, "请求次数过多，请稍后再试")
			return
		}
		c.Next()
		r.Add(c, kind)
	}
}

// parseInt64 解析十进制整数字符串，非法输入返回 0。
func parseInt64(s string) int64 {
	var n int64
	for _, ch := range s {
		if ch < '0' || ch > '9' {
			return 0
		}
		n = n*10 + int64(ch-'0')
	}
	return n
}

func itoa64(n int64) string {
	if n == 0 {
		return "0"
	}
	neg := n < 0
	if neg {
		n = -n
	}
	var buf [21]byte
	i := len(buf)
	for n > 0 {
		i--
		buf[i] = byte('0' + n%10)
		n /= 10
	}
	if neg {
		i--
		buf[i] = '-'
	}
	return string(buf[i:])
}