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

optimize code

Browse Source
This commit is contained in:
xing 2025-03-16 22:11:28 +08:00
parent 2ac3d82542
commit 3c2896e06d
4 changed files with 33 additions and 31 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
app/middleware/flowLimit.go
View File

@ -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
}

35
app/middleware/iplimit.go
View File

@ -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
}

4
app/middleware/searchlimit.go
View File

@ -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")

7
app/route/route.go
View File

@ -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)