mirror of
https://github.com/yangjian102621/geekai.git
synced 2026-04-12 22:24:28 +08:00
44 lines
1.1 KiB
Go
44 lines
1.1 KiB
Go
package middleware
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
"geekai/core/types"
|
||
"geekai/utils"
|
||
"net/http"
|
||
"time"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
"github.com/go-redis/redis/v8"
|
||
)
|
||
|
||
// RateLimitEvery 使用 Redis 做固定间隔限流：在 interval 内仅允许一次请求
|
||
// Key 优先使用登录用户ID，若没有则退化为 route + IP
|
||
func RateLimitEvery(redisClient *redis.Client, interval time.Duration) gin.HandlerFunc {
|
||
return func(c *gin.Context) {
|
||
keyID := ""
|
||
if userID, ok := c.Get(types.LoginUserID); ok {
|
||
keyID = fmt.Sprintf("user:%s", utils.InterfaceToString(userID))
|
||
} else {
|
||
keyID = fmt.Sprintf("ip:%s", c.ClientIP())
|
||
}
|
||
|
||
fullPath := c.FullPath()
|
||
if fullPath == "" {
|
||
fullPath = c.Request.URL.Path
|
||
}
|
||
key := fmt.Sprintf("rl:%s:%s", fullPath, keyID)
|
||
|
||
okSet, err := redisClient.SetNX(context.Background(), key, 1, interval).Result()
|
||
if err != nil {
|
||
// Redis 异常时放行，避免误伤可用性
|
||
return
|
||
}
|
||
if !okSet {
|
||
c.JSON(http.StatusTooManyRequests, types.BizVo{Code: types.Failed, Message: "请求过于频繁，请稍后重试"})
|
||
c.Abort()
|
||
return
|
||
}
|
||
}
|
||
}
|