From 38dfd08332dda5809191bc54ef06b781b1d81e44 Mon Sep 17 00:00:00 2001
From: Loyalsoldier <10487845+Loyalsoldier@users.noreply.github.com>
Date: Fri, 25 Oct 2024 07:24:30 +0800
Subject: Refine: make lookup command more specific

---
 lib/container.go | 50 +++++++++++++++++++-------------------------------
 1 file changed, 19 insertions(+), 31 deletions(-)

(limited to 'lib')

diff --git a/lib/container.go b/lib/container.go
index 3ecb5b0e..b08cbc13 100644
--- a/lib/container.go
+++ b/lib/container.go
@@ -232,41 +232,29 @@ func (c *container) lookup(addrOrPrefix any, iptype IPType, searchList ...string
 			continue
 		}
 
+		var ipset *netipx.IPSet
+		var err error
 		switch iptype {
 		case IPv4:
-			ipset, err := entry.GetIPv4Set()
-			if err != nil {
-				return nil, false, err
-			}
-			switch addrOrPrefix := addrOrPrefix.(type) {
-			case netip.Prefix:
-				if found := ipset.ContainsPrefix(addrOrPrefix); found {
-					isfound = true
-					result = append(result, entry.GetName())
-				}
-			case netip.Addr:
-				if found := ipset.Contains(addrOrPrefix); found {
-					isfound = true
-					result = append(result, entry.GetName())
-				}
-			}
-
+			ipset, err = entry.GetIPv4Set()
 		case IPv6:
-			ipset, err := entry.GetIPv6Set()
-			if err != nil {
-				return nil, false, err
+			ipset, err = entry.GetIPv6Set()
+		}
+
+		if err != nil {
+			return nil, false, err
+		}
+
+		switch addrOrPrefix := addrOrPrefix.(type) {
+		case netip.Prefix:
+			if found := ipset.ContainsPrefix(addrOrPrefix); found {
+				isfound = true
+				result = append(result, entry.GetName())
 			}
-			switch addrOrPrefix := addrOrPrefix.(type) {
-			case netip.Prefix:
-				if found := ipset.ContainsPrefix(addrOrPrefix); found {
-					isfound = true
-					result = append(result, entry.GetName())
-				}
-			case netip.Addr:
-				if found := ipset.Contains(addrOrPrefix); found {
-					isfound = true
-					result = append(result, entry.GetName())
-				}
+		case netip.Addr:
+			if found := ipset.Contains(addrOrPrefix); found {
+				isfound = true
+				result = append(result, entry.GetName())
 			}
 		}
 	}
-- 
cgit v1.3.1