mirror of
https://github.com/TeaOSLab/EdgeNode.git
synced 2025-11-04 16:00:25 +08:00
103 lines
2.7 KiB
Go
103 lines
2.7 KiB
Go
package utils
|
||
|
||
import (
|
||
teaconst "github.com/TeaOSLab/EdgeNode/internal/const"
|
||
"github.com/TeaOSLab/EdgeNode/internal/re"
|
||
"github.com/TeaOSLab/EdgeNode/internal/utils/cachehits"
|
||
"github.com/TeaOSLab/EdgeNode/internal/utils/fasttime"
|
||
"github.com/TeaOSLab/EdgeNode/internal/utils/ttlcache"
|
||
"github.com/TeaOSLab/EdgeNode/internal/waf/requests"
|
||
"github.com/cespare/xxhash/v2"
|
||
"github.com/iwind/TeaGo/types"
|
||
stringutil "github.com/iwind/TeaGo/utils/string"
|
||
"strconv"
|
||
)
|
||
|
||
var SharedCache = ttlcache.NewCache[int8]()
|
||
var cacheHits *cachehits.Stat
|
||
|
||
func init() {
|
||
if !teaconst.IsMain {
|
||
return
|
||
}
|
||
cacheHits = cachehits.NewStat(5)
|
||
}
|
||
|
||
const (
|
||
MaxCacheDataSize = 1024
|
||
)
|
||
|
||
type CacheLife = int64
|
||
|
||
const (
|
||
CacheDisabled CacheLife = 0
|
||
CacheShortLife CacheLife = 600
|
||
CacheMiddleLife CacheLife = 1800
|
||
CacheLongLife CacheLife = 7200
|
||
)
|
||
|
||
// MatchStringCache 正则表达式匹配字符串，并缓存结果
|
||
func MatchStringCache(regex *re.Regexp, s string, cacheLife CacheLife) bool {
|
||
if regex == nil {
|
||
return false
|
||
}
|
||
|
||
var regIdString = regex.IdString()
|
||
|
||
// 如果长度超过一定数量，大概率是不能重用的
|
||
if cacheLife <= 0 || len(s) > MaxCacheDataSize || !cacheHits.IsGood(regIdString) {
|
||
return regex.MatchString(s)
|
||
}
|
||
|
||
var hash = xxhash.Sum64String(s)
|
||
var key = regIdString + "@" + strconv.FormatUint(hash, 10)
|
||
var item = SharedCache.Read(key)
|
||
if item != nil {
|
||
cacheHits.IncreaseHit(regIdString)
|
||
return item.Value == 1
|
||
}
|
||
var b = regex.MatchString(s)
|
||
if b {
|
||
SharedCache.Write(key, 1, fasttime.Now().Unix()+cacheLife)
|
||
} else {
|
||
SharedCache.Write(key, 0, fasttime.Now().Unix()+cacheLife)
|
||
}
|
||
cacheHits.IncreaseCached(regIdString)
|
||
return b
|
||
}
|
||
|
||
// MatchBytesCache 正则表达式匹配字节slice，并缓存结果
|
||
func MatchBytesCache(regex *re.Regexp, byteSlice []byte, cacheLife CacheLife) bool {
|
||
if regex == nil {
|
||
return false
|
||
}
|
||
|
||
var regIdString = regex.IdString()
|
||
|
||
// 如果长度超过一定数量，大概率是不能重用的
|
||
if cacheLife <= 0 || len(byteSlice) > MaxCacheDataSize || !cacheHits.IsGood(regIdString) {
|
||
return regex.Match(byteSlice)
|
||
}
|
||
|
||
var hash = xxhash.Sum64(byteSlice)
|
||
var key = regIdString + "@" + strconv.FormatUint(hash, 10)
|
||
var item = SharedCache.Read(key)
|
||
if item != nil {
|
||
cacheHits.IncreaseHit(regIdString)
|
||
return item.Value == 1
|
||
}
|
||
var b = regex.Match(byteSlice)
|
||
if b {
|
||
SharedCache.Write(key, 1, fasttime.Now().Unix()+cacheLife)
|
||
} else {
|
||
SharedCache.Write(key, 0, fasttime.Now().Unix()+cacheLife)
|
||
}
|
||
cacheHits.IncreaseCached(regIdString)
|
||
return b
|
||
}
|
||
|
||
// ComposeIPType 组合IP类型
|
||
func ComposeIPType(setId int64, req requests.Request) string {
|
||
return "set:" + types.String(setId) + "@" + stringutil.Md5(req.WAFRaw().UserAgent())
|
||
}
|