chore: ipcidr direct using go4.org/netipx

This commit is contained in:
wwqgtxx 2024-01-11 09:33:59 +08:00
parent 425bc692ad
commit 7c54775ddc
4 changed files with 28 additions and 65 deletions

View File

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

View File

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

2
go.mod
View File

@ -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

View File

@ -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 {