130 lines
3.6 KiB
Go
130 lines
3.6 KiB
Go
package middleware
|
||
|
||
import (
|
||
"fmt"
|
||
"strings"
|
||
|
||
"gitea.com/red-future/common/redis"
|
||
"gitea.com/red-future/common/utils"
|
||
"github.com/gogf/gf/v2/frame/g"
|
||
"github.com/gogf/gf/v2/net/ghttp"
|
||
"github.com/gogf/gf/v2/util/gconv"
|
||
)
|
||
|
||
// GlobalLimiter 全局限流中间件（使用Redis分布式控制）
|
||
func GlobalLimiter(r *ghttp.Request) {
|
||
// 从配置文件读取全局限流参数
|
||
globalLimit := g.Cfg().MustGet(r.GetCtx(), "rate.limit", 800).Int64()
|
||
|
||
key := redis.RateLimitKeyGlobal
|
||
|
||
// 使用Redis计数器进行全局限流
|
||
count, err := redis.IncrRateLimit(r.GetCtx(), key, 1) // 1秒窗口
|
||
if err != nil {
|
||
g.Log().Errorf(r.GetCtx(), "全局限流Redis错误: %v", err)
|
||
r.Middleware.Next()
|
||
return
|
||
}
|
||
|
||
if count > globalLimit {
|
||
g.Log().Warningf(r.GetCtx(), "全局限流触发: count: %d, limit: %d", count, globalLimit)
|
||
r.Response.WriteStatusExit(429, "系统当前繁忙,请稍后再试")
|
||
return
|
||
}
|
||
|
||
r.Middleware.Next()
|
||
}
|
||
|
||
// IPLimiter IP限流中间件（防DDoS）
|
||
func IPLimiter(r *ghttp.Request) {
|
||
ip := r.GetClientIp()
|
||
key := fmt.Sprintf(redis.RateLimitKeyIP, ip)
|
||
|
||
// 从配置文件读取IP限流参数
|
||
ipLimit := g.Cfg().MustGet(r.GetCtx(), "rate.ip.limit", 100).Int64()
|
||
|
||
// 使用Redis计数器
|
||
count, err := redis.IncrRateLimit(r.GetCtx(), key, 1) // 1秒窗口
|
||
if err != nil {
|
||
g.Log().Errorf(r.GetCtx(), "IP限流Redis错误: %v", err)
|
||
r.Middleware.Next()
|
||
return
|
||
}
|
||
|
||
if count > ipLimit {
|
||
g.Log().Warningf(r.GetCtx(), "IP限流触发: %s, count: %d, limit: %d", ip, count, ipLimit)
|
||
r.Response.WriteStatusExit(429, "请求过于频繁,请稍后再试")
|
||
return
|
||
}
|
||
|
||
r.Middleware.Next()
|
||
}
|
||
|
||
// UserLimiter 用户维度限流中间件（防止单用户滥用）
|
||
func UserLimiter(r *ghttp.Request) {
|
||
if r.RequestURI == "/swagger" || r.RequestURI == "/admin-go/api/v1/pub/captcha/get" || r.RequestURI == "/admin-go/api/v1/system/login" {
|
||
r.Middleware.Next()
|
||
return
|
||
}
|
||
var userName string
|
||
user, err := utils.GetUserInfo(r.GetCtx())
|
||
if err != nil {
|
||
r.Response.WriteStatusExit(429, err.Error())
|
||
return
|
||
}
|
||
userName = gconv.String(user.UserName)
|
||
// 从配置文件读取用户限流参数
|
||
userLimit := g.Cfg().MustGet(r.GetCtx(), "rate.user.limit", 50).Int64()
|
||
key := fmt.Sprintf(redis.RateLimitKeyUser, userName)
|
||
count, err := redis.IncrRateLimit(r.GetCtx(), key, 1)
|
||
if err != nil {
|
||
g.Log().Errorf(r.GetCtx(), "用户限流Redis错误: %v", err)
|
||
return
|
||
}
|
||
|
||
if count > userLimit {
|
||
r.Response.WriteStatusExit(429, "您的请求过于频繁,请稍后再试")
|
||
return
|
||
}
|
||
|
||
r.Middleware.Next()
|
||
}
|
||
|
||
// ServiceLimiter 服务维度限流中间件（保护微服务）
|
||
func ServiceLimiter(r *ghttp.Request) {
|
||
// 从URL路径提取服务名: /customerService/xxx -> customerService
|
||
pathParts := strings.Split(strings.Trim(r.URL.Path, "/"), "/")
|
||
if len(pathParts) == 0 {
|
||
r.Middleware.Next()
|
||
return
|
||
}
|
||
|
||
serverName := pathParts[0]
|
||
|
||
// 从配置文件读取服务限流参数
|
||
serviceLimitKey := fmt.Sprintf("rate.services.%s.limit", serverName)
|
||
limit := g.Cfg().MustGet(r.GetCtx(), serviceLimitKey, 0).Int64()
|
||
|
||
// 如果配置为0，说明该服务没有限流配置，跳过限流
|
||
if limit == 0 {
|
||
r.Middleware.Next()
|
||
return
|
||
}
|
||
|
||
key := fmt.Sprintf(redis.RateLimitKeyService, serverName)
|
||
count, err := redis.IncrRateLimit(r.GetCtx(), key, 1)
|
||
if err != nil {
|
||
g.Log().Errorf(r.GetCtx(), "服务限流Redis错误: %v", err)
|
||
r.Middleware.Next()
|
||
return
|
||
}
|
||
|
||
if count > limit {
|
||
g.Log().Warningf(r.GetCtx(), "服务限流触发: %s, count: %d, limit: %d", serverName, count, limit)
|
||
r.Response.WriteStatusExit(429, fmt.Sprintf("服务 '%s' 当前繁忙,请稍后再试", serverName))
|
||
return
|
||
}
|
||
|
||
r.Middleware.Next()
|
||
}
|