github.com/TeaOSLab/EdgeNode@v1.3.8/internal/waf/utils/utils.go (about)

     1  package utils
     2  
     3  import (
     4  	teaconst "github.com/TeaOSLab/EdgeNode/internal/const"
     5  	"github.com/TeaOSLab/EdgeNode/internal/re"
     6  	"github.com/TeaOSLab/EdgeNode/internal/ttlcache"
     7  	"github.com/TeaOSLab/EdgeNode/internal/utils/cachehits"
     8  	"github.com/TeaOSLab/EdgeNode/internal/utils/fasttime"
     9  	"github.com/TeaOSLab/EdgeNode/internal/waf/requests"
    10  	"github.com/cespare/xxhash/v2"
    11  	"github.com/iwind/TeaGo/types"
    12  	stringutil "github.com/iwind/TeaGo/utils/string"
    13  	"strconv"
    14  )
    15  
    16  var SharedCache = ttlcache.NewCache[int8]()
    17  var cacheHits *cachehits.Stat
    18  
    19  func init() {
    20  	if !teaconst.IsMain {
    21  		return
    22  	}
    23  	cacheHits = cachehits.NewStat(5)
    24  }
    25  
    26  const (
    27  	MaxCacheDataSize = 1024
    28  )
    29  
    30  type CacheLife = int64
    31  
    32  const (
    33  	CacheDisabled   CacheLife = 0
    34  	CacheShortLife  CacheLife = 600
    35  	CacheMiddleLife CacheLife = 1800
    36  	CacheLongLife   CacheLife = 7200
    37  )
    38  
    39  // MatchStringCache 正则表达式匹配字符串，并缓存结果
    40  func MatchStringCache(regex *re.Regexp, s string, cacheLife CacheLife) bool {
    41  	if regex == nil {
    42  		return false
    43  	}
    44  
    45  	var regIdString = regex.IdString()
    46  
    47  	// 如果长度超过一定数量，大概率是不能重用的
    48  	if cacheLife <= 0 || len(s) > MaxCacheDataSize || !cacheHits.IsGood(regIdString) {
    49  		return regex.MatchString(s)
    50  	}
    51  
    52  	var hash = xxhash.Sum64String(s)
    53  	var key = regIdString + "@" + strconv.FormatUint(hash, 10)
    54  	var item = SharedCache.Read(key)
    55  	if item != nil {
    56  		cacheHits.IncreaseHit(regIdString)
    57  		return item.Value == 1
    58  	}
    59  	var b = regex.MatchString(s)
    60  	if b {
    61  		SharedCache.Write(key, 1, fasttime.Now().Unix()+cacheLife)
    62  	} else {
    63  		SharedCache.Write(key, 0, fasttime.Now().Unix()+cacheLife)
    64  	}
    65  	cacheHits.IncreaseCached(regIdString)
    66  	return b
    67  }
    68  
    69  // MatchBytesCache 正则表达式匹配字节slice，并缓存结果
    70  func MatchBytesCache(regex *re.Regexp, byteSlice []byte, cacheLife CacheLife) bool {
    71  	if regex == nil {
    72  		return false
    73  	}
    74  
    75  	var regIdString = regex.IdString()
    76  
    77  	// 如果长度超过一定数量，大概率是不能重用的
    78  	if cacheLife <= 0 || len(byteSlice) > MaxCacheDataSize || !cacheHits.IsGood(regIdString) {
    79  		return regex.Match(byteSlice)
    80  	}
    81  
    82  	var hash = xxhash.Sum64(byteSlice)
    83  	var key = regIdString + "@" + strconv.FormatUint(hash, 10)
    84  	var item = SharedCache.Read(key)
    85  	if item != nil {
    86  		cacheHits.IncreaseHit(regIdString)
    87  		return item.Value == 1
    88  	}
    89  	var b = regex.Match(byteSlice)
    90  	if b {
    91  		SharedCache.Write(key, 1, fasttime.Now().Unix()+cacheLife)
    92  	} else {
    93  		SharedCache.Write(key, 0, fasttime.Now().Unix()+cacheLife)
    94  	}
    95  	cacheHits.IncreaseCached(regIdString)
    96  	return b
    97  }
    98  
    99  // ComposeIPType 组合IP类型
   100  func ComposeIPType(setId int64, req requests.Request) string {
   101  	return "set:" + types.String(setId) + "@" + stringutil.Md5(req.WAFRaw().UserAgent())
   102  }