90CNServer/server/internal/platform/ratelimit/redis_limiter.go

package ratelimit

import (
	"context"
	"encoding/base64"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"longheng.io/server/internal/platform/db"
)

// ErrRedisExecutorRequired 在 RedisLimiter 没有配置 executor 时返回。
var ErrRedisExecutorRequired = errors.New("ratelimit redis executor is required")

// FailPolicy 描述 Redis 不可用时的行为。
type FailPolicy string

const (
	// FailOpen 在 Redis 出错或超时时放行请求（适合非关键 HTTP 路径）。
	FailOpen FailPolicy = "open"
	// FailClosed 在 Redis 出错或超时时拒绝请求（适合管理接口、危险操作）。
	FailClosed FailPolicy = "closed"
)

// 默认的 token bucket Lua 脚本：所有限流计算在 Redis 内部一次完成，避免
// "GET → 计算 → SET" 类竞争。返回 {allowed, retry_after_ms}。
//
// KEYS[1]   bucket key（每个 rule + identity 一个）
// ARGV[1]   capacity（桶容量，浮点）
// ARGV[2]   refill_per_second（每秒补充令牌数，浮点）
// ARGV[3]   now_ms（当前时间，毫秒）
// ARGV[4]   ttl_ms（key 过期时间，毫秒，应该 >= bucket 满后多保留一段）
// #nosec G101 -- Lua 脚本里的 token 表示限流令牌，不是硬编码凭证。
const redisBucketScript = `
local key = KEYS[1]
local capacity = tonumber(ARGV[1])
local refill_per_second = tonumber(ARGV[2])
local now_ms = tonumber(ARGV[3])
local ttl_ms = tonumber(ARGV[4])
if capacity <= 0 then return {0, ttl_ms} end
if refill_per_second <= 0 then return {0, ttl_ms} end
local data = redis.call("HMGET", key, "tokens", "updated_at_ms")
local tokens = tonumber(data[1])
local last_ms = tonumber(data[2])
if tokens == nil then
  tokens = capacity
  last_ms = now_ms
end
local elapsed_ms = now_ms - last_ms
if elapsed_ms < 0 then elapsed_ms = 0 end
tokens = tokens + (elapsed_ms / 1000.0) * refill_per_second
if tokens > capacity then tokens = capacity end
local allowed = 0
local retry_after_ms = 0
if tokens >= 1 then
  tokens = tokens - 1
  allowed = 1
else
  local missing = 1 - tokens
  retry_after_ms = math.ceil(missing / refill_per_second * 1000.0)
  if retry_after_ms < 1 then retry_after_ms = 1 end
end
redis.call("HMSET", key, "tokens", tostring(tokens), "updated_at_ms", tostring(now_ms))
redis.call("PEXPIRE", key, ttl_ms)
return {allowed, retry_after_ms}
`

// RedisConfig 用来构造 RedisLimiter。Rules / Window / MaxRequests 语义与单机
// Limiter 完全一致，便于同一套配置在 single-node 演示和 production 之间切换。
type RedisConfig struct {
	Enabled     bool
	Executor    db.RedisExecutor
	KeyPrefix   string
	Window      time.Duration
	MaxRequests int
	Rules       []Rule
	FailPolicy  FailPolicy
	Now         func() time.Time
}

// RedisLimiter 用 Redis Lua token bucket 提供跨节点一致的限流。
// 多个 gateway / logic / adminops 实例共用一组 Redis 即可看到同一个桶。
//
// 与 Limiter 的关键差异：
//   - 失败需要返回 error，HTTP 包装路径会根据 FailPolicy 决定放行或拒绝。
//   - 不再维护进程内 buckets / cleanup；TTL 交给 Redis。
//   - 不支持 trusted-proxy header 识别身份，调用方传 identity；HTTP 包装
//     仍可复用 *Limiter 的 identity 提取逻辑，但 RedisLimiter 自身保持
//     "传 key + path 进来" 的纯函数语义。
type RedisLimiter struct {
	enabled     bool
	executor    db.RedisExecutor
	keyPrefix   string
	defaultRule Rule
	rules       []Rule
	failPolicy  FailPolicy
	now         func() time.Time
}

