106 lines
2.8 KiB
Go
106 lines
2.8 KiB
Go
package middleware
|
||
|
||
import (
|
||
"fmt"
|
||
"strings"
|
||
"time"
|
||
|
||
"ems.agt/src/framework/constants/cachekey"
|
||
"ems.agt/src/framework/i18n"
|
||
"ems.agt/src/framework/redis"
|
||
"ems.agt/src/framework/utils/ctx"
|
||
"ems.agt/src/framework/utils/ip2region"
|
||
"ems.agt/src/framework/vo/result"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
)
|
||
|
||
const (
|
||
// 默认策略全局限流
|
||
LIMIT_GLOBAL = 1
|
||
|
||
// 根据请求者IP进行限流
|
||
LIMIT_IP = 2
|
||
|
||
// 根据用户ID进行限流
|
||
LIMIT_USER = 3
|
||
)
|
||
|
||
// LimitOption 请求限流参数
|
||
type LimitOption struct {
|
||
Time int64 `json:"time"` // 限流时间,单位秒
|
||
Count int64 `json:"count"` // 限流次数
|
||
Type int64 `json:"type"` // 限流条件类型,默认LIMIT_GLOBAL
|
||
}
|
||
|
||
// RateLimit 请求限流
|
||
//
|
||
// 示例参数：middleware.LimitOption{ Time: 5, Count: 10, Type: middleware.LIMIT_IP }
|
||
//
|
||
// 参数表示：5秒内，最多请求10次，限制类型为 IP
|
||
//
|
||
// 使用 USER 时，请在用户身份授权认证校验后使用
|
||
// 以便获取登录用户信息，无用户信息时默认为 GLOBAL
|
||
func RateLimit(option LimitOption) gin.HandlerFunc {
|
||
return func(c *gin.Context) {
|
||
language := ctx.AcceptLanguage(c)
|
||
|
||
// 初始可选参数数据
|
||
if option.Time < 5 {
|
||
option.Time = 5
|
||
}
|
||
if option.Count < 10 {
|
||
option.Count = 10
|
||
}
|
||
if option.Type == 0 {
|
||
option.Type = LIMIT_GLOBAL
|
||
}
|
||
|
||
// 获取执行函数名称
|
||
funcName := c.HandlerName()
|
||
lastDotIndex := strings.LastIndex(funcName, "/")
|
||
funcName = funcName[lastDotIndex+1:]
|
||
// 生成限流key
|
||
var limitKey string = cachekey.RATE_LIMIT_KEY + funcName
|
||
|
||
// 用户
|
||
if option.Type == LIMIT_USER {
|
||
loginUser, err := ctx.LoginUser(c)
|
||
if err != nil {
|
||
c.JSON(401, result.Err(map[string]any{
|
||
"code": 401,
|
||
"msg": i18n.TKey(language, err.Error()),
|
||
}))
|
||
c.Abort() // 停止执行后续的处理函数
|
||
return
|
||
}
|
||
limitKey = cachekey.RATE_LIMIT_KEY + loginUser.UserID + ":" + funcName
|
||
}
|
||
|
||
// IP
|
||
if option.Type == LIMIT_IP {
|
||
clientIP := ip2region.ClientIP(c.ClientIP())
|
||
limitKey = cachekey.RATE_LIMIT_KEY + clientIP + ":" + funcName
|
||
}
|
||
|
||
// 在Redis查询并记录请求次数
|
||
rateCount, _ := redis.RateLimit("", limitKey, option.Time, option.Count)
|
||
rateTime, _ := redis.GetExpire("", limitKey)
|
||
|
||
// 设置响应头中的限流声明字段
|
||
c.Header("X-RateLimit-Limit", fmt.Sprintf("%d", option.Count)) // 总请求数限制
|
||
c.Header("X-RateLimit-Remaining", fmt.Sprintf("%d", option.Count-rateCount)) // 剩余可用请求数
|
||
c.Header("X-RateLimit-Reset", fmt.Sprintf("%d", time.Now().Unix()+int64(rateTime))) // 重置时间戳
|
||
|
||
if rateCount >= option.Count {
|
||
// 访问过于频繁，请稍候再试
|
||
c.JSON(200, i18n.TKey(language, "app.common.rateLimitTip"))
|
||
c.Abort() // 停止执行后续的处理函数
|
||
return
|
||
}
|
||
|
||
// 调用下一个处理程序
|
||
c.Next()
|
||
}
|
||
}
|