xing/wp-go
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

小优化

Browse Source
This commit is contained in:
xing 2022-09-29 17:25:15 +08:00
parent 5254dd1f46
commit 8efddea822
1 changed files with 6 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
middleware/flowLimit.go
View File

@ -39,7 +39,8 @@ func FlowLimit() func(ctx *gin.Context) {
c.Next() c.Next()
return return
} }
if m.searchLimit(true, c, f) { ip := c.ClientIP()
if m.searchLimit(true, c, ip, f) {
c.Abort() c.Abort()
return return
} }
@ -51,12 +52,12 @@ func FlowLimit() func(ctx *gin.Context) {
c.String(http.StatusForbidden, "请求太多了，服务器君压力山大中==!, 请稍后访问") c.String(http.StatusForbidden, "请求太多了，服务器君压力山大中==!, 请稍后访问")
c.Abort() c.Abort()
atomic.AddInt64(&flow, -1) atomic.AddInt64(&flow, -1)
m.searchLimit(false, c, f) m.searchLimit(false, c, ip, f)
return return
} }
c.Next() c.Next()
m.searchLimit(false, c, f) m.searchLimit(false, c, ip, f)
atomic.AddInt64(&flow, -1) atomic.AddInt64(&flow, -1)
} }
} }
@ -67,8 +68,8 @@ func (m *IpLimitMap) set(k string, n int) {
m.m[k] = n m.m[k] = n
} }
func (m *IpLimitMap) searchLimit(a bool, c *gin.Context, f []string) (isForbid bool) { func (m *IpLimitMap) searchLimit(a bool, c *gin.Context, ip string, f []string) (isForbid bool) {
ip := c.ClientIP()
if f[0] == "" && c.Query("s") != "" { if f[0] == "" && c.Query("s") != "" {
if a { if a {
i, ok := m.m[ip] i, ok := m.m[ip]