[WAF]规则支持导入导出

internal/ttlcache/cache.go

@@ -0,0 +1,151 @@
+package ttlcache
+
+import (
+	"github.com/TeaOSLab/EdgeAdmin/internal/utils"
+	"time"
+)
+
+var DefaultCache = NewCache()
+
+// TTL缓存
+// 最大的缓存时间为30 * 86400
+// Piece数据结构：
+//      Piece1          |  Piece2 | Piece3 | ...
+//  [ Item1, Item2, ... |   ...
+// KeyMap列表数据结构
+// { timestamp1 => [key1, key2, ...] }, ...
+type Cache struct {
+	isDestroyed bool
+	pieces      []*Piece
+	countPieces uint64
+	maxItems    int
+
+	gcPieceIndex int
+	ticker       *utils.Ticker
+}
+
+func NewCache(opt ...OptionInterface) *Cache {
+	countPieces := 128
+	maxItems := 1_000_000
+	for _, option := range opt {
+		if option == nil {
+			continue
+		}
+		switch o := option.(type) {
+		case *PiecesOption:
+			if o.Count > 0 {
+				countPieces = o.Count
+			}
+		case *MaxItemsOption:
+			if o.Count > 0 {
+				maxItems = o.Count
+			}
+		}
+	}
+
+	cache := &Cache{
+		countPieces: uint64(countPieces),
+		maxItems:    maxItems,
+	}
+
+	for i := 0; i < countPieces; i++ {
+		cache.pieces = append(cache.pieces, NewPiece(maxItems/countPieces))
+	}
+
+	// start timer
+	go func() {
+		cache.ticker = utils.NewTicker(5 * time.Second)
+		for cache.ticker.Next() {
+			cache.GC()
+		}
+	}()
+
+	return cache
+}
+
+func (this *Cache) Write(key string, value interface{}, expiredAt int64) {
+	if this.isDestroyed {
+		return
+	}
+
+	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) IncreaseInt64(key string, delta int64, expiredAt int64) int64 {
+	if this.isDestroyed {
+		return 0
+	}
+
+	currentTimestamp := time.Now().Unix()
+	if expiredAt <= currentTimestamp {
+		return 0
+	}
+
+	maxExpiredAt := currentTimestamp + 30*86400
+	if expiredAt > maxExpiredAt {
+		expiredAt = maxExpiredAt
+	}
+	uint64Key := HashKey([]byte(key))
+	pieceIndex := uint64Key % this.countPieces
+	return this.pieces[pieceIndex].IncreaseInt64(uint64Key, delta, expiredAt)
+}
+
+func (this *Cache) Read(key string) (item *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
+}
+
+func (this *Cache) Destroy() {
+	this.isDestroyed = true
+
+	if this.ticker != nil {
+		this.ticker.Stop()
+		this.ticker = nil
+	}
+	for _, piece := range this.pieces {
+		piece.Destroy()
+	}
+}

internal/ttlcache/cache_test.go

