diff --git a/app/middleware/flowLimit.go b/app/middleware/flowLimit.go index 811a2ee..9a2edc7 100644 --- a/app/middleware/flowLimit.go +++ b/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 } diff --git a/app/middleware/iplimit.go b/app/middleware/iplimit.go index 48804aa..12ef767 100644 --- a/app/middleware/iplimit.go +++ b/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)) { - m := IpLimitMap{ +func IpLimit(num int64, clearNum ...int64) (func(ctx *gin.Context), func(int64, ...int64)) { + 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) - if ii > 0 { - atomic.AddInt64(i, -1) - if atomic.LoadInt64(i) == 0 { + atomic.AddInt64(i, -1) + if atomic.LoadInt64(i) <= 0 { + cNum := int(atomic.LoadInt64(m.clearNum)) + 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 } diff --git a/app/middleware/searchlimit.go b/app/middleware/searchlimit.go index 8c6257b..cdd9677 100644 --- a/app/middleware/searchlimit.go +++ b/app/middleware/searchlimit.go @@ -6,8 +6,8 @@ import ( "github.com/gin-gonic/gin" ) -func SearchLimit(num int64) func(ctx *gin.Context) { - fn, reFn := IpLimit(num) +func SearchLimit(num int64, clearNum ...int64) func(ctx *gin.Context) { + fn, reFn := IpLimit(num, clearNum...) reload.Append(func() { reFn(config.GetConfig().SingleIpSearchNum) }, "search-ip-limit-number") diff --git a/app/route/route.go b/app/route/route.go index b7fe701..8c0878c 100644 --- a/app/route/route.go +++ b/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) - r.POST("/comment", commentMiddleWare, actions.PostComment) + commentMiddleWare, _ := middleware.FlowLimit(5, c.SingleIpSearchNum, c.CacheTime.SleepTime) + commentIpMiddleware, _ := middleware.IpLimit(5, 5) + r.POST("/comment", commentMiddleWare, commentIpMiddleware, actions.PostComment) r.NoRoute(actions.ThemeHook(constraints.NoRoute)) }, 84.6)