优化删除IP名单时操作

internal/iplibrary/ip_list.go

@@ -4,6 +4,7 @@ import (
 	"github.com/TeaOSLab/EdgeNode/internal/utils"
 	"github.com/TeaOSLab/EdgeNode/internal/utils/expires"
 	"github.com/TeaOSLab/EdgeNode/internal/utils/fasttime"
+	"github.com/iwind/TeaGo/logs"
 	"sort"
 	"sync"
 )
@@ -14,6 +15,8 @@ var GlobalWhiteIPList = NewIPList()
 // IPList IP名单
 // TODO IP名单可以分片关闭，这样让每一片的数据量减少，查询更快
 type IPList struct {
+	isDeleted bool
+
 	itemsMap    map[uint64]*IPItem // id => item
 	sortedItems []*IPItem
 	allItemsMap map[uint64]*IPItem // id => item
@@ -38,11 +41,19 @@ func NewIPList() *IPList {
 }
 
 func (this *IPList) Add(item *IPItem) {
+	if this.isDeleted {
+		return
+	}
+
 	this.addItem(item, true)
 }
 
 // AddDelay 延迟添加，需要手工调用Sort()函数
 func (this *IPList) AddDelay(item *IPItem) {
+	if this.isDeleted {
+		return
+	}
+
 	this.addItem(item, false)
 }
 
@@ -60,6 +71,10 @@ func (this *IPList) Delete(itemId uint64) {
 
 // Contains 判断是否包含某个IP
 func (this *IPList) Contains(ip uint64) bool {
+	if this.isDeleted {
+		return false
+	}
+
 	this.locker.RLock()
 	if len(this.allItemsMap) > 0 {
 		this.locker.RUnlock()
@@ -75,6 +90,10 @@ func (this *IPList) Contains(ip uint64) bool {
 
 // ContainsExpires 判断是否包含某个IP
 func (this *IPList) ContainsExpires(ip uint64) (expiresAt int64, ok bool) {
+	if this.isDeleted {
+		return
+	}
+
 	this.locker.RLock()
 	if len(this.allItemsMap) > 0 {
 		this.locker.RUnlock()
@@ -94,6 +113,10 @@ func (this *IPList) ContainsExpires(ip uint64) (expiresAt int64, ok bool) {
 
 // ContainsIPStrings 是否包含一组IP中的任意一个，并返回匹配的第一个Item
 func (this *IPList) ContainsIPStrings(ipStrings []string) (item *IPItem, found bool) {
+	if this.isDeleted {
+		return
+	}
+
 	if len(ipStrings) == 0 {
 		return
 	}
@@ -125,6 +148,11 @@ func (this *IPList) ContainsIPStrings(ipStrings []string) (item *IPItem, found b
 	return
 }
 
+func (this *IPList) SetDeleted() {
+	this.isDeleted = true
+	logs.Println("set deleted:", this.isDeleted) // TODO
+}
+
 func (this *IPList) addItem(item *IPItem, sortable bool) {
 	if item == nil {
 		return

internal/nodes/node_tasks.go

@@ -4,6 +4,7 @@ package nodes
 
 import (
 	"encoding/json"
+	"errors"
 	"fmt"
 	"github.com/TeaOSLab/EdgeCommon/pkg/nodeconfigs"
 	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
@@ -16,9 +17,12 @@ import (
 	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
 	"github.com/TeaOSLab/EdgeNode/internal/rpc"
 	"github.com/TeaOSLab/EdgeNode/internal/trackers"
+	"github.com/TeaOSLab/EdgeNode/internal/waf"
 	"github.com/iwind/TeaGo/Tea"
+	"github.com/iwind/TeaGo/maps"
 	"github.com/iwind/TeaGo/types"
 	"os"
+	"strings"
 	"time"
 )
 
@@ -94,8 +98,13 @@ func (this *Node) execTask(rpcClient *rpc.RPCClient, task *pb.NodeTask) error {
 	case "toaChanged":
 		err = this.execTOAChangedTask()
 	default:
+		// 特殊任务
+		if strings.HasPrefix(task.Type, "ipListDeleted") { // 删除IP名单
+			err = this.execDeleteIPList(task.Type)
+		} else { // 未处理的任务
 			remotelogs.Error("NODE", "task '"+types.String(task.Id)+"', type '"+task.Type+"' has not been handled")
 		}
+	}
 
 	return err
 }
@@ -285,6 +294,34 @@ func (this *Node) execUpdatingServersTask(rpcClient *rpc.RPCClient) error {
 	return nil
 }
 
+// 删除IP名单
+func (this *Node) execDeleteIPList(taskType string) error {
+	optionsString, ok := strings.CutPrefix(taskType, "ipListDeleted@")
+	if !ok {
+		return errors.New("invalid task type '" + taskType + "'")
+	}
+	var optionMap = maps.Map{}
+	err := json.Unmarshal([]byte(optionsString), &optionMap)
+	if err != nil {
+		return fmt.Errorf("decode options failed: %w, options: %s", err, optionsString)
+	}
+	var listId = optionMap.GetInt64("listId")
+	if listId <= 0 {
+		return nil
+	}
+
+	// 标记已被删除
+	waf.AddDeletedIPList(listId)
+
+	var list = iplibrary.SharedIPListManager.FindList(listId)
+
+	if list != nil {
+		list.SetDeleted()
+	}
+
+	return nil
+}
+
 // 标记任务完成
 func (this *Node) finishTask(taskId int64, taskVersion int64, taskErr error) (success bool) {
 	if taskId <= 0 {

internal/waf/action_record_ip.go

@@ -110,6 +110,7 @@ type RecordIPAction struct {
 
 	Type            string `yaml:"type" json:"type"`
 	IPListId        int64  `yaml:"ipListId" json:"ipListId"`
+	IPListIsDeleted bool   `yaml:"ipListIsDeleted" json:"ipListIsDeleted"`
 	Level           string `yaml:"level" json:"level"`
 	Timeout         int32  `yaml:"timeout" json:"timeout"`
 	Scope           string `yaml:"scope" json:"scope"`
@@ -132,6 +133,9 @@ func (this *RecordIPAction) WillChange() bool {
 }
 
 func (this *RecordIPAction) Perform(waf *WAF, group *RuleGroup, set *RuleSet, request requests.Request, writer http.ResponseWriter) (continueRequest bool, goNextSet bool) {
+	// 是否已删除
+	var ipListIsAvailable = this.IPListId > 0 && !this.IPListIsDeleted && !ExistDeletedIPList(this.IPListId)
+
 	// 是否在本地白名单中
 	if SharedIPWhiteList.Contains("set:"+types.String(set.Id), this.Scope, request.WAFServerId(), request.WAFRemoteIP()) {
 		return true, false
@@ -152,14 +156,18 @@ func (this *RecordIPAction) Perform(waf *WAF, group *RuleGroup, set *RuleSet, re
 		request.WAFClose()
 
 		// 先加入本地的黑名单
+		if ipListIsAvailable {
 			SharedIPBlackList.Add(IPTypeAll, this.Scope, request.WAFServerId(), request.WAFRemoteIP(), expiresAt)
+		}
 	} else {
 		// 加入本地白名单
+		if ipListIsAvailable {
 			SharedIPWhiteList.Add("set:"+types.String(set.Id), this.Scope, request.WAFServerId(), request.WAFRemoteIP(), expiresAt)
 		}
+	}
 
 	// 上报
-	if this.IPListId > 0 {
+	if this.IPListId > 0 && ipListIsAvailable {
 		var serverId int64
 		if this.Scope == firewallconfigs.FirewallScopeService {
 			serverId = request.WAFServerId()

internal/waf/ip_lists_deleted.go

@@ -0,0 +1,30 @@
+// Copyright 2023 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package waf
+
+import (
+	"github.com/TeaOSLab/EdgeNode/internal/zero"
+	"sync"
+)
+
+var deletedIPListIdMap = map[int64]zero.Zero{} // listId => Zero
+var deletedIPListLocker = sync.RWMutex{}
+
+// AddDeletedIPList add deleted ip list
+func AddDeletedIPList(ipListId int64) {
+	if ipListId <= 0 {
+		return
+	}
+
+	deletedIPListLocker.Lock()
+	deletedIPListIdMap[ipListId] = zero.Zero{}
+	deletedIPListLocker.Unlock()
+}
+
+// ExistDeletedIPList check if ip list has been deleted
+func ExistDeletedIPList(ipListId int64) bool {
+	deletedIPListLocker.RLock()
+	_, ok := deletedIPListIdMap[ipListId]
+	deletedIPListLocker.RUnlock()
+	return ok
+}