From 9c94c1df5f1890a183f5af7cbbc3254e3fab8e72 Mon Sep 17 00:00:00 2001
From: GoEdgeLab <goedgecloud@gmail.com>
Date: Tue, 13 Jul 2021 15:49:16 +0800
Subject: [PATCH] =?UTF-8?q?IP=E6=B5=8B=E8=AF=95=E6=97=B6=E5=90=8C=E6=97=B6?=
 =?UTF-8?q?=E4=B9=9F=E6=A3=80=E6=9F=A5=E7=BB=91=E5=AE=9A=E7=9A=84IP?=
 =?UTF-8?q?=E5=90=8D=E5=8D=95?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../services/service_http_firewall_policy.go  | 136 ++++++++++++------
 1 file changed, 92 insertions(+), 44 deletions(-)

diff --git a/internal/rpc/services/service_http_firewall_policy.go b/internal/rpc/services/service_http_firewall_policy.go
index 3f99a2e5..1a62ce2b 100644
--- a/internal/rpc/services/service_http_firewall_policy.go
+++ b/internal/rpc/services/service_http_firewall_policy.go
@@ -667,29 +667,53 @@ func (this *HTTPFirewallPolicyService) CheckHTTPFirewallPolicyIPStatus(ctx conte
 		firewallPolicy.Inbound.AllowListRef != nil &&
 		firewallPolicy.Inbound.AllowListRef.IsOn &&
 		firewallPolicy.Inbound.AllowListRef.ListId > 0 {
-		item, err := models.SharedIPItemDAO.FindEnabledItemContainsIP(tx, firewallPolicy.Inbound.AllowListRef.ListId, ipLong)
-		if err != nil {
-			return nil, err
+
+		var listIds = []int64{}
+		if firewallPolicy.Inbound.AllowListRef.ListId > 0 {
+			listIds = append(listIds, firewallPolicy.Inbound.AllowListRef.ListId)
 		}
-		if item != nil {
-			return &pb.CheckHTTPFirewallPolicyIPStatusResponse{
-				IsOk:      true,
-				Error:     "",
-				IsFound:   true,
-				IsAllowed: true,
-				IpList:    &pb.IPList{Name: "白名单", Id: firewallPolicy.Inbound.AllowListRef.ListId},
-				IpItem: &pb.IPItem{
-					Id:         int64(item.Id),
-					IpFrom:     item.IpFrom,
-					IpTo:       item.IpTo,
-					ExpiredAt:  int64(item.ExpiredAt),
-					Reason:     item.Reason,
-					Type:       item.Type,
-					EventLevel: item.EventLevel,
-				},
-				RegionCountry:  nil,
-				RegionProvince: nil,
-			}, nil
+		if len(firewallPolicy.Inbound.PublicAllowListRefs) > 0 {
+			for _, ref := range firewallPolicy.Inbound.PublicAllowListRefs {
+				if !ref.IsOn {
+					continue
+				}
+
+				listIds = append(listIds, ref.ListId)
+			}
+		}
+
+		for _, listId := range listIds {
+			item, err := models.SharedIPItemDAO.FindEnabledItemContainsIP(tx, listId, ipLong)
+			if err != nil {
+				return nil, err
+			}
+			if item != nil {
+				listName, err := models.SharedIPListDAO.FindIPListName(tx, listId)
+				if err != nil {
+					return nil, err
+				}
+				if len(listName) == 0 {
+					listName = "白名单"
+				}
+				return &pb.CheckHTTPFirewallPolicyIPStatusResponse{
+					IsOk:      true,
+					Error:     "",
+					IsFound:   true,
+					IsAllowed: true,
+					IpList:    &pb.IPList{Name: listName, Id: listId},
+					IpItem: &pb.IPItem{
+						Id:         int64(item.Id),
+						IpFrom:     item.IpFrom,
+						IpTo:       item.IpTo,
+						ExpiredAt:  int64(item.ExpiredAt),
+						Reason:     item.Reason,
+						Type:       item.Type,
+						EventLevel: item.EventLevel,
+					},
+					RegionCountry:  nil,
+					RegionProvince: nil,
+				}, nil
+			}
 		}
 	}
 
@@ -699,29 +723,53 @@ func (this *HTTPFirewallPolicyService) CheckHTTPFirewallPolicyIPStatus(ctx conte
 		firewallPolicy.Inbound.AllowListRef != nil &&
 		firewallPolicy.Inbound.AllowListRef.IsOn &&
 		firewallPolicy.Inbound.AllowListRef.ListId > 0 {
-		item, err := models.SharedIPItemDAO.FindEnabledItemContainsIP(tx, firewallPolicy.Inbound.DenyListRef.ListId, ipLong)
-		if err != nil {
-			return nil, err
+
+		var listIds = []int64{}
+		if firewallPolicy.Inbound.DenyListRef.ListId > 0 {
+			listIds = append(listIds, firewallPolicy.Inbound.DenyListRef.ListId)
 		}
-		if item != nil {
-			return &pb.CheckHTTPFirewallPolicyIPStatusResponse{
-				IsOk:      true,
-				Error:     "",
-				IsFound:   true,
-				IsAllowed: false,
-				IpList:    &pb.IPList{Name: "黑名单", Id: firewallPolicy.Inbound.DenyListRef.ListId},
-				IpItem: &pb.IPItem{
-					Id:         int64(item.Id),
-					IpFrom:     item.IpFrom,
-					IpTo:       item.IpTo,
-					ExpiredAt:  int64(item.ExpiredAt),
-					Reason:     item.Reason,
-					Type:       item.Type,
-					EventLevel: item.EventLevel,
-				},
-				RegionCountry:  nil,
-				RegionProvince: nil,
-			}, nil
+		if len(firewallPolicy.Inbound.PublicDenyListRefs) > 0 {
+			for _, ref := range firewallPolicy.Inbound.PublicDenyListRefs {
+				if !ref.IsOn {
+					continue
+				}
+
+				listIds = append(listIds, ref.ListId)
+			}
+		}
+
+		for _, listId := range listIds {
+			item, err := models.SharedIPItemDAO.FindEnabledItemContainsIP(tx, listId, ipLong)
+			if err != nil {
+				return nil, err
+			}
+			if item != nil {
+				listName, err := models.SharedIPListDAO.FindIPListName(tx, listId)
+				if err != nil {
+					return nil, err
+				}
+				if len(listName) == 0 {
+					listName = "黑名单"
+				}
+				return &pb.CheckHTTPFirewallPolicyIPStatusResponse{
+					IsOk:      true,
+					Error:     "",
+					IsFound:   true,
+					IsAllowed: false,
+					IpList:    &pb.IPList{Name: listName, Id: listId},
+					IpItem: &pb.IPItem{
+						Id:         int64(item.Id),
+						IpFrom:     item.IpFrom,
+						IpTo:       item.IpTo,
+						ExpiredAt:  int64(item.ExpiredAt),
+						Reason:     item.Reason,
+						Type:       item.Type,
+						EventLevel: item.EventLevel,
+					},
+					RegionCountry:  nil,
+					RegionProvince: nil,
+				}, nil
+			}
 		}
 	}