app/middleware/flowLimit.go

@@ -5,7 +5,6 @@ import (
 	"github.com/fthvgb1/wp-go/safety"
 	"github.com/gin-gonic/gin"
 	"net/http"
-	"strings"
 	"sync/atomic"
 	"time"
 )
@@ -37,21 +36,3 @@ func FlowLimit(maxRequestSleepNum, maxRequestNum int64, sleepTime []time.Duratio
 		c.Next()
 	}, fn
 }
-
-func WPFlowLimit(maxRequestSleepNum, maxRequestNum int64, sleepTime []time.Duration) (func(ctx *gin.Context), func(int64, int64, []time.Duration)) {
-	statPath := map[string]struct{}{
-		"wp-includes": {},
-		"wp-content":  {},
-		"favicon.ico": {},
-	}
-	next, reFn := FlowLimit(maxRequestSleepNum, maxRequestNum, sleepTime)
-	return func(c *gin.Context) {
-		f := strings.Split(strings.TrimLeft(c.FullPath(), "/"), "/")
-		_, ok := statPath[f[0]]
-		if len(f) > 0 && ok {
-			c.Next()
-			return
-		}
-		next(c)
-	}, reFn
-}

app/middleware/iplimit.go

@@ -1,192 +1,79 @@
 package middleware
 
 import (
-	"github.com/fthvgb1/wp-go/helper/number"
-	str "github.com/fthvgb1/wp-go/helper/strings"
 	"github.com/gin-gonic/gin"
 	"net/http"
-	"strings"
 	"sync"
 	"sync/atomic"
-	"time"
 )
 
