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

optimize code and fix bug

Browse Source
This commit is contained in:
xing 2025-03-19 15:26:00 +08:00
parent 64556d69a0
commit 092f30b644
2 changed files with 30 additions and 28 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

52
app/middleware/iplimit.go
View File

@ -20,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, currentTotalFlow *int64) DeferClearFn func(c *gin.Context, m LimitMap[K], k K, v *int64, currentTotalFlow int64)
AddFn func(c *gin.Context, m LimitMap[K], k K, v, currentTotalFlow *int64) AddFn func(c *gin.Context, m LimitMap[K], k K, v, currentTotalFlow *int64) (current int64, currentTotal 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, currentTotalFlow *int64), deferClearFn func(c *gin.Context, m LimitMap[K], k K, v *int64, currentTotalFlow int64),
addFns ...func(c *gin.Context, m LimitMap[K], k K, v, currentTotalFlow *int64), addFns ...func(c *gin.Context, m LimitMap[K], k K, v, currentTotalFlow *int64) (current int64, currentTotal int64),
) *FlowLimits[K] { ) *FlowLimits[K] {
f := &FlowLimits[K]{ f := &FlowLimits[K]{
@ -50,24 +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, currentTotalFlow *int64) { func (f FlowLimits[K]) DeferClear(c *gin.Context, m LimitMap[K], k K, v *int64, currentTotalFlow int64) {
f.DeferClearFn(c, m, k, v, currentTotalFlow) f.DeferClearFn(c, m, k, v, currentTotalFlow)
} }
func (f FlowLimits[K]) Add(c *gin.Context, m LimitMap[K], k K, v, currentTotalFlow *int64) { func (f FlowLimits[K]) Add(c *gin.Context, m LimitMap[K], k K, v, currentTotalFlow *int64) (current int64, currentTotal int64) {
f.AddFn(c, m, k, v, currentTotalFlow) return f.AddFn(c, m, k, v, currentTotalFlow)
} }
func (f FlowLimits[K]) Adds(_ *gin.Context, _ LimitMap[K], _ K, v, currentTotalFlow *int64) { func (f FlowLimits[K]) Adds(_ *gin.Context, _ LimitMap[K], _ K, v, currentTotalFlow *int64) (current int64, currentTotal int64) {
atomic.AddInt64(v, 1)
atomic.AddInt64(currentTotalFlow, 1) return 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, currentTotalFlow *int64) DeferClear(c *gin.Context, m LimitMap[K], k K, v *int64, currentTotalFlow int64)
Add(c *gin.Context, m LimitMap[K], k K, v, currentTotalFlow *int64) Add(c *gin.Context, m LimitMap[K], k K, v, currentTotalFlow *int64) (current int64, currentTotal 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)) {
@ -104,10 +104,12 @@ 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, currentTotalFlow) ii, total := a.Add(c, m, key, i, currentTotalFlow)
defer a.DeferClear(c, m, key, i, currentTotalFlow) defer func() {
defer atomic.AddInt64(currentTotalFlow, -1) total = atomic.AddInt64(currentTotalFlow, -1)
if atomic.LoadInt64(i) > atomic.LoadInt64(m.LimitNum) { a.DeferClear(c, m, key, i, total)
}()
if ii > atomic.LoadInt64(m.LimitNum) {
a.Limit(c) a.Limit(c)
return return
} }
@ -115,18 +117,18 @@ func CustomFlowLimit[K comparable](a MapFlowLimit[K], maxRequestNum int64, clear
}, fn }, fn
} }
func IpLimitClear[K comparable](_ *gin.Context, m LimitMap[K], key K, i, currentTotalFlow *int64) { func IpLimitClear[K comparable](_ *gin.Context, m LimitMap[K], key K, i *int64, currentTotalFlow int64) {
atomic.AddInt64(i, -1) ii := atomic.AddInt64(i, -1)
if atomic.LoadInt64(i) <= 0 { if ii <= 0 {
cNum := int(atomic.LoadInt64(m.ClearNum)) cNum := atomic.LoadInt64(m.ClearNum)
if cNum <= 0 { if cNum < 0 {
m.Mux.Lock() m.Mux.Lock()
delete(m.Map, key) delete(m.Map, key)
m.Mux.Unlock() m.Mux.Unlock()
return return
} }
if int(atomic.LoadInt64(currentTotalFlow)) <= cNum { if currentTotalFlow < 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 {
@ -165,10 +167,10 @@ func IpMinuteLimit(num int64, clearNum ...int64) (func(ctx *gin.Context), func(i
return CustomFlowLimit(a, num, clearNum...) return CustomFlowLimit(a, num, clearNum...)
} }
func IpMinuteLimitDeferFn(_ *gin.Context, m LimitMap[string], k string, _, currentTotalFlow *int64) { func IpMinuteLimitDeferFn(_ *gin.Context, m LimitMap[string], k string, _ *int64, currentTotalFlow int64) {
cNum := min(int(atomic.LoadInt64(m.ClearNum)), 1) cNum := atomic.LoadInt64(m.ClearNum)
minu := strings.Split(k, "|")[1] minu := strings.Split(k, "|")[1]
if int(atomic.LoadInt64(currentTotalFlow)) < cNum { if 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]

6
app/route/route.go
View File

@ -102,7 +102,7 @@ func SetupRouter() *gin.Engine {
}, 85.1) }, 85.1)
SetGinAction("setRoute", func(r *gin.Engine) { SetGinAction("setRoute", func(r *gin.Engine) {
r.GET("/", actions.Feed, middleware.SearchLimit(c.SingleIpSearchNum, 5), r.GET("/", actions.Feed, middleware.SearchLimit(c.SingleIpSearchNum),
actions.ThemeHook(constraints.Home)) actions.ThemeHook(constraints.Home))
r.GET("/page/:page", actions.ThemeHook(constraints.Home)) r.GET("/page/:page", actions.ThemeHook(constraints.Home))
r.GET("/p/category/:category", actions.ThemeHook(constraints.Category)) r.GET("/p/category/:category", actions.ThemeHook(constraints.Category))
@ -119,8 +119,8 @@ func SetupRouter() *gin.Engine {
r.GET("/p/:id/feed", actions.PostFeed) r.GET("/p/:id/feed", actions.PostFeed)
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(10, c.SingleIpSearchNum, c.CacheTime.SleepTime)
commentIpMiddleware, _ := middleware.IpLimit(5, 2) commentIpMiddleware, _ := middleware.IpLimit(5)
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))