optimize code

app/middleware/flowLimit.go

@@ -5,18 +5,12 @@ import (
 	"github.com/fthvgb1/wp-go/safety"
 	"github.com/gin-gonic/gin"
 	"net/http"
-	"strings"
 	"sync/atomic"
 	"time"
 )
 
 func FlowLimit(maxRequestSleepNum, maxRequestNum int64, sleepTime []time.Duration) (func(ctx *gin.Context), func(int64, int64, []time.Duration)) {
 	var flow int64
-	statPath := map[string]struct{}{
-		"wp-includes": {},
-		"wp-content":  {},
-		"favicon.ico": {},
-	}
 	s := safety.Var[[]time.Duration]{}
 	s.Store(sleepTime)
 	fn := func(sleepNum, maxNum int64, st []time.Duration) {
@@ -25,19 +19,12 @@ func FlowLimit(maxRequestSleepNum, maxRequestNum int64, sleepTime []time.Duratio
 		s.Store(st)
 	}
 	return func(c *gin.Context) {
-		f := strings.Split(strings.TrimLeft(c.FullPath(), "/"), "/")
-		_, ok := statPath[f[0]]
-		if len(f) > 0 && ok {
-			c.Next()
-			return
-		}
-
 		n := atomic.LoadInt64(&flow)
-		if n >= atomic.LoadInt64(&maxRequestSleepNum) && n <= atomic.LoadInt64(&maxRequestNum) {
+		if n >= atomic.LoadInt64(&maxRequestSleepNum) && n < atomic.LoadInt64(&maxRequestNum) {
 			ss := s.Load()
 			t := number.Rand(ss[0], ss[1])
 			time.Sleep(t)
-		} else if n > atomic.LoadInt64(&maxRequestNum) {
+		} else if n >= atomic.LoadInt64(&maxRequestNum) {
 			c.String(http.StatusForbidden, "请求太多了，服务器君表示压力山大==!, 请稍后访问")
 			c.Abort()
 			return
@@ -46,7 +33,6 @@ func FlowLimit(maxRequestSleepNum, maxRequestNum int64, sleepTime []time.Duratio
 		defer func() {
 			atomic.AddInt64(&flow, -1)
 		}()
-
 		c.Next()
 	}, fn
 }

app/middleware/iplimit.go

@@ -7,22 +7,27 @@ import (
 	"sync/atomic"
 )
 
-type IpLimitMap struct {
+type ipLimitMap struct {
 	mux      *sync.RWMutex
 	m        map[string]*int64
 	limitNum *int64
+	clearNum *int64
 }
 
-func IpLimit(num int64) (func(ctx *gin.Context), func(int64)) {
+func IpLimit(num int64, clearNum ...int64) (func(ctx *gin.Context), func(int64, ...int64)) {
-	m := IpLimitMap{
+	m := ipLimitMap{
 		mux:      &sync.RWMutex{},
 		m:        make(map[string]*int64),
 		limitNum: new(int64),
+		clearNum: new(int64),
 	}
-	fn := func(num int64) {
+	fn := func(num int64, clearNum ...int64) {
 		atomic.StoreInt64(m.limitNum, num)
+		if len(clearNum) > 0 {
+			atomic.StoreInt64(m.clearNum, clearNum[0])
+		}
 	}
-	fn(num)
+	fn(num, clearNum...)
 
 	return func(c *gin.Context) {
 		if atomic.LoadInt64(m.limitNum) <= 0 {
@@ -42,10 +47,20 @@ func IpLimit(num int64) (func(ctx *gin.Context), func(int64)) {
 		}
 
 		defer func() {
-			ii := atomic.LoadInt64(i)
+			atomic.AddInt64(i, -1)
-			if ii > 0 {
+			if atomic.LoadInt64(i) <= 0 {
-				atomic.AddInt64(i, -1)
+				cNum := int(atomic.LoadInt64(m.clearNum))
-				if atomic.LoadInt64(i) == 0 {
+				if cNum <= 0 {
+					m.mux.Lock()
+					delete(m.m, ip)
+					m.mux.Unlock()
+					return
+				}
+
+				m.mux.RLock()
+				l := len(m.m)
+				m.mux.RUnlock()
+				if l < cNum {
 					m.mux.Lock()
 					delete(m.m, ip)
 					m.mux.Unlock()
@@ -54,7 +69,7 @@ func IpLimit(num int64) (func(ctx *gin.Context), func(int64)) {
 		}()
 
 		if atomic.LoadInt64(i) >= atomic.LoadInt64(m.limitNum) {
-			c.Status(http.StatusForbidden)
+			c.String(http.StatusForbidden, "请求太多了，服务器君表示压力山大==!, 请稍后访问")
 			c.Abort()
 			return
 		}

app/middleware/searchlimit.go

@@ -6,8 +6,8 @@ import (
 	"github.com/gin-gonic/gin"
 )
 
-func SearchLimit(num int64) func(ctx *gin.Context) {
+func SearchLimit(num int64, clearNum ...int64) func(ctx *gin.Context) {
-	fn, reFn := IpLimit(num)
+	fn, reFn := IpLimit(num, clearNum...)
 	reload.Append(func() {
 		reFn(config.GetConfig().SingleIpSearchNum)
 	}, "search-ip-limit-number")

app/route/route.go

@@ -102,7 +102,7 @@ func SetupRouter() *gin.Engine {
 	}, 85.1)
 
 	SetGinAction("setRoute", func(r *gin.Engine) {
-		r.GET("/", actions.Feed, middleware.SearchLimit(c.SingleIpSearchNum),
+		r.GET("/", actions.Feed, middleware.SearchLimit(c.SingleIpSearchNum, 5),
 			actions.ThemeHook(constraints.Home))
 		r.GET("/page/:page", actions.ThemeHook(constraints.Home))
 		r.GET("/p/category/:category", actions.ThemeHook(constraints.Category))
@@ -119,8 +119,9 @@ func SetupRouter() *gin.Engine {
 		r.GET("/p/:id/feed", actions.PostFeed)
 		r.GET("/feed", actions.SiteFeed)
 		r.GET("/comments/feed", actions.CommentsFeed)
-		commentMiddleWare, _ := middleware.FlowLimit(c.MaxRequestSleepNum, 5, c.CacheTime.SleepTime)
+		commentMiddleWare, _ := middleware.FlowLimit(5, c.SingleIpSearchNum, c.CacheTime.SleepTime)
-		r.POST("/comment", commentMiddleWare, actions.PostComment)
+		commentIpMiddleware, _ := middleware.IpLimit(5, 5)
+		r.POST("/comment", commentMiddleWare, commentIpMiddleware, actions.PostComment)
 
 		r.NoRoute(actions.ThemeHook(constraints.NoRoute))
 	}, 84.6)