golang,go,博客,开源,编程
在 Gin 框架中,Rate Limiting(速率限制)是限制客户端在一定时间内发起请求次数的一种技术。常见的应用场景包括防止暴力破解、保护 API 免受滥用、减轻服务器负载等。
Rate Limiting 主要通过限制在指定时间内可以接受的请求数量来防止过度的请求。一般的策略包括:
在实际开发中,最常见的策略是 固定窗口 和 令牌桶。
一个常见的方式是利用 Redis 的 INCR
操作和过期时间来实现速率限制。Redis 非常适合做这种场景,因为它的单线程特性可以确保操作的原子性,并且支持高效的计数。
下面是一个简单的基于 Redis 的 Rate Limiting 中间件实现。
首先需要安装 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")
}
init()
函数中,配置了 Redis 客户端。使用 redis.NewClient
创建 Redis 客户端并连接到本地的 Redis 服务。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
错误。令牌桶算法是另一种常见的实现方式,下面是一个基于令牌桶算法的 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")
}
Rate Limiting 是提高服务稳定性、保证公平使用的重要手段,尤其在面对大规模请求时非常有效。