// NewRedis 构造 RedisLimiter；当 executor 为 nil 时返回错误，避免静默退化为单机。
func NewRedis(cfg RedisConfig) (*RedisLimiter, error) {
	if cfg.Executor == nil {
		return nil, ErrRedisExecutorRequired
	}
	if cfg.Window <= 0 {
		cfg.Window = time.Second
	}
	if cfg.MaxRequests <= 0 {
		cfg.MaxRequests = 60
	}
	prefix := strings.Trim(strings.TrimSpace(cfg.KeyPrefix), ":")
	if prefix == "" {
		prefix = "longheng:ratelimit"
	}
	now := cfg.Now
	if now == nil {
		now = time.Now
	}
	policy := normalizeFailPolicy(cfg.FailPolicy)
	rules := normalizeRules(cfg.Rules, cfg.Window, cfg.MaxRequests)
	return &RedisLimiter{
		enabled:   cfg.Enabled,
		executor:  cfg.Executor,
		keyPrefix: prefix,
		defaultRule: Rule{
			Path:        "*",
			Window:      cfg.Window,
			WindowSec:   int64(cfg.Window / time.Second),
			MaxRequests: cfg.MaxRequests,
			Prefix:      true,
		},
		rules:      rules,
		failPolicy: policy,
		now:        now,
	}, nil
}

// Enabled 让 RedisLimiter 与 *Limiter 在调用上保持兼容。
func (r *RedisLimiter) Enabled() bool {
	if r == nil {
		return false
	}
	return r.enabled
}

// FailPolicy 暴露当前的失败策略，便于上层在调用 AllowKeyCtx 出错时做决策。
func (r *RedisLimiter) FailPolicy() FailPolicy {
	if r == nil {
		return FailOpen
	}
	return normalizeFailPolicy(r.failPolicy)
}

// AllowKeyCtx 是核心入口。返回 allowed / retry-after / error。
// error != nil 时调用方应配合 FailPolicy 决定放行还是拒绝。
func (r *RedisLimiter) AllowKeyCtx(ctx context.Context, key string, path string) (bool, time.Duration, error) {
	if r == nil {
		return true, 0, nil
	}
	if !r.enabled {
		return true, 0, nil
	}
	if r.executor == nil {
		return false, 0, ErrRedisExecutorRequired
	}
	if ctx == nil {
		ctx = context.Background()
	}
	if err := ctx.Err(); err != nil {
		return false, 0, err
	}
	key = strings.TrimSpace(key)
	if key == "" {
		key = "unknown"
	}
	rule := r.matchRule(path)
	window := rule.Window
	if window <= 0 {
		window = time.Second
	}
	capacity := float64(rule.MaxRequests)
	if capacity <= 0 {
		capacity = 1
	}
	refillPerSecond := capacity / window.Seconds()
	if refillPerSecond <= 0 {
		refillPerSecond = capacity
	}
	nowFunc := r.now
	if nowFunc == nil {
		nowFunc = time.Now
	}
	now := nowFunc()
	bucketKey := r.bucketKey(key, rule)
	// TTL 给 bucket 至少留两个 window 的余量，避免活跃用户被频繁重建。
	ttlMS := int64(window/time.Millisecond)*2 + 1000

	raw, err := r.executor.Do(ctx, "EVAL", redisBucketScript,
		1, bucketKey,
		strconv.FormatFloat(capacity, 'f', -1, 64),
		strconv.FormatFloat(refillPerSecond, 'f', -1, 64),
		strconv.FormatInt(now.UnixNano()/int64(time.Millisecond), 10),
		strconv.FormatInt(ttlMS, 10),
	)
	if err != nil {
		return false, 0, fmt.Errorf("ratelimit redis EVAL failed: %w", err)
	}
	allowed, retryAfter, err := parseBucketResp(raw)
	if err != nil {
		return false, 0, err
	}
	return allowed, retryAfter, nil
}

// Allow 提供与 *Limiter 一致的便捷形式，使用 Background context。
// 注意：失败语义由 FailPolicy 决定 —— FailOpen 放行、FailClosed 拒绝。
// 需要保留 error 的调用方应直接使用 AllowKeyCtx。
func (r *RedisLimiter) Allow(identity, path string) (bool, time.Duration) {
	allowed, retry, err := r.AllowKeyCtx(context.Background(), identity, path)
	if err != nil {
		return r.FailPolicy() == FailOpen, retry
	}
	return allowed, retry
}

