rebuild ipcidr

wwqgtxx · Jan 13, 2024 · fa104ca · fa104ca
1 parent a7f4034
commit fa104ca
Show file tree

Hide file tree

Showing 6 changed files with 190 additions and 19 deletions.
diff --git a/component/cidr/ipcidr_set.go b/component/cidr/ipcidr_set.go
@@ -0,0 +1,66 @@
+package cidr
+
+import (
+	"fmt"
+	"net/netip"
+	"unsafe"
+
+	"go4.org/netipx"
+)
+
+type IpCidrSet struct {
+	// must same with netipx.IPSet
+	rr []netipx.IPRange
+}
+
+func NewIpCidrSet() *IpCidrSet {
+	return &IpCidrSet{}
+}
+
+func (set *IpCidrSet) AddIpCidrForString(ipCidr string) error {
+	prefix, err := netip.ParsePrefix(ipCidr)
+	if err != nil {
+		return err
+	}
+	return set.AddIpCidr(prefix)
+}
+
+func (set *IpCidrSet) AddIpCidr(ipCidr netip.Prefix) (err error) {
+	if r := netipx.RangeOfPrefix(ipCidr); r.IsValid() {
+		set.rr = append(set.rr, r)
+	} else {
+		err = fmt.Errorf("not valid ipcidr range: %s", ipCidr)
+	}
+	return
+}
+
+func (set *IpCidrSet) IsContainForString(ipString string) bool {
+	ip, err := netip.ParseAddr(ipString)
+	if err != nil {
+		return false
+	}
+	return set.IsContain(ip)
+}
+
+func (set *IpCidrSet) IsContain(ip netip.Addr) bool {
+	return set.toIPSet().Contains(ip.WithZone(""))
+}
+
+func (set *IpCidrSet) Merge() error {
+	var b netipx.IPSetBuilder
+	b.AddSet(set.toIPSet())
+	i, err := b.IPSet()
+	if err != nil {
+		return err
+	}
+	set.fromIPSet(i)
+	return nil
+}
+
+func (set *IpCidrSet) toIPSet() *netipx.IPSet {
+	return (*netipx.IPSet)(unsafe.Pointer(set))
+}
+
+func (set *IpCidrSet) fromIPSet(i *netipx.IPSet) {
+	*set = *(*IpCidrSet)(unsafe.Pointer(i))
+}
diff --git a/component/cidr/ipcidr_set_test.go b/component/cidr/ipcidr_set_test.go
@@ -0,0 +1,107 @@
+package cidr
+
+import (
+	"testing"
+)
+
+func TestIpv4(t *testing.T) {
+	tests := []struct {
+		name     string
+		ipCidr   string
+		ip       string
+		expected bool
+	}{
+		{
+			name:     "Test Case 1",
+			ipCidr:   "149.154.160.0/20",
+			ip:       "149.154.160.0",
+			expected: true,
+		},
+		{
+			name:     "Test Case 2",
+			ipCidr:   "192.168.0.0/16",
+			ip:       "10.0.0.1",
+			expected: false,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			set := &IpCidrSet{}
+			set.AddIpCidrForString(test.ipCidr)
+
+			result := set.IsContainForString(test.ip)
+			if result != test.expected {
+				t.Errorf("Expected result: %v, got: %v", test.expected, result)
+			}
+		})
+	}
+}
+
+func TestIpv6(t *testing.T) {
+	tests := []struct {
+		name     string
+		ipCidr   string
+		ip       string
+		expected bool
+	}{
+		{
+			name:     "Test Case 1",
+			ipCidr:   "2409:8000::/20",
+			ip:       "2409:8087:1e03:21::27",
+			expected: true,
+		},
+		{
+			name:     "Test Case 2",
+			ipCidr:   "240e::/16",
+			ip:       "240e:964:ea02:100:1800::71",
+			expected: true,
+		},
+	}
+	// Add more test cases as needed
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			set := &IpCidrSet{}
+			set.AddIpCidrForString(test.ipCidr)
+
+			result := set.IsContainForString(test.ip)
+			if result != test.expected {
+				t.Errorf("Expected result: %v, got: %v", test.expected, result)
+			}
+		})
+	}
+}
+
+func TestMerge(t *testing.T) {
+	tests := []struct {
+		name        string
+		ipCidr1     string
+		ipCidr2     string
+		ipCidr3     string
+		expectedLen int
+	}{
+		{
+			name:        "Test Case 1",
+			ipCidr1:     "2409:8000::/20",
+			ipCidr2:     "2409:8000::/21",
+			ipCidr3:     "2409:8000::/48",
+			expectedLen: 1,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			set := &IpCidrSet{}
+			set.AddIpCidrForString(test.ipCidr1)
+			set.AddIpCidrForString(test.ipCidr2)
+			set.Merge()
+
+			rangesLen := len(set.rr)
+
+			if rangesLen != test.expectedLen {
+				t.Errorf("Expected len: %v, got: %v", test.expectedLen, rangesLen)
+			}
+		})
+	}
+}
diff --git a/rule/domain_tree.go b/rule/domain_tree.go
@@ -33,9 +33,10 @@ func (d *DomainTree) Insert(domain string) error {
 	return nil
 }
 
