diff --git a/internal/nodes/client_conn.go b/internal/nodes/client_conn.go index 769918a..ce2708c 100644 --- a/internal/nodes/client_conn.go +++ b/internal/nodes/client_conn.go @@ -137,7 +137,7 @@ func (this *ClientConn) increaseSYNFlood(synFloodConfig *firewallconfigs.SYNFloo var ip = this.RawIP() if len(ip) > 0 && !iplibrary.IsInWhiteList(ip) && (!synFloodConfig.IgnoreLocal || !utils.IsLocalIP(ip)) { var timestamp = utils.NextMinuteUnixTime() - var result = ttlcache.SharedCache.IncreaseInt64("SYN_FLOOD:"+ip, 1, timestamp) + var result = ttlcache.SharedCache.IncreaseInt64("SYN_FLOOD:"+ip, 1, timestamp, true) var minAttempts = synFloodConfig.MinAttempts if minAttempts < 5 { minAttempts = 5 diff --git a/internal/ttlcache/cache.go b/internal/ttlcache/cache.go index 564e4c8..76d572f 100644 --- a/internal/ttlcache/cache.go +++ b/internal/ttlcache/cache.go @@ -91,7 +91,7 @@ func (this *Cache) Write(key string, value interface{}, expiredAt int64) (ok boo }) } -func (this *Cache) IncreaseInt64(key string, delta int64, expiredAt int64) int64 { +func (this *Cache) IncreaseInt64(key string, delta int64, expiredAt int64, extend bool) int64 { if this.isDestroyed { return 0 } @@ -107,7 +107,7 @@ func (this *Cache) IncreaseInt64(key string, delta int64, expiredAt int64) int64 } uint64Key := HashKey([]byte(key)) pieceIndex := uint64Key % this.countPieces - return this.pieces[pieceIndex].IncreaseInt64(uint64Key, delta, expiredAt) + return this.pieces[pieceIndex].IncreaseInt64(uint64Key, delta, expiredAt, extend) } func (this *Cache) Read(key string) (item *Item) { diff --git a/internal/ttlcache/cache_test.go b/internal/ttlcache/cache_test.go index 71b80ee..6234fbc 100644 --- a/internal/ttlcache/cache_test.go +++ b/internal/ttlcache/cache_test.go @@ -65,14 +65,14 @@ func TestCache_IncreaseInt64(t *testing.T) { var unixTime = time.Now().Unix() { - cache.IncreaseInt64("a", 1, unixTime+3600) + cache.IncreaseInt64("a", 1, unixTime+3600, false) var item = cache.Read("a") t.Log(item) a.IsTrue(item.Value == int64(1)) a.IsTrue(item.expiredAt == unixTime+3600) } { - cache.IncreaseInt64("a", 1, unixTime+3600+1) + cache.IncreaseInt64("a", 1, unixTime+3600+1, true) var item = cache.Read("a") t.Log(item) a.IsTrue(item.Value == int64(2)) @@ -83,7 +83,7 @@ func TestCache_IncreaseInt64(t *testing.T) { t.Log(cache.Read("b")) } { - cache.IncreaseInt64("b", 1, time.Now().Unix()+3600+3) + cache.IncreaseInt64("b", 1, time.Now().Unix()+3600+3, false) t.Log(cache.Read("b")) } } diff --git a/internal/ttlcache/piece.go b/internal/ttlcache/piece.go index 5110bb5..2244e68 100644 --- a/internal/ttlcache/piece.go +++ b/internal/ttlcache/piece.go @@ -39,13 +39,15 @@ func (this *Piece) Add(key uint64, item *Item) (ok bool) { return true } -func (this *Piece) IncreaseInt64(key uint64, delta int64, expiredAt int64) (result int64) { +func (this *Piece) IncreaseInt64(key uint64, delta int64, expiredAt int64, extend bool) (result int64) { this.locker.Lock() item, ok := this.m[key] if ok && item.expiredAt > time.Now().Unix() { result = types.Int64(item.Value) + delta item.Value = result - item.expiredAt = expiredAt + if extend { + item.expiredAt = expiredAt + } this.expiresList.Add(key, expiredAt) } else { if len(this.m) < this.maxItems { diff --git a/internal/waf/captcha_validator.go b/internal/waf/captcha_validator.go index ae3b905..d98da03 100644 --- a/internal/waf/captcha_validator.go +++ b/internal/waf/captcha_validator.go @@ -167,7 +167,7 @@ func (this *CaptchaValidator) validate(actionConfig *CaptchaAction, maxFails int } else { // 增加计数 if maxFails > 0 && failBlockTimeout > 0 { - var countFails = ttlcache.SharedCache.IncreaseInt64("CAPTCHA:FAILS:"+request.WAFRemoteIP(), 1, time.Now().Unix()+300) + var countFails = ttlcache.SharedCache.IncreaseInt64("CAPTCHA:FAILS:"+request.WAFRemoteIP(), 1, time.Now().Unix()+300, true) if int(countFails) >= maxFails { SharedIPBlackList.RecordIP(IPTypeAll, firewallconfigs.FirewallScopeService, request.WAFServerId(), request.WAFRemoteIP(), time.Now().Unix()+int64(failBlockTimeout), policyId, false, groupId, setId, "CAPTCHA验证连续失败") return false diff --git a/internal/waf/checkpoints/cc.go b/internal/waf/checkpoints/cc.go index de6c3c9..f1b9722 100644 --- a/internal/waf/checkpoints/cc.go +++ b/internal/waf/checkpoints/cc.go @@ -114,7 +114,7 @@ func (this *CCCheckpoint) RequestValue(req requests.Request, param string, optio if len(key) == 0 { key = req.WAFRemoteIP() } - value = this.cache.IncreaseInt64(key, int64(1), time.Now().Unix()+period) + value = this.cache.IncreaseInt64(key, int64(1), time.Now().Unix()+period, false) } return diff --git a/internal/waf/checkpoints/cc2.go b/internal/waf/checkpoints/cc2.go index d7980a6..35d1011 100644 --- a/internal/waf/checkpoints/cc2.go +++ b/internal/waf/checkpoints/cc2.go @@ -38,7 +38,7 @@ func (this *CC2Checkpoint) RequestValue(req requests.Request, param string, opti threshold = 1000 } - value = ccCache.IncreaseInt64("WAF-CC-"+strings.Join(keyValues, "@"), 1, time.Now().Unix()+period) + value = ccCache.IncreaseInt64("WAF-CC-"+strings.Join(keyValues, "@"), 1, time.Now().Unix()+period, false) return }