// AllowKey 与 *Limiter.AllowKey 同名，便于在 wiring 层用 Allower 接口互换实现。
func (r *RedisLimiter) AllowKey(key, path string) (bool, time.Duration) {
	return r.Allow(key, path)
}

func (r *RedisLimiter) matchRule(path string) Rule {
	path = strings.TrimSpace(path)
	best := Rule{}
	for _, rule := range r.rules {
		if rule.Prefix {
			if strings.HasPrefix(path, rule.Path) && len(rule.Path) > len(best.Path) {
				best = rule
			}
			continue
		}
		if path == rule.Path && len(rule.Path) > len(best.Path) {
			best = rule
		}
	}
	if best.Path != "" {
		return best
	}
	return r.fallbackRule()
}

func (r *RedisLimiter) bucketKey(identity string, rule Rule) string {
	rulePath := rule.Path
	if rulePath == "" {
		rulePath = "*"
	}
	prefix := strings.Trim(strings.TrimSpace(r.keyPrefix), ":")
	if prefix == "" {
		prefix = "longheng:ratelimit"
	}
	return prefix + ":" + encodeRedisPart(identity) + ":" + encodeRedisPart(rulePath)
}

func (r *RedisLimiter) fallbackRule() Rule {
	rule := r.defaultRule
	if rule.Path == "" {
		rule.Path = "*"
		rule.Prefix = true
	}
	if rule.Window <= 0 {
		rule.Window = time.Second
	}
	if rule.MaxRequests <= 0 {
		rule.MaxRequests = 60
	}
	rule.WindowSec = int64(rule.Window / time.Second)
	return rule
}

func encodeRedisPart(value string) string {
	if value == "" {
		return "_"
	}
	return base64.RawURLEncoding.EncodeToString([]byte(value))
}

func normalizeFailPolicy(v FailPolicy) FailPolicy {
	switch FailPolicy(strings.ToLower(strings.TrimSpace(string(v)))) {
	case FailOpen, "":
		return FailOpen
	case FailClosed:
		return FailClosed
	default:
		return FailOpen
	}
}

// parseBucketResp 解析 EVAL 的返回值。RedisExecutor.Do 返回 any，
// 实际类型取决于驱动：redigo 返回 []interface{} + 元素为 int64；其他驱动
// 可能返回 []any 或 nested 类型。这里兼容常见情况，遇到未知形态时给出
// 明确错误，避免静默放行。
func parseBucketResp(raw any) (bool, time.Duration, error) {
	switch v := raw.(type) {
	case []any:
		return parseBucketArray(v)
	default:
		return false, 0, fmt.Errorf("ratelimit redis returned unexpected type %T", raw)
	}
}

func parseBucketArray(values []any) (bool, time.Duration, error) {
	if len(values) < 2 {
		return false, 0, fmt.Errorf("ratelimit redis returned %d fields, want 2", len(values))
	}
	allowed, err := toInt64(values[0])
	if err != nil {
		return false, 0, fmt.Errorf("ratelimit redis allowed field: %w", err)
	}
	retryMS, err := toInt64(values[1])
	if err != nil {
		return false, 0, fmt.Errorf("ratelimit redis retry_after field: %w", err)
	}
	if retryMS < 0 {
		retryMS = 0
	}
	return allowed == 1, time.Duration(retryMS) * time.Millisecond, nil
}

func toInt64(v any) (int64, error) {
	switch n := v.(type) {
	case int64:
		return n, nil
	case int:
		return int64(n), nil
	case int32:
		return int64(n), nil
	case uint64:
		if n > uint64(1<<63-1) {
			return 0, fmt.Errorf("uint64 value overflows int64: %d", n)
		}
		return int64(n), nil
	case []byte:
		return strconv.ParseInt(string(n), 10, 64)
	case string:
		return strconv.ParseInt(n, 10, 64)
	default:
		return 0, fmt.Errorf("unsupported numeric type %T", v)
	}
}