[waf]支持包含二进制、不支持二进制等操作符；支持对参数值编解码

internal/cache/cache.go

@@ -0,0 +1,104 @@
+package cache
+
+import (
+	"time"
+)
+
+// TTL缓存
+// 最大的缓存时间为30 * 86400
+// Piece数据结构：
+//      Piece1          |  Piece2 | Piece3 | ...
+//  [ Item1, Item2, ... |   ...
+// KeyMap列表数据结构
+// { timestamp1 => [key1, key2, ...] }, ...
+type Cache struct {
+	pieces      []*Piece
+	countPieces uint64
+
+	gcPieceIndex int
+}
+
+func NewCache(opt ...OptionInterface) *Cache {
+	countPieces := 128
+	for _, option := range opt {
+		if option == nil {
+			continue
+		}
+		switch o := option.(type) {
+		case *PiecesOption:
+			if o.Count > 0 {
+				countPieces = o.Count
+			}
+		}
+	}
+
+	cache := &Cache{
+		countPieces: uint64(countPieces),
+	}
+
+	for i := 0; i < countPieces; i++ {
+		cache.pieces = append(cache.pieces, NewPiece())
+	}
+
+	// start timer
+	go func() {
+		ticker := time.NewTicker(1 * time.Second)
+		for range ticker.C {
+			cache.GC()
+		}
+	}()
+
+	return cache
+}
+
+func (this *Cache) Add(key string, value interface{}, expiredAt int64) {
+	currentTimestamp := time.Now().Unix()
+	if expiredAt <= currentTimestamp {
+		return
+	}
+
+	maxExpiredAt := currentTimestamp + 30*86400
+	if expiredAt > maxExpiredAt {
+		expiredAt = maxExpiredAt
+	}
+	uint64Key := HashKey([]byte(key))
+	pieceIndex := uint64Key % this.countPieces
+	this.pieces[pieceIndex].Add(uint64Key, &Item{
+		value:     value,
+		expiredAt: expiredAt,
+	})
+}
+
+func (this *Cache) Read(key string) (value *Item) {
+	uint64Key := HashKey([]byte(key))
+	return this.pieces[uint64Key%this.countPieces].Read(uint64Key)
+}
+
+func (this *Cache) readIntKey(key uint64) (value *Item) {
+	return this.pieces[key%this.countPieces].Read(key)
+}
+
+func (this *Cache) Delete(key string) {
+	uint64Key := HashKey([]byte(key))
+	this.pieces[uint64Key%this.countPieces].Delete(uint64Key)
+}
+
+func (this *Cache) deleteIntKey(key uint64) {
+	this.pieces[key%this.countPieces].Delete(key)
+}
+
+func (this *Cache) Count() (count int) {
+	for _, piece := range this.pieces {
+		count += piece.Count()
+	}
+	return
+}
+
+func (this *Cache) GC() {
+	this.pieces[this.gcPieceIndex].GC()
+	newIndex := this.gcPieceIndex + 1
+	if newIndex >= int(this.countPieces) {
+		newIndex = 0
+	}
+	this.gcPieceIndex = newIndex
+}

internal/cache/cache_test.go

