103 lines
2.9 KiB
Go
103 lines
2.9 KiB
Go
package middleware
|
||
|
||
import (
|
||
"fmt"
|
||
"strconv"
|
||
"strings"
|
||
"time"
|
||
|
||
"github.com/gogf/gf/v2/errors/gerror"
|
||
"github.com/gogf/gf/v2/net/ghttp"
|
||
|
||
"easygoadmin/internal/consts"
|
||
"easygoadmin/utility/cache"
|
||
"easygoadmin/utility/response"
|
||
)
|
||
|
||
// ============================== API 限流中间件 ==============================
|
||
//
|
||
// 固定窗口计数：键按窗口起点对齐（ratelimit:{ip}:{path}:{窗口起点}），
|
||
// 计数存 utility/cache（Redis 优先、内存降级），多实例部署时共享计数。
|
||
//
|
||
// 规则：
|
||
// - 仅限 /api/ 前缀接口；文档（/api.json）、静态资源、健康检查不受限
|
||
// - 认证入口（登录/验证码/令牌兑换）按 rateLimit.authLimit 单独更严阈值
|
||
// - 计数存储异常时放行（限流不阻断业务）
|
||
|
||
// RateLimit API 限流
|
||
func RateLimit(r *ghttp.Request) {
|
||
if !cfgBool(r.Context(), "rateLimit.enable", true) {
|
||
r.Middleware.Next()
|
||
return
|
||
}
|
||
path := r.URL.Path
|
||
if !isRateLimitedPath(path) {
|
||
r.Middleware.Next()
|
||
return
|
||
}
|
||
window := cfgInt(r.Context(), "rateLimit.windowSeconds", 60)
|
||
limit := cfgInt(r.Context(), "rateLimit.limit", 120)
|
||
if isAuthPath(path) {
|
||
limit = cfgInt(r.Context(), "rateLimit.authLimit", 10)
|
||
}
|
||
if limit <= 0 || window <= 0 {
|
||
r.Middleware.Next()
|
||
return
|
||
}
|
||
|
||
// 窗口起点对齐：同一窗口内共享计数键
|
||
windowStart := time.Now().Unix() / int64(window) * int64(window)
|
||
key := fmt.Sprintf("ratelimit:%s:%s:%d", realClientIP(r), path, windowStart)
|
||
n, err := incrCounter(r, key, window)
|
||
if err != nil {
|
||
r.Middleware.Next()
|
||
return
|
||
}
|
||
if n > int64(limit) {
|
||
response.Write(r, response.Fail(gerror.NewCode(consts.CodeTooManyRequests)))
|
||
return
|
||
}
|
||
r.Middleware.Next()
|
||
}
|
||
|
||
// isRateLimitedPath 判断是否为参与限流的接口路径
|
||
// /api.json 与 /swagger 不带 /api/ 前缀，自然排除；健康检查显式豁免
|
||
func isRateLimitedPath(path string) bool {
|
||
if path == "/api/health" {
|
||
return false
|
||
}
|
||
return path == "/api" || strings.HasPrefix(path, "/api/")
|
||
}
|
||
|
||
// isAuthPath 判断是否为认证入口（适用更严格的独立阈值）
|
||
func isAuthPath(path string) bool {
|
||
switch path {
|
||
case "/api/login", "/api/captcha", "/api/oauth2/token":
|
||
return true
|
||
}
|
||
return false
|
||
}
|
||
|
||
// incrCounter 固定窗口计数自增
|
||
// 键已含窗口起点，TTL 仅用于回收存储，续期不影响窗口归属
|
||
func incrCounter(r *ghttp.Request, key string, windowSeconds int) (int64, error) {
|
||
ttl := windowSeconds + 5
|
||
if raw, ok := cache.Get(r.Context(), key); ok {
|
||
n, _ := strconv.ParseInt(raw, 10, 64)
|
||
n++
|
||
_ = cache.Set(r.Context(), key, strconv.FormatInt(n, 10), ttl)
|
||
return n, nil
|
||
}
|
||
if cache.SetNX(r.Context(), key, "1", ttl) {
|
||
return 1, nil
|
||
}
|
||
// 并发下 SetNX 未抢到：读现值再加一
|
||
if raw, ok := cache.Get(r.Context(), key); ok {
|
||
n, _ := strconv.ParseInt(raw, 10, 64)
|
||
n++
|
||
_ = cache.Set(r.Context(), key, strconv.FormatInt(n, 10), ttl)
|
||
return n, nil
|
||
}
|
||
return 1, nil
|
||
}
|