From 7c54775ddcec7c9dd2e7f2688c69bcbcd51f68cb Mon Sep 17 00:00:00 2001
From: wwqgtxx <wwqgtxx@gmail.com>
Date: Thu, 11 Jan 2024 09:33:59 +0800
Subject: [PATCH] chore: ipcidr direct using go4.org/netipx

---
 component/cidr/ipcidr_set.go      | 81 ++++++++-----------------------
 component/cidr/ipcidr_set_test.go |  8 ++-
 go.mod                            |  2 +-
 rules/provider/ipcidr_strategy.go |  2 +-
 4 files changed, 28 insertions(+), 65 deletions(-)

diff --git a/component/cidr/ipcidr_set.go b/component/cidr/ipcidr_set.go
index b8dec0ee..dac2d38a 100644
--- a/component/cidr/ipcidr_set.go
+++ b/component/cidr/ipcidr_set.go
@@ -1,89 +1,48 @@
 package cidr
 
 import (
-	"math/big"
-	"net"
-	"sort"
+	"go4.org/netipx"
+	"net/netip"
 )
 
-type Range struct {
-	Start *big.Int
-	End   *big.Int
-}
-
 type IpCidrSet struct {
-	Ranges []Range
+	Ranges *netipx.IPSet
 }
 
 func NewIpCidrSet() *IpCidrSet {
 	return &IpCidrSet{}
 }
 
-func ipToBigInt(ip net.IP) *big.Int {
-	ipBigInt := big.NewInt(0)
-	ipBigInt.SetBytes(ip.To16())
-	return ipBigInt
-}
-
-func cidrToRange(cidr string) (Range, error) {
-	_, ipNet, err := net.ParseCIDR(cidr)
-	if err != nil {
-		return Range{}, err
-	}
-	firstIP, lastIP := networkRange(ipNet)
-	return Range{Start: ipToBigInt(firstIP), End: ipToBigInt(lastIP)}, nil
-}
-
-func networkRange(network *net.IPNet) (net.IP, net.IP) {
-	firstIP := network.IP
-	lastIP := make(net.IP, len(firstIP))
-	copy(lastIP, firstIP)
-	for i := range firstIP {
-		lastIP[i] |= ^network.Mask[i]
-	}
-	return firstIP, lastIP
-}
-
 func (set *IpCidrSet) AddIpCidrForString(ipCidr string) error {
-	ipRange, err := cidrToRange(ipCidr)
+	prefix, err := netip.ParsePrefix(ipCidr)
 	if err != nil {
 		return err
 	}
-	set.Ranges = append(set.Ranges, ipRange)
-	sort.Slice(set.Ranges, func(i, j int) bool {
-		return set.Ranges[i].Start.Cmp(set.Ranges[j].Start) < 0
-	})
+	err = set.AddIpCidr(prefix)
 	return nil
 }
 
