gin限流中间件的一个bug

年关在即，各个平台会出相应的活动促一下活跃，有了活动，就会有脚本来刷，前几天运维查了prometheus说请求没有被限制住
由于是配合元旦活动顺便加的限流功能，QA专注于测业务并没有测限流器的功能,果然还是翻车了。
于是构造了一段最简化的代码压测看看用的中间件¹是不是确实有问题:

package main

import (
	"github.com/Salvatore-Giordano/gin-redis-ip-limiter"
	"github.com/gin-gonic/gin"
	"github.com/go-redis/redis"
	"time"
)

func main() {
	r := gin.Default()
	r.Use(iplimiter.NewRateLimiterMiddleware(redis.NewClient(&redis.Options{
		Addr:     "localhost:6379",
		Password: "",
		DB:       1,
	}), "general", 10, time.Second))
	r.GET("/ping", func(c *gin.Context) {
		c.String(200, "pong")
	})
	r.Run(":8000")
}

压测单客户端发100条请求

➜  ~ oha -n 100 http://localhost:8000/ping
Summary:
  Success rate:	100.00%
  Total:	0.0316 secs
  Slowest:	0.0215 secs
  Fastest:	0.0039 secs
  Average:	0.0125 secs
  Requests/sec:	3166.8020

  Total data:	2.71 KiB
  Size/request:	27 B
  Size/sec:	85.94 KiB

Response time histogram:
  0.004 [1]  |■
  0.006 [30] |■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■■
  0.007 [19] |■■■■■■■■■■■■■■■■■■■■
  0.009 [0]  |
  0.011 [0]  |
  0.013 [0]  |
  0.014 [0]  |
  0.016 [0]  |
  0.018 [12] |■■■■■■■■■■■■
  0.020 [13] |■■■■■■■■■■■■■
  0.021 [25] |■■■■■■■■■■■■■■■■■■■■■■■■■■

Response time distribution:
  10.00% in 0.0047 secs
  25.00% in 0.0053 secs
  50.00% in 0.0174 secs
  75.00% in 0.0202 secs
  90.00% in 0.0206 secs
  95.00% in 0.0212 secs
  99.00% in 0.0215 secs
  99.90% in 0.0215 secs
  99.99% in 0.0215 secs


Details (average, fastest, slowest):
  DNS+dialup:	0.0024 secs, 0.0014 secs, 0.0033 secs
  DNS-lookup:	0.0000 secs, 0.0000 secs, 0.0005 secs

Status code distribution:
  [429] 61 responses
  [200] 39 responses

结果显示 0.03 秒发完了100个 Request，有 39 个请求被放行了，果然限流中间件有问题。中间件是在 Gin Middleware Github² 仓库找的，看完README就拿来直接用了。

中间件代码非常少，很容易分析出问题的原因

func NewRateLimiterMiddleware(redisClient *redis.Client, key string, limit int, slidingWindow time.Duration) gin.HandlerFunc {
	...
	return func(c *gin.Context) {
		now := time.Now().UnixNano()
		userCntKey := fmt.Sprint(c.ClientIP(), ":", key)

		redisClient.ZRemRangeByScore(userCntKey,
			"0",
			fmt.Sprint(now-(slidingWindow.Nanoseconds()))).Result()

		reqs, _ := redisClient.ZRange(userCntKey, 0, -1).Result()

		if len(reqs) >= limit {
			c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
				"status":  http.StatusTooManyRequests,
				"message": "too many request",
			})
			return
		}

		c.Next()
		redisClient.ZAddNX(userCntKey, redis.Z{Score: float64(now), Member: float64(now)})
		redisClient.Expire(userCntKey, slidingWindow)
	}

}

实现原理是将 Request 的客户端 IP 当作 key，当前时间纳秒当作 value 存入 redis 有序集合中，每次请求过来了，先删除窗口时间之前的所有数据, 然后统计当前时间减去设定的窗口时间之内集合中有多少条数据，如果没有达到阈值，则放行处理此请求，HTTP请求处理完成后，将请求时间存入集合。
#### 这段实现存在两个问题 1. 正在处理HTTP请求时，未将此次请求提前加入到Redis集合中，导致处理过程中的相同IP请求也会被放行，也是BUG的主要原因。 2. 即使查完之后发现不满足立即加入集合，也会有问题。因为这两步操作依然是分开的，可能会稍有缓解，限流误差不会有刚才那么大而已。

修复：

既然问题找到了，修复也很简单，将先查后增加封装成一个原子操作即可，例如利用 Redis SETNX 做一个简单的锁或者使用lua脚本来包装操作，由于 SETNX 可能存在锁未释放的风险，所以选择简单可靠一点的lua脚本
将上面的Go限流判断逻辑翻译成lua，参考了 gozero³

local limit = tonumber(ARGV[1])
local window = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local min = now - window
local current = redis.call("ZREMRANGEBYSCORE", key, '-inf', min)
local cnt = redis.call('ZCOUNT', key, '-inf', '+inf')
if cnt >= limit then
	return 1
else
	redis.call('ZADD', key, now, now)
	redis.call('PEXPIRE', key, window)
	return 0
end

参考资料