golang接口IP限流,IP黑名单,IP白名单

增加中间件

可以选择普通模式和LUA脚本模式,建议选择普通模式,实际上不需要控制的那么精确。

package Middlewares

import (
	"github.com/gin-gonic/gin"
	"strconv"
	"time"
	"voteapi/pkg/app/response"
	"voteapi/pkg/gredis"
	"voteapi/pkg/util"
)

const IP_LIMIT_NUM_KEY = "ipLimit:ipLimitNum"
const IP_BLACK_LIST_KEY = "ipLimit:ipBlackList"

var prefix = "{gateway}"
var delaySeconds int64 = 60   // 观察时间跨度,秒
var maxAttempts int64 = 10000 // 限制请求数
var blackSeconds int64 = 0   // 封禁时长,秒,0-不封禁

func GateWayPlus() gin.HandlerFunc {
	return func(c *gin.Context) {
		path := c.FullPath()
		clientIp := c.ClientIP()

		// redis配置集群时必须
		param := make(map[string]string)
		param["path"] = path
		param["clientIp"] = clientIp
		if !main(param) {
			c.Abort()
			response.JsonResponseError(c, "当前IP请求过于频繁,暂时被封禁~")
		}
	}
}

func main(param map[string]string) bool {
	// 预知的IP黑名单
	var blackList []string
	if util.InStringArray(param["clientIp"], blackList) {
		return false
	}

	// 预知的IP白名单
	var whiteList []string
	if util.InStringArray(param["clientIp"], whiteList) {
		return false
	}

	blackKey := prefix + ":" + IP_BLACK_LIST_KEY
	limitKey := prefix + ":" + IP_LIMIT_NUM_KEY

	curr := time.Now().Unix()
	item := util.Md5(param["path"] + "|" + param["clientIp"])

	return normal(blackKey, limitKey, item, curr)
}

// 普通模式
func normal(blackKey string, limitKey string, item string, time int64) (res bool) {
	if blackSeconds > 0 {
		timeout, _ := gredis.RawCommand("HGET", blackKey, item)
		if timeout != nil {
			to, _ := strconv.Atoi(string(timeout.([]uint8)))
			if int64(to) > time {
				// 未解封
				return false
			}
			// 已解封,移除黑名单
			gredis.RawCommand("HDEL", blackKey, item)
		}
	}

	l, _ := gredis.RawCommand("HGET", limitKey, item)
	if l != nil {
		last, _ := strconv.Atoi(string(l.([]uint8)))
		if int64(last) >= maxAttempts {
			return false
		}
	}

	num, _ := gredis.RawCommand("HINCRBY", limitKey, item, 1)
	if ttl, _ := gredis.TTLKey(limitKey); ttl == int64(-1) {
		gredis.Expire(limitKey, int64(delaySeconds))
	}

	if num.(int64) >= maxAttempts && blackSeconds > 0 {
		// 加入黑名单
		gredis.RawCommand("HSET", blackKey, item, time+blackSeconds)
		// 删除记录
		gredis.RawCommand("HDEL", limitKey, item)
	}

	return true
}

// LUA脚本模式
// 支持redis集群部署
func luaScript(blackKey string, limitKey string, item string, time int64) (res bool) {
	script := `
local blackSeconds = tonumber(ARGV[5])
if(blackSeconds > 0)
then
    local timeout = redis.call('hget', KEYS[1], ARGV[1])
    if(timeout ~= false)
    then
        if(tonumber(timeout) > tonumber(ARGV[2]))
        then
            return false
        end
        redis.call('hdel', KEYS[1], ARGV[1])
    end
end

local last = redis.call('hget', KEYS[2], ARGV[1])
if(last ~= false and tonumber(last) >= tonumber(ARGV[3]))
then
    return false
end

local num = redis.call('hincrby', KEYS[2], ARGV[1], 1)
local ttl = redis.call('ttl', KEYS[2])
if(ttl == -1)
then
    redis.call('expire', KEYS[2], ARGV[4])
end

if(tonumber(num) >= tonumber(ARGV[3]) and blackSeconds > 0)
then 
    redis.call('hset', KEYS[1], ARGV[1], ARGV[2] + ARGV[5])
    redis.call('hdel', KEYS[2], ARGV[1])
end
return true
`
	result, err := gredis.RawCommand("EVAL", script, 2, blackKey, limitKey, item, time, maxAttempts, delaySeconds, blackSeconds)
	if err != nil {
		return false
	}

	if result == int64(1) {
		return true
	} else {
		return false
	}
}
  • 0
    点赞
  • 8
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值