From 5d49b1a2fc72d367d3a36d743c8ad54d5202dbf0 Mon Sep 17 00:00:00 2001
From: GoEdgeLab <goedgecloud@gmail.com>
Date: Sun, 21 Apr 2024 10:32:04 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8DIP=E5=90=8D=E5=8D=95=E4=B8=AD?=
 =?UTF-8?q?=E2=80=9C=E6=89=80=E6=9C=89IP=E2=80=9D=E7=B1=BB=E5=9E=8B?=
 =?UTF-8?q?=E7=9A=84IP=E4=B8=8D=E8=B5=B7=E4=BD=9C=E7=94=A8=E7=9A=84?=
 =?UTF-8?q?=E9=97=AE=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/iplibrary/ip_list.go      | 19 +++++++-------
 internal/iplibrary/ip_list_test.go | 40 ++++++++++++++++++++++--------
 2 files changed, 38 insertions(+), 21 deletions(-)

diff --git a/internal/iplibrary/ip_list.go b/internal/iplibrary/ip_list.go
index 15786ae..6ecd99a 100644
--- a/internal/iplibrary/ip_list.go
+++ b/internal/iplibrary/ip_list.go
@@ -221,15 +221,15 @@ func (this *IPList) addItem(item *IPItem, lock bool, sortable bool) {
 	this.itemsMap[item.Id] = item
 
 	// 展开
-	if !IsZero(item.IPFrom) {
+	if item.Type == IPItemTypeAll {
+		this.allItemsMap[item.Id] = item
+	} else if !IsZero(item.IPFrom) {
 		if !IsZero(item.IPTo) {
 			this.sortedRangeItems = append(this.sortedRangeItems, item)
 			shouldSort = true
 		} else {
 			this.ipMap[ToHex(item.IPFrom)] = item
 		}
-	} else {
-		this.allItemsMap[item.Id] = item
 	}
 
 	if item.ExpiredAt > 0 {
@@ -310,6 +310,12 @@ func (this *IPList) deleteItem(itemId uint64) {
 	// 从buffer中删除
 	delete(this.bufferItemsMap, itemId)
 
+	// 从all items中删除
+	_, ok := this.allItemsMap[itemId]
+	if ok {
+		delete(this.allItemsMap, itemId)
+	}
+
 	// 检查是否存在
 	oldItem, existsOld := this.itemsMap[itemId]
 	if !existsOld {
@@ -327,13 +333,6 @@ func (this *IPList) deleteItem(itemId uint64) {
 
 	delete(this.itemsMap, itemId)
 
-	// 是否为All Item
-	_, ok := this.allItemsMap[itemId]
-	if ok {
-		delete(this.allItemsMap, itemId)
-		return
-	}
-
 	// 删除排序中的Item
 	if !IsZero(oldItem.IPTo) {
 		var index = -1
diff --git a/internal/iplibrary/ip_list_test.go b/internal/iplibrary/ip_list_test.go
index 5475e28..af39d58 100644
--- a/internal/iplibrary/ip_list_test.go
+++ b/internal/iplibrary/ip_list_test.go
@@ -245,19 +245,37 @@ func TestIPList_Contains_Many(t *testing.T) {
 func TestIPList_ContainsAll(t *testing.T) {
 	var a = assert.NewAssertion(t)
 
-	var list = iplibrary.NewIPList()
-	list.Add(&iplibrary.IPItem{
-		Id:     1,
-		Type:   "all",
-		IPFrom: nil,
-	})
-	var b = list.Contains(iputils.ToBytes("192.168.1.1"))
-	a.IsTrue(b)
+	{
+		var list = iplibrary.NewIPList()
+		list.Add(&iplibrary.IPItem{
+			Id:     1,
+			Type:   "all",
+			IPFrom: nil,
+		})
+		var b = list.Contains(iputils.ToBytes("192.168.1.1"))
+		a.IsTrue(b)
 
-	list.Delete(1)
+		list.Delete(1)
 
-	b = list.Contains(iputils.ToBytes("192.168.1.1"))
-	a.IsFalse(b)
+		b = list.Contains(iputils.ToBytes("192.168.1.1"))
+		a.IsFalse(b)
+	}
+
+	{
+		var list = iplibrary.NewIPList()
+		list.Add(&iplibrary.IPItem{
+			Id:     1,
+			Type:   "all",
+			IPFrom: iputils.ToBytes("0.0.0.0"),
+		})
+		var b = list.Contains(iputils.ToBytes("192.168.1.1"))
+		a.IsTrue(b)
+
+		list.Delete(1)
+
+		b = list.Contains(iputils.ToBytes("192.168.1.1"))
+		a.IsFalse(b)
+	}
 }
 
 func TestIPList_ContainsIPStrings(t *testing.T) {