// Package ratelimit 内存限流与配额(MVP 起步,Redis 后置)。 // 覆盖:密钥级每日请求数 / 每日 token 数配额、用户级每秒速率。 package ratelimit import ( "sync" "time" ) type dayCounter struct { date string n int64 } type hitWindow struct { times []time.Time limit int window time.Duration } // Limiter 内存计数器。并发安全。 type Limiter struct { mu sync.Mutex reqDaily map[uint64]*dayCounter // 密钥每日请求数 tokDaily map[uint64]*dayCounter // 密钥每日 token 用量 userHits map[uint64]*hitWindow // 用户速率窗口 authHits map[string]*hitWindow // 认证端点速率窗口(按 IP / IP+账号) } func New() *Limiter { return &Limiter{ reqDaily: map[uint64]*dayCounter{}, tokDaily: map[uint64]*dayCounter{}, userHits: map[uint64]*hitWindow{}, authHits: map[string]*hitWindow{}, } } func today() string { return time.Now().UTC().Format("2006-01-02") } // AllowRequestDaily 检查并计数密钥每日请求配额;无配额(limit<=0)时仅计数。 // 返回 false 表示超过配额。 func (l *Limiter) AllowRequestDaily(keyID uint64, limit int) bool { l.mu.Lock() defer l.mu.Unlock() d := today() c, ok := l.reqDaily[keyID] if !ok || c.date != d { c = &dayCounter{date: d} l.reqDaily[keyID] = c } c.n++ if limit > 0 && c.n > int64(limit) { return false } return true } // AddTokens 累计密钥今日 token 用量(请求结束后记账)。 func (l *Limiter) AddTokens(keyID uint64, n int64) { if n <= 0 { return } l.mu.Lock() defer l.mu.Unlock() d := today() c, ok := l.tokDaily[keyID] if !ok || c.date != d { c = &dayCounter{date: d} l.tokDaily[keyID] = c } c.n += n } // TokensUsed 返回密钥今日已用 token。 func (l *Limiter) TokensUsed(keyID uint64) int64 { l.mu.Lock() defer l.mu.Unlock() c, ok := l.tokDaily[keyID] if !ok || c.date != today() { return 0 } return c.n } // AllowUserRate 用户级每秒请求速率限制(滑动窗口);limit<=0 不限制。 func (l *Limiter) AllowUserRate(userID uint64, limit int) bool { if limit <= 0 { return true } l.mu.Lock() defer l.mu.Unlock() now := time.Now() win := time.Second w, ok := l.userHits[userID] if !ok || w.limit != limit || w.window != win { w = &hitWindow{limit: limit, window: win} l.userHits[userID] = w } // 清理窗口外的时间戳 cutoff := now.Add(-win) keep := w.times[:0] for _, t := range w.times { if t.After(cutoff) { keep = append(keep, t) } } w.times = keep if len(w.times) >= limit { return false } w.times = append(w.times, now) return true } // AllowWindow 字符串键滑动窗口限速;limit<=0 表示不限制,返回 false 表示已超限。 // 用于认证端点(登录爆破、批量注册),键由调用方拼出(IP、IP+用户名等)。 // 窗口内已无命中时删除条目,避免陌生键把 map 撑大。 func (l *Limiter) AllowWindow(key string, limit int, window time.Duration) bool { if limit <= 0 { return true } if window <= 0 { window = time.Minute } l.mu.Lock() defer l.mu.Unlock() now := time.Now() w, ok := l.authHits[key] if !ok { w = &hitWindow{limit: limit, window: window} l.authHits[key] = w } cutoff := now.Add(-w.window) keep := w.times[:0] for _, t := range w.times { if t.After(cutoff) { keep = append(keep, t) } } w.times = keep if len(w.times) >= w.limit { return false } w.times = append(w.times, now) // 键由外部拼出(IP 等),条目只增不减会长期堆积:规模到阈值就整体扫一遍过期键 if len(l.authHits) > 4096 { for k, v := range l.authHits { vcut := now.Add(-v.window) alive := v.times[:0] for _, t := range v.times { if t.After(vcut) { alive = append(alive, t) } } v.times = alive if len(v.times) == 0 { delete(l.authHits, k) } } } return true }