65 lines
1.4 KiB
Go
65 lines
1.4 KiB
Go
package middleware
|
|
|
|
import (
|
|
"net/http"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/ulule/limiter/v3"
|
|
"github.com/ulule/limiter/v3/drivers/store/memory"
|
|
)
|
|
|
|
// RateLimitConfig 限流配置
|
|
type RateLimitConfig struct {
|
|
GlobalRPS float64
|
|
PerIPRPS float64
|
|
PerIPBurst int
|
|
}
|
|
|
|
// RateLimiter 双维度限流：全局令牌桶 + 单 IP 令牌桶
|
|
// 开发/测试环境可用内存存储，生产环境建议换 Redis 存储
|
|
func RateLimiter(cfg RateLimitConfig) gin.HandlerFunc {
|
|
store := memory.NewStore()
|
|
|
|
// 全局限流
|
|
globalRate := limiter.Rate{
|
|
Period: time.Second,
|
|
Limit: int64(cfg.GlobalRPS),
|
|
}
|
|
globalLimiter := limiter.New(store, globalRate)
|
|
|
|
// 单 IP 限流
|
|
perIPRate := limiter.Rate{
|
|
Period: time.Second,
|
|
Limit: int64(cfg.PerIPRPS),
|
|
}
|
|
perIPLimiter := limiter.New(store, perIPRate)
|
|
|
|
return func(c *gin.Context) {
|
|
// 1. 检查全局限流
|
|
globalCtx, err := globalLimiter.Get(c.Request.Context(), "global")
|
|
if err != nil || globalCtx.Reached {
|
|
c.JSON(http.StatusTooManyRequests, gin.H{
|
|
"code": 4290,
|
|
"message": "server busy, please retry later",
|
|
})
|
|
c.Abort()
|
|
return
|
|
}
|
|
|
|
// 2. 检查单 IP 限流
|
|
ip := c.ClientIP()
|
|
ipCtx, err := perIPLimiter.Get(c.Request.Context(), "ip:"+ip)
|
|
if err != nil || ipCtx.Reached {
|
|
c.JSON(http.StatusTooManyRequests, gin.H{
|
|
"code": 4291,
|
|
"message": "too many requests from this IP",
|
|
})
|
|
c.Abort()
|
|
return
|
|
}
|
|
|
|
c.Next()
|
|
}
|
|
}
|