-func (set *IpCidrSet) AddIpCidr(ipCidr *net.IPNet) error {
-	return set.AddIpCidrForString(ipCidr.String())
+func (set *IpCidrSet) AddIpCidr(ipCidr netip.Prefix) (err error) {
+	var b netipx.IPSetBuilder
+	b.AddSet(set.Ranges)
+	b.AddPrefix(ipCidr)
+	set.Ranges, err = b.IPSet()
+	return
 }
 
 func (set *IpCidrSet) IsContainForString(ipString string) bool {
-	ip := ipToBigInt(net.ParseIP(ipString))
-	idx := sort.Search(len(set.Ranges), func(i int) bool {
-		return set.Ranges[i].End.Cmp(ip) >= 0
-	})
-	if idx < len(set.Ranges) && set.Ranges[idx].Start.Cmp(ip) <= 0 && set.Ranges[idx].End.Cmp(ip) >= 0 {
-		return true
-	}
-	return false
-}
-
-func (set *IpCidrSet) IsContain(ip net.IP) bool {
-	if ip == nil {
+	ip, err := netip.ParseAddr(ipString)
+	if err != nil {
 		return false
 	}
-	return set.IsContainForString(ip.String())
+	return set.IsContain(ip)
 }
 
-func (set *IpCidrSet) Merge() {
-	for i := 0; i < len(set.Ranges)-1; i++ {
-		if set.Ranges[i].End.Cmp(set.Ranges[i+1].Start) >= 0 {
-			set.Ranges[i].End = set.Ranges[i+1].End
-			set.Ranges = append(set.Ranges[:i+1], set.Ranges[i+2:]...)
-			i--
-		}
+func (set *IpCidrSet) IsContain(ip netip.Addr) bool {
+	if set.Ranges == nil {
+		return false
 	}
+	return set.Ranges.Contains(ip.WithZone(""))
 }
+
+func (set *IpCidrSet) Merge() {}
diff --git a/component/cidr/ipcidr_set_test.go b/component/cidr/ipcidr_set_test.go
index a8cfef7f..a6eaec84 100644
--- a/component/cidr/ipcidr_set_test.go
+++ b/component/cidr/ipcidr_set_test.go
@@ -1,7 +1,9 @@
 package cidr
 
 import (
+	"go4.org/netipx"
 	"testing"
+	"unsafe"
 )
 
 func TestIpv4(t *testing.T) {
@@ -97,8 +99,10 @@ func TestMerge(t *testing.T) {
 			set.AddIpCidrForString(test.ipCidr2)
 			set.Merge()
 
-			if len(set.Ranges) != test.expectedLen {
-				t.Errorf("Expected len: %v, got: %v", test.expectedLen, len(set.Ranges))
+			rangesLen := len(*(*[]netipx.IPRange)(unsafe.Pointer(set.Ranges)))
+
+			if rangesLen != test.expectedLen {
+				t.Errorf("Expected len: %v, got: %v", test.expectedLen, rangesLen)
 			}
 		})
 	}
diff --git a/go.mod b/go.mod
index 866ac035..07458af8 100644
--- a/go.mod
+++ b/go.mod
@@ -47,6 +47,7 @@ require (
 	github.com/wk8/go-ordered-map/v2 v2.1.8
 	github.com/zhangyunhao116/fastrand v0.3.0
 	go.uber.org/automaxprocs v1.5.3
+	go4.org/netipx v0.0.0-20231129151722-fdeea329fbba
 	golang.org/x/crypto v0.16.0
 	golang.org/x/exp v0.0.0-20231206192017-f3f8817b8deb
 	golang.org/x/net v0.19.0
@@ -105,7 +106,6 @@ require (
 	github.com/yusufpapurcu/wmi v1.2.3 // indirect
 	gitlab.com/yawning/bsaes.git v0.0.0-20190805113838-0a714cd429ec // indirect
 	go.uber.org/mock v0.3.0 // indirect
-	go4.org/netipx v0.0.0-20231129151722-fdeea329fbba // indirect
 	golang.org/x/mod v0.14.0 // indirect
 	golang.org/x/text v0.14.0 // indirect
 	golang.org/x/time v0.5.0 // indirect
diff --git a/rules/provider/ipcidr_strategy.go b/rules/provider/ipcidr_strategy.go
index dad48305..c93facd9 100644
--- a/rules/provider/ipcidr_strategy.go
+++ b/rules/provider/ipcidr_strategy.go
@@ -19,7 +19,7 @@ func (i *ipcidrStrategy) ShouldFindProcess() bool {
 
 func (i *ipcidrStrategy) Match(metadata *C.Metadata) bool {
 	// return i.trie != nil && i.trie.IsContain(metadata.DstIP.AsSlice())
-	return i.cidrSet != nil && i.cidrSet.IsContain(metadata.DstIP.AsSlice())
+	return i.cidrSet != nil && i.cidrSet.IsContain(metadata.DstIP)
 }
 
 func (i *ipcidrStrategy) Count() int {