package middleware import ( "fmt" "strconv" "strings" "time" "github.com/gogf/gf/v2/errors/gerror" "github.com/gogf/gf/v2/net/ghttp" "easygoadmin/internal/consts" "easygoadmin/utility/cache" "easygoadmin/utility/response" ) // ============================== API 限流中间件 ============================== // // 固定窗口计数:键按窗口起点对齐(ratelimit:{ip}:{path}:{窗口起点}), // 计数存 utility/cache(Redis 优先、内存降级),多实例部署时共享计数。 // // 规则: // - 仅限 /api/ 前缀接口;文档(/api.json)、静态资源、健康检查不受限 // - 认证入口(登录/验证码/令牌兑换)按 rateLimit.authLimit 单独更严阈值 // - 计数存储异常时放行(限流不阻断业务) // RateLimit API 限流 func RateLimit(r *ghttp.Request) { if !cfgBool(r.Context(), "rateLimit.enable", true) { r.Middleware.Next() return } path := r.URL.Path if !isRateLimitedPath(path) { r.Middleware.Next() return } window := cfgInt(r.Context(), "rateLimit.windowSeconds", 60) limit := cfgInt(r.Context(), "rateLimit.limit", 120) if isAuthPath(path) { limit = cfgInt(r.Context(), "rateLimit.authLimit", 10) } if limit <= 0 || window <= 0 { r.Middleware.Next() return } // 窗口起点对齐:同一窗口内共享计数键 windowStart := time.Now().Unix() / int64(window) * int64(window) key := fmt.Sprintf("ratelimit:%s:%s:%d", realClientIP(r), path, windowStart) n, err := incrCounter(r, key, window) if err != nil { r.Middleware.Next() return } if n > int64(limit) { response.Write(r, response.Fail(gerror.NewCode(consts.CodeTooManyRequests))) return } r.Middleware.Next() } // isRateLimitedPath 判断是否为参与限流的接口路径 // /api.json 与 /swagger 不带 /api/ 前缀,自然排除;健康检查显式豁免 func isRateLimitedPath(path string) bool { if path == "/api/health" { return false } return path == "/api" || strings.HasPrefix(path, "/api/") } // isAuthPath 判断是否为认证入口(适用更严格的独立阈值) func isAuthPath(path string) bool { switch path { case "/api/login", "/api/captcha", "/api/oauth2/token": return true } return false } // incrCounter 固定窗口计数自增 // 键已含窗口起点,TTL 仅用于回收存储,续期不影响窗口归属 func incrCounter(r *ghttp.Request, key string, windowSeconds int) (int64, error) { ttl := windowSeconds + 5 if raw, ok := cache.Get(r.Context(), key); ok { n, _ := strconv.ParseInt(raw, 10, 64) n++ _ = cache.Set(r.Context(), key, strconv.FormatInt(n, 10), ttl) return n, nil } if cache.SetNX(r.Context(), key, "1", ttl) { return 1, nil } // 并发下 SetNX 未抢到:读现值再加一 if raw, ok := cache.Get(r.Context(), key); ok { n, _ := strconv.ParseInt(raw, 10, 64) n++ _ = cache.Set(r.Context(), key, strconv.FormatInt(n, 10), ttl) return n, nil } return 1, nil }