package middleware import ( "math" "net/http" "sync" "time" "github.com/gin-gonic/gin" "vision-tool/internal/model" ) // CORS 跨域中间件(开发环境放开,生产环境应限制为具体域名)。 func CORS() gin.HandlerFunc { return func(c *gin.Context) { c.Header("Access-Control-Allow-Origin", "*") c.Header("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS") c.Header("Access-Control-Allow-Headers", "Content-Type, Authorization") if c.Request.Method == http.MethodOptions { c.AbortWithStatus(http.StatusNoContent) return } c.Next() } } // ============================================================================ // 内存令牌桶限流(全局限流,防匿名刷接口) // ============================================================================ // RateLimiter 基于时间的令牌桶。 type RateLimiter struct { mu sync.Mutex rate float64 // 每秒补充令牌数 burst float64 // 桶容量 tokens float64 lastTime time.Time } func NewRateLimiter(rps, burst int) *RateLimiter { return &RateLimiter{ rate: float64(rps), burst: float64(burst), tokens: float64(burst), lastTime: time.Now(), } } func (rl *RateLimiter) Allow() bool { rl.mu.Lock() defer rl.mu.Unlock() now := time.Now() rl.tokens = math.Min(rl.burst, rl.tokens+now.Sub(rl.lastTime).Seconds()*rl.rate) rl.lastTime = now if rl.tokens >= 1 { rl.tokens-- return true } return false } // Middleware 返回限流 gin 中间件。 func (rl *RateLimiter) Middleware() gin.HandlerFunc { return func(c *gin.Context) { if !rl.Allow() { c.AbortWithStatusJSON(http.StatusTooManyRequests, model.Err(429, "请求过于频繁,请稍后再试")) return } c.Next() } }