WAF增加多个动作

internal/utils/encrypt.go

@@ -0,0 +1,159 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package utils
+
+import (
+	"bytes"
+	"crypto/aes"
+	"crypto/cipher"
+	"encoding/base64"
+	"encoding/json"
+	"errors"
+	"github.com/TeaOSLab/EdgeCommon/pkg/nodeconfigs"
+	"github.com/TeaOSLab/EdgeNode/internal/events"
+	"github.com/iwind/TeaGo/logs"
+	"github.com/iwind/TeaGo/maps"
+	"github.com/iwind/TeaGo/rands"
+	stringutil "github.com/iwind/TeaGo/utils/string"
+)
+
+var (
+	simpleEncryptMagicKey = rands.HexString(32)
+)
+
+func init() {
+	events.On(events.EventReload, func() {
+		nodeConfig, _ := nodeconfigs.SharedNodeConfig()
+		if nodeConfig != nil {
+			simpleEncryptMagicKey = stringutil.Md5(nodeConfig.NodeId + "@" + nodeConfig.Secret)
+		}
+	})
+}
+
+// SimpleEncrypt 加密特殊信息
+func SimpleEncrypt(data []byte) []byte {
+	var method = &AES256CFBMethod{}
+	err := method.Init([]byte(simpleEncryptMagicKey), []byte(simpleEncryptMagicKey[:16]))
+	if err != nil {
+		logs.Println("[SimpleEncrypt]" + err.Error())
+		return data
+	}
+
+	dst, err := method.Encrypt(data)
+	if err != nil {
+		logs.Println("[SimpleEncrypt]" + err.Error())
+		return data
+	}
+	return dst
+}
+
+// SimpleDecrypt 解密特殊信息
+func SimpleDecrypt(data []byte) []byte {
+	var method = &AES256CFBMethod{}
+	err := method.Init([]byte(simpleEncryptMagicKey), []byte(simpleEncryptMagicKey[:16]))
+	if err != nil {
+		logs.Println("[MagicKeyEncode]" + err.Error())
+		return data
+	}
+
+	src, err := method.Decrypt(data)
+	if err != nil {
+		logs.Println("[MagicKeyEncode]" + err.Error())
+		return data
+	}
+	return src
+}
+
+func SimpleEncryptMap(m maps.Map) (base64String string, err error) {
+	mJSON, err := json.Marshal(m)
+	if err != nil {
+		return "", err
+	}
+	data := SimpleEncrypt(mJSON)
+	return base64.StdEncoding.EncodeToString(data), nil
+}
+
+func SimpleDecryptMap(base64String string) (maps.Map, error) {
+	data, err := base64.StdEncoding.DecodeString(base64String)
+	if err != nil {
+		return nil, err
+	}
+	mJSON := SimpleDecrypt(data)
+	var result = maps.Map{}
+	err = json.Unmarshal(mJSON, &result)
+	if err != nil {
+		return nil, err
+	}
+	return result, nil
+}
+
+type AES256CFBMethod struct {
+	block cipher.Block
+	iv    []byte
+}
+
+func (this *AES256CFBMethod) Init(key, iv []byte) error {
+	// 判断key是否为32长度
+	l := len(key)
+	if l > 32 {
+		key = key[:32]
+	} else if l < 32 {
+		key = append(key, bytes.Repeat([]byte{' '}, 32-l)...)
+	}
+
+	block, err := aes.NewCipher(key)
+	if err != nil {
+		return err
+	}
+	this.block = block
+
+	// 判断iv长度
+	l2 := len(iv)
+	if l2 > aes.BlockSize {
+		iv = iv[:aes.BlockSize]
+	} else if l2 < aes.BlockSize {
+		iv = append(iv, bytes.Repeat([]byte{' '}, aes.BlockSize-l2)...)
+	}
+	this.iv = iv
+
+	return nil
+}
+
+func (this *AES256CFBMethod) Encrypt(src []byte) (dst []byte, err error) {
+	if len(src) == 0 {
+		return
+	}
+
+	defer func() {
+		r := recover()
+		if r != nil {
+			err = errors.New("encrypt failed")
+		}
+	}()
+
+	dst = make([]byte, len(src))
+
+	encrypter := cipher.NewCFBEncrypter(this.block, this.iv)
+	encrypter.XORKeyStream(dst, src)
+
+	return
+}
+
+func (this *AES256CFBMethod) Decrypt(dst []byte) (src []byte, err error) {
+	if len(dst) == 0 {
+		return
+	}
+
+	defer func() {
+		r := recover()
+		if r != nil {
+			err = errors.New("decrypt failed")
+		}
+	}()
+
+	src = make([]byte, len(dst))
+	decrypter := cipher.NewCFBDecrypter(this.block, this.iv)
+	decrypter.XORKeyStream(src, dst)
+
+	return
+}

