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) }