d5f0b8da63
feat: 🚦 新增基于 Redis 令牌桶的限流中间件 - 使用 Redis 实现令牌桶算法进行限流 🪣 - 覆盖除登录、注册、刷新 token 以外的所有接口 🔒 fix: 🐛 修复任务块添加到日程接口可修改已安排任务时间的问题 - 禁止通过该接口直接修改已安排任务块的时间 - 修正不合理的业务逻辑，保证数据一致性 ✅
39 lines
1.1 KiB
Go
39 lines
1.1 KiB
Go
package middleware
|
||
|
||
import (
|
||
"fmt"
|
||
"log"
|
||
|
||
"github.com/LoveLosita/smartflow/backend/pkg"
|
||
"github.com/LoveLosita/smartflow/backend/respond"
|
||
"github.com/gin-gonic/gin"
|
||
)
|
||
|
||
func RateLimitMiddleware(limiter *pkg.RateLimiter, capacity, rate int) gin.HandlerFunc {
|
||
return func(c *gin.Context) {
|
||
// 1. 确定限流对象：可以用 UserID，也可以用 IP
|
||
// 这里建议用 UserID，防止某个用户换 IP 疯狂刷
|
||
userID := c.GetInt("user_id") // 假设你之前的 JWT 已经塞进去了
|
||
key := fmt.Sprintf("rate_limit:user:%d", userID)
|
||
|
||
// 2. 执行限流检查
|
||
allowed, err := limiter.Allow(c.Request.Context(), key, capacity, rate)
|
||
if err != nil {
|
||
// 如果 Redis 挂了，为了保证业务可用，通常选择“放行”并记录日志
|
||
log.Printf("Redis limiter error: %v", err)
|
||
c.Next()
|
||
return
|
||
}
|
||
|
||
if !allowed {
|
||
// 3. 触发限流：直接调你写好的 DealWithError
|
||
// 可以在 respond 里定义一个新错误：TooManyRequests
|
||
respond.DealWithError(c, respond.TooManyRequests)
|
||
c.Abort() // 拦截，不执行后续 Handler
|
||
return
|
||
}
|
||
|
||
c.Next()
|
||
}
|
||
}
|