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-19 10:33:35 +08:00
parent 274610be8b
commit 64556d69a0
2 changed files with 30 additions and 41 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

55
app/middleware/iplimit.go
View File

@ -1,7 +1,6 @@
package middleware package middleware
import ( import (
"github.com/fthvgb1/wp-go/helper/number"
str "github.com/fthvgb1/wp-go/helper/strings" str "github.com/fthvgb1/wp-go/helper/strings"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"net/http" "net/http"
@ -21,14 +20,14 @@ type LimitMap[K comparable] struct {
type FlowLimits[K comparable] struct { type FlowLimits[K comparable] struct {
GetKeyFn func(ctx *gin.Context) K GetKeyFn func(ctx *gin.Context) K
LimitedFns func(ctx *gin.Context) LimitedFns func(ctx *gin.Context)
DeferClearFn func(c *gin.Context, m LimitMap[K], k K, v *int64) DeferClearFn func(c *gin.Context, m LimitMap[K], k K, v, currentTotalFlow *int64)
AddFn func(c *gin.Context, m LimitMap[K], k K, v *int64) AddFn func(c *gin.Context, m LimitMap[K], k K, v, currentTotalFlow *int64)
} }
func NewFlowLimits[K comparable](getKeyFn func(ctx *gin.Context) K, func NewFlowLimits[K comparable](getKeyFn func(ctx *gin.Context) K,
limitedFns func(ctx *gin.Context), limitedFns func(ctx *gin.Context),
deferClearFn func(c *gin.Context, m LimitMap[K], k K, v *int64), deferClearFn func(c *gin.Context, m LimitMap[K], k K, v, currentTotalFlow *int64),
addFns ...func(c *gin.Context, m LimitMap[K], k K, v *int64), addFns ...func(c *gin.Context, m LimitMap[K], k K, v, currentTotalFlow *int64),
) *FlowLimits[K] { ) *FlowLimits[K] {
f := &FlowLimits[K]{ f := &FlowLimits[K]{
@ -51,23 +50,24 @@ func (f FlowLimits[K]) GetKey(c *gin.Context) K {
func (f FlowLimits[K]) Limit(c *gin.Context) { func (f FlowLimits[K]) Limit(c *gin.Context) {
f.LimitedFns(c) f.LimitedFns(c)
} }
func (f FlowLimits[K]) DeferClear(c *gin.Context, m LimitMap[K], k K, v *int64) { func (f FlowLimits[K]) DeferClear(c *gin.Context, m LimitMap[K], k K, v, currentTotalFlow *int64) {
f.DeferClearFn(c, m, k, v) f.DeferClearFn(c, m, k, v, currentTotalFlow)
} }
func (f FlowLimits[K]) Add(c *gin.Context, m LimitMap[K], k K, v *int64) { func (f FlowLimits[K]) Add(c *gin.Context, m LimitMap[K], k K, v, currentTotalFlow *int64) {
f.AddFn(c, m, k, v) f.AddFn(c, m, k, v, currentTotalFlow)
} }
func (f FlowLimits[K]) Adds(_ *gin.Context, _ LimitMap[K], _ K, v *int64) { func (f FlowLimits[K]) Adds(_ *gin.Context, _ LimitMap[K], _ K, v, currentTotalFlow *int64) {
atomic.AddInt64(v, 1) atomic.AddInt64(v, 1)
atomic.AddInt64(currentTotalFlow, 1)
} }
type MapFlowLimit[K comparable] interface { type MapFlowLimit[K comparable] interface {
GetKey(c *gin.Context) K GetKey(c *gin.Context) K
Limit(c *gin.Context) Limit(c *gin.Context)
DeferClear(c *gin.Context, m LimitMap[K], k K, v *int64) DeferClear(c *gin.Context, m LimitMap[K], k K, v, currentTotalFlow *int64)
Add(c *gin.Context, m LimitMap[K], k K, v *int64) Add(c *gin.Context, m LimitMap[K], k K, v, currentTotalFlow *int64)
} }
func CustomFlowLimit[K comparable](a MapFlowLimit[K], maxRequestNum int64, clearNum ...int64) (func(ctx *gin.Context), func(int64, ...int64)) { func CustomFlowLimit[K comparable](a MapFlowLimit[K], maxRequestNum int64, clearNum ...int64) (func(ctx *gin.Context), func(int64, ...int64)) {
@ -88,6 +88,7 @@ func CustomFlowLimit[K comparable](a MapFlowLimit[K], maxRequestNum int64, clear
atomic.StoreInt64(m.ClearNum, clearNum[0]) atomic.StoreInt64(m.ClearNum, clearNum[0])
} }
} }
currentTotalFlow := new(int64)
return func(c *gin.Context) { return func(c *gin.Context) {
if atomic.LoadInt64(m.LimitNum) <= 0 { if atomic.LoadInt64(m.LimitNum) <= 0 {
c.Next() c.Next()
@ -103,8 +104,9 @@ func CustomFlowLimit[K comparable](a MapFlowLimit[K], maxRequestNum int64, clear
m.Map[key] = i m.Map[key] = i
m.Mux.Unlock() m.Mux.Unlock()
} }
a.Add(c, m, key, i) a.Add(c, m, key, i, currentTotalFlow)
defer a.DeferClear(c, m, key, i) defer a.DeferClear(c, m, key, i, currentTotalFlow)
defer atomic.AddInt64(currentTotalFlow, -1)
if atomic.LoadInt64(i) > atomic.LoadInt64(m.LimitNum) { if atomic.LoadInt64(i) > atomic.LoadInt64(m.LimitNum) {
a.Limit(c) a.Limit(c)
return return
@ -113,7 +115,7 @@ func CustomFlowLimit[K comparable](a MapFlowLimit[K], maxRequestNum int64, clear
}, fn }, fn
} }
func IpLimitClear[K comparable](_ *gin.Context, m LimitMap[K], key K, i *int64) { func IpLimitClear[K comparable](_ *gin.Context, m LimitMap[K], key K, i, currentTotalFlow *int64) {
atomic.AddInt64(i, -1) atomic.AddInt64(i, -1)
if atomic.LoadInt64(i) <= 0 { if atomic.LoadInt64(i) <= 0 {
cNum := int(atomic.LoadInt64(m.ClearNum)) cNum := int(atomic.LoadInt64(m.ClearNum))
@ -124,10 +126,7 @@ func IpLimitClear[K comparable](_ *gin.Context, m LimitMap[K], key K, i *int64)
return return
} }
m.Mux.RLock() if int(atomic.LoadInt64(currentTotalFlow)) <= cNum {
l := len(m.Map)
m.Mux.RUnlock()
if l < cNum {
m.Mux.Lock() m.Mux.Lock()
for k, v := range m.Map { for k, v := range m.Map {
if atomic.LoadInt64(v) < 1 { if atomic.LoadInt64(v) < 1 {
@ -158,27 +157,18 @@ func IpLimit(num int64, clearNum ...int64) (func(ctx *gin.Context), func(int64,
} }
func IpMinuteLimit(num int64, clearNum ...int64) (func(ctx *gin.Context), func(int64, ...int64)) { func IpMinuteLimit(num int64, clearNum ...int64) (func(ctx *gin.Context), func(int64, ...int64)) {
total := new(int64)
a := NewFlowLimits(func(c *gin.Context) string { a := NewFlowLimits(func(c *gin.Context) string {
return str.Join(c.ClientIP(), "|", time.Now().Format("2006-01-02 15:04")) return str.Join(c.ClientIP(), "|", time.Now().Format("2006-01-02 15:04"))
}, }, ToManyRequest(), IpMinuteLimitDeferFn,
ToManyRequest(),
IpMinuteLimitDeferFn(total),
func(c *gin.Context, m LimitMap[string], k string, v *int64) {
atomic.AddInt64(v, 1)
atomic.AddInt64(total, 1)
},
) )
return CustomFlowLimit(a, num, clearNum...) return CustomFlowLimit(a, num, clearNum...)
} }
func IpMinuteLimitDeferFn(total *int64) func(_ *gin.Context, m LimitMap[string], k string, _ *int64) { func IpMinuteLimitDeferFn(_ *gin.Context, m LimitMap[string], k string, _, currentTotalFlow *int64) {
return func(_ *gin.Context, m LimitMap[string], k string, _ *int64) { cNum := min(int(atomic.LoadInt64(m.ClearNum)), 1)
atomic.AddInt64(total, -1)
cNum := number.Min(int(atomic.LoadInt64(m.ClearNum)), 1)
minu := strings.Split(k, "|")[1] minu := strings.Split(k, "|")[1]
if int(atomic.LoadInt64(total)) < cNum { if int(atomic.LoadInt64(currentTotalFlow)) < cNum {
m.Mux.Lock() m.Mux.Lock()
for key := range m.Map { for key := range m.Map {
t := strings.Split(key, "|")[1] t := strings.Split(key, "|")[1]
@ -189,4 +179,3 @@ func IpMinuteLimitDeferFn(total *int64) func(_ *gin.Context, m LimitMap[string],
m.Mux.Unlock() m.Mux.Unlock()
} }
} }
}

2
app/route/route.go
View File

@ -120,7 +120,7 @@ func SetupRouter() *gin.Engine {
r.GET("/feed", actions.SiteFeed) r.GET("/feed", actions.SiteFeed)
r.GET("/comments/feed", actions.CommentsFeed) r.GET("/comments/feed", actions.CommentsFeed)
commentMiddleWare, _ := middleware.FlowLimit(5, c.SingleIpSearchNum, c.CacheTime.SleepTime) commentMiddleWare, _ := middleware.FlowLimit(5, c.SingleIpSearchNum, c.CacheTime.SleepTime)
commentIpMiddleware, _ := middleware.IpLimit(5, 5) commentIpMiddleware, _ := middleware.IpLimit(5, 2)
r.POST("/comment", commentMiddleWare, commentIpMiddleware, actions.PostComment) r.POST("/comment", commentMiddleWare, commentIpMiddleware, actions.PostComment)
r.NoRoute(actions.ThemeHook(constraints.NoRoute)) r.NoRoute(actions.ThemeHook(constraints.NoRoute))