73 lines
2.3 KiB
Go
73 lines
2.3 KiB
Go
// Package middleware ...
|
|
//
|
|
// Description : middleware ...
|
|
//
|
|
// Author : go_developer@163.com<白茶清欢>
|
|
//
|
|
// Date : 2026-01-04 14:17
|
|
package middleware
|
|
|
|
import (
|
|
"fmt"
|
|
"net/http"
|
|
|
|
"git.zhangdeman.cn/zhangdeman/gin/define"
|
|
"git.zhangdeman.cn/zhangdeman/gin/logger"
|
|
"git.zhangdeman.cn/zhangdeman/gin/util"
|
|
loggerPkg "git.zhangdeman.cn/zhangdeman/logger"
|
|
"git.zhangdeman.cn/zhangdeman/rate_limit/abstract"
|
|
rateLimitDefine "git.zhangdeman.cn/zhangdeman/rate_limit/define"
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
// RateLimit 接口流控
|
|
func RateLimit(rl abstract.IRateLimit, uriConfigTable map[string]define.UriConfig) gin.HandlerFunc {
|
|
return func(ctx *gin.Context) {
|
|
if nil == rl {
|
|
ctx.Next()
|
|
return
|
|
}
|
|
requestUri := ctx.Request.URL.Path
|
|
if _, exist := uriConfigTable[requestUri]; !exist {
|
|
ctx.Next()
|
|
return
|
|
}
|
|
if uriConfigTable[requestUri].RateLimitConfig.Total <= 0 || uriConfigTable[requestUri].RateLimitConfig.TimeInterval <= 0 {
|
|
// 配置异常, 不做处理
|
|
ctx.Next()
|
|
return
|
|
}
|
|
var (
|
|
allow bool
|
|
err error
|
|
)
|
|
|
|
if allow, err = rl.AllowN(ctx, &rateLimitDefine.LimitConfig{
|
|
Key: fmt.Sprintf("__gin_pkg_%v", requestUri), // TODO: 支持配置key的格式
|
|
Total: uriConfigTable[requestUri].RateLimitConfig.Total,
|
|
Rate: uriConfigTable[requestUri].RateLimitConfig.Use,
|
|
TimeInterval: uriConfigTable[requestUri].RateLimitConfig.TimeInterval,
|
|
}, uriConfigTable[requestUri].RateLimitConfig.Use); nil != err {
|
|
// 留空处理出现异常, 默认放行, 避免放大问题, 但是记录 Error 日志
|
|
logger.Instance.Error("无法正确执行流量控制程序", loggerPkg.NewLogData(util.GinCtxToContext(ctx), logger.RecordType, logger.CodeServiceRateLimitFailure, map[string]any{
|
|
"err_msg": err.Error(),
|
|
}).ToFieldList()...)
|
|
ctx.Next()
|
|
return
|
|
}
|
|
if !allow {
|
|
logger.Instance.Error("流量负载已饱和, 请稍后再试", loggerPkg.NewLogData(util.GinCtxToContext(ctx), logger.RecordType, logger.CodeServiceRateLimitDisable, map[string]any{
|
|
"rate_limit_config": uriConfigTable[requestUri].RateLimitConfig,
|
|
}).ToFieldList()...)
|
|
// 429 : Too Many Requests
|
|
ctx.JSON(http.StatusTooManyRequests, gin.H{
|
|
"err_msg": "流量负载已饱和, 请稍后再试",
|
|
})
|
|
ctx.Abort()
|
|
return
|
|
}
|
|
// 流控校验通过
|
|
ctx.Next()
|
|
}
|
|
}
|