@@ -0,0 +1,103 @@
+package cache
+
+import (
+	"github.com/iwind/TeaGo/rands"
+	"runtime"
+	"strconv"
+	"testing"
+	"time"
+)
+
+func TestNewCache(t *testing.T) {
+	cache := NewCache()
+	cache.Add("a", 1, time.Now().Unix()+3600)
+	cache.Add("b", 2, time.Now().Unix()+3601)
+	cache.Add("a", 1, time.Now().Unix()+3602)
+	cache.Add("d", 1, time.Now().Unix()+1)
+
+	for _, piece := range cache.pieces {
+		if len(piece.m) > 0 {
+			for k, item := range piece.m {
+				t.Log(k, "=>", item.value, item.expiredAt)
+			}
+		}
+	}
+	t.Log(cache.Read("a"))
+	time.Sleep(2 * time.Second)
+	t.Log(cache.Read("d"))
+}
+
+func BenchmarkCache_Add(b *testing.B) {
+	runtime.GOMAXPROCS(1)
+
+	cache := NewCache()
+	for i := 0; i < b.N; i++ {
+		cache.Add(strconv.Itoa(i), i, time.Now().Unix()+int64(i%1024))
+	}
+}
+
+func TestCache_Read(t *testing.T) {
+	runtime.GOMAXPROCS(1)
+
+	var cache = NewCache(PiecesOption{Count: 512})
+
+	for i := 0; i < 10_000_000; i++ {
+		cache.Add("HELLO_WORLD_"+strconv.Itoa(i), i, time.Now().Unix()+int64(i%10240)+1)
+	}
+
+	/**total := 0
+	for _, piece := range cache.pieces {
+		t.Log(len(piece.m), "keys")
+		total += len(piece.m)
+	}
+	t.Log(total, "total keys")**/
+
+	before := time.Now()
+	for i := 0; i < 10_240; i++ {
+		_ = cache.Read("HELLO_WORLD_" + strconv.Itoa(i))
+	}
+	t.Log(time.Since(before).Seconds()*1000, "ms")
+}
+
+func TestCache_GC(t *testing.T) {
+	var cache = NewCache(&PiecesOption{Count: 5})
+	cache.Add("a", 1, time.Now().Unix()+1)
+	cache.Add("b", 2, time.Now().Unix()+2)
+	cache.Add("c", 3, time.Now().Unix()+3)
+	cache.Add("d", 4, time.Now().Unix()+4)
+	cache.Add("e", 5, time.Now().Unix()+10)
+
+	go func() {
+		for i := 0; i < 1000; i++ {
+			cache.Add("f", 1, time.Now().Unix()+1)
+			time.Sleep(10 * time.Millisecond)
+		}
+	}()
+
+	for i := 0; i < 20; i++ {
+		cache.GC()
+		t.Log("items:", cache.Count())
+		time.Sleep(1 * time.Second)
+	}
+
+	t.Log("now:", time.Now().Unix())
+	for _, p := range cache.pieces {
+		for k, v := range p.m {
+			t.Log(k, v.value, v.expiredAt)
+		}
+	}
+}
+
+func TestCache_GC2(t *testing.T) {
+	runtime.GOMAXPROCS(1)
+
+	cache := NewCache()
+	for i := 0; i < 1_000_000; i++ {
+		cache.Add(strconv.Itoa(i), i, time.Now().Unix()+int64(rands.Int(0, 100)))
+	}
+
+	for i := 0; i < 100; i++ {
+		t.Log(cache.Count(), "items")
+		time.Sleep(1 * time.Second)
+	}
+}

internal/cache/item.go

@@ -0,0 +1,6 @@
+package cache
+
+type Item struct {
+	value     interface{}
+	expiredAt int64
+}

internal/cache/option.go

@@ -0,0 +1,12 @@
+package cache
+
+type OptionInterface interface {
+}
+
+type PiecesOption struct {
+	Count int
+}
+
+func NewPiecesOption(count int) *PiecesOption {
+	return &PiecesOption{Count: count}
+}

internal/cache/piece.go

@@ -0,0 +1,57 @@
+package cache
+
+import (
+	"github.com/TeaOSLab/EdgeNode/internal/utils"
+	"sync"
+	"time"
+)
+
+type Piece struct {
+	m      map[uint64]*Item
+	locker sync.RWMutex
+}
+
+func NewPiece() *Piece {
+	return &Piece{m: map[uint64]*Item{}}
+}
+
+func (this *Piece) Add(key uint64, item *Item) () {
+	this.locker.Lock()
+	this.m[key] = item
+	this.locker.Unlock()
+}
+
+func (this *Piece) Delete(key uint64) {
+	this.locker.Lock()
+	delete(this.m, key)
+	this.locker.Unlock()
+}
+
+func (this *Piece) Read(key uint64) (item *Item) {
+	this.locker.RLock()
+	item = this.m[key]
+	if item != nil && item.expiredAt < utils.UnixTime() {
+		item = nil
+	}
+	this.locker.RUnlock()
+
+	return
+}
+
+func (this *Piece) Count() (count int) {
+	this.locker.RLock()
+	count = len(this.m)
+	this.locker.RUnlock()
+	return
+}
+
+func (this *Piece) GC() {
+	this.locker.Lock()
+	timestamp := time.Now().Unix()
+	for k, item := range this.m {
+		if item.expiredAt <= timestamp {
+			delete(this.m, k)
+		}
+	}
+	this.locker.Unlock()
+}