-func (d *DomainTree) FinishInsert() {
+func (d *DomainTree) FinishInsert() error {
 	d.ds = d.dt.NewDomainSet()
 	d.dt = nil
+	return nil
 }
 
 func NewDomainTree() *DomainTree {

diff --git a/rule/ipcidr_tree.go b/rule/ipcidr_tree.go
@@ -1,14 +1,13 @@
 package rules
 
 import (
+	"github.com/metacubex/mihomo/component/cidr"
 	C "github.com/metacubex/mihomo/constant"
-	"go4.org/netipx"
-	"net/netip"
 )
 
 type IpCidrTree struct {
 	IPCIDR    // for C.Rule interface
-	ipSet     *netipx.IPSet
+	cidrSet   *cidr.IpCidrSet
 	ruleCount int
 }
 
@@ -28,36 +27,29 @@ func (i *IpCidrTree) Match(metadata *C.Metadata) (bool, string) {
 	if !ip.IsValid() {
 		return false, ""
 	}
-	if i.ipSet == nil {
-		return false, ""
-	}
-	if i.ipSet.Contains(ip) {
+	if i.cidrSet.IsContain(ip) {
 		return true, i.adapter
 	}
 	return false, ""
 }
 
 func (i *IpCidrTree) Insert(ipCidr string) error {
-	prefix, err := netip.ParsePrefix(ipCidr)
-	if err != nil {
-		return err
-	}
-	var b netipx.IPSetBuilder
-	b.AddSet(i.ipSet)
-	b.AddPrefix(prefix)
-	i.ipSet, err = b.IPSet()
+	err := i.cidrSet.AddIpCidrForString(ipCidr)
 	if err != nil {
 		return err
 	}
 	i.ruleCount++
 	return nil
 }
 
-func (i *IpCidrTree) FinishInsert() {}
+func (i *IpCidrTree) FinishInsert() error {
+	return i.cidrSet.Merge()
+}
 
 func NewIPCIDRTree() *IpCidrTree {
 	return &IpCidrTree{
 		IPCIDR:    IPCIDR{},
+		cidrSet:   cidr.NewIpCidrSet(),
 		ruleCount: 0,
 	}
 }
diff --git a/rule/ipcidr_tree_test.go b/rule/ipcidr_tree_test.go
@@ -42,6 +42,7 @@ func TestIpv4Match(t *testing.T) {
 	assert.NoError(t, tree.Insert("11.2.13.2/26"))
 	assert.NoError(t, tree.Insert("55.5.6.3/8"))
 	assert.NoError(t, tree.Insert("66.23.25.4/6"))
+	assert.NoError(t, tree.FinishInsert())
 	assert.Equal(t, true, match(tree, "129.2.3.65"))
 	assert.Equal(t, false, match(tree, "15.2.3.1"))
 	assert.Equal(t, true, match(tree, "11.2.13.1"))
@@ -89,6 +90,7 @@ func TestIpv6Match(t *testing.T) {
 	assert.NoError(t, tree.Insert("2001:67c:4e8:f002::a/32"))
 	assert.NoError(t, tree.Insert("2001:67c:4e8:f004::a/60"))
 	assert.NoError(t, tree.Insert("2001:b28:f23f:f005::a/64"))
+	assert.NoError(t, tree.FinishInsert())
 	assert.Equal(t, true, match(tree, "2001:b28:f23d:f001::e"))
 	assert.Equal(t, false, match(tree, "2222::fff2"))
 	assert.Equal(t, true, match(tree, "2000::ffa0"))

diff --git a/rule/provider.go b/rule/provider.go
@@ -85,7 +85,7 @@ type RuleTree interface {
 	C.Rule
 	RuleCount() int
 	Insert(string) error
-	FinishInsert()
+	FinishInsert() error
 }
 
 var ErrNoPayload = errors.New("file must have a `payload` field")
@@ -229,7 +229,10 @@ func rulesParse(buf []byte, behavior P.RuleBehavior, format P.RuleFormat) (any,
 	}
 
 	if rt != nil {
-		rt.FinishInsert()
+		err := rt.FinishInsert()
+		if err != nil {
+			return nil, err
+		}
 	}
 
 	if len(rules) == 0 {