Gin: 设计一个中间件,使用Redis进行不同IP的频繁请求限制

通过记录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()
}
暂无评论

发送评论 编辑评论


				
上一篇
下一篇