通过记录IP的访问次数到Redis,来限制频繁请求/访问接口。可以动态地在响应内容中显示实际的限制频率。需要注意当前的响应内容message是以分钟为单位作提醒,由于
rl.Expires.Minutes()
可能会返回一个浮点数,所以使用int()
来将其转换为整数。如果需要更精确的时间单位表示(比如秒)还需自行修改。
调用示例
全局使用和单独使用
func RunServer(redisClient *redis.Client) *gin.Engine {
r := gin.Default()
// 全局限制:每60分钟最多请求40次
r.Use(middleware.NewRateLimiter(redisClient, 40, 60*time.Minute))
// 单独限制:每10分钟最多请求1次
r.POST("/register",middleware.NewRateLimiter(redisClient, 1, 10*time.Minute), handler.UserRegister)
}
中间件具体实现
package middleware
import (
"context"
"fmt"
"net/http"
"time"
"github.com/gin-gonic/gin"
"github.com/redis/go-redis/v9"
)
// 定义结构
type RateLimiter struct {
Client *redis.Client
Max int
Expires time.Duration
}
// 创建中间件
func NewRateLimiter(client *redis.Client, max int, expires time.Duration) gin.HandlerFunc {
return (&RateLimiter{
Client: client,
Max: max,
Expires: expires,
}).Limit
}
// 请求频率限制
func (rl *RateLimiter) Limit(c *gin.Context) {
ip := c.ClientIP()
key := fmt.Sprintf("rate_limit:%s:%s", ip, c.Request.URL.Path)
ctx := context.Background()
// 尝试设置key,如果key不存在则设置成功并返回true
allowed, err := rl.Client.SetNX(ctx, key, 0, rl.Expires).Result()
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"code": 500, "message": "Redis服务错误: " + err.Error()})
c.Abort()
return
}
if !allowed {
// 获取当前请求次数
current, err := rl.Client.Get(ctx, key).Int()
if err != nil {
if err == redis.Nil {
// Redis中没有该键,可能是由于某种原因键过期了,继续处理请求
c.Next()
return
}
c.JSON(http.StatusInternalServerError, gin.H{"code": 500, "message": "Redis查询错误: " + err.Error()})
c.Abort()
return
}
if current >= rl.Max {
// 如果当前请求次数超过最大限制,返回429状态码
c.JSON(http.StatusTooManyRequests, gin.H{
"code": 429,
"message": fmt.Sprintf("触发频繁请求限制,每%d分钟内最多允许请求%d次", int(rl.Expires.Minutes()), rl.Max),
})
c.Abort()
return
}
// 当前请求次数未超过限制,增加请求次数
if _, err := rl.Client.Incr(ctx, key).Result(); err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"code": 500, "message": "Redis计数错误: " + err.Error()})
c.Abort()
return
}
}
// 继续处理请求
c.Next()
}