golang,go,博客,开源,编程

gin中间件之ratelimit

Published on with 0 views and 0 comments

在 Gin 框架中,Rate Limiting(速率限制)是限制客户端在一定时间内发起请求次数的一种技术。常见的应用场景包括防止暴力破解、保护 API 免受滥用、减轻服务器负载等。

1. Rate Limiting 的基本原理

Rate Limiting 主要通过限制在指定时间内可以接受的请求数量来防止过度的请求。一般的策略包括:

  • 固定窗口(Fixed Window):在一个固定时间窗口内,限制请求的数量。
  • 滑动窗口(Sliding Window):使用一个时间滑动窗口,每个请求都在时间窗口内限制。
  • 漏桶算法(Leaky Bucket):请求按照固定速率处理,超出速率的请求会被丢弃。
  • 令牌桶算法(Token Bucket):令牌以固定速率加入桶中,只有当请求能获得令牌时才允许通过。

在实际开发中,最常见的策略是 固定窗口令牌桶

2. 使用 Redis 实现 Rate Limiting

一个常见的方式是利用 Redis 的 INCR 操作和过期时间来实现速率限制。Redis 非常适合做这种场景,因为它的单线程特性可以确保操作的原子性,并且支持高效的计数。

下面是一个简单的基于 Redis 的 Rate Limiting 中间件实现。

3. 实现 Rate Limit 中间件(Redis)

首先需要安装 Redis 的 Go 客户端库:

go get github.com/go-redis/redis/v8

然后在 Gin 中实现一个基于 Redis 的 Rate Limit 中间件:

package main

import (
	"context"
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/go-redis/redis/v8"
	"net/http"
	"time"
)

var rdb *redis.Client

// 初始化 Redis 客户端
func init() {
	rdb = redis.NewClient(&redis.Options{
		Addr: "localhost:6379", // Redis 地址
		DB:   0,                // 默认 DB
	})
}

// RateLimit 中间件
func RateLimitMiddleware(limit int, window time.Duration) gin.HandlerFunc {
	return func(c *gin.Context) {
		// 获取客户端 IP 地址作为 Rate Limiting 的 Key
		clientIP := c.ClientIP()
		// 使用 Redis 的 INCR 命令来实现请求计数
		key := fmt.Sprintf("rate_limit:%s", clientIP)
		count, err := rdb.Incr(context.Background(), key).Result()
		if err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": "Rate limit error"})
			c.Abort()
			return
		}

		// 设置过期时间,确保每个时间窗口重置计数
		if count == 1 {
			rdb.Expire(context.Background(), key, window)
		}

		// 如果超出限制,拒绝请求
		if count > int64(limit) {
			c.JSON(http.StatusTooManyRequests, gin.H{
				"message": fmt.Sprintf("Too many requests, try again later. Allowed requests: %d", limit),
			})
			c.Abort()
			return
		}

		// 请求未超限,继续执行
		c.Next()
	}
}

func main() {
	r := gin.Default()

	// 设置 Rate Limit 中间件,限制每个 IP 每 10 秒钟只能请求 5 次
	r.Use(RateLimitMiddleware(5, 10*time.Second))

	r.GET("/ping", func(c *gin.Context) {
		c.JSON(http.StatusOK, gin.H{"message": "pong"})
	})

	r.Run(":8080")
}

解析代码

  1. Redis 配置:在 init() 函数中,配置了 Redis 客户端。使用 redis.NewClient 创建 Redis 客户端并连接到本地的 Redis 服务。
  2. RateLimitMiddleware
    • limit: 每个 IP 地址允许的最大请求次数。
    • window: 限制的时间窗口,比如 10 * time.Second 表示 10 秒内最多允许 limit 次请求。
    • key: 以客户端 IP 地址作为唯一标识,生成 Redis 键名。每个 IP 的请求计数都存储在 Redis 中。
    • Incr: 使用 Redis 的 INCR 命令进行计数。Redis 会返回当前计数,如果当前计数超过限制则拒绝请求。
    • Expire: 使用 Expire 设置 Redis 键的过期时间,即每个时间窗口结束时,计数会自动重置。
    • 如果请求超过限制,则返回 HTTP 429 Too Many Requests 错误。
  3. 客户端请求计数:每次请求都会通过 Redis 增加计数。如果计数超过了限制,则立即返回错误响应。
  4. 请求通过:如果请求没有超过限制,继续执行后续的处理(如返回 "pong" 响应)。

4. 使用令牌桶(Token Bucket)算法实现 Rate Limiting

令牌桶算法是另一种常见的实现方式,下面是一个基于令牌桶算法的 Rate Limit 中间件实现。我们可以使用 Redis 来模拟令牌桶的行为,按固定速率将令牌放入桶中,只有当请求获取到令牌时才允许继续。

package main

import (
	"context"
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/go-redis/redis/v8"
	"net/http"
	"time"
)

var rdb *redis.Client

// 初始化 Redis 客户端
func init() {
	rdb = redis.NewClient(&redis.Options{
		Addr: "localhost:6379",
		DB:   0,
	})
}

// TokenBucketRateLimit 中间件
func TokenBucketRateLimitMiddleware(limit int, rate time.Duration) gin.HandlerFunc {
	return func(c *gin.Context) {
		clientIP := c.ClientIP()
		key := fmt.Sprintf("token_bucket:%s", clientIP)

		// 获取当前令牌数
		tokens, err := rdb.Get(context.Background(), key).Int64()
		if err != nil && err != redis.Nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": "Rate limit error"})
			c.Abort()
			return
		}

		// 如果令牌数为负,表示限流
		if tokens <= 0 {
			c.JSON(http.StatusTooManyRequests, gin.H{
				"message": "Too many requests, try again later",
			})
			c.Abort()
			return
		}

		// 允许通过后,扣除令牌数
		rdb.Decr(context.Background(), key)

		// 如果当前令牌数为 0,则设置令牌重新生成的时间
		if tokens == 1 {
			rdb.Expire(context.Background(), key, rate)
		}

		// 继续处理请求
		c.Next()
	}
}

func main() {
	r := gin.Default()

	// 使用令牌桶限流中间件,令牌数限制为 5 个,每秒生成一个令牌
	r.Use(TokenBucketRateLimitMiddleware(5, 1*time.Second))

	r.GET("/ping", func(c *gin.Context) {
		c.JSON(http.StatusOK, gin.H{"message": "pong"})
	})

	r.Run(":8080")
}

5. 总结

  • 基于 Redis 的 Rate Limiting:通过 Redis 实现计数,结合过期时间,能够有效地实现固定窗口的 Rate Limiting。
  • 令牌桶算法:通过 Redis 来模拟令牌桶算法,通过令牌控制请求速率。
  • Gin 中间件的实现:通过 Gin 的中间件功能,可以很容易地集成 Rate Limiting 到整个应用中,保护 API 免受滥用。

Rate Limiting 是提高服务稳定性、保证公平使用的重要手段,尤其在面对大规模请求时非常有效。


标题:gin中间件之ratelimit
作者:mooncakeee
地址:http://blog.dd95828.com/articles/2025/01/07/1736219466263.html
联系:scotttu@163.com