backend/internal/ratelimit/bucket.go

@@ -0,0 +1,172 @@
+package ratelimit
+
+import (
+	"context"
+	"sync"
+	"time"
+
+	"github.com/hhs/camtalk/internal/config"
+)
+
+// TokenBucket 内存令牌桶，适用于单实例部署。
+type TokenBucket struct {
+	capacity   int       // 桶容量
+	rate       float64   // 每秒填充令牌数
+	tokens     float64   // 当前令牌数
+	lastRefill time.Time // 上次填充时间
+	mu         sync.Mutex
+}
+
+// newTokenBucket 创建令牌桶。
+func newTokenBucket(capacity int, rate float64) *TokenBucket {
+	return &TokenBucket{
+		capacity:   capacity,
+		rate:       rate,
+		tokens:     float64(capacity), // 初始满桶
+		lastRefill: time.Now(),
+	}
+}
+
+// allow 尝试消耗一个令牌。
+func (b *TokenBucket) allow() (bool, time.Duration) {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+
+	now := time.Now()
+	elapsed := now.Sub(b.lastRefill).Seconds()
+
+	// 补充令牌
+	newTokens := elapsed * b.rate
+	b.tokens = min(float64(b.capacity), b.tokens+newTokens)
+	b.lastRefill = now
+
+	// 尝试消耗一个令牌
+	if b.tokens >= 1 {
+		b.tokens -= 1
+		return true, 0
+	}
+
+	// 计算需要等待的时间
+	if b.rate == 0 {
+		// rate=0 时永远无法补充令牌
+		return false, 24 * time.Hour // 返回一个很大的值
+	}
+	retryAfter := time.Duration((1-b.tokens)/b.rate*1000) * time.Millisecond
+	return false, retryAfter
+}
+
+// MemoryLimiter 管理多个用户的令牌桶。
+type MemoryLimiter struct {
+	buckets  map[string]*TokenBucket
+	config   config.RateLimitConfig
+	mu       sync.RWMutex
+	stopOnce sync.Once
+	done     chan struct{}
+}
+
+// NewMemoryLimiter 创建内存限流器。
+func NewMemoryLimiter(cfg config.RateLimitConfig) *MemoryLimiter {
+	limiter := &MemoryLimiter{
+		buckets: make(map[string]*TokenBucket),
+		config:  cfg,
+		done:    make(chan struct{}),
+	}
+
+	// 启动后台清理 goroutine
+	go limiter.cleanup()
+
+	return limiter
+}
+
+// Allow 实现 Limiter 接口。
+func (l *MemoryLimiter) Allow(ctx context.Context, key string) (bool, time.Duration) {
+	bucket := l.getOrCreateBucket(key)
+	return bucket.allow()
+}
+
+// Stop 实现 Limiter 接口。
+func (l *MemoryLimiter) Stop() {
+	l.stopOnce.Do(func() {
+		close(l.done)
+	})
+}
+
+// getOrCreateBucket 获取或创建令牌桶。
+func (l *MemoryLimiter) getOrCreateBucket(key string) *TokenBucket {
+	// 先尝试读锁
+	l.mu.RLock()
+	bucket, exists := l.buckets[key]
+	l.mu.RUnlock()
+
+	if exists {
+		return bucket
+	}
+
+	// 需要创建新桶，升级为写锁
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	// 双重检查（可能其他 goroutine 已创建）
+	bucket, exists = l.buckets[key]
+	if exists {
+		return bucket
+	}
+
+	// 根据 key 确定配置（简化版：假设 key 格式为 "userID:action"）
+	cfg := l.getBucketConfig(key)
+	bucket = newTokenBucket(cfg.Capacity, cfg.Rate)
+	l.buckets[key] = bucket
+
+	return bucket
+}
+
+// getBucketConfig 根据 key 获取桶配置。
+func (l *MemoryLimiter) getBucketConfig(key string) config.BucketConfig {
+	// 简化实现：从 key 后缀判断动作类型
+	// 实际使用时调用方会传递正确的 key
+	// 默认使用 query 配置
+	return l.config.Query
+}
+
+// cleanup 定期清理不活跃的桶。
+func (l *MemoryLimiter) cleanup() {
+	ticker := time.NewTicker(10 * time.Minute)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ticker.C:
+			l.removeInactiveBuckets()
+		case <-l.done:
+			return
+		}
+	}
+}
+
+// removeInactiveBuckets 移除超过 10 分钟无活动的桶。
+func (l *MemoryLimiter) removeInactiveBuckets() {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	now := time.Now()
+	for key, bucket := range l.buckets {
+		bucket.mu.Lock()
+		inactive := now.Sub(bucket.lastRefill) > 10*time.Minute
+		bucket.mu.Unlock()
+
+		if inactive {
+			delete(l.buckets, key)
+		}
+	}
+}
+
+// min 返回两个 float64 中的较小值。
+func min(a, b float64) float64 {
+	if a < b {
+		return a
+	}
+	return b
+}
+
+// 编译期接口检查
+var _ Limiter = (*MemoryLimiter)(nil)

