package cache import ( "context" "errors" "fmt" "strconv" "strings" "time" "hyapp-admin-server/internal/config" "github.com/redis/go-redis/v9" ) type Redis struct { client *redis.Client prefix string } type UnlockFunc func(context.Context) error var ErrRedisUnavailable = errors.New("redis is unavailable") type FixedWindowRule struct { Key string Limit uint64 } type FixedWindowDecision struct { Allowed bool RetryAfter time.Duration } // fixedWindowScript consumes every rule in one Redis-side operation. All rate-limit // keys use one Cluster hash tag at the caller, so EVAL remains atomic without CROSSSLOT. // Rejected attempts still increment every layer; expiry is assigned only on the first // increment and is never renewed, preserving true fixed-window behavior. const fixedWindowScript = ` local allowed = 1 local retry_ms = 0 local window_ms = tonumber(ARGV[1]) for i = 1, #KEYS do local count = redis.call("INCR", KEYS[i]) if count == 1 then redis.call("PEXPIRE", KEYS[i], window_ms) end local limit = tonumber(ARGV[i + 1]) if count > limit then allowed = 0 local ttl = redis.call("PTTL", KEYS[i]) if ttl > retry_ms then retry_ms = ttl end end end if allowed == 0 and retry_ms <= 0 then retry_ms = window_ms end return {allowed, retry_ms}` func NewRedis(ctx context.Context, cfg config.RedisConfig) (*Redis, error) { if !cfg.Enabled { return nil, nil } client := redis.NewClient(&redis.Options{ Addr: cfg.Addr, Password: cfg.Password, DB: cfg.DB, ContextTimeoutEnabled: true, }) if err := client.Ping(ctx).Err(); err != nil { _ = client.Close() return nil, err } return &Redis{client: client, prefix: cfg.KeyPrefix}, nil } func (r *Redis) Close() error { if r == nil || r.client == nil { return nil } return r.client.Close() } func (r *Redis) Ping(ctx context.Context) error { if r == nil || r.client == nil { return nil } return r.client.Ping(ctx).Err() } func (r *Redis) TryLock(ctx context.Context, key string, owner string, ttl time.Duration) (bool, UnlockFunc, error) { if r == nil || r.client == nil { return true, func(context.Context) error { return nil }, nil } if strings.TrimSpace(owner) == "" { return false, nil, errors.New("lock owner is required") } redisKey := r.key(key) ok, err := r.client.SetNX(ctx, redisKey, owner, ttl).Result() if err != nil || !ok { return ok, nil, err } return true, func(unlockCtx context.Context) error { const script = ` if redis.call("GET", KEYS[1]) == ARGV[1] then return redis.call("DEL", KEYS[1]) end return 0` return r.client.Eval(unlockCtx, script, []string{redisKey}, owner).Err() }, nil } // RefreshLock 只延长仍由当前 owner 持有的锁;任务 worker 续租 DB lease 时同步续租 Redis, // 防止大导出超过初始 TTL 后被另一实例重复执行并覆盖产物。 func (r *Redis) RefreshLock(ctx context.Context, key string, owner string, ttl time.Duration) (bool, error) { if r == nil || r.client == nil { return true, nil } if strings.TrimSpace(owner) == "" || ttl <= 0 { return false, errors.New("lock owner and positive ttl are required") } const script = ` if redis.call("GET", KEYS[1]) == ARGV[1] then return redis.call("PEXPIRE", KEYS[1], ARGV[2]) end return 0` result, err := r.client.Eval(ctx, script, []string{r.key(key)}, owner, ttl.Milliseconds()).Int64() return result == 1, err } // ConsumeFixedWindow atomically increments all supplied counters and decides whether // every limit still allows the request. A nil/disabled Redis is an availability error, // not an implicit allow: security callers must fail closed instead of falling back to // per-process memory that attackers could bypass by switching replicas. func (r *Redis) ConsumeFixedWindow(ctx context.Context, window time.Duration, rules []FixedWindowRule) (FixedWindowDecision, error) { if r == nil || r.client == nil { return FixedWindowDecision{}, ErrRedisUnavailable } if window < time.Millisecond || len(rules) == 0 { return FixedWindowDecision{}, errors.New("positive fixed window and at least one rule are required") } keys := make([]string, 0, len(rules)) args := make([]any, 0, len(rules)+1) args = append(args, window.Milliseconds()) seen := make(map[string]struct{}, len(rules)) for _, rule := range rules { key := strings.TrimSpace(rule.Key) if key == "" || rule.Limit == 0 { return FixedWindowDecision{}, errors.New("fixed-window keys and limits must be non-empty") } key = r.key(key) if _, exists := seen[key]; exists { return FixedWindowDecision{}, fmt.Errorf("duplicate fixed-window key %q", key) } seen[key] = struct{}{} keys = append(keys, key) args = append(args, strconv.FormatUint(rule.Limit, 10)) } result, err := r.client.Eval(ctx, fixedWindowScript, keys, args...).Int64Slice() if err != nil { return FixedWindowDecision{}, fmt.Errorf("%w: %v", ErrRedisUnavailable, err) } if len(result) != 2 || (result[0] != 0 && result[0] != 1) { return FixedWindowDecision{}, fmt.Errorf("%w: invalid fixed-window response", ErrRedisUnavailable) } retryAfter := time.Duration(result[1]) * time.Millisecond if result[0] == 0 && retryAfter <= 0 { retryAfter = window } return FixedWindowDecision{Allowed: result[0] == 1, RetryAfter: retryAfter}, nil } func (r *Redis) key(key string) string { if r.prefix == "" { return key } return r.prefix + strings.TrimPrefix(key, r.prefix) }