hyapp-server/server/admin/internal/cache/redis.go

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