openteam/server/internal/pkg/ratelimit/ratelimit.go

// 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
	}
	// 代理热路径顺带做回收：停用密钥的旧日期条目和过期的用户窗口不再堆积
	if len(l.reqDaily) > dailySweepThreshold {
		l.sweepLocked(d)
	}
	return true
}

const dailySweepThreshold = 4096

// sweepLocked 回收不再活跃的条目（调用方持锁）：
//   - 日计数器：日期不是今天的条目下次使用本来就会清零，直接删除
//   - 用户速率窗口：窗口内无命中的条目删除
func (l *Limiter) sweepLocked(today string) {
	now := time.Now()
	for k, c := range l.reqDaily {
		if c.date != today {
			delete(l.reqDaily, k)
		}
	}
	for k, c := range l.tokDaily {
		if c.date != today {
			delete(l.tokDaily, k)
		}
	}
	for k, w := range l.userHits {
		cutoff := now.Add(-w.window)
		alive := 0
		for _, t := range w.times {
			if t.After(cutoff) {
				alive++
			}
		}
		if alive == 0 {
			delete(l.userHits, k)
		}
	}
}

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