60 lines
1.3 KiB
Go
60 lines
1.3 KiB
Go
package middleware
|
||
|
||
import (
|
||
"fmt"
|
||
"nl-video-api/utility/response"
|
||
"time"
|
||
|
||
"github.com/gogf/gf/v2/database/gredis"
|
||
"github.com/gogf/gf/v2/frame/g"
|
||
"github.com/gogf/gf/v2/net/ghttp"
|
||
)
|
||
|
||
// RateLimit 限流中间件
|
||
func RateLimit(maxRequests int, window time.Duration) func(r *ghttp.Request) {
|
||
return func(r *ghttp.Request) {
|
||
var (
|
||
ctx = r.Context()
|
||
client = g.Redis()
|
||
key = fmt.Sprintf("rate_limit:%s", r.GetClientIp())
|
||
)
|
||
|
||
// 获取当前请求次数
|
||
count, err := client.Get(ctx, key)
|
||
if err != nil {
|
||
g.Log().Error(ctx, "Redis获取失败:", err)
|
||
r.Middleware.Next()
|
||
return
|
||
}
|
||
|
||
// 检查是否超过限制
|
||
if count.Int() >= maxRequests {
|
||
response.Error(r, response.CodeError, "请求过于频繁，请稍后再试")
|
||
return
|
||
}
|
||
|
||
// 增加计数
|
||
if count.Int() == 0 {
|
||
// 第一次请求，设置过期时间
|
||
seconds := int64(window.Seconds())
|
||
client.Set(ctx, key, 1, gredis.SetOption{
|
||
TTLOption: gredis.TTLOption{EX: &seconds},
|
||
})
|
||
} else {
|
||
// 增加计数
|
||
client.Incr(ctx, key)
|
||
}
|
||
|
||
r.Middleware.Next()
|
||
}
|
||
}
|
||
|
||
// APIRateLimit API接口限流（每分钟60次）
|
||
func APIRateLimit(r *ghttp.Request) {
|
||
RateLimit(60, time.Minute)(r)
|
||
}
|
||
|
||
// LoginRateLimit 登录接口限流（每分钟5次）
|
||
func LoginRateLimit(r *ghttp.Request) {
|
||
RateLimit(5, time.Minute)(r)
|
||
} |