user-system/internal/security/ratelimit.go

package security

import (
	"sync"
	"time"
)

// RateLimitAlgorithm 限流算法类型
type RateLimitAlgorithm string

const (
	AlgorithmTokenBucket   RateLimitAlgorithm = "token_bucket"
	AlgorithmLeakyBucket   RateLimitAlgorithm = "leaky_bucket"
	AlgorithmSlidingWindow RateLimitAlgorithm = "sliding_window"
	AlgorithmFixedWindow   RateLimitAlgorithm = "fixed_window"
)

// TokenBucket 令牌桶算法
type TokenBucket struct {
	capacity    int64
	tokens      int64
	rate        int64         // 每秒产生的令牌数
	lastRefill  time.Time
	mu          sync.Mutex
}

// NewTokenBucket 创建令牌桶
func NewTokenBucket(capacity, rate int64) *TokenBucket {
	return &TokenBucket{
		capacity:   capacity,
		tokens:     capacity,
		rate:       rate,
		lastRefill: time.Now(),
	}
}

// Allow 检查是否允许访问
func (tb *TokenBucket) Allow() bool {
	tb.mu.Lock()
	defer tb.mu.Unlock()

	now := time.Now()
	elapsed := now.Sub(tb.lastRefill).Seconds()

	// 计算需要补充的令牌数
	refillTokens := int64(elapsed * float64(tb.rate))
	tb.tokens += refillTokens
	if tb.tokens > tb.capacity {
		tb.tokens = tb.capacity
	}
	tb.lastRefill = now

	// 检查是否有足够的令牌
	if tb.tokens > 0 {
		tb.tokens--
		return true
	}

	return false
}

// LeakyBucket 漏桶算法
type LeakyBucket struct {
	capacity int64
	water    int64
	rate     int64         // 每秒漏出的水量
	lastLeak time.Time
	mu       sync.Mutex
}

// NewLeakyBucket 创建漏桶
func NewLeakyBucket(capacity, rate int64) *LeakyBucket {
	return &LeakyBucket{
		capacity: capacity,
		water:    0,
		rate:     rate,
		lastLeak: time.Now(),
	}
}

// Allow 检查是否允许访问
func (lb *LeakyBucket) Allow() bool {
	lb.mu.Lock()
	defer lb.mu.Unlock()

	now := time.Now()
	elapsed := now.Sub(lb.lastLeak).Seconds()

	// 计算漏出的水量
	leakWater := int64(elapsed * float64(lb.rate))
	lb.water -= leakWater
	if lb.water < 0 {
		lb.water = 0
	}
	lb.lastLeak = now

	// 检查桶是否已满
	if lb.water < lb.capacity {
		lb.water++
		return true
	}

	return false
}

// SlidingWindow 滑动窗口算法
type SlidingWindow struct {
	window time.Duration
	capacity int64
	requests []time.Time
	mu       sync.Mutex
}

// NewSlidingWindow 创建滑动窗口
func NewSlidingWindow(window time.Duration, capacity int64) *SlidingWindow {
	return &SlidingWindow{
		window:   window,
		capacity: capacity,
		requests: make([]time.Time, 0),
	}
}

// Allow 检查是否允许访问
func (sw *SlidingWindow) Allow() bool {
	sw.mu.Lock()
	defer sw.mu.Unlock()

	now := time.Now()
	
	// 移除窗口外的请求
	validRequests := make([]time.Time, 0)
	for _, req := range sw.requests {
		if now.Sub(req) < sw.window {
			validRequests = append(validRequests, req)
		}
	}
	sw.requests = validRequests

	// 检查是否超过容量
	if int64(len(sw.requests)) < sw.capacity {
		sw.requests = append(sw.requests, now)
		return true
	}

	return false
}

// RateLimiter 限流器
type RateLimiter struct {
	algorithm RateLimitAlgorithm
	limiter   interface{}
}

// NewRateLimiter 创建限流器
func NewRateLimiter(algorithm RateLimitAlgorithm, capacity, rate int64, window time.Duration) *RateLimiter {
	limiter := &RateLimiter{algorithm: algorithm}

	switch algorithm {
	case AlgorithmTokenBucket:
		limiter.limiter = NewTokenBucket(capacity, rate)
	case AlgorithmLeakyBucket:
		limiter.limiter = NewLeakyBucket(capacity, rate)
	case AlgorithmSlidingWindow:
		limiter.limiter = NewSlidingWindow(window, capacity)
	default:
		limiter.limiter = NewSlidingWindow(window, capacity)
	}

	return limiter
}

// Allow 检查是否允许访问
func (rl *RateLimiter) Allow() bool {
	switch rl.algorithm {
	case AlgorithmTokenBucket:
		return rl.limiter.(*TokenBucket).Allow()
	case AlgorithmLeakyBucket:
		return rl.limiter.(*LeakyBucket).Allow()
	case AlgorithmSlidingWindow:
		return rl.limiter.(*SlidingWindow).Allow()
	default:
		return rl.limiter.(*SlidingWindow).Allow()
	}
}