Initial commit: Basic project structure and dependencies

internal/shared/middleware/ratelimit.go

@@ -0,0 +1,166 @@
+package middleware
+
+import (
+	"fmt"
+	"net/http"
+	"sync"
+	"time"
+
+	"tyapi-server/internal/config"
+
+	"github.com/gin-gonic/gin"
+	"golang.org/x/time/rate"
+)
+
+// RateLimitMiddleware 限流中间件
+type RateLimitMiddleware struct {
+	config   *config.Config
+	limiters map[string]*rate.Limiter
+	mutex    sync.RWMutex
+}
+
+// NewRateLimitMiddleware 创建限流中间件
+func NewRateLimitMiddleware(cfg *config.Config) *RateLimitMiddleware {
+	return &RateLimitMiddleware{
+		config:   cfg,
+		limiters: make(map[string]*rate.Limiter),
+	}
+}
+
+// GetName 返回中间件名称
+func (m *RateLimitMiddleware) GetName() string {
+	return "ratelimit"
+}
+
+// GetPriority 返回中间件优先级
+func (m *RateLimitMiddleware) GetPriority() int {
+	return 90 // 高优先级
+}
+
+// Handle 返回中间件处理函数
+func (m *RateLimitMiddleware) Handle() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		// 获取客户端标识（IP地址）
+		clientID := m.getClientID(c)
+
+		// 获取或创建限流器
+		limiter := m.getLimiter(clientID)
+
+		// 检查是否允许请求
+		if !limiter.Allow() {
+			c.Header("X-RateLimit-Limit", fmt.Sprintf("%d", m.config.RateLimit.Requests))
+			c.Header("X-RateLimit-Window", m.config.RateLimit.Window.String())
+			c.Header("Retry-After", "60")
+
+			c.JSON(http.StatusTooManyRequests, gin.H{
+				"success": false,
+				"message": "Rate limit exceeded",
+				"error":   "Too many requests",
+			})
+			c.Abort()
+			return
+		}
+
+		// 添加限流头部信息
+		c.Header("X-RateLimit-Limit", fmt.Sprintf("%d", m.config.RateLimit.Requests))
+		c.Header("X-RateLimit-Window", m.config.RateLimit.Window.String())
+
+		c.Next()
+	}
+}
+
+// IsGlobal 是否为全局中间件
+func (m *RateLimitMiddleware) IsGlobal() bool {
+	return true
+}
+
+// getClientID 获取客户端标识
+func (m *RateLimitMiddleware) getClientID(c *gin.Context) string {
+	// 优先使用X-Forwarded-For头部
+	if xff := c.GetHeader("X-Forwarded-For"); xff != "" {
+		return xff
+	}
+
+	// 使用X-Real-IP头部
+	if xri := c.GetHeader("X-Real-IP"); xri != "" {
+		return xri
+	}
+
+	// 使用RemoteAddr
+	return c.ClientIP()
+}
+
+// getLimiter 获取或创建限流器
+func (m *RateLimitMiddleware) getLimiter(clientID string) *rate.Limiter {
+	m.mutex.RLock()
+	limiter, exists := m.limiters[clientID]
+	m.mutex.RUnlock()
+
+	if exists {
+		return limiter
+	}
+
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+
+	// 双重检查
+	if limiter, exists := m.limiters[clientID]; exists {
+		return limiter
+	}
+
+	// 创建新的限流器
+	// rate.Every计算每个请求之间的间隔
+	rateLimit := rate.Every(m.config.RateLimit.Window / time.Duration(m.config.RateLimit.Requests))
+	limiter = rate.NewLimiter(rateLimit, m.config.RateLimit.Burst)
+
+	m.limiters[clientID] = limiter
+
+	// 启动清理协程（仅第一次创建时）
+	if len(m.limiters) == 1 {
+		go m.cleanupRoutine()
+	}
+
+	return limiter
+}
+
+// cleanupRoutine 定期清理不活跃的限流器
+func (m *RateLimitMiddleware) cleanupRoutine() {
+	ticker := time.NewTicker(10 * time.Minute) // 每10分钟清理一次
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ticker.C:
+			m.cleanup()
+		}
+	}
+}
+
+// cleanup 清理不活跃的限流器
+func (m *RateLimitMiddleware) cleanup() {
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+
+	now := time.Now()
+	for clientID, limiter := range m.limiters {
+		// 如果限流器在过去1小时内没有被使用，则删除它
+		if limiter.Reserve().Delay() == 0 && now.Sub(time.Now()) > time.Hour {
+			delete(m.limiters, clientID)
+		}
+	}
+}
+
+// GetStats 获取限流统计
+func (m *RateLimitMiddleware) GetStats() map[string]interface{} {
+	m.mutex.RLock()
+	defer m.mutex.RUnlock()
+
+	return map[string]interface{}{
+		"active_limiters": len(m.limiters),
+		"rate_limit": map[string]interface{}{
+			"requests": m.config.RateLimit.Requests,
+			"window":   m.config.RateLimit.Window,
+			"burst":    m.config.RateLimit.Burst,
+		},
+	}
+}