internal/cache/piece_test.go

@@ -0,0 +1,52 @@
+package cache
+
+import (
+	"github.com/iwind/TeaGo/rands"
+	"testing"
+	"time"
+)
+
+func TestPiece_Add(t *testing.T) {
+	piece := NewPiece()
+	piece.Add(1, &Item{expiredAt: time.Now().Unix() + 3600})
+	piece.Add(2, &Item{})
+	piece.Add(3, &Item{})
+	piece.Delete(3)
+	for key, item := range piece.m {
+		t.Log(key, item.value)
+	}
+	t.Log(piece.Read(1))
+}
+
+func TestPiece_GC(t *testing.T) {
+	piece := NewPiece()
+	piece.Add(1, &Item{value: 1, expiredAt: time.Now().Unix() + 1})
+	piece.Add(2, &Item{value: 2, expiredAt: time.Now().Unix() + 1})
+	piece.Add(3, &Item{value: 3, expiredAt: time.Now().Unix() + 1})
+	t.Log("before gc ===")
+	for key, item := range piece.m {
+		t.Log(key, item.value)
+	}
+
+	time.Sleep(1 * time.Second)
+	piece.GC()
+
+	t.Log("after gc ===")
+	for key, item := range piece.m {
+		t.Log(key, item.value)
+	}
+}
+
+func TestPiece_GC2(t *testing.T) {
+	piece := NewPiece()
+	for i := 0; i < 10_000; i++ {
+		piece.Add(uint64(i), &Item{value: 1, expiredAt: time.Now().Unix() + int64(rands.Int(1, 10))})
+	}
+
+	time.Sleep(1 * time.Second)
+
+	before := time.Now()
+	piece.GC()
+	t.Log(time.Since(before).Seconds()*1000, "ms")
+	t.Log(piece.Count())
+}

internal/cache/utils.go

@@ -0,0 +1,7 @@
+package cache
+
+import "github.com/dchest/siphash"
+
+func HashKey(key []byte) uint64 {
+	return siphash.Hash(0, 0, key)
+}

internal/cache/utils_test.go

@@ -0,0 +1,13 @@
+package cache
+
+import (
+	"runtime"
+	"testing"
+)
+
+func BenchmarkHashKey(b *testing.B) {
+	runtime.GOMAXPROCS(1)
+	for i := 0; i < b.N; i++ {
+		HashKey([]byte("HELLO,WORLDHELLO,WORLDHELLO,WORLDHELLO,WORLDHELLO,WORLDHELLO,WORLD"))
+	}
+}

internal/nodes/waf_manager.go

@@ -93,11 +93,20 @@ func (this *WAFManager) convertWAF(policy *firewallconfigs.HTTPFirewallPolicy) (
 					r := &waf.Rule{
 						Description:       rule.Description,
 						Param:             rule.Param,
+						ParamFilters:      []*waf.ParamFilter{},
 						Operator:          rule.Operator,
 						Value:             rule.Value,
 						IsCaseInsensitive: rule.IsCaseInsensitive,
 						CheckpointOptions: rule.CheckpointOptions,
 					}
+
+					for _, paramFilter := range rule.ParamFilters {
+						r.ParamFilters = append(r.ParamFilters, &waf.ParamFilter{
+							Code:    paramFilter.Code,
+							Options: paramFilter.Options,
+						})
+					}
+
 					s.Rules = append(s.Rules, r)
 				}
 

internal/waf/param_filter.go

@@ -0,0 +1,8 @@
+package waf
+
+import "github.com/iwind/TeaGo/maps"
+
+type ParamFilter struct {
+	Code    string   `yaml:"code" json:"code"`
+	Options maps.Map `yaml:"options" json:"options"`
+}

internal/waf/rule.go

