// Package giftlimit 为 gateway 送礼入口提供跨节点容量保护。 // // 送礼请求的最坏下游成本与 gift_count*去重 target_count 成正比:幸运礼物需要逐份开奖, // 每个目标也会形成独立账务和事实。因此本包不按 HTTP 次数计费,而按服务端已经校验过 // 的 total_gift_units 在 App、房间、用户三层扣令牌,并额外约束单用户在途请求数。 package giftlimit import ( "context" "crypto/sha256" "encoding/hex" "errors" "fmt" "math" "strconv" "strings" "time" "github.com/redis/go-redis/v9" ) const ( defaultKeyPrefix = "gateway:gift:capacity:" // 一个合法请求最多 5000 gift units,因此用户桶必须至少容纳一个完整请求;500/s // 让极端 999 数量请求形成自然冷却,同时普通 300ms 连击批次不会感到逐击等待。 defaultUserBucketCapacity = 5_000 defaultUserRefillTokensPerSecond = 500 // 房间和 App 层允许多个真实用户同时连击,但把 5 人叠加造成的无界增长转换成 429 // 背压。数值均为可灰度配置的保护起点,不冒充未经压测的 Broker 容量结论。 defaultRoomBucketCapacity = 20_000 defaultRoomRefillTokensPerSecond = 5_000 defaultAppBucketCapacity = 100_000 defaultAppRefillTokensPerSecond = 25_000 defaultMaxInFlight = 2 defaultInFlightLease = 10 * time.Second defaultRedisOperationTimeout = 500 * time.Millisecond redisReleaseTimeout = time.Second redisDialTimeout = time.Second redisReadTimeout = 300 * time.Millisecond redisWriteTimeout = 300 * time.Millisecond redisPoolTimeout = 300 * time.Millisecond redisMinRetryBackoff = 20 * time.Millisecond redisMaxRetryBackoff = 50 * time.Millisecond ) // Config 描述三层令牌桶、单用户在途租约和 Redis 操作预算。 type Config struct { Enabled bool KeyPrefix string UserBucketCapacity int UserRefillTokensPerSecond float64 RoomBucketCapacity int RoomRefillTokensPerSecond float64 AppBucketCapacity int AppRefillTokensPerSecond float64 MaxInFlight int InFlightLease time.Duration RedisOperationTimeout time.Duration } // Input 是一次送礼请求在容量维度上的稳定身份。 type Input struct { AppCode string UserID int64 RoomID string RequestID string Weight int } // Lease 标识已占用的分布式在途槽。字段只供 Limiter.Release 使用。 type Lease struct { inFlightKey string member string } // Decision 返回是否获准以及客户端下一次可重试的最短等待时间。 type Decision struct { Allowed bool RetryAfter time.Duration Lease Lease } // Limiter 是 room HTTP adapter 所需的最小容量保护边界。 type Limiter interface { Acquire(ctx context.Context, input Input) (Decision, error) Release(ctx context.Context, lease Lease) error } type redisScriptRunner interface { evalSlice(ctx context.Context, script string, keys []string, args ...any) ([]any, error) } type goRedisScriptRunner struct { client *redis.Client } type bucketPolicy struct { capacity int refillPerSecond float64 } type redisScopeKeys struct { appBucket string roomBucket string userBucket string inFlight string leaseMember string } // RedisLimiter 按 user→room→app 顺序执行三个独立 Lua。 // // Redis Cluster 只允许一次脚本访问同一 hash slot。App、房间、用户是不同聚合域,不能 // 伪装成一个跨 slot 原子事务。用户层先检查 in-flight,可防止已占满并发槽的账号用必拒 // 请求持续消耗房间/App 额度;后层拒绝时释放用户 lease,但保留用户令牌作为请求方成本。 // App 拒绝时不回滚房间令牌,少量保守损耗比并发回滚造成超发或绕过容量上限更安全。 type RedisLimiter struct { runner redisScriptRunner config Config } // redisTokenBucketScript 在一个聚合维度内原子补充并扣减令牌,返回 // {allowed, retry_after_ms}。Redis TIME 是多 gateway 节点的统一时间源。 const redisTokenBucketScript = ` local redis_time = redis.call("TIME") local now_ms = tonumber(redis_time[1]) * 1000 + math.floor(tonumber(redis_time[2]) / 1000) local capacity = tonumber(ARGV[1]) local refill_per_ms = tonumber(ARGV[2]) local weight = tonumber(ARGV[3]) local bucket_ttl_ms = tonumber(ARGV[4]) local tokens = tonumber(redis.call("HGET", KEYS[1], "tokens")) local last_ms = tonumber(redis.call("HGET", KEYS[1], "last_ms")) if not tokens or not last_ms then tokens = capacity last_ms = now_ms else tokens = math.min(capacity, tokens + math.max(0, now_ms - last_ms) * refill_per_ms) end if tokens < weight then redis.call("HSET", KEYS[1], "tokens", tokens, "last_ms", now_ms) redis.call("PEXPIRE", KEYS[1], bucket_ttl_ms) local retry_ms = math.ceil((weight - tokens) / refill_per_ms) return {0, math.max(1, retry_ms)} end redis.call("HSET", KEYS[1], "tokens", tokens - weight, "last_ms", now_ms) redis.call("PEXPIRE", KEYS[1], bucket_ttl_ms) return {1, 0} ` // redisUserAcquireScript 把用户桶和 app:user 在途 ZSET 放进同一个 hash slot,因此用户 // 额度检查、并发检查和租约写入在单个 Redis 节点内保持原子。 const redisUserAcquireScript = ` local redis_time = redis.call("TIME") local now_ms = tonumber(redis_time[1]) * 1000 + math.floor(tonumber(redis_time[2]) / 1000) local capacity = tonumber(ARGV[1]) local refill_per_ms = tonumber(ARGV[2]) local weight = tonumber(ARGV[3]) local max_in_flight = tonumber(ARGV[4]) local lease_ms = tonumber(ARGV[5]) local member = ARGV[6] local bucket_ttl_ms = tonumber(ARGV[7]) -- gateway 异常退出不会执行 Release;每次申请前清理过期租约,确保 lease_ms 后自动恢复。 redis.call("ZREMRANGEBYSCORE", KEYS[2], "-inf", now_ms) local in_flight = redis.call("ZCARD", KEYS[2]) if in_flight >= max_in_flight then local earliest = redis.call("ZRANGE", KEYS[2], 0, 0, "WITHSCORES") local retry_ms = 1 if earliest[2] then retry_ms = math.max(1, tonumber(earliest[2]) - now_ms) end return {0, retry_ms} end local tokens = tonumber(redis.call("HGET", KEYS[1], "tokens")) local last_ms = tonumber(redis.call("HGET", KEYS[1], "last_ms")) if not tokens or not last_ms then tokens = capacity last_ms = now_ms else tokens = math.min(capacity, tokens + math.max(0, now_ms - last_ms) * refill_per_ms) end if tokens < weight then redis.call("HSET", KEYS[1], "tokens", tokens, "last_ms", now_ms) redis.call("PEXPIRE", KEYS[1], bucket_ttl_ms) local retry_ms = math.ceil((weight - tokens) / refill_per_ms) return {0, math.max(1, retry_ms)} end redis.call("HSET", KEYS[1], "tokens", tokens - weight, "last_ms", now_ms) redis.call("PEXPIRE", KEYS[1], bucket_ttl_ms) redis.call("ZADD", KEYS[2], now_ms + lease_ms, member) redis.call("PEXPIRE", KEYS[2], lease_ms * 2) return {1, 0} ` const redisReleaseScript = ` redis.call("ZREM", KEYS[1], ARGV[1]) if redis.call("ZCARD", KEYS[1]) == 0 then redis.call("DEL", KEYS[1]) end return {1} ` // Normalize 收敛与 10k MQ TPS 和当前送礼行为匹配的安全默认值。 func Normalize(config Config) Config { config.KeyPrefix = strings.TrimSpace(config.KeyPrefix) if config.KeyPrefix == "" { config.KeyPrefix = defaultKeyPrefix } if config.UserBucketCapacity <= 0 { config.UserBucketCapacity = defaultUserBucketCapacity } if config.UserRefillTokensPerSecond <= 0 { config.UserRefillTokensPerSecond = defaultUserRefillTokensPerSecond } if config.RoomBucketCapacity <= 0 { config.RoomBucketCapacity = defaultRoomBucketCapacity } if config.RoomRefillTokensPerSecond <= 0 { config.RoomRefillTokensPerSecond = defaultRoomRefillTokensPerSecond } if config.AppBucketCapacity <= 0 { config.AppBucketCapacity = defaultAppBucketCapacity } if config.AppRefillTokensPerSecond <= 0 { config.AppRefillTokensPerSecond = defaultAppRefillTokensPerSecond } if config.MaxInFlight <= 0 { config.MaxInFlight = defaultMaxInFlight } if config.InFlightLease <= 0 { config.InFlightLease = defaultInFlightLease } if config.RedisOperationTimeout <= 0 { config.RedisOperationTimeout = defaultRedisOperationTimeout } return config } // NewRedisClient 建立并探测容量保护使用的 Redis 连接。 // ContextTimeoutEnabled 必须开启,否则 go-redis 可能只依赖较长 socket timeout,使 handler // 的 500ms fail-closed 预算和 Release 的 1s 回收预算失效。 func NewRedisClient(ctx context.Context, addr string, password string, db int) (*redis.Client, error) { client := redis.NewClient(redisClientOptions(addr, password, db)) if err := client.Ping(ctx).Err(); err != nil { _ = client.Close() return nil, err } return client, nil } func redisClientOptions(addr string, password string, db int) *redis.Options { return &redis.Options{ Addr: strings.TrimSpace(addr), Password: password, DB: db, MaxRetries: 1, MinRetryBackoff: redisMinRetryBackoff, MaxRetryBackoff: redisMaxRetryBackoff, DialTimeout: redisDialTimeout, DialerRetries: 1, ReadTimeout: redisReadTimeout, WriteTimeout: redisWriteTimeout, PoolTimeout: redisPoolTimeout, ContextTimeoutEnabled: true, } } // NewRedisLimiter 创建生产使用的分布式 limiter。 func NewRedisLimiter(client *redis.Client, config Config) *RedisLimiter { return newRedisLimiter(goRedisScriptRunner{client: client}, config) } func newRedisLimiter(runner redisScriptRunner, config Config) *RedisLimiter { return &RedisLimiter{runner: runner, config: Normalize(config)} } // Acquire 在一个独立短 deadline 内按 user→room→app 顺序保守扣减。 // 用户并发先行可挡住低成本全局 DoS;后层失败会立即释放用户 lease,但不会返还已扣令牌。 func (l *RedisLimiter) Acquire(ctx context.Context, input Input) (Decision, error) { if l == nil || l.runner == nil { return Decision{}, errors.New("gift capacity limiter is not configured") } appCode := strings.ToLower(strings.TrimSpace(input.AppCode)) if appCode == "" || input.UserID <= 0 || strings.TrimSpace(input.RoomID) == "" || strings.TrimSpace(input.RequestID) == "" || input.Weight <= 0 { return Decision{}, errors.New("gift capacity limiter input is invalid") } config := Normalize(l.config) if input.Weight > config.UserBucketCapacity || input.Weight > config.RoomBucketCapacity || input.Weight > config.AppBucketCapacity { // handler 会先用 max_targets_per_request 返回 4xx;到达这里说明容量配置与协议上限漂移, // 必须 fail-closed,而不是返回一个永远无法成功的 429 重试循环。 return Decision{}, errors.New("gift target weight exceeds configured bucket capacity") } if ctx == nil { ctx = context.Background() } operationCtx, cancel := context.WithTimeout(ctx, config.RedisOperationTimeout) defer cancel() keys := buildRedisScopeKeys(config.KeyPrefix, appCode, input) userPolicy := bucketPolicy{capacity: config.UserBucketCapacity, refillPerSecond: config.UserRefillTokensPerSecond} result, err := l.runner.evalSlice(operationCtx, redisUserAcquireScript, []string{keys.userBucket, keys.inFlight}, userPolicy.capacity, refillPerMillisecond(userPolicy.refillPerSecond), input.Weight, config.MaxInFlight, config.InFlightLease.Milliseconds(), keys.leaseMember, bucketTTL(userPolicy, config.InFlightLease).Milliseconds(), ) if err != nil { return Decision{}, err } userDecision, err := decodeDecision(result) if err != nil { // Lua 已返回但协议解码失败时执行幂等 ZREM;脚本可能已经写入 lease,不能等到 TTL 才恢复。 l.releaseAcquiredLease(operationCtx, Lease{inFlightKey: keys.inFlight, member: keys.leaseMember}) return Decision{}, err } if !userDecision.Allowed { // 并发已满或用户桶不足时尚未触达 room/app key,因此攻击者不能靠必拒请求耗尽共享额度。 return userDecision, nil } lease := Lease{inFlightKey: keys.inFlight, member: keys.leaseMember} layers := []struct { key string policy bucketPolicy }{ {key: keys.roomBucket, policy: bucketPolicy{capacity: config.RoomBucketCapacity, refillPerSecond: config.RoomRefillTokensPerSecond}}, {key: keys.appBucket, policy: bucketPolicy{capacity: config.AppBucketCapacity, refillPerSecond: config.AppRefillTokensPerSecond}}, } for _, layer := range layers { decision, err := l.consumeBucket(operationCtx, layer.key, layer.policy, input.Weight) if err != nil || !decision.Allowed { // 使用同一个短 deadline 立即释放用户 lease;deadline 已耗尽时释放可能失败, // 但 ZSET score 的 InFlightLease 仍会兜底回收,绝不会永久锁死用户。 l.releaseAcquiredLease(operationCtx, lease) return decision, err } } userDecision.Lease = lease return userDecision, nil } func (l *RedisLimiter) consumeBucket(ctx context.Context, key string, policy bucketPolicy, weight int) (Decision, error) { result, err := l.runner.evalSlice(ctx, redisTokenBucketScript, []string{key}, policy.capacity, refillPerMillisecond(policy.refillPerSecond), weight, bucketTTL(policy, 0).Milliseconds(), ) if err != nil { return Decision{}, err } return decodeDecision(result) } func (l *RedisLimiter) releaseAcquiredLease(ctx context.Context, lease Lease) { if lease.inFlightKey == "" || lease.member == "" { return } // 不调用公开 Release,避免给一次 Acquire 额外再增加 1s;补偿动作必须服从 Acquire 的 // 共同短预算。失败时租约 TTL 是安全兜底,用户 token 则有意不返还。 _, _ = l.runner.evalSlice(ctx, redisReleaseScript, []string{lease.inFlightKey}, lease.member) } // Release 始终再包一层 1s deadline。handler 使用与客户端取消解耦的 background context, // go-redis 又启用 ContextTimeoutEnabled,因此慢 Redis 不会把请求收尾永久挂住。 func (l *RedisLimiter) Release(ctx context.Context, lease Lease) error { if l == nil || l.runner == nil || lease.inFlightKey == "" || lease.member == "" { return nil } if ctx == nil { ctx = context.Background() } releaseCtx, cancel := context.WithTimeout(ctx, redisReleaseTimeout) defer cancel() _, err := l.runner.evalSlice(releaseCtx, redisReleaseScript, []string{lease.inFlightKey}, lease.member) return err } func (r goRedisScriptRunner) evalSlice(ctx context.Context, script string, keys []string, args ...any) ([]any, error) { return r.client.Eval(ctx, script, keys, args...).Slice() } func buildRedisScopeKeys(prefix string, appCode string, input Input) redisScopeKeys { appPart := boundedHash(appCode) roomPart := boundedHash(appCode + ":" + strings.TrimSpace(input.RoomID)) userPart := boundedHash(appCode + ":" + strconv.FormatInt(input.UserID, 10)) member := boundedHash(input.RequestID) // 每个聚合域拥有自己的 hash tag;只有用户桶与 in-flight ZSET 共享 user tag,因而只有 // 这两个 key 会放入同一 Lua。App/room/user 三层绝不跨 slot 伪原子。 return redisScopeKeys{ appBucket: prefix + "{app:" + appPart + "}:bucket", roomBucket: prefix + "{room:" + roomPart + "}:bucket", userBucket: prefix + "{user:" + userPart + "}:bucket", inFlight: prefix + "{user:" + userPart + "}:inflight", leaseMember: member, } } func bucketTTL(policy bucketPolicy, minimum time.Duration) time.Duration { ttl := time.Duration(math.Ceil(float64(policy.capacity)/policy.refillPerSecond*2)) * time.Second if ttl < time.Second { ttl = time.Second } if ttl < minimum { ttl = minimum } return ttl } func refillPerMillisecond(tokensPerSecond float64) string { return strconv.FormatFloat(tokensPerSecond/1000, 'f', -1, 64) } func decodeDecision(result []any) (Decision, error) { if len(result) != 2 { return Decision{}, fmt.Errorf("gift capacity limiter returned %d values", len(result)) } allowed, err := redisInt64(result[0]) if err != nil { return Decision{}, fmt.Errorf("decode gift capacity allowed: %w", err) } retryAfterMS, err := redisInt64(result[1]) if err != nil { return Decision{}, fmt.Errorf("decode gift capacity retry_after_ms: %w", err) } if allowed == 1 { return Decision{Allowed: true}, nil } if retryAfterMS < 1 { retryAfterMS = 1 } return Decision{RetryAfter: time.Duration(retryAfterMS) * time.Millisecond}, nil } func boundedHash(value string) string { sum := sha256.Sum256([]byte(strings.TrimSpace(value))) return hex.EncodeToString(sum[:16]) } func redisInt64(value any) (int64, error) { switch typed := value.(type) { case int64: return typed, nil case int: return int64(typed), nil case string: return strconv.ParseInt(typed, 10, 64) case []byte: return strconv.ParseInt(string(typed), 10, 64) default: return 0, fmt.Errorf("unsupported redis integer type %T", value) } }