Gin 限流中间件

func limiting() gin.HandlerFunc {
    return func(c *gin.Context) {
        // limitingConfig 限流配置
        type limitingConfig struct {
            Time  time.Duration // 时间范围
            Value uint16        // 时间范围内的最大次数
        }

        // limitingList 限速列表
        // key 是 method 和 c.FullPath() 用 `_` 来拼接
        // 用 c.FullPath() 是因为获取的是路由的规则参数
        // 比如 /user/:name 获取的是 /user/:name 而不是 /user/hongfs
        var limitingList = make(map[string]limitingConfig)

        limitingList["GET_/user"] = limitingConfig{
            Time:  time.Minute * 1,
            Value: 10,
        }

        routeKey := fmt.Sprintf("%s_%s", c.Request.Method, c.FullPath())

        limitingInfo, exists := limitingList[routeKey]

        if exists == false {
            // 获取不到当前路由的限流配置信息
            // 可以工具业务来决定抛出 AbortWithStatus 结束还是继续
            c.Next()

            return
        }

        // 获取当前请求来源的用户标识
        identification := func(c *gin.Context) string {
            // 客户端 IP
            return c.ClientIP()

            // 用户 UID
            // return c.MustGet("uid").(string)

            // 共用？
            // return "1"
        }(c)

        // 计算出当前是第几个时间范围了
        timeIndex := float64(time.Now().Unix()) / limitingInfo.Time.Seconds()

        // 缓存键值
        cacheKey := fmt.Sprintf("%s:%s:%d", c.Request.Method, sha1Value(c.FullPath()), int64(math.Floor(timeIndex)))

        rdb := redis.NewClient(&redis.Options{
            Addr: "redis.hongfs.cn",
        })

        // 通过 Redis Hash 自增返回值来判断是不是超过限制了
        res, err := rdb.HIncrBy(context.Background(), cacheKey, identification, 1).Result()

        if err != nil {
            log.Println("Redis 处理异常：" + err.Error())

            c.AbortWithStatus(http.StatusInternalServerError)

            return
        }

        // 设置过期时间，因为采用 Hash，所以会存在重复性设置
        if res == 1 {
            rdb.Expire(context.Background(), cacheKey, time.Duration(limitingInfo.Time.Seconds()+1))
        }

        // X-Ratelimit-Limit 当前时间范围内的最大请求数
        c.Writer.Header().Set("X-Ratelimit-Limit", strconv.Itoa(int(limitingInfo.Value)))
        
        // X-Ratelimit-Reset 下一个时间范围的开始时间戳
        c.Writer.Header().Set("X-Ratelimit-Reset", strconv.FormatInt(int64(timeIndex+1)*int64(limitingInfo.Time.Seconds()), 10))

        if res > int64(limitingInfo.Value) {
            c.Writer.Header().Set("X-Ratelimit-Remaining", "0")

            // 429 代表已达限制
            c.AbortWithStatus(http.StatusTooManyRequests)

            return
        }

        // X-Ratelimit-Remaining 当前时间范围还剩下的请求次数
        c.Writer.Header().Set("X-Ratelimit-Remaining", strconv.FormatInt(int64(limitingInfo.Value)-res, 10))

        c.Next()
    }
}