internal/utils/encrypt_test.go

@@ -0,0 +1,52 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package utils
+
+import (
+	"github.com/iwind/TeaGo/maps"
+	"sync"
+	"testing"
+)
+
+func TestSimpleEncrypt(t *testing.T) {
+	var arr = []string{"Hello", "World", "People"}
+	for _, s := range arr {
+		var value = []byte(s)
+		encoded := SimpleEncrypt(value)
+		t.Log(encoded, string(encoded))
+		decoded := SimpleDecrypt(encoded)
+		t.Log(decoded, string(decoded))
+	}
+}
+
+func TestSimpleEncrypt_Concurrent(t *testing.T) {
+	wg := sync.WaitGroup{}
+	var arr = []string{"Hello", "World", "People"}
+	wg.Add(len(arr))
+	for _, s := range arr {
+		go func(s string) {
+			defer wg.Done()
+			t.Log(string(SimpleDecrypt(SimpleEncrypt([]byte(s)))))
+		}(s)
+	}
+	wg.Wait()
+}
+
+func TestSimpleEncryptMap(t *testing.T) {
+	var m = maps.Map{
+		"s": "Hello",
+		"i": 20,
+		"b": true,
+	}
+	encodedResult, err := SimpleEncryptMap(m)
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Log("result:", encodedResult)
+
+	decodedResult, err := SimpleDecryptMap(encodedResult)
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Log(decodedResult)
+}

internal/utils/expires/list.go

@@ -12,6 +12,7 @@ type List struct {
 	itemsMap  map[int64]int64   // itemId => timestamp
 
 	locker sync.Mutex
+	ticker *time.Ticker
 }
 
 func NewList() *List {
@@ -21,10 +22,7 @@ func NewList() *List {
 	}
 }
 
