66 lines
1.3 KiB
Go
66 lines
1.3 KiB
Go
package middleware
|
||
|
||
import (
|
||
"net/http"
|
||
"sync"
|
||
"time"
|
||
|
||
"golang.org/x/time/rate"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
)
|
||
|
||
// RateLimiter 按 IP 的限流器
|
||
type RateLimiter struct {
|
||
mu sync.Mutex
|
||
clients map[string]*rate.Limiter
|
||
r rate.Limit
|
||
b int
|
||
}
|
||
|
||
// NewRateLimiter 创建限流中间件，r 每秒请求数，b 突发容量
|
||
func NewRateLimiter(r rate.Limit, b int) *RateLimiter {
|
||
return &RateLimiter{
|
||
clients: make(map[string]*rate.Limiter),
|
||
r: r,
|
||
b: b,
|
||
}
|
||
}
|
||
|
||
// getLimiter 获取或创建该 key 的 limiter
|
||
func (rl *RateLimiter) getLimiter(key string) *rate.Limiter {
|
||
rl.mu.Lock()
|
||
defer rl.mu.Unlock()
|
||
if lim, ok := rl.clients[key]; ok {
|
||
return lim
|
||
}
|
||
lim := rate.NewLimiter(rl.r, rl.b)
|
||
rl.clients[key] = lim
|
||
return lim
|
||
}
|
||
|
||
// Middleware 返回 Gin 限流中间件（按客户端 IP）
|
||
func (rl *RateLimiter) Middleware() gin.HandlerFunc {
|
||
return func(c *gin.Context) {
|
||
key := c.ClientIP()
|
||
lim := rl.getLimiter(key)
|
||
if !lim.Allow() {
|
||
c.AbortWithStatus(http.StatusTooManyRequests)
|
||
return
|
||
}
|
||
c.Next()
|
||
}
|
||
}
|
||
|
||
// Cleanup 定期清理过期 limiter（可选，避免 map 无限增长）
|
||
func (rl *RateLimiter) Cleanup(interval time.Duration) {
|
||
ticker := time.NewTicker(interval)
|
||
go func() {
|
||
for range ticker.C {
|
||
rl.mu.Lock()
|
||
rl.clients = make(map[string]*rate.Limiter)
|
||
rl.mu.Unlock()
|
||
}
|
||
}()
|
||
}
|