From f9db161ecf9f6e8fbf9d7e04e99052cf9c0490e0 Mon Sep 17 00:00:00 2001
From: xxxsen <naivesen@gmail.com>
Date: Thu, 14 Mar 2024 21:01:13 +0800
Subject: [PATCH] chore: ip check logic

---
 cleaner/rule.go      | 11 ++++++++---
 cleaner/rule_test.go | 15 +++++++++++++++
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/cleaner/rule.go b/cleaner/rule.go
index d329824..6250bfc 100644
--- a/cleaner/rule.go
+++ b/cleaner/rule.go
@@ -66,10 +66,11 @@ type ipRuleSet struct {
 }
 
 func (r *ipRuleSet) isMatch(strip string) bool {
-	if _, ok := r.ip[strip]; ok {
+	ip := net.ParseIP(strip)
+	normalizedIP := ip.String()
+	if _, ok := r.ip[normalizedIP]; ok {
 		return true
 	}
-	ip := net.ParseIP(strip)
 	if ip == nil {
 		return false
 	}
@@ -96,7 +97,11 @@ func makeIPRuleSet(rs []string) (*ipRuleSet, error) {
 			set.cidr = append(set.cidr, cidr)
 			continue
 		}
-		set.ip[item] = struct{}{}
+		ip := net.ParseIP(item)
+		if ip == nil {
+			return nil, fmt.Errorf("invalid ip:%s", item)
+		}
+		set.ip[ip.String()] = struct{}{}
 	}
 	return set, nil
 }
diff --git a/cleaner/rule_test.go b/cleaner/rule_test.go
index c642ea8..f05e14d 100644
--- a/cleaner/rule_test.go
+++ b/cleaner/rule_test.go
@@ -10,15 +10,30 @@ func TestIP(t *testing.T) {
 		"127.0.0.1",
 		"192.168.0.0/16",
 		"10.0.0.0/8",
+		"2001:db8::/66",
+		"2001:db8:0:0:8000::5",
 	})
 	assert.NoError(t, err)
 	assert.True(t, set.isMatch("127.0.0.1"))
 	assert.True(t, set.isMatch("192.168.50.6"))
 	assert.True(t, set.isMatch("10.1.2.3"))
+	assert.True(t, set.isMatch("2001:db8::1"))
+	assert.True(t, set.isMatch("2001:db8:0:0:8000::5"))
 	assert.False(t, set.isMatch("11.1.2.3"))
 	assert.False(t, set.isMatch("192.167.50.1"))
 	assert.False(t, set.isMatch("127.0.0.2"))
+	assert.False(t, set.isMatch("2001:db8:0:0:8000::4"))
+}
 
+func TestInvalidIP(t *testing.T) {
+	_, err := makeIPRuleSet([]string{
+		"2001:db8.3.4::/66",
+	})
+	assert.Error(t, err)
+	_, err = makeIPRuleSet([]string{
+		"1.2.3.4.5",
+	})
+	assert.Error(t, err)
 }
 
 func TestStrSet(t *testing.T) {