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