package http import ( "context" "crypto/sha256" "encoding/hex" "hyapp/services/gateway-service/internal/transport/http/httpkit" "net/http" "strings" "sync" "time" "github.com/redis/go-redis/v9" ) const ( authOperationPassword = "password" authOperationThirdParty = "third_party" authOperationRefresh = "refresh" authOperationLogout = "logout" ) // AuthRateLimitConfig 描述 gateway public auth 入口的固定窗口频控策略。 type AuthRateLimitConfig struct { // Enabled 为 false 时关闭 gateway 频控。 Enabled bool // KeyPrefix 是 Redis key 前缀,用于区分环境和服务。 KeyPrefix string // Window 是固定计数窗口。 Window time.Duration // IPLimit 限制同一入口 IP 对所有 public auth API 的请求数。 IPLimit int // DeviceIDLimit 限制同一 device_id 对登录、注册和 refresh 的请求数。 DeviceIDLimit int // ProviderIPLimit 限制同一 IP 对同一三方 provider 的注册/登录请求数。 ProviderIPLimit int // DisplayUserIDIPLimit 限制同一 IP 对同一 display_user_id 的密码尝试。 DisplayUserIDIPLimit int // DisplayUserIDLimit 限制同一 display_user_id 被跨 IP 喷射尝试。 DisplayUserIDLimit int // SessionIDIPLimit 限制同一 IP 对同一 session_id 的 logout 请求数。 SessionIDIPLimit int } type authRateInput struct { operation string deviceID string provider string displayUserID string sessionID string } type authRateKey struct { key string limit int } type authRateWindow struct { startedAt time.Time count int } type publicAuthRateLimiter interface { allow(ctx context.Context, keys []authRateKey) (bool, error) } type memoryAuthRateLimiter struct { mu sync.Mutex now func() time.Time window time.Duration windows map[string]authRateWindow } type redisAuthRateLimiter struct { runner authRateRedisRunner window time.Duration } type authRateRedisRunner interface { evalInt(ctx context.Context, script string, keys []string, args ...any) (int64, error) } type goRedisAuthRateRunner struct { client *redis.Client } const redisAuthRateLimitScript = ` for i = 1, #KEYS do local current = tonumber(redis.call("GET", KEYS[i]) or "0") local limit = tonumber(ARGV[i + 1]) if current >= limit then return 0 end end for i = 1, #KEYS do local next = redis.call("INCR", KEYS[i]) if next == 1 then redis.call("PEXPIRE", KEYS[i], ARGV[1]) end end return 1 ` func newMemoryAuthRateLimiter(config AuthRateLimitConfig) *memoryAuthRateLimiter { config = normalizeAuthRateLimitConfig(config) return &memoryAuthRateLimiter{ now: time.Now, window: config.Window, windows: make(map[string]authRateWindow), } } func NewAuthRateLimitRedisClient(ctx context.Context, addr string, password string, db int) (*redis.Client, error) { client := redis.NewClient(&redis.Options{ Addr: strings.TrimSpace(addr), Password: password, DB: db, }) if err := client.Ping(ctx).Err(); err != nil { _ = client.Close() return nil, err } return client, nil } func NewRedisAuthRateLimiter(client *redis.Client, config AuthRateLimitConfig) publicAuthRateLimiter { return newRedisAuthRateLimiter(goRedisAuthRateRunner{client: client}, config) } func newRedisAuthRateLimiter(runner authRateRedisRunner, config AuthRateLimitConfig) *redisAuthRateLimiter { config = normalizeAuthRateLimitConfig(config) return &redisAuthRateLimiter{runner: runner, window: config.Window} } func defaultAuthRateLimitConfig() AuthRateLimitConfig { return AuthRateLimitConfig{ Enabled: true, KeyPrefix: "gateway:auth:rate:", Window: time.Minute, IPLimit: 120, DeviceIDLimit: 30, ProviderIPLimit: 20, DisplayUserIDIPLimit: 10, DisplayUserIDLimit: 50, SessionIDIPLimit: 30, } } func normalizeAuthRateLimitConfig(config AuthRateLimitConfig) AuthRateLimitConfig { defaults := defaultAuthRateLimitConfig() if strings.TrimSpace(config.KeyPrefix) == "" { config.KeyPrefix = defaults.KeyPrefix } if config.Window <= 0 { config.Window = defaults.Window } if config.IPLimit <= 0 { config.IPLimit = defaults.IPLimit } if config.DeviceIDLimit <= 0 { config.DeviceIDLimit = defaults.DeviceIDLimit } if config.ProviderIPLimit <= 0 { config.ProviderIPLimit = defaults.ProviderIPLimit } if config.DisplayUserIDIPLimit <= 0 { config.DisplayUserIDIPLimit = defaults.DisplayUserIDIPLimit } if config.DisplayUserIDLimit <= 0 { config.DisplayUserIDLimit = defaults.DisplayUserIDLimit } if config.SessionIDIPLimit <= 0 { config.SessionIDIPLimit = defaults.SessionIDIPLimit } return config } func (l *memoryAuthRateLimiter) allow(_ context.Context, keys []authRateKey) (bool, error) { now := l.now() l.mu.Lock() defer l.mu.Unlock() l.pruneLocked(now) for _, key := range keys { if key.key == "" || key.limit <= 0 { continue } window := l.currentWindowLocked(now, key.key) if window.count >= key.limit { return false, nil } } for _, key := range keys { if key.key == "" || key.limit <= 0 { continue } window := l.currentWindowLocked(now, key.key) window.count++ l.windows[key.key] = window } return true, nil } func (l *memoryAuthRateLimiter) currentWindowLocked(now time.Time, key string) authRateWindow { window, exists := l.windows[key] if !exists || now.Sub(window.startedAt) >= l.window { return authRateWindow{startedAt: now} } return window } func (l *memoryAuthRateLimiter) pruneLocked(now time.Time) { if len(l.windows) < 4096 { return } for key, window := range l.windows { if now.Sub(window.startedAt) >= 2*l.window { delete(l.windows, key) } } } func (l *redisAuthRateLimiter) allow(ctx context.Context, keys []authRateKey) (bool, error) { if len(keys) == 0 { return true, nil } redisKeys := make([]string, 0, len(keys)) args := make([]any, 0, len(keys)+1) args = append(args, l.window.Milliseconds()) for _, key := range keys { if key.key == "" || key.limit <= 0 { continue } redisKeys = append(redisKeys, key.key) args = append(args, key.limit) } if len(redisKeys) == 0 { return true, nil } result, err := l.runner.evalInt(ctx, redisAuthRateLimitScript, redisKeys, args...) if err != nil { return false, err } return result == 1, nil } func (r goRedisAuthRateRunner) evalInt(ctx context.Context, script string, keys []string, args ...any) (int64, error) { return r.client.Eval(ctx, script, keys, args...).Int64() } // SetAuthRateLimit 允许测试或本地构造覆盖 gateway public auth 入口频控策略,使用内存 backend。 func (h *Handler) SetAuthRateLimit(config AuthRateLimitConfig) { config = normalizeAuthRateLimitConfig(config) h.authRateLimitConfig = config h.authRateLimiter = newMemoryAuthRateLimiter(config) } // SetAuthRateLimitBackend 允许 app 层注入 Redis backend,避免多 gateway 节点各自计数。 func (h *Handler) SetAuthRateLimitBackend(config AuthRateLimitConfig, limiter publicAuthRateLimiter) { config = normalizeAuthRateLimitConfig(config) h.authRateLimitConfig = config if limiter == nil { h.authRateLimiter = newMemoryAuthRateLimiter(config) return } h.authRateLimiter = limiter } func (h *Handler) allowPublicAuthRequest(writer http.ResponseWriter, request *http.Request, input authRateInput) bool { config := normalizeAuthRateLimitConfig(h.authRateLimitConfig) if !config.Enabled { return true } if h.authRateLimiter == nil { h.SetAuthRateLimit(config) } allowed, err := h.authRateLimiter.allow(request.Context(), h.authRateKeys(request, input)) if err != nil { httpkit.WriteError(writer, request, http.StatusBadGateway, httpkit.CodeUpstreamError, "upstream service error") return false } if allowed { return true } httpkit.WriteError(writer, request, http.StatusTooManyRequests, httpkit.CodeRateLimited, "rate limited") return false } func (h *Handler) authRateKeys(request *http.Request, input authRateInput) []authRateKey { config := normalizeAuthRateLimitConfig(h.authRateLimitConfig) ip := clientIP(request) ipPart := authRateKeyPart(ip) keyPrefix := strings.TrimSpace(config.KeyPrefix) keys := []authRateKey{{key: keyPrefix + "ip:" + ipPart, limit: config.IPLimit}} deviceID := strings.TrimSpace(input.deviceID) if deviceID != "" { keys = append(keys, authRateKey{key: keyPrefix + "device:" + authRateKeyPart(deviceID), limit: config.DeviceIDLimit}) } switch input.operation { case authOperationThirdParty: provider := strings.ToLower(strings.TrimSpace(input.provider)) if provider != "" { keys = append(keys, authRateKey{key: keyPrefix + "third_party:provider_ip:" + authRateKeyPart(provider) + ":" + ipPart, limit: config.ProviderIPLimit}) } case authOperationPassword: displayUserID := strings.TrimSpace(input.displayUserID) if displayUserID != "" { displayPart := authRateKeyPart(displayUserID) keys = append(keys, authRateKey{key: keyPrefix + "password:display_ip:" + displayPart + ":" + ipPart, limit: config.DisplayUserIDIPLimit}, authRateKey{key: keyPrefix + "password:display:" + displayPart, limit: config.DisplayUserIDLimit}, ) } case authOperationLogout: sessionID := strings.TrimSpace(input.sessionID) if sessionID != "" { keys = append(keys, authRateKey{key: keyPrefix + "logout:session_ip:" + authRateKeyPart(sessionID) + ":" + ipPart, limit: config.SessionIDIPLimit}) } } return keys } func authRateKeyPart(value string) string { sum := sha256.Sum256([]byte(strings.TrimSpace(value))) // Public auth 字段和 X-Forwarded-For 都来自入口流量,Redis key 只保留固定长度摘要。 return hex.EncodeToString(sum[:16]) }