package middleware import ( "net/http" "sync" "time" "github.com/gin-gonic/gin" ) // RateLimiter is a simple per-IP sliding-window rate limiter that runs // entirely in-process. It is intended for auth endpoints (login, register) // where brute-force protection is critical. type RateLimiter struct { mu sync.Mutex window time.Duration max int requests map[string][]time.Time stop chan struct{} stopped sync.Once } // NewRateLimiter creates a rate limiter allowing max requests per window // per client IP. func NewRateLimiter(max int, window time.Duration) *RateLimiter { rl := &RateLimiter{ window: window, max: max, requests: make(map[string][]time.Time), stop: make(chan struct{}), } go rl.cleanup() return rl } // Close 停止后台清理 goroutine:清理循环此前无停止机制,每建一个实例 // 就永久滞留一条 goroutine(测试场景会随实例创建不断累积)。 func (rl *RateLimiter) Close() { rl.stopped.Do(func() { close(rl.stop) }) } func (rl *RateLimiter) cleanup() { ticker := time.NewTicker(5 * time.Minute) defer ticker.Stop() for { select { case <-rl.stop: return case <-ticker.C: } rl.mu.Lock() now := time.Now() for ip, times := range rl.requests { var valid []time.Time for _, t := range times { if now.Sub(t) <= rl.window { valid = append(valid, t) } } if len(valid) == 0 { delete(rl.requests, ip) } else { rl.requests[ip] = valid } } rl.mu.Unlock() } } // Allow returns true if the request from ip is within the rate limit. func (rl *RateLimiter) Allow(ip string) bool { rl.mu.Lock() defer rl.mu.Unlock() now := time.Now() times := rl.requests[ip] var valid []time.Time for _, t := range times { if now.Sub(t) <= rl.window { valid = append(valid, t) } } if len(valid) >= rl.max { rl.requests[ip] = valid return false } rl.requests[ip] = append(valid, now) return true } // RateLimit returns a Gin middleware that rejects requests exceeding the // per-IP rate limit with 429 Too Many Requests. func RateLimit(limiter *RateLimiter) gin.HandlerFunc { return func(c *gin.Context) { if !limiter.Allow(c.ClientIP()) { c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{ "code": 42901, "message": "too many requests, please try again later", }) return } c.Next() } }