@@ -0,0 +1,124 @@
+package ttlcache
+
+import (
+	"github.com/iwind/TeaGo/rands"
+	"runtime"
+	"strconv"
+	"testing"
+	"time"
+)
+
+func TestNewCache(t *testing.T) {
+	cache := NewCache()
+	cache.Write("a", 1, time.Now().Unix()+3600)
+	cache.Write("b", 2, time.Now().Unix()+3601)
+	cache.Write("a", 1, time.Now().Unix()+3602)
+	cache.Write("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.Write(strconv.Itoa(i), i, time.Now().Unix()+int64(i%1024))
+	}
+}
+
+func TestCache_IncreaseInt64(t *testing.T) {
+	var cache = NewCache()
+
+	{
+		cache.IncreaseInt64("a", 1, time.Now().Unix()+3600)
+		t.Log(cache.Read("a"))
+	}
+	{
+		cache.IncreaseInt64("a", 1, time.Now().Unix()+3600+1)
+		t.Log(cache.Read("a"))
+	}
+	{
+		cache.Write("b", 1, time.Now().Unix()+3600+2)
+		t.Log(cache.Read("b"))
+	}
+	{
+		cache.IncreaseInt64("b", 1, time.Now().Unix()+3600+3)
+		t.Log(cache.Read("b"))
+	}
+}
+
+func TestCache_Read(t *testing.T) {
+	runtime.GOMAXPROCS(1)
+
+	var cache = NewCache(PiecesOption{Count: 32})
+
+	for i := 0; i < 10_000_000; i++ {
+		cache.Write("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.Write("a", 1, time.Now().Unix()+1)
+	cache.Write("b", 2, time.Now().Unix()+2)
+	cache.Write("c", 3, time.Now().Unix()+3)
+	cache.Write("d", 4, time.Now().Unix()+4)
+	cache.Write("e", 5, time.Now().Unix()+10)
+
+	go func() {
+		for i := 0; i < 1000; i++ {
+			cache.Write("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.Write(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/ttlcache/item.go

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

internal/ttlcache/option.go

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

internal/ttlcache/piece.go

@@ -0,0 +1,88 @@
+package ttlcache
+
+import (
+	"github.com/iwind/TeaGo/types"
+	"sync"
+	"time"
+)
+
+type Piece struct {
+	m        map[uint64]*Item
+	maxItems int
+	locker   sync.RWMutex
+}
+
+func NewPiece(maxItems int) *Piece {
+	return &Piece{m: map[uint64]*Item{}, maxItems: maxItems}
+}
+
+func (this *Piece) Add(key uint64, item *Item) () {
+	this.locker.Lock()
+	if len(this.m) >= this.maxItems {
+		this.locker.Unlock()
+		return
+	}
+	this.m[key] = item
+	this.locker.Unlock()
+}
+
+func (this *Piece) IncreaseInt64(key uint64, delta int64, expiredAt int64) (result int64) {
+	this.locker.Lock()
+	item, ok := this.m[key]
+	if ok {
+		result := types.Int64(item.Value) + delta
+		item.Value = result
+		item.expiredAt = expiredAt
+	} else {
+		if len(this.m) < this.maxItems {
+			result = delta
+			this.m[key] = &Item{
+				Value:     delta,
+				expiredAt: expiredAt,
+			}
+		}
+	}
+	this.locker.Unlock()
+	return
+}
+
+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 < time.Now().Unix() {
+		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()
+}
+
+func (this *Piece) Destroy() {
+	this.locker.Lock()
+	this.m = nil
+	this.locker.Unlock()
+}

internal/ttlcache/piece_test.go

@@ -0,0 +1,60 @@
+package ttlcache
+
+import (
+	"github.com/iwind/TeaGo/rands"
+	"testing"
+	"time"
+)
+
+func TestPiece_Add(t *testing.T) {
+	piece := NewPiece(10)
+	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_MaxItems(t *testing.T) {
+	piece := NewPiece(10)
+	for i := 0; i < 1000; i++ {
+		piece.Add(uint64(i), &Item{expiredAt: time.Now().Unix() + 3600})
+	}
+	t.Log(len(piece.m))
+}
+
+func TestPiece_GC(t *testing.T) {
+	piece := NewPiece(10)
+	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(10)
+	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/ttlcache/utils.go

@@ -0,0 +1,7 @@
+package ttlcache
+
+import "github.com/cespare/xxhash"
+
+func HashKey(key []byte) uint64 {
+	return xxhash.Sum64(key)
+}

internal/ttlcache/utils_test.go

@@ -0,0 +1,13 @@
+package ttlcache
+
+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/utils/ticker.go

@@ -0,0 +1,47 @@
+package utils
+
+import (
+	"time"
+)
+
+// 类似于time.Ticker，但能够真正地停止
+type Ticker struct {
+	raw *time.Ticker
+
+	S chan bool
+	C <-chan time.Time
+
+	isStopped bool
+}
+
+// 创建新Ticker
+func NewTicker(duration time.Duration) *Ticker {
+	raw := time.NewTicker(duration)
+	return &Ticker{
+		raw: raw,
+		C:   raw.C,
+		S:   make(chan bool, 1),
+	}
+}
+
+// 查找下一个Tick
+func (this *Ticker) Next() bool {
+	select {
+	case <-this.raw.C:
+		return true
+	case <-this.S:
+		return false
+	}
+}
+
+// 停止
+func (this *Ticker) Stop() {
+	if this.isStopped {
+		return
+	}
+
+	this.isStopped = true
+
+	this.raw.Stop()
+	this.S <- true
+}

internal/web/actions/default/servers/components/waf/export.go

@@ -1,6 +1,15 @@
 package waf
 
-import "github.com/TeaOSLab/EdgeAdmin/internal/web/actions/actionutils"
+import (
+	"encoding/json"
+	"github.com/TeaOSLab/EdgeAdmin/internal/ttlcache"
+	"github.com/TeaOSLab/EdgeAdmin/internal/web/actions/actionutils"
+	"github.com/TeaOSLab/EdgeAdmin/internal/web/models"
+	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs/firewallconfigs"
+	"github.com/iwind/TeaGo/actions"
+	"github.com/iwind/TeaGo/rands"
+	"time"
+)
 
 type ExportAction struct {
 	actionutils.ParentAction
@@ -10,6 +19,102 @@ func (this *ExportAction) Init() {
 	this.Nav("", "", "export")
 }
 
-func (this *ExportAction) RunGet(params struct{}) {
+func (this *ExportAction) RunGet(params struct {
+	FirewallPolicyId int64
+}) {
+	policy, err := models.SharedHTTPFirewallPolicyDAO.FindEnabledPolicyConfig(this.AdminContext(), params.FirewallPolicyId)
+	if err != nil {
+		this.ErrorPage(err)
+		return
+	}
+	if policy == nil {
+		this.NotFound("firewallPolicy", policy.Id)
+		return
+	}
+
+	inboundGroups := []*firewallconfigs.HTTPFirewallRuleGroup{}
+	outboundGroups := []*firewallconfigs.HTTPFirewallRuleGroup{}
+	if policy.Inbound != nil {
+		for _, g := range policy.Inbound.Groups {
+			if g.IsOn {
+				inboundGroups = append(inboundGroups, g)
+			}
+		}
+	}
+	if policy.Outbound != nil {
+		for _, g := range policy.Outbound.Groups {
+			if g.IsOn {
+				outboundGroups = append(outboundGroups, g)
+			}
+		}
+	}
+	this.Data["inboundGroups"] = inboundGroups
+	this.Data["outboundGroups"] = outboundGroups
+
 	this.Show()
 }
+
+func (this *ExportAction) RunPost(params struct {
+	FirewallPolicyId int64
+	InboundGroupIds  []int64
+	OutboundGroupIds []int64
+
+	Must *actions.Must
+	CSRF *actionutils.CSRF
+}) {
+	defer this.CreateLogInfo("导出WAF策略 %d", params.FirewallPolicyId)
+
+	policy, err := models.SharedHTTPFirewallPolicyDAO.FindEnabledPolicyConfig(this.AdminContext(), params.FirewallPolicyId)
+	if err != nil {
+		this.ErrorPage(err)
+		return
+	}
+	if policy == nil {
+		this.NotFound("firewallPolicy", policy.Id)
+		return
+	}
+
+	// inbound
+	newInboundGroups := []*firewallconfigs.HTTPFirewallRuleGroup{}
+	for _, inboundGroupId := range params.InboundGroupIds {
+		group := policy.FindRuleGroup(inboundGroupId)
+		if group != nil {
+			newInboundGroups = append(newInboundGroups, group)
+		}
+	}
+	if policy.Inbound == nil {
+		policy.Inbound = &firewallconfigs.HTTPFirewallInboundConfig{
+			IsOn: true,
+		}
+	}
+	policy.Inbound.Groups = newInboundGroups
+	policy.Inbound.GroupRefs = nil
+
+	// outbound
+	newOutboundGroups := []*firewallconfigs.HTTPFirewallRuleGroup{}
+	for _, outboundGroupId := range params.OutboundGroupIds {
+		group := policy.FindRuleGroup(outboundGroupId)
+		if group != nil {
+			newOutboundGroups = append(newOutboundGroups, group)
+		}
+	}
+	if policy.Outbound == nil {
+		policy.Outbound = &firewallconfigs.HTTPFirewallOutboundConfig{
+			IsOn: true,
+		}
+	}
+	policy.Outbound.Groups = newOutboundGroups
+	policy.Outbound.GroupRefs = nil
+
+	configJSON, err := json.Marshal(policy)
+	if err != nil {
+		this.ErrorPage(err)
+		return
+	}
+
+	key := "waf." + rands.HexString(32)
+	ttlcache.DefaultCache.Write(key, configJSON, time.Now().Unix()+600)
+
+	this.Data["key"] = key
+	this.Success()
+}

internal/web/actions/default/servers/components/waf/exportDownload.go

@@ -0,0 +1,37 @@
+package waf
+
+import (
+	"github.com/TeaOSLab/EdgeAdmin/internal/ttlcache"
+	"github.com/TeaOSLab/EdgeAdmin/internal/web/actions/actionutils"
+	"strconv"
+)
+
+type ExportDownloadAction struct {
+	actionutils.ParentAction
+}
+
+func (this *ExportDownloadAction) Init() {
+	this.Nav("", "", "")
+}
+
+func (this *ExportDownloadAction) RunGet(params struct {
+	Key string
+}) {
+	item := ttlcache.DefaultCache.Read(params.Key)
+	if item == nil || item.Value == nil {
+		this.WriteString("找不到要导出的内容")
+		return
+	}
+
+	ttlcache.DefaultCache.Delete(params.Key)
+
+	data, ok := item.Value.([]byte)
+	if ok {
+		this.AddHeader("Content-Disposition", "attachment; filename=\"WAF.json\";")
+		this.AddHeader("Content-Length", strconv.Itoa(len(data)))
+		this.Write(data)
+	} else {
+		this.WriteString("找不到要导出的内容")
+		return
+	}
+}

internal/web/actions/default/servers/components/waf/import.go

@@ -1,6 +1,12 @@
 package waf
 
-import "github.com/TeaOSLab/EdgeAdmin/internal/web/actions/actionutils"
+import (
+	"encoding/json"
+	"github.com/TeaOSLab/EdgeAdmin/internal/web/actions/actionutils"
+	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
+	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs/firewallconfigs"
+	"github.com/iwind/TeaGo/actions"
+)
 
 type ImportAction struct {
 	actionutils.ParentAction
@@ -13,3 +19,41 @@ func (this *ImportAction) Init() {
 func (this *ImportAction) RunGet(params struct{}) {
 	this.Show()
 }
+
+func (this *ImportAction) RunPost(params struct {
+	FirewallPolicyId int64
+	File             *actions.File
+
+	Must *actions.Must
+	CSRF *actionutils.CSRF
+}) {
+	defer this.CreateLogInfo("从文件中导入规则到WAF策略 %d", params.FirewallPolicyId)
+
+	if params.File == nil {
+		this.Fail("请上传要导入的文件")
+	}
+	if params.File.Ext != ".json" {
+		this.Fail("规则文件的扩展名只能是.json")
+	}
+
+	data, err := params.File.Read()
+	if err != nil {
+		this.Fail("读取文件时发生错误：" + err.Error())
+	}
+
+	config := &firewallconfigs.HTTPFirewallPolicy{}
+	err = json.Unmarshal(data, config)
+	if err != nil {
+		this.Fail("解析文件时发生错误：" + err.Error())
+	}
+
+	_, err = this.RPC().HTTPFirewallPolicyRPC().ImportHTTPFirewallPolicy(this.AdminContext(), &pb.ImportHTTPFirewallPolicyRequest{
+		FirewallPolicyId:   params.FirewallPolicyId,
+		FirewallPolicyJSON: data,
+	})
+	if err != nil {
+		this.Fail("导入失败：" + err.Error())
+	}
+
+	this.Success()
+}

internal/web/actions/default/servers/components/waf/init.go

@@ -24,6 +24,7 @@ func init() {
 			GetPost("/update", new(UpdateAction)).
 			GetPost("/test", new(TestAction)).
 			GetPost("/export", new(ExportAction)).
+			Get("/exportDownload", new(ExportDownloadAction)).
 			GetPost("/import", new(ImportAction)).
 			Post("/updateGroupOn", new(UpdateGroupOnAction)).
 			Post("/deleteGroup", new(DeleteGroupAction)).