chore: ip check logic

xxxsen · Mar 14, 2024 · f9db161 · f9db161
1 parent abb5757
commit f9db161
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 3 deletions.
diff --git 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
@@ -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) {