feat: backend core - auth, user, role, permission, device, webhook, monitoring, cache, repository, service, middleware, API handlers

internal/security/ratelimit.go

@@ -0,0 +1,184 @@
+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()
+	}
+}