From 55b161deab1ce8b8968edb80f0e08938fc88e8d4 Mon Sep 17 00:00:00 2001
From: GoEdgeLab <goedgecloud@gmail.com>
Date: Fri, 26 Nov 2021 10:39:04 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8DWAF=E7=9A=84=E4=B8=B4?=
 =?UTF-8?q?=E6=97=B6=E7=99=BD=E5=90=8D=E5=8D=95=E8=A2=AB=E5=BD=93=E5=81=9A?=
 =?UTF-8?q?=E9=BB=91=E5=90=8D=E5=8D=95=E4=BD=BF=E7=94=A8=E7=9A=84Bug?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/waf/ip_list.go | 47 +++++++++++++++++++++++++----------------
 1 file changed, 29 insertions(+), 18 deletions(-)

diff --git a/internal/waf/ip_list.go b/internal/waf/ip_list.go
index 6716a91..9c2120b 100644
--- a/internal/waf/ip_list.go
+++ b/internal/waf/ip_list.go
@@ -10,8 +10,15 @@ import (
 	"sync/atomic"
 )
 
-var SharedIPWhiteList = NewIPList()
-var SharedIPBlackList = NewIPList()
+var SharedIPWhiteList = NewIPList(IPListTypeAllow)
+var SharedIPBlackList = NewIPList(IPListTypeDeny)
+
+type IPListType = string
+
+const (
+	IPListTypeAllow IPListType = "allow"
+	IPListTypeDeny  IPListType = "deny"
+)
 
 const IPTypeAll = "*"
 
@@ -20,16 +27,18 @@ type IPList struct {
 	expireList *expires.List
 	ipMap      map[string]int64 // ip => id
 	idMap      map[int64]string // id => ip
+	listType   IPListType
 
 	id     int64
 	locker sync.RWMutex
 }
 
 // NewIPList 获取新对象
-func NewIPList() *IPList {
+func NewIPList(listType IPListType) *IPList {
 	var list = &IPList{
-		ipMap: map[string]int64{},
-		idMap: map[int64]string{},
+		ipMap:    map[string]int64{},
+		idMap:    map[int64]string{},
+		listType: listType,
 	}
 
 	e := expires.NewList()
@@ -67,20 +76,22 @@ func (this *IPList) Add(ipType string, scope firewallconfigs.FirewallScope, serv
 func (this *IPList) RecordIP(ipType string, scope firewallconfigs.FirewallScope, serverId int64, ip string, expiresAt int64, policyId int64, groupId int64, setId int64) {
 	this.Add(ipType, scope, serverId, ip, expiresAt)
 
-	select {
-	case recordIPTaskChan <- &recordIPTask{
-		ip:                            ip,
-		listId:                        firewallconfigs.GlobalListId,
-		expiredAt:                     expiresAt,
-		level:                         firewallconfigs.DefaultEventLevel,
-		serverId:                      serverId,
-		sourceServerId:                serverId,
-		sourceHTTPFirewallPolicyId:    policyId,
-		sourceHTTPFirewallRuleGroupId: groupId,
-		sourceHTTPFirewallRuleSetId:   setId,
-	}:
-	default:
+	if this.listType == IPListTypeDeny {
+		select {
+		case recordIPTaskChan <- &recordIPTask{
+			ip:                            ip,
+			listId:                        firewallconfigs.GlobalListId,
+			expiredAt:                     expiresAt,
+			level:                         firewallconfigs.DefaultEventLevel,
+			serverId:                      serverId,
+			sourceServerId:                serverId,
+			sourceHTTPFirewallPolicyId:    policyId,
+			sourceHTTPFirewallRuleGroupId: groupId,
+			sourceHTTPFirewallRuleSetId:   setId,
+		}:
+		default:
 
+		}
 	}
 }