84 lines
1.9 KiB
Go
84 lines
1.9 KiB
Go
// Package middleware ...
|
|
//
|
|
// Description : 基于redis的流控管理
|
|
//
|
|
// Author : go_developer@163.com<白茶清欢>
|
|
//
|
|
// Date : 2021-07-30 20:52 下午
|
|
package middleware
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"time"
|
|
|
|
"github.com/pkg/errors"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/go-redis/redis/v8"
|
|
"github.com/go-redis/redis_rate/v9"
|
|
)
|
|
|
|
// BuildRateLimitConfig 构建流控配置
|
|
//
|
|
// Author : go_developer@163.com<白茶清欢>
|
|
//
|
|
// Date : 10:35 下午 2021/7/30
|
|
func BuildRateLimitConfig(rate int, burst int, period time.Duration) redis_rate.Limit {
|
|
return redis_rate.Limit{
|
|
Rate: rate,
|
|
Burst: burst,
|
|
Period: period,
|
|
}
|
|
}
|
|
|
|
// RateLimit gin 可以直接引用的中间件
|
|
//
|
|
// Author : go_developer@163.com<白茶清欢>
|
|
//
|
|
// Date : 10:35 下午 2021/7/30
|
|
func RateLimit(redisInstance *redis.Client, key string, limitRule redis_rate.Limit, responseData map[string]interface{}) gin.HandlerFunc {
|
|
return func(ctx *gin.Context) {
|
|
var (
|
|
err error
|
|
)
|
|
|
|
if _, err = Allow(redisInstance, key, limitRule); nil != err {
|
|
if nil == responseData {
|
|
responseData = make(map[string]interface{})
|
|
responseData["code"] = -1
|
|
responseData["message"] = err.Error()
|
|
responseData["data"] = make(map[string]interface{})
|
|
}
|
|
ctx.AbortWithStatusJSON(http.StatusOK, responseData)
|
|
return
|
|
}
|
|
|
|
ctx.Next()
|
|
}
|
|
}
|
|
|
|
// Allow 是否允许访问
|
|
//
|
|
// Author : go_developer@163.com<白茶清欢>
|
|
//
|
|
// Date : 11:52 下午 2021/7/30
|
|
func Allow(redisInstance *redis.Client, key string, limitRule redis_rate.Limit) (bool, error) {
|
|
var (
|
|
limiter *redis_rate.Limiter
|
|
result *redis_rate.Result
|
|
err error
|
|
)
|
|
|
|
limiter = redis_rate.NewLimiter(redisInstance)
|
|
if result, err = limiter.Allow(context.Background(), key, limitRule); nil != err {
|
|
return false, err
|
|
}
|
|
|
|
if result.Allowed > 0 {
|
|
return true, nil
|
|
}
|
|
|
|
return false, errors.New("rate limit! current request is not allowed!")
|
|
}
|