DNF90 服务端与启动器完整源码（含 GM 工具功能增强）
服务端（Go）： - GM 塞背包/带属性装备发放：在线角色推送 op14 增量物品行，背包实时刷新，无需回角色选择界面 - GM 改金币/改等级：去掉在线拒绝，落库后推送 op14 钱包/op37 等级状态，在线实时刷新 - 等级端点统一切换到新版 GMSetCharacterLevel（自带在线刷新） - 无限/锁死疲劳模式 + 燃烧疲劳：每房间消耗 2 点、打怪与通关经验双倍 - 燃烧疲劳开关时给所有在线角色推送系统提示（op12 mode 0） - 深渊史诗掉率：GM 全局倍率 + 城镇增益药剂窗口倍率，真实接入掉落结算 启动器（Python/JS）： - 标题栏自绘、窗口 1204x717、禁止拖拽改大小、仅最小化/关闭 - GM 工具：疲劳与燃烧、深渊史诗掉率、拍卖行上架（单一定价对齐） - 拍卖行上架前端对齐服务端单一价格语义，修复 400 报错
This commit is contained in:
hanfuit committed
commit
645b64fe21
4388 files changed
+1405457
No files matched your search
@@ -0,0 +1,545 @@
|
||||
package ratelimit
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Config 定义单机限流器的开关、默认规则和可信代理来源。
|
||||
type Config struct {
|
||||
Enabled bool
|
||||
Algorithm string
|
||||
Window time.Duration
|
||||
MaxRequests int
|
||||
Rules []Rule
|
||||
CleanupEvery time.Duration
|
||||
TrustedHeader string
|
||||
TrustedProxyCIDRs []string
|
||||
}
|
||||
|
||||
// Rule 定义一个 path 的窗口、最大请求数和前缀匹配方式。
|
||||
type Rule struct {
|
||||
Path string `json:"path"`
|
||||
Window time.Duration `json:"-"`
|
||||
WindowSec int64 `json:"window_seconds"`
|
||||
MaxRequests int `json:"max_requests"`
|
||||
Prefix bool `json:"prefix"`
|
||||
}
|
||||
|
||||
// Snapshot 描述限流器当前配置和内存桶数量。
|
||||
type Snapshot struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
Algorithm string `json:"algorithm"`
|
||||
WindowSeconds int64 `json:"window_seconds"`
|
||||
MaxRequests int `json:"max_requests"`
|
||||
Rules []Rule `json:"rules"`
|
||||
Buckets int `json:"buckets"`
|
||||
}
|
||||
|
||||
// Limiter 提供单进程内存限流，适合单节点或 Redis 限流不可用时的降级场景。
|
||||
type Limiter struct {
|
||||
enabled bool
|
||||
algorithm string
|
||||
defaultRule Rule
|
||||
rules []Rule
|
||||
cleanupEvery time.Duration
|
||||
trustedHeader string
|
||||
trustedProxyCIDRs []*net.IPNet
|
||||
|
||||
mu sync.Mutex
|
||||
buckets map[string]bucket
|
||||
lastCleanup time.Time
|
||||
now func() time.Time
|
||||
}
|
||||
|
||||
type bucket struct {
|
||||
windowStart time.Time
|
||||
count int
|
||||
window time.Duration
|
||||
tokens float64
|
||||
updatedAt time.Time
|
||||
}
|
||||
|
||||
const (
|
||||
AlgorithmWindow = "window"
|
||||
AlgorithmTokenBucket = "token_bucket"
|
||||
)
|
||||
|
||||
// New 根据配置创建单机限流器，并补齐安全默认值。
|
||||
func New(cfg Config) *Limiter {
|
||||
if cfg.Window <= 0 {
|
||||
cfg.Window = time.Second
|
||||
}
|
||||
if cfg.MaxRequests <= 0 {
|
||||
cfg.MaxRequests = 60
|
||||
}
|
||||
if cfg.CleanupEvery <= 0 {
|
||||
cfg.CleanupEvery = time.Minute
|
||||
}
|
||||
cfg.Algorithm = normalizeAlgorithm(cfg.Algorithm)
|
||||
if !ValidAlgorithm(cfg.Algorithm) {
|
||||
cfg.Algorithm = AlgorithmWindow
|
||||
}
|
||||
rules := normalizeRules(cfg.Rules, cfg.Window, cfg.MaxRequests)
|
||||
return &Limiter{
|
||||
enabled: cfg.Enabled,
|
||||
algorithm: cfg.Algorithm,
|
||||
defaultRule: Rule{
|
||||
Path: "*",
|
||||
Window: cfg.Window,
|
||||
WindowSec: int64(cfg.Window / time.Second),
|
||||
MaxRequests: cfg.MaxRequests,
|
||||
Prefix: true,
|
||||
},
|
||||
rules: rules,
|
||||
cleanupEvery: cfg.CleanupEvery,
|
||||
trustedHeader: strings.TrimSpace(cfg.TrustedHeader),
|
||||
trustedProxyCIDRs: parseProxyCIDRs(cfg.TrustedProxyCIDRs),
|
||||
buckets: make(map[string]bucket),
|
||||
now: time.Now,
|
||||
}
|
||||
}
|
||||
|
||||
// Enabled 返回当前限流开关状态。
|
||||
func (l *Limiter) Enabled() bool {
|
||||
if l == nil {
|
||||
return false
|
||||
}
|
||||
l.mu.Lock()
|
||||
l.ensureReadyLocked()
|
||||
enabled := l.enabled
|
||||
l.mu.Unlock()
|
||||
return enabled
|
||||
}
|
||||
|
||||
// Wrap 将限流器包装到 HTTP handler 前面。
|
||||
func (l *Limiter) Wrap(next http.Handler) http.Handler {
|
||||
if l == nil {
|
||||
return next
|
||||
}
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
key := l.identity(r)
|
||||
path := r.URL.Path
|
||||
allowed, retryAfter := l.Allow(key, path)
|
||||
if !allowed {
|
||||
if retryAfter < time.Second {
|
||||
retryAfter = time.Second
|
||||
}
|
||||
w.Header().Set("Retry-After", strconv.Itoa(int(retryAfter/time.Second)))
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusTooManyRequests)
|
||||
_, _ = w.Write([]byte(`{"error":"rate limit exceeded"}` + "\n"))
|
||||
return
|
||||
}
|
||||
next.ServeHTTP(w, r)
|
||||
})
|
||||
}
|
||||
|
||||
// WrapFunc 将限流器包装到 HTTP handler 函数前面。
|
||||
func (l *Limiter) WrapFunc(next http.HandlerFunc) http.HandlerFunc {
|
||||
if l == nil {
|
||||
return next
|
||||
}
|
||||
return l.Wrap(next).ServeHTTP
|
||||
}
|
||||
|
||||
// Allow 使用 identity/path 判断一次请求是否允许通过。
|
||||
func (l *Limiter) Allow(identity, path string) (bool, time.Duration) {
|
||||
return l.AllowKey(identity, path)
|
||||
}
|
||||
|
||||
// AllowComposite 将多个业务维度组成安全 key 后执行限流判断。
|
||||
func (l *Limiter) AllowComposite(path string, parts ...string) (bool, time.Duration) {
|
||||
return l.AllowKey(CompositeKey(parts...), path)
|
||||
}
|
||||
|
||||
// AllowKey 使用调用方给出的稳定 key 执行限流判断。
|
||||
func (l *Limiter) AllowKey(key, path string) (bool, time.Duration) {
|
||||
if l == nil {
|
||||
return true, 0
|
||||
}
|
||||
key = strings.TrimSpace(key)
|
||||
if key == "" {
|
||||
key = "unknown"
|
||||
}
|
||||
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
l.ensureReadyLocked()
|
||||
if !l.enabled {
|
||||
return true, 0
|
||||
}
|
||||
now := l.now()
|
||||
rule := l.matchLocked(path)
|
||||
window := rule.Window
|
||||
if window <= 0 {
|
||||
window = time.Second
|
||||
}
|
||||
l.cleanupLocked(now)
|
||||
|
||||
bucketKey := key + "\x00" + rule.Path
|
||||
if l.algorithm == AlgorithmTokenBucket {
|
||||
return l.allowBucketLocked(bucketKey, rule, window, now)
|
||||
}
|
||||
return l.allowWindowLocked(bucketKey, rule, window, now)
|
||||
}
|
||||
|
||||
func (l *Limiter) allowWindowLocked(key string, rule Rule, window time.Duration, now time.Time) (bool, time.Duration) {
|
||||
b := l.buckets[key]
|
||||
if b.windowStart.IsZero() || now.Sub(b.windowStart) >= window {
|
||||
b = bucket{windowStart: now, window: window}
|
||||
}
|
||||
b.window = window
|
||||
b.count++
|
||||
l.buckets[key] = b
|
||||
if b.count <= rule.MaxRequests {
|
||||
return true, 0
|
||||
}
|
||||
return false, window - now.Sub(b.windowStart)
|
||||
}
|
||||
|
||||
func (l *Limiter) allowBucketLocked(key string, rule Rule, window time.Duration, now time.Time) (bool, time.Duration) {
|
||||
capacity := float64(rule.MaxRequests)
|
||||
if capacity <= 0 {
|
||||
capacity = 1
|
||||
}
|
||||
refillPerSecond := capacity / window.Seconds()
|
||||
if refillPerSecond <= 0 {
|
||||
refillPerSecond = capacity
|
||||
}
|
||||
b := l.buckets[key]
|
||||
if b.updatedAt.IsZero() {
|
||||
b = bucket{windowStart: now, window: window, tokens: capacity, updatedAt: now}
|
||||
}
|
||||
if elapsed := now.Sub(b.updatedAt).Seconds(); elapsed > 0 {
|
||||
b.tokens += elapsed * refillPerSecond
|
||||
if b.tokens > capacity {
|
||||
b.tokens = capacity
|
||||
}
|
||||
}
|
||||
b.window = window
|
||||
b.updatedAt = now
|
||||
if b.tokens >= 1 {
|
||||
b.tokens--
|
||||
b.count++
|
||||
l.buckets[key] = b
|
||||
return true, 0
|
||||
}
|
||||
l.buckets[key] = b
|
||||
missing := 1 - b.tokens
|
||||
retry := time.Duration(missing/refillPerSecond*float64(time.Second) + 0.5)
|
||||
if retry <= 0 {
|
||||
retry = time.Nanosecond
|
||||
}
|
||||
return false, retry
|
||||
}
|
||||
|
||||
// Configure 原子替换限流配置，并清空旧桶避免不同规则间互相污染。
|
||||
func (l *Limiter) Configure(cfg Config) {
|
||||
if l == nil {
|
||||
return
|
||||
}
|
||||
if cfg.Window <= 0 {
|
||||
cfg.Window = time.Second
|
||||
}
|
||||
if cfg.MaxRequests <= 0 {
|
||||
cfg.MaxRequests = 60
|
||||
}
|
||||
if cfg.CleanupEvery <= 0 {
|
||||
cfg.CleanupEvery = time.Minute
|
||||
}
|
||||
cfg.Algorithm = normalizeAlgorithm(cfg.Algorithm)
|
||||
if !ValidAlgorithm(cfg.Algorithm) {
|
||||
cfg.Algorithm = AlgorithmWindow
|
||||
}
|
||||
rules := normalizeRules(cfg.Rules, cfg.Window, cfg.MaxRequests)
|
||||
l.mu.Lock()
|
||||
l.enabled = cfg.Enabled
|
||||
l.algorithm = cfg.Algorithm
|
||||
l.defaultRule = Rule{
|
||||
Path: "*",
|
||||
Window: cfg.Window,
|
||||
WindowSec: int64(cfg.Window / time.Second),
|
||||
MaxRequests: cfg.MaxRequests,
|
||||
Prefix: true,
|
||||
}
|
||||
l.rules = rules
|
||||
l.cleanupEvery = cfg.CleanupEvery
|
||||
l.trustedHeader = strings.TrimSpace(cfg.TrustedHeader)
|
||||
l.trustedProxyCIDRs = parseProxyCIDRs(cfg.TrustedProxyCIDRs)
|
||||
l.buckets = make(map[string]bucket)
|
||||
l.lastCleanup = time.Time{}
|
||||
l.ensureReadyLocked()
|
||||
l.mu.Unlock()
|
||||
}
|
||||
|
||||
// Snapshot 返回当前限流配置和桶数量。
|
||||
func (l *Limiter) Snapshot() Snapshot {
|
||||
if l == nil {
|
||||
return Snapshot{}
|
||||
}
|
||||
l.mu.Lock()
|
||||
l.ensureReadyLocked()
|
||||
buckets := len(l.buckets)
|
||||
rules := make([]Rule, len(l.rules))
|
||||
copy(rules, l.rules)
|
||||
snapshot := Snapshot{
|
||||
Enabled: l.enabled,
|
||||
Algorithm: l.algorithm,
|
||||
WindowSeconds: int64(l.defaultRule.Window / time.Second),
|
||||
MaxRequests: l.defaultRule.MaxRequests,
|
||||
Rules: rules,
|
||||
Buckets: buckets,
|
||||
}
|
||||
l.mu.Unlock()
|
||||
return snapshot
|
||||
}
|
||||
|
||||
func (l *Limiter) ensureReadyLocked() {
|
||||
if l.algorithm == "" {
|
||||
l.algorithm = AlgorithmWindow
|
||||
}
|
||||
if !ValidAlgorithm(l.algorithm) {
|
||||
l.algorithm = AlgorithmWindow
|
||||
}
|
||||
if l.defaultRule.Path == "" {
|
||||
l.defaultRule.Path = "*"
|
||||
l.defaultRule.Prefix = true
|
||||
}
|
||||
if l.defaultRule.Window <= 0 {
|
||||
l.defaultRule.Window = time.Second
|
||||
}
|
||||
if l.defaultRule.MaxRequests <= 0 {
|
||||
l.defaultRule.MaxRequests = 60
|
||||
}
|
||||
l.defaultRule.WindowSec = int64(l.defaultRule.Window / time.Second)
|
||||
if l.cleanupEvery <= 0 {
|
||||
l.cleanupEvery = time.Minute
|
||||
}
|
||||
if l.buckets == nil {
|
||||
l.buckets = make(map[string]bucket)
|
||||
}
|
||||
if l.now == nil {
|
||||
l.now = time.Now
|
||||
}
|
||||
}
|
||||
|
||||
func (l *Limiter) matchLocked(path string) Rule {
|
||||
path = strings.TrimSpace(path)
|
||||
best := Rule{}
|
||||
for _, rule := range l.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 l.defaultRule
|
||||
}
|
||||
|
||||
func (l *Limiter) cleanupLocked(now time.Time) {
|
||||
if l.lastCleanup.IsZero() {
|
||||
l.lastCleanup = now
|
||||
return
|
||||
}
|
||||
if now.Sub(l.lastCleanup) < l.cleanupEvery {
|
||||
return
|
||||
}
|
||||
for key, b := range l.buckets {
|
||||
window := b.window
|
||||
if window <= 0 {
|
||||
window = l.defaultRule.Window
|
||||
}
|
||||
base := b.windowStart
|
||||
if l.algorithm == AlgorithmTokenBucket && !b.updatedAt.IsZero() {
|
||||
base = b.updatedAt
|
||||
}
|
||||
if now.Sub(base) > 2*window+l.cleanupEvery {
|
||||
delete(l.buckets, key)
|
||||
}
|
||||
}
|
||||
l.lastCleanup = now
|
||||
}
|
||||
|
||||
func (l *Limiter) identity(r *http.Request) string {
|
||||
l.mu.Lock()
|
||||
trustedHeader := l.trustedHeader
|
||||
trustedProxyCIDRs := append([]*net.IPNet(nil), l.trustedProxyCIDRs...)
|
||||
l.mu.Unlock()
|
||||
if trustedHeader != "" && trustedProxyAllowed(r.RemoteAddr, trustedProxyCIDRs) {
|
||||
if value := strings.TrimSpace(r.Header.Get(trustedHeader)); value != "" {
|
||||
if idx := strings.IndexByte(value, ','); idx >= 0 {
|
||||
value = strings.TrimSpace(value[:idx])
|
||||
}
|
||||
return value
|
||||
}
|
||||
}
|
||||
host, _, err := net.SplitHostPort(r.RemoteAddr)
|
||||
if err == nil && host != "" {
|
||||
return host
|
||||
}
|
||||
return strings.TrimSpace(r.RemoteAddr)
|
||||
}
|
||||
|
||||
func trustedProxyAllowed(remoteAddr string, trustedProxyCIDRs []*net.IPNet) bool {
|
||||
if len(trustedProxyCIDRs) == 0 {
|
||||
return true
|
||||
}
|
||||
host, _, err := net.SplitHostPort(strings.TrimSpace(remoteAddr))
|
||||
if err != nil {
|
||||
host = strings.TrimSpace(remoteAddr)
|
||||
}
|
||||
ip := net.ParseIP(host)
|
||||
if ip == nil {
|
||||
return false
|
||||
}
|
||||
for _, cidr := range trustedProxyCIDRs {
|
||||
if cidr.Contains(ip) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func parseProxyCIDRs(values []string) []*net.IPNet {
|
||||
out := make([]*net.IPNet, 0, len(values))
|
||||
for _, value := range values {
|
||||
value = strings.TrimSpace(value)
|
||||
if value == "" {
|
||||
continue
|
||||
}
|
||||
if ip := net.ParseIP(value); ip != nil {
|
||||
bits := 32
|
||||
if ip.To4() == nil {
|
||||
bits = 128
|
||||
}
|
||||
_, network, _ := net.ParseCIDR(ip.String() + "/" + strconv.Itoa(bits))
|
||||
out = append(out, network)
|
||||
continue
|
||||
}
|
||||
if _, network, err := net.ParseCIDR(value); err == nil {
|
||||
out = append(out, network)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// ParseRules 解析 path,window_seconds,max_requests 格式的分号分隔规则。
|
||||
func ParseRules(raw string) ([]Rule, error) {
|
||||
raw = strings.TrimSpace(raw)
|
||||
if raw == "" {
|
||||
return nil, nil
|
||||
}
|
||||
parts := strings.Split(raw, ";")
|
||||
rules := make([]Rule, 0, len(parts))
|
||||
for _, part := range parts {
|
||||
part = strings.TrimSpace(part)
|
||||
if part == "" {
|
||||
continue
|
||||
}
|
||||
fields := strings.Split(part, ",")
|
||||
if len(fields) != 3 {
|
||||
return nil, fmt.Errorf("rate limit rule %q must be path,window_seconds,max_requests", part)
|
||||
}
|
||||
path := strings.TrimSpace(fields[0])
|
||||
if path == "" {
|
||||
return nil, fmt.Errorf("rate limit rule %q has empty path", part)
|
||||
}
|
||||
prefix := false
|
||||
if strings.HasSuffix(path, "*") {
|
||||
prefix = true
|
||||
path = strings.TrimSuffix(path, "*")
|
||||
}
|
||||
if strings.HasSuffix(path, "/") {
|
||||
prefix = true
|
||||
}
|
||||
windowSec, err := strconv.Atoi(strings.TrimSpace(fields[1]))
|
||||
if err != nil || windowSec <= 0 {
|
||||
return nil, fmt.Errorf("rate limit rule %q has invalid window_seconds", part)
|
||||
}
|
||||
maxRequests, err := strconv.Atoi(strings.TrimSpace(fields[2]))
|
||||
if err != nil || maxRequests <= 0 {
|
||||
return nil, fmt.Errorf("rate limit rule %q has invalid max_requests", part)
|
||||
}
|
||||
rules = append(rules, Rule{
|
||||
Path: path,
|
||||
Window: time.Duration(windowSec) * time.Second,
|
||||
WindowSec: int64(windowSec),
|
||||
MaxRequests: maxRequests,
|
||||
Prefix: prefix,
|
||||
})
|
||||
}
|
||||
return rules, nil
|
||||
}
|
||||
|
||||
func normalizeRules(rules []Rule, defaultWindow time.Duration, defaultMax int) []Rule {
|
||||
out := make([]Rule, 0, len(rules))
|
||||
for _, rule := range rules {
|
||||
rule.Path = strings.TrimSpace(rule.Path)
|
||||
if rule.Path == "" {
|
||||
continue
|
||||
}
|
||||
if rule.Window <= 0 {
|
||||
rule.Window = defaultWindow
|
||||
}
|
||||
if rule.MaxRequests <= 0 {
|
||||
rule.MaxRequests = defaultMax
|
||||
}
|
||||
rule.WindowSec = int64(rule.Window / time.Second)
|
||||
out = append(out, rule)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func normalizeAlgorithm(value string) string {
|
||||
value = strings.ToLower(strings.TrimSpace(value))
|
||||
switch value {
|
||||
case "", AlgorithmWindow:
|
||||
return AlgorithmWindow
|
||||
case "tokenbucket", "token-bucket", AlgorithmTokenBucket:
|
||||
return AlgorithmTokenBucket
|
||||
default:
|
||||
return value
|
||||
}
|
||||
}
|
||||
|
||||
// ValidAlgorithm 判断限流算法是否被当前实现支持。
|
||||
func ValidAlgorithm(value string) bool {
|
||||
switch normalizeAlgorithm(value) {
|
||||
case AlgorithmWindow, AlgorithmTokenBucket:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// CompositeKey 使用长度前缀组合多个业务维度，避免分隔符碰撞。
|
||||
func CompositeKey(parts ...string) string {
|
||||
if len(parts) == 0 {
|
||||
return "unknown"
|
||||
}
|
||||
var b strings.Builder
|
||||
for _, part := range parts {
|
||||
part = strings.TrimSpace(part)
|
||||
b.WriteString(strconv.Itoa(len(part)))
|
||||
b.WriteByte(':')
|
||||
b.WriteString(part)
|
||||
b.WriteByte('|')
|
||||
}
|
||||
if b.Len() == 0 {
|
||||
return "unknown"
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
@@ -0,0 +1,351 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user