@@ -2,9 +2,12 @@ package waf
 
 import (
 	"bytes"
+	"encoding/base64"
 	"encoding/binary"
 	"errors"
 	"github.com/TeaOSLab/EdgeCommon/pkg/configutils"
+	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs/filterconfigs"
+	"github.com/TeaOSLab/EdgeNode/internal/logs"
 	"github.com/TeaOSLab/EdgeNode/internal/waf/checkpoints"
 	"github.com/TeaOSLab/EdgeNode/internal/waf/requests"
 	"github.com/TeaOSLab/EdgeNode/internal/waf/utils"
@@ -23,7 +26,8 @@ var singleParamRegexp = regexp.MustCompile("^\\${[\\w.-]+}$")
 // rule
 type Rule struct {
 	Description       string                 `yaml:"description" json:"description"`
-	Param             string                 `yaml:"param" json:"param"`       // such as ${arg.name} or ${args}, can be composite as ${arg.firstName}${arg.lastName}
+	Param             string                 `yaml:"param" json:"param"` // such as ${arg.name} or ${args}, can be composite as ${arg.firstName}${arg.lastName}
+	ParamFilters      []*ParamFilter         `yaml:"paramFilters" json:"paramFilters"`
 	Operator          RuleOperator           `yaml:"operator" json:"operator"` // such as contains, gt,  ...
 	Value             string                 `yaml:"value" json:"value"`       // compared value
 	IsCaseInsensitive bool                   `yaml:"isCaseInsensitive" json:"isCaseInsensitive"`
@@ -122,7 +126,6 @@ func (this *Rule) Init() error {
 		} else {
 			return errors.New("invalid ip range")
 		}
-
 	}
 
 	if singleParamRegexp.MatchString(this.Param) {
@@ -187,6 +190,11 @@ func (this *Rule) MatchRequest(req *requests.Request) (b bool, err error) {
 			return false, err
 		}
 
+		// execute filters
+		if len(this.ParamFilters) > 0 {
+			value = this.execFilter(value)
+		}
+
 		// if is composed checkpoint, we just returns true or false
 		if this.singleCheckpoint.IsComposed() {
 			return types.Bool(value), nil
@@ -233,6 +241,12 @@ func (this *Rule) MatchResponse(req *requests.Request, resp *requests.Response)
 			if err != nil {
 				return false, err
 			}
+
+			// execute filters
+			if len(this.ParamFilters) > 0 {
+				value = this.execFilter(value)
+			}
+
 			return this.Test(value), nil
 		}
 
@@ -420,6 +434,20 @@ func (this *Rule) Test(value interface{}) bool {
 		} else {
 			return strings.HasSuffix(types.String(value), this.Value)
 		}
+	case RuleOperatorContainsBinary:
+		data, _ := base64.StdEncoding.DecodeString(types.String(this.Value))
+		if this.IsCaseInsensitive {
+			return bytes.Contains(bytes.ToUpper([]byte(types.String(value))), bytes.ToUpper(data))
+		} else {
+			return bytes.Contains([]byte(types.String(value)), data)
+		}
+	case RuleOperatorNotContainsBinary:
+		data, _ := base64.StdEncoding.DecodeString(types.String(this.Value))
+		if this.IsCaseInsensitive {
+			return !bytes.Contains(bytes.ToUpper([]byte(types.String(value))), bytes.ToUpper(data))
+		} else {
+			return !bytes.Contains([]byte(types.String(value)), data)
+		}
 	case RuleOperatorHasKey:
 		if types.IsSlice(value) {
 			index := types.Int(this.Value)
@@ -594,3 +622,24 @@ func (this *Rule) ipToInt64(ip net.IP) int64 {
 	}
 	return int64(binary.BigEndian.Uint32(ip))
 }
+
+func (this *Rule) execFilter(value interface{}) interface{} {
+	var goNext bool
+	var err error
+
+	for _, filter := range this.ParamFilters {
+		filterInstance := filterconfigs.FindFilter(filter.Code)
+		if filterInstance == nil {
+			continue
+		}
+		value, goNext, err = filterInstance.Do(value, filter.Options)
+		if err != nil {
+			logs.Println("WAF", "filter error: "+err.Error())
+			break
+		}
+		if !goNext {
+			break
+		}
+	}
+	return value
+}

internal/waf/rule_operator.go

@@ -23,6 +23,9 @@ const (
 	RuleOperatorVersionLt    RuleOperator = "version lt"
 	RuleOperatorVersionRange RuleOperator = "version range"
 
+	RuleOperatorContainsBinary    RuleOperator = "contains binary"     // contains binary
+	RuleOperatorNotContainsBinary RuleOperator = "not contains binary" // not contains binary
+
 	// ip
 	RuleOperatorEqIP       RuleOperator = "eq ip"
 	RuleOperatorGtIP       RuleOperator = "gt ip"