egg_app/app/mw/mw_limiter.go

package mw

import (
	"applet/app/e"
	"applet/app/md"
	"applet/app/svc"
	"applet/app/utils"
	"applet/app/utils/cache"
	"bytes"
	"fmt"
	"github.com/gin-gonic/gin"
	"io/ioutil"
)

func Limiter(c *gin.Context) {
	limit := 500 // 限流次数
	ttl := 2     // 限流过期时间
	ip := utils.GetIP(c.Request)
	//if ip != "221.4.210.167" && c.Request.Host != "127.0.0.1:4000" {
	//	e.OutErr(c, e.ERR, "系统维护中~")
	//	return
	//}
	// 读取token或者ip
	token := c.GetHeader("Authorization")
	// 判断是否已经超出限额次数
	method := c.Request.Method
	host := c.Request.Host
	uri := c.Request.URL.String()

	buf := make([]byte, 1024*1024*5) //TODO::5M
	num, _ := c.Request.Body.Read(buf)
	body := buf[:num]
	//body, _ := ioutil.ReadAll(c.Request.Body)
	// Write body back
	c.Request.Body = ioutil.NopCloser(bytes.NewBuffer(body))
	utils.FilePutContents("Limiter", utils.SerializeStr(map[string]interface{}{
		"uri":  c.Request.URL.Path,
		"body": string(bytes.NewBuffer(body).Bytes()),
	}))
	//queryValue := utils.SerializeStr(c.Request.URL.Query()) //不建议开启，失去限流的意义

	//TODO::分布式锁阻拦(保证原子性)
	requestIdPrefix := fmt.Sprintf(md.DealAppLimiterRequestIdPrefix, ip)
	cb, err := svc.HandleLimiterDistributedLock(ip, requestIdPrefix)
	if err != nil {
		e.OutErr(c, e.ERR, err.Error())
		return
	}
	if cb != nil {
		defer cb() // 释放锁
	}

	Md5 := utils.Md5(ip + token + method + host + uri + string(body))
	//Md5 := utils.Md5(ip + token + method + host + uri + string(body) + queryValue)
	//if cache.Exists(Md5) {
	//	c.AbortWithStatusJSON(428, gin.H{
	//		"code": 428,
	//		"msg":  "don't repeat the request",
	//		"data": struct{}{},
	//	})
	//	return
	//}

	// 2s后没返回自动释放
	go cache.SetEx(Md5, "0", ttl)

	key := "LIMITER_APP_EGG_" + ip
	reqs, _ := cache.GetInt(key)
	if reqs >= limit {
		c.AbortWithStatusJSON(429, gin.H{
			"code": 429,
			"msg":  "too many requests",
			"data": struct{}{},
		})
		return
	}
	if reqs > 0 {
		//go cache.Incr(key)
		go cache.SetEx(key, reqs+1, ttl)
	} else {
		go cache.SetEx(key, 1, ttl)
	}
	c.Next()
	go cache.Del(Md5)
}