请求限制

config.example.yaml

@@ -39,3 +39,11 @@ postDataCacheTime: 1h
 commentsCacheTime: 5m
 # 定时清理缓存周期时间
 crontabClearCacheTime: 5m
+# 到达指定并发请求数时随机sleep
+maxRequestSleepNum: 100
+# 随机sleep时间
+sleepTime: [1s,3s]
+# 全局最大请求数，超过直接403
+maxRequestNum: 500
+# 单ip同时最大搜索请求数
+singleIpSearchNum: 10

middleware/flowLimit.go

@@ -0,0 +1,98 @@
+package middleware
+
+import (
+	"github.com/gin-gonic/gin"
+	"github/fthvgb1/wp-go/vars"
+	"math/rand"
+	"net/http"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+type IpLimitMap struct {
+	mux *sync.Mutex
+	m   map[string]int
+}
+
+func FlowLimit() func(ctx *gin.Context) {
+	var flow int64
+	rand.Seed(time.Now().UnixNano())
+	randFn := func(start, end time.Duration) time.Duration {
+		end++
+		return time.Duration(rand.Intn(int(end-start)) + int(start))
+	}
+	m := IpLimitMap{
+		mux: &sync.Mutex{},
+		m:   make(map[string]int),
+	}
+	statPath := map[string]struct{}{
+		"wp-includes": {},
+		"wp-content":  {},
+		"favicon.ico": {},
+	}
+	return func(c *gin.Context) {
+		f := strings.Split(strings.TrimLeft(c.FullPath(), "/"), "/")
+		_, ok := statPath[f[0]]
+		if len(f) > 0 && ok {
+			c.Next()
+			return
+		}
+		if m.searchLimit(true, c, f) {
+			c.Abort()
+			return
+		}
+		atomic.AddInt64(&flow, 1)
+		if flow >= vars.Conf.MaxRequestSleepNum && flow <= vars.Conf.MaxRequestNum {
+			t := randFn(vars.Conf.SleepTime[0], vars.Conf.SleepTime[1])
+			time.Sleep(t)
+		} else if flow > vars.Conf.MaxRequestNum {
+			c.String(http.StatusForbidden, "请求太多了，服务器君压力山大中==!, 请稍后访问")
+			c.Abort()
+			atomic.AddInt64(&flow, -1)
+			m.searchLimit(false, c, f)
+			return
+		}
+
+		c.Next()
+		m.searchLimit(false, c, f)
+		atomic.AddInt64(&flow, -1)
+	}
+}
+
+func (m *IpLimitMap) set(k string, n int) {
+	m.mux.Lock()
+	defer m.mux.Unlock()
+	m.m[k] = n
+}
+
+func (m *IpLimitMap) searchLimit(a bool, c *gin.Context, f []string) (isForbid bool) {
+	ip := c.ClientIP()
+	if f[0] == "" && c.Query("s") != "" {
+		if a {
+			i, ok := m.m[ip]
+			if ok {
+				num := vars.Conf.SingleIpSearchNum
+				if num < 1 {
+					num = 10
+				}
+				if i > num {
+					return true
+				}
+			} else {
+				i = 0
+			}
+			i++
+			m.set(ip, i)
+		} else {
+			m.set(ip, m.m[ip]-1)
+			if m.m[ip] == 0 {
+				m.mux.Lock()
+				delete(m.m, ip)
+				m.mux.Unlock()
+			}
+		}
+	}
+	return
+}

plugins/digest.go

@@ -113,4 +113,5 @@ func Digest(p *Plugin[models.WpPosts], c *gin.Context, post *models.WpPosts, sce
 		return
 	}
 	post.PostContent = DigestCache(c, post.Id, post.PostContent)
+	p.Next()
 }

route/route.go

@@ -6,6 +6,7 @@ import (
 	"github.com/gin-contrib/sessions/cookie"
 	"github.com/gin-gonic/gin"
 	"github/fthvgb1/wp-go/actions"
+	"github/fthvgb1/wp-go/helper"
 	"github/fthvgb1/wp-go/middleware"
 	"github/fthvgb1/wp-go/static"
 	"github/fthvgb1/wp-go/templates"
@@ -17,7 +18,7 @@ import (
 func SetupRouter() *gin.Engine {
 	// Disable Console Color
 	// gin.DisableConsoleColor()
-	r := gin.Default()
+	r := gin.New()
 	r.HTMLRender = templates.NewFsTemplate(template.FuncMap{
 		"unescaped": func(s string) any {
 			return template.HTML(s)
@@ -26,7 +27,7 @@ func SetupRouter() *gin.Engine {
 			return t.Format("2006年 01月 02日")
 		},
 	}).SetTemplate()
-	r.Use(middleware.SetStaticFileCache)
+	r.Use(gin.Logger(), gin.Recovery(), middleware.FlowLimit(), middleware.SetStaticFileCache)
 	//gzip 因为一般会用nginx做反代时自动使用gzip,所以go这边本身可以不用
 	/*r.Use(gzip.Gzip(gzip.DefaultCompression, gzip.WithExcludedPaths([]string{
 		"/wp-includes/", "/wp-content/",
@@ -51,6 +52,8 @@ func SetupRouter() *gin.Engine {
 	r.GET("/p/date/:year/:month/page/:page", actions.Index)
 	r.POST("/login", actions.Login)
 	r.GET("/p/:id", actions.Detail)
+	if helper.IsContainInArr(gin.Mode(), []string{gin.DebugMode, gin.TestMode}) {
 		pprof.Register(r, "dev/pprof")
+	}
 	return r
 }

vars/config.go

@@ -24,6 +24,10 @@ type Config struct {
 	PostDataCacheTime       time.Duration   `yaml:"postDataCacheTime"`
 	CommentsCacheTime       time.Duration   `yaml:"commentsCacheTime"`
 	CrontabClearCacheTime   time.Duration   `yaml:"crontabClearCacheTime"`
+	MaxRequestSleepNum      int64           `yaml:"maxRequestSleepNum"`
+	SleepTime               []time.Duration `yaml:"sleepTime"`
+	MaxRequestNum           int64           `yaml:"maxRequestNum"`
+	SingleIpSearchNum       int             `yaml:"singleIpSearchNum"`
 }
 
 type Mysql struct {