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 package cidr
import ( import (
"math/big" "go4.org/netipx"
"net" "net/netip"
"sort"
) )
type Range struct {
Start *big.Int
End *big.Int
}
type IpCidrSet struct { type IpCidrSet struct {
Ranges []Range Ranges *netipx.IPSet
} }
func NewIpCidrSet() *IpCidrSet { func NewIpCidrSet() *IpCidrSet {
return &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 { func (set *IpCidrSet) AddIpCidrForString(ipCidr string) error {
ipRange, err := cidrToRange(ipCidr) prefix, err := netip.ParsePrefix(ipCidr)
if err != nil { if err != nil {
return err return err
} }
set.Ranges = append(set.Ranges, ipRange) err = set.AddIpCidr(prefix)
sort.Slice(set.Ranges, func(i, j int) bool {
return set.Ranges[i].Start.Cmp(set.Ranges[j].Start) < 0
})
return nil return nil
} }
func (set *IpCidrSet) AddIpCidr(ipCidr *net.IPNet) error { func (set *IpCidrSet) AddIpCidr(ipCidr netip.Prefix) (err error) {
return set.AddIpCidrForString(ipCidr.String()) var b netipx.IPSetBuilder
b.AddSet(set.Ranges)
b.AddPrefix(ipCidr)
set.Ranges, err = b.IPSet()
return
} }
func (set *IpCidrSet) IsContainForString(ipString string) bool { func (set *IpCidrSet) IsContainForString(ipString string) bool {
ip := ipToBigInt(net.ParseIP(ipString)) ip, err := netip.ParseAddr(ipString)
idx := sort.Search(len(set.Ranges), func(i int) bool { if err != nil {
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 {
return false return false
} }
return set.IsContainForString(ip.String()) return set.IsContain(ip)
} }
func (set *IpCidrSet) Merge() { func (set *IpCidrSet) IsContain(ip netip.Addr) bool {
for i := 0; i < len(set.Ranges)-1; i++ { if set.Ranges == nil {
if set.Ranges[i].End.Cmp(set.Ranges[i+1].Start) >= 0 { return false
set.Ranges[i].End = set.Ranges[i+1].End
set.Ranges = append(set.Ranges[:i+1], set.Ranges[i+2:]...)
i--
}
} }
return set.Ranges.Contains(ip.WithZone(""))
} }
func (set *IpCidrSet) Merge() {}

View File

@ -1,7 +1,9 @@
package cidr package cidr
import ( import (
"go4.org/netipx"
"testing" "testing"
"unsafe"
) )
func TestIpv4(t *testing.T) { func TestIpv4(t *testing.T) {
@ -97,8 +99,10 @@ func TestMerge(t *testing.T) {
set.AddIpCidrForString(test.ipCidr2) set.AddIpCidrForString(test.ipCidr2)
set.Merge() set.Merge()
if len(set.Ranges) != test.expectedLen { rangesLen := len(*(*[]netipx.IPRange)(unsafe.Pointer(set.Ranges)))
t.Errorf("Expected len: %v, got: %v", test.expectedLen, len(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/wk8/go-ordered-map/v2 v2.1.8
github.com/zhangyunhao116/fastrand v0.3.0 github.com/zhangyunhao116/fastrand v0.3.0
go.uber.org/automaxprocs v1.5.3 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/crypto v0.16.0
golang.org/x/exp v0.0.0-20231206192017-f3f8817b8deb golang.org/x/exp v0.0.0-20231206192017-f3f8817b8deb
golang.org/x/net v0.19.0 golang.org/x/net v0.19.0
@ -105,7 +106,6 @@ require (
github.com/yusufpapurcu/wmi v1.2.3 // indirect github.com/yusufpapurcu/wmi v1.2.3 // indirect
gitlab.com/yawning/bsaes.git v0.0.0-20190805113838-0a714cd429ec // indirect gitlab.com/yawning/bsaes.git v0.0.0-20190805113838-0a714cd429ec // indirect
go.uber.org/mock v0.3.0 // 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/mod v0.14.0 // indirect
golang.org/x/text v0.14.0 // indirect golang.org/x/text v0.14.0 // indirect
golang.org/x/time v0.5.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 { func (i *ipcidrStrategy) Match(metadata *C.Metadata) bool {
// return i.trie != nil && i.trie.IsContain(metadata.DstIP.AsSlice()) // 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 { func (i *ipcidrStrategy) Count() int {