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)
|
|||
|
|
}
|