-type LimitMap[K comparable] struct {
+type ipLimitMap struct {
-	Mux      *sync.RWMutex
+	mux      *sync.RWMutex
-	Map      map[K]*int64
+	m        map[string]*int64
-	LimitNum *int64
+	limitNum *int64
-	ClearNum *int64
+	clearNum *int64
-}
-
-type FlowLimits[K comparable] struct {
-	GetKeyFn     func(ctx *gin.Context) K
-	LimitedFns   func(ctx *gin.Context)
-	DeferClearFn func(c *gin.Context, m LimitMap[K], k K, v *int64)
-	AddFn        func(c *gin.Context, m LimitMap[K], k K, v *int64)
-}
-
-func NewFlowLimits[K comparable](getKeyFn func(ctx *gin.Context) K,
-	limitedFns func(ctx *gin.Context),
-	deferClearFn func(c *gin.Context, m LimitMap[K], k K, v *int64),
-	addFns ...func(c *gin.Context, m LimitMap[K], k K, v *int64),
-) *FlowLimits[K] {
-
-	f := &FlowLimits[K]{
-		GetKeyFn:     getKeyFn,
-		LimitedFns:   limitedFns,
-		DeferClearFn: deferClearFn,
-	}
-	fn := f.Adds
-	if len(addFns) > 0 {
-		fn = addFns[0]
-	}
-	f.AddFn = fn
-	return f
-}
-
-func (f FlowLimits[K]) GetKey(c *gin.Context) K {
-	return f.GetKeyFn(c)
-}
-
-func (f FlowLimits[K]) Limit(c *gin.Context) {
-	f.LimitedFns(c)
-}
-func (f FlowLimits[K]) DeferClear(c *gin.Context, m LimitMap[K], k K, v *int64) {
-	f.DeferClearFn(c, m, k, v)
-}
-
-func (f FlowLimits[K]) Add(c *gin.Context, m LimitMap[K], k K, v *int64) {
-	f.AddFn(c, m, k, v)
-}
-
-func (f FlowLimits[K]) Adds(_ *gin.Context, _ LimitMap[K], _ K, v *int64) {
-	atomic.AddInt64(v, 1)
-}
-
-type MapFlowLimit[K comparable] interface {
-	GetKey(c *gin.Context) K
-	Limit(c *gin.Context)
-	DeferClear(c *gin.Context, m LimitMap[K], k K, v *int64)
-	Add(c *gin.Context, m LimitMap[K], k K, v *int64)
-}
-
-func CustomFlowLimit[K comparable](a MapFlowLimit[K], maxRequestNum int64, clearNum ...int64) (func(ctx *gin.Context), func(int64, ...int64)) {
-	m := LimitMap[K]{
-		Mux:      &sync.RWMutex{},
-		Map:      make(map[K]*int64),
-		LimitNum: new(int64),
-		ClearNum: new(int64),
-	}
-	atomic.StoreInt64(m.LimitNum, maxRequestNum)
-	if len(clearNum) > 0 {
-		atomic.StoreInt64(m.ClearNum, clearNum[0])
-	}
-
-	fn := func(num int64, clearNum ...int64) {
-		atomic.StoreInt64(m.LimitNum, num)
-		if len(clearNum) > 0 {
-			atomic.StoreInt64(m.ClearNum, clearNum[0])
-		}
-	}
-	return func(c *gin.Context) {
-		if atomic.LoadInt64(m.LimitNum) <= 0 {
-			c.Next()
-			return
-		}
-		key := a.GetKey(c)
-		m.Mux.RLock()
-		i, ok := m.Map[key]
-		m.Mux.RUnlock()
-		if !ok {
-			m.Mux.Lock()
-			i = new(int64)
-			m.Map[key] = i
-			m.Mux.Unlock()
-		}
-		a.Add(c, m, key, i)
-		defer a.DeferClear(c, m, key, i)
-		if atomic.LoadInt64(i) > atomic.LoadInt64(m.LimitNum) {
-			a.Limit(c)
-			return
-		}
-		c.Next()
-	}, fn
-}
-
-func IpLimitClear[K comparable](_ *gin.Context, m LimitMap[K], key K, i *int64) {
-	atomic.AddInt64(i, -1)
-	if atomic.LoadInt64(i) <= 0 {
-		cNum := int(atomic.LoadInt64(m.ClearNum))
-		if cNum <= 0 {
-			m.Mux.Lock()
-			delete(m.Map, key)
-			m.Mux.Unlock()
-			return
-		}
-
-		m.Mux.RLock()
-		l := len(m.Map)
-		m.Mux.RUnlock()
-		if l < cNum {
-			m.Mux.Lock()
-			for k, v := range m.Map {
-				if atomic.LoadInt64(v) < 1 {
-					delete(m.Map, k)
-				}
-			}
-			m.Mux.Unlock()
-		}
-	}
-}
-
-func ToManyRequest(messages ...string) func(c *gin.Context) {
-	message := "请求太多了，服务器君表示压力山大==!, 请稍后访问"
-	if len(messages) > 0 {
-		message = messages[0]
-	}
-	return func(c *gin.Context) {
-		c.String(http.StatusForbidden, message)
-		c.Abort()
-	}
 }
 
 func IpLimit(num int64, clearNum ...int64) (func(ctx *gin.Context), func(int64, ...int64)) {
-	a := NewFlowLimits(func(c *gin.Context) string {
+	m := ipLimitMap{
-		return c.ClientIP()
+		mux:      &sync.RWMutex{},
-	}, ToManyRequest(), IpLimitClear)
+		m:        make(map[string]*int64),
-	return CustomFlowLimit[string](a, num, clearNum...)
+		limitNum: new(int64),
+		clearNum: new(int64),
+	}
+	fn := func(num int64, clearNum ...int64) {
+		atomic.StoreInt64(m.limitNum, num)
+		if len(clearNum) > 0 {
+			atomic.StoreInt64(m.clearNum, clearNum[0])
+		}
+	}
+	fn(num, clearNum...)
+
+	return func(c *gin.Context) {
+		if atomic.LoadInt64(m.limitNum) <= 0 {
+			c.Next()
+			return
+		}
+		ip := c.ClientIP()
+		m.mux.RLock()
+		i, ok := m.m[ip]
+		m.mux.RUnlock()
+
+		if !ok {
+			m.mux.Lock()
+			i = new(int64)
+			m.m[ip] = i
+			m.mux.Unlock()
 		}
 
-func IpMinuteLimit(num int64, clearNum ...int64) (func(ctx *gin.Context), func(int64, ...int64)) {
+		defer func() {
-	total := new(int64)
+			atomic.AddInt64(i, -1)
-	a := NewFlowLimits(func(c *gin.Context) string {
+			if atomic.LoadInt64(i) <= 0 {
-		return str.Join(c.ClientIP(), "|", time.Now().Format("2006-01-02 15:04"))
+				cNum := int(atomic.LoadInt64(m.clearNum))
-	},
+				if cNum <= 0 {
-		ToManyRequest(),
+					m.mux.Lock()
-		IpMinuteLimitDeferFn(total),
+					delete(m.m, ip)
-		func(c *gin.Context, m LimitMap[string], k string, v *int64) {
+					m.mux.Unlock()
-			atomic.AddInt64(v, 1)
+					return
-			atomic.AddInt64(total, 1)
-		},
-	)
-
-	return CustomFlowLimit(a, num, clearNum...)
 				}
 
-func IpMinuteLimitDeferFn(total *int64) func(_ *gin.Context, m LimitMap[string], k string, _ *int64) {
+				m.mux.RLock()
-	return func(_ *gin.Context, m LimitMap[string], k string, _ *int64) {
+				l := len(m.m)
-		atomic.AddInt64(total, -1)
+				m.mux.RUnlock()
-		cNum := number.Min(int(atomic.LoadInt64(m.ClearNum)), 1)
+				if l < cNum {
-		minu := strings.Split(k, "|")[1]
+					m.mux.Lock()
-		if int(atomic.LoadInt64(total)) < cNum {
+					delete(m.m, ip)
-			m.Mux.Lock()
+					m.mux.Unlock()
-			for key := range m.Map {
-				t := strings.Split(key, "|")[1]
-				if minu != t {
-					delete(m.Map, key)
 				}
 			}
-			m.Mux.Unlock()
+		}()
-		}
+
+		if atomic.LoadInt64(i) >= atomic.LoadInt64(m.limitNum) {
+			c.String(http.StatusForbidden, "请求太多了，服务器君表示压力山大==!, 请稍后访问")
+			c.Abort()
+			return
 		}
+		atomic.AddInt64(i, 1)
+		c.Next()
+	}, fn
 }

app/plugins/devexample/plugintt/make.sh

@@ -1,4 +1,4 @@
-#!/bin/bash
+#/bin/bash
 
 # copy plugintt to other dir and remove .dev suffix
 # note the go version and build tool flag must same to server build

app/route/route.go

@@ -59,7 +59,7 @@ func SetupRouter() *gin.Engine {
 		wpconfig.SetTemplateFs(theme.TemplateFs)
 	}, 90.5)
 
-	siteFlowLimitMiddleware, siteFlow := middleware.WPFlowLimit(c.MaxRequestSleepNum, c.MaxRequestNum, c.CacheTime.SleepTime)
+	siteFlowLimitMiddleware, siteFlow := middleware.FlowLimit(c.MaxRequestSleepNum, c.MaxRequestNum, c.CacheTime.SleepTime)
 	reload.Append(func() {
 		c = config.GetConfig()
 		siteFlow(c.MaxRequestSleepNum, c.MaxRequestNum, c.CacheTime.SleepTime)