backend/internal/ratelimit/bucket_test.go

@@ -0,0 +1,203 @@
+package ratelimit
+
+import (
+	"context"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/hhs/camtalk/internal/config"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestTokenBucket_Allow_FirstRequest(t *testing.T) {
+	bucket := newTokenBucket(5, 0.2)
+
+	allowed, retryAfter := bucket.allow()
+
+	assert.True(t, allowed)
+	assert.Equal(t, time.Duration(0), retryAfter)
+}
+
+func TestTokenBucket_Allow_ConsumeUntilEmpty(t *testing.T) {
+	bucket := newTokenBucket(3, 0.2)
+
+	// 连续消耗 3 个令牌
+	for i := 0; i < 3; i++ {
+		allowed, _ := bucket.allow()
+		assert.True(t, allowed, "request %d should be allowed", i+1)
+	}
+
+	// 第 4 个请求应被拒绝
+	allowed, retryAfter := bucket.allow()
+	assert.False(t, allowed)
+	assert.Greater(t, retryAfter, time.Duration(0))
+}
+
+func TestTokenBucket_Allow_RetryAfterCorrect(t *testing.T) {
+	bucket := newTokenBucket(1, 1.0) // 每秒 1 个令牌
+
+	// 消耗唯一的令牌
+	allowed, _ := bucket.allow()
+	require.True(t, allowed)
+
+	// 立即再次请求应被拒绝
+	allowed, retryAfter := bucket.allow()
+	assert.False(t, allowed)
+	// retryAfter 应约为 1 秒（允许一定误差）
+	assert.InDelta(t, 1000, retryAfter.Milliseconds(), 100)
+}
+
+func TestTokenBucket_Allow_RefillAfterWait(t *testing.T) {
+	bucket := newTokenBucket(2, 10.0) // 每秒 10 个令牌（每 100ms 一个）
+
+	// 消耗 2 个令牌
+	bucket.allow()
+	bucket.allow()
+
+	// 等待 150ms，应补充至少 1 个令牌
+	time.Sleep(150 * time.Millisecond)
+
+	allowed, _ := bucket.allow()
+	assert.True(t, allowed)
+}
+
+func TestTokenBucket_Allow_CapacityLimit(t *testing.T) {
+	bucket := newTokenBucket(3, 1.0)
+
+	// 等待足够长时间让桶"溢出"
+	time.Sleep(100 * time.Millisecond)
+
+	// 但最多只能消耗 capacity 个令牌
+	for i := 0; i < 3; i++ {
+		allowed, _ := bucket.allow()
+		assert.True(t, allowed, "request %d should be allowed", i+1)
+	}
+
+	// 第 4 个应被拒绝
+	allowed, _ := bucket.allow()
+	assert.False(t, allowed)
+}
+
+func TestTokenBucket_Allow_ConcurrentSafe(t *testing.T) {
+	bucket := newTokenBucket(100, 10.0)
+	var wg sync.WaitGroup
+	successCount := 0
+	var mu sync.Mutex
+
+	// 100 个并发请求
+	for i := 0; i < 100; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			allowed, _ := bucket.allow()
+			if allowed {
+				mu.Lock()
+				successCount++
+				mu.Unlock()
+			}
+		}()
+	}
+
+	wg.Wait()
+
+	// 应该正好 100 个成功（桶容量为 100）
+	assert.Equal(t, 100, successCount)
+}
+
+func TestTokenBucket_Allow_ZeroCapacity(t *testing.T) {
+	bucket := newTokenBucket(0, 1.0)
+
+	allowed, retryAfter := bucket.allow()
+	assert.False(t, allowed)
+	assert.Greater(t, retryAfter, time.Duration(0))
+}
+
+func TestTokenBucket_Allow_ZeroRate(t *testing.T) {
+	bucket := newTokenBucket(1, 0.0)
+
+	// 第一个通过
+	allowed, _ := bucket.allow()
+	assert.True(t, allowed)
+
+	// 第二个被拒绝，且 retryAfter 应为无限大（实际上会很大）
+	allowed, retryAfter := bucket.allow()
+	assert.False(t, allowed)
+	// rate=0 时，retryAfter 理论上无限大，实际会是一个很大的值
+	assert.Greater(t, retryAfter, 1*time.Hour)
+}
+
+func TestMemoryLimiter_Allow_DifferentKeys(t *testing.T) {
+	cfg := config.RateLimitConfig{
+		Enabled: true,
+		Query:   config.BucketConfig{Capacity: 2, Rate: 1.0},
+	}
+	limiter := NewMemoryLimiter(cfg)
+	defer limiter.Stop()
+
+	ctx := context.Background()
+
+	// user1 消耗 2 个令牌
+	allowed, _ := limiter.Allow(ctx, "user1:query")
+	assert.True(t, allowed)
+	allowed, _ = limiter.Allow(ctx, "user1:query")
+	assert.True(t, allowed)
+
+	// user1 第 3 个被拒绝
+	allowed, _ = limiter.Allow(ctx, "user1:query")
+	assert.False(t, allowed)
+
+	// user2 应该不受影响
+	allowed, _ = limiter.Allow(ctx, "user2:query")
+	assert.True(t, allowed)
+}
+
+func TestMemoryLimiter_Cleanup(t *testing.T) {
+	cfg := config.RateLimitConfig{
+		Enabled: true,
+		Query:   config.BucketConfig{Capacity: 1, Rate: 1.0},
+	}
+	limiter := NewMemoryLimiter(cfg)
+	defer limiter.Stop()
+
+	ctx := context.Background()
+
+	// 创建一个桶
+	limiter.Allow(ctx, "user1:query")
+
+	// 验证桶已创建
+	limiter.mu.RLock()
+	initialCount := len(limiter.buckets)
+	limiter.mu.RUnlock()
+	assert.Equal(t, 1, initialCount)
+
+	// 手动触发清理（模拟 10 分钟后）
+	limiter.mu.Lock()
+	for _, bucket := range limiter.buckets {
+		bucket.mu.Lock()
+		bucket.lastRefill = time.Now().Add(-11 * time.Minute)
+		bucket.mu.Unlock()
+	}
+	limiter.mu.Unlock()
+
+	limiter.removeInactiveBuckets()
+
+	// 验证桶已被清理
+	limiter.mu.RLock()
+	finalCount := len(limiter.buckets)
+	limiter.mu.RUnlock()
+	assert.Equal(t, 0, finalCount)
+}
+
+func TestMemoryLimiter_Stop(t *testing.T) {
+	cfg := config.RateLimitConfig{
+		Enabled: true,
+		Query:   config.BucketConfig{Capacity: 1, Rate: 1.0},
+	}
+	limiter := NewMemoryLimiter(cfg)
+
+	// 多次调用 Stop 不应 panic
+	limiter.Stop()
+	limiter.Stop()
+}

backend/internal/ratelimit/limiter.go

@@ -0,0 +1,17 @@
+package ratelimit
+
+import (
+	"context"
+	"time"
+)
+
+// Limiter 速率限制器接口。
+type Limiter interface {
+	// Allow 判断 key 是否允许执行一次操作。
+	// key 通常为 "userID:action" 格式。
+	// 返回 (allowed, retryAfter)。retryAfter 表示需要等待的时间。
+	Allow(ctx context.Context, key string) (bool, time.Duration)
+
+	// Stop 停止限流器，清理资源（如后台 goroutine）。
+	Stop()
+}