chore: ipcidr direct using go4.org/netipx

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() {}

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)
 			}
 		})
 	}