-func (this *List) Add(itemId int64, expiredAt int64) {
-	if expiredAt <= time.Now().Unix() {
-		return
-	}
+func (this *List) Add(itemId int64, expiresAt int64) {
 	this.locker.Lock()
 	defer this.locker.Unlock()
 
@@ -34,17 +32,17 @@ func (this *List) Add(itemId int64, expiredAt int64) {
 		this.removeItem(itemId)
 	}
 
-	expireItemMap, ok := this.expireMap[expiredAt]
+	expireItemMap, ok := this.expireMap[expiresAt]
 	if ok {
 		expireItemMap[itemId] = true
 	} else {
 		expireItemMap = ItemMap{
 			itemId: true,
 		}
-		this.expireMap[expiredAt] = expireItemMap
+		this.expireMap[expiresAt] = expireItemMap
 	}
 
-	this.itemsMap[itemId] = expiredAt
+	this.itemsMap[itemId] = expiresAt
 }
 
 func (this *List) Remove(itemId int64) {
@@ -64,21 +62,22 @@ func (this *List) GC(timestamp int64, callback func(itemId int64)) {
 }
 
 func (this *List) StartGC(callback func(itemId int64)) {
-	ticker := time.NewTicker(1 * time.Second)
+	this.ticker = time.NewTicker(1 * time.Second)
 	lastTimestamp := int64(0)
-	for range ticker.C {
+	for range this.ticker.C {
 		timestamp := time.Now().Unix()
 		if lastTimestamp == 0 {
 			lastTimestamp = timestamp - 3600
 		}
 
-		// 防止死循环
-		if lastTimestamp > timestamp {
-			continue
-		}
-
-		for i := lastTimestamp; i <= timestamp; i++ {
-			this.GC(timestamp, callback)
+		if timestamp >= lastTimestamp {
+			for i := lastTimestamp; i <= timestamp; i++ {
+				this.GC(i, callback)
+			}
+		} else {
+			for i := timestamp; i <= lastTimestamp; i++ {
+				this.GC(i, callback)
+			}
 		}
 
 		// 这样做是为了防止系统时钟突变

internal/utils/expires/list_test.go

@@ -58,6 +58,10 @@ func TestList_Start_GC(t *testing.T) {
 	list.Add(2, time.Now().Unix()+1)
 	list.Add(3, time.Now().Unix()+2)
 	list.Add(4, time.Now().Unix()+5)
+	list.Add(5, time.Now().Unix()+5)
+	list.Add(6, time.Now().Unix()+6)
+	list.Add(7, time.Now().Unix()+6)
+	list.Add(8, time.Now().Unix()+6)
 
 	go func() {
 		list.StartGC(func(itemId int64) {
@@ -66,7 +70,7 @@ func TestList_Start_GC(t *testing.T) {
 		})
 	}()
 
-	time.Sleep(10 * time.Second)
+	time.Sleep(20 * time.Second)
 }
 
 func TestList_ManyItems(t *testing.T) {

internal/utils/jsonutils/map.go

@@ -0,0 +1,35 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package jsonutils
+
+import (
+	"encoding/json"
+	"github.com/iwind/TeaGo/maps"
+)
+
+func MapToObject(m maps.Map, ptr interface{}) error {
+	if m == nil {
+		return nil
+	}
+	mJSON, err := json.Marshal(m)
+	if err != nil {
+		return err
+	}
+	return json.Unmarshal(mJSON, ptr)
+}
+
+func ObjectToMap(ptr interface{}) (maps.Map, error) {
+	if ptr == nil {
+		return maps.Map{}, nil
+	}
+	ptrJSON, err := json.Marshal(ptr)
+	if err != nil {
+		return nil, err
+	}
+	var result = maps.Map{}
+	err = json.Unmarshal(ptrJSON, &result)
+	if err != nil {
+		return nil, err
+	}
+	return result, nil
+}

internal/utils/jsonutils/map_test.go

@@ -0,0 +1,46 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package jsonutils
+
+import (
+	"github.com/iwind/TeaGo/assert"
+	"github.com/iwind/TeaGo/maps"
+	"testing"
+)
+
+func TestMapToObject(t *testing.T) {
+	a := assert.NewAssertion(t)
+
+	type typeA struct {
+		B int  `json:"b"`
+		C bool `json:"c"`
+	}
+
+	{
+		var obj = &typeA{B: 1, C: true}
+		m, err := ObjectToMap(obj)
+		if err != nil {
+			t.Fatal(err)
+		}
+		PrintT(m, t)
+		a.IsTrue(m.GetInt("b") == 1)
+		a.IsTrue(m.GetBool("c") == true)
+	}
+
+	{
+		var obj = &typeA{}
+		err := MapToObject(maps.Map{
+			"b": 1024,
+			"c": true,
+		}, obj)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if obj == nil {
+			t.Fatal("obj should not be nil")
+		}
+		a.IsTrue(obj.B == 1024)
+		a.IsTrue(obj.C == true)
+		PrintT(obj, t)
+	}
+}

internal/utils/jsonutils/utils.go

@@ -0,0 +1,17 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package jsonutils
+
+import (
+	"encoding/json"
+	"testing"
+)
+
+func PrintT(obj interface{}, t *testing.T) {
+	data, err := json.MarshalIndent(obj, "", "  ")
+	if err != nil {
+		t.Log(err)
+	} else {
+		t.Log(string(data))
+	}
+}