54 lines
1.2 KiB
Go
54 lines
1.2 KiB
Go
package middleware
|
|
|
|
import (
|
|
"net/http"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
type rateWindow struct {
|
|
start time.Time
|
|
count int
|
|
}
|
|
|
|
var rateWindows sync.Map
|
|
var rateMu sync.Mutex
|
|
|
|
// RateLimitByIP applies a fixed-window limit to public endpoints.
|
|
// A shared gateway limiter is still needed when running multiple replicas.
|
|
func RateLimitByIP(max int, window time.Duration) gin.HandlerFunc {
|
|
if max < 1 {
|
|
max = 1
|
|
}
|
|
if window <= 0 {
|
|
window = time.Minute
|
|
}
|
|
return func(c *gin.Context) {
|
|
key := c.Request.Method + ":" + c.FullPath() + ":" + c.ClientIP()
|
|
now := time.Now()
|
|
rateMu.Lock()
|
|
allowed := false
|
|
var current rateWindow
|
|
if raw, ok := rateWindows.Load(key); ok {
|
|
current, _ = raw.(rateWindow)
|
|
}
|
|
if current.start.IsZero() || now.Sub(current.start) >= window {
|
|
current = rateWindow{start: now, count: 1}
|
|
allowed = true
|
|
} else if current.count < max {
|
|
current.count++
|
|
allowed = true
|
|
}
|
|
rateWindows.Store(key, current)
|
|
rateMu.Unlock()
|
|
if !allowed {
|
|
c.Header("Retry-After", "60")
|
|
c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{"code": 429, "message": "请求过于频繁，请稍后重试"})
|
|
return
|
|
}
|
|
c.Next()
|
|
}
|
|
}
|