From e860497c0c92bfc765fd0caa4c4b61b85c23ee81 Mon Sep 17 00:00:00 2001 From: wwqgtxx Date: Sat, 13 Jan 2024 11:44:02 +0800 Subject: [PATCH] chore: cleanup IPSet code --- component/geodata/router/condition.go | 206 +++----------------------- dns/filters.go | 2 +- rules/common/geoip.go | 2 +- 3 files changed, 20 insertions(+), 190 deletions(-) diff --git a/component/geodata/router/condition.go b/component/geodata/router/condition.go index 73bc88d4f..c2ac80710 100644 --- a/component/geodata/router/condition.go +++ b/component/geodata/router/condition.go @@ -1,12 +1,11 @@ package router import ( - "encoding/binary" "fmt" - "net" - "sort" + "net/netip" "strings" + "github.com/metacubex/mihomo/component/cidr" "github.com/metacubex/mihomo/component/geodata/strmatcher" "github.com/metacubex/mihomo/component/trie" ) @@ -121,115 +120,24 @@ func (m *v2rayDomainMatcher) ApplyDomain(domain string) bool { return isMatched } -// CIDRList is an alias of []*CIDR to provide sort.Interface. -type CIDRList []*CIDR - -// Len implements sort.Interface. -func (l *CIDRList) Len() int { - return len(*l) -} - -// Less implements sort.Interface. -func (l *CIDRList) Less(i int, j int) bool { - ci := (*l)[i] - cj := (*l)[j] - - if len(ci.Ip) < len(cj.Ip) { - return true - } - - if len(ci.Ip) > len(cj.Ip) { - return false - } - - for k := 0; k < len(ci.Ip); k++ { - if ci.Ip[k] < cj.Ip[k] { - return true - } - - if ci.Ip[k] > cj.Ip[k] { - return false - } - } - - return ci.Prefix < cj.Prefix -} - -// Swap implements sort.Interface. -func (l *CIDRList) Swap(i int, j int) { - (*l)[i], (*l)[j] = (*l)[j], (*l)[i] -} - -type ipv6 struct { - a uint64 - b uint64 -} - type GeoIPMatcher struct { countryCode string reverseMatch bool - ip4 []uint32 - prefix4 []uint8 - ip6 []ipv6 - prefix6 []uint8 -} - -func normalize4(ip uint32, prefix uint8) uint32 { - return (ip >> (32 - prefix)) << (32 - prefix) -} - -func normalize6(ip ipv6, prefix uint8) ipv6 { - if prefix <= 64 { - ip.a = (ip.a >> (64 - prefix)) << (64 - prefix) - ip.b = 0 - } else { - ip.b = (ip.b >> (128 - prefix)) << (128 - prefix) - } - return ip + cidrSet *cidr.IpCidrSet } func (m *GeoIPMatcher) Init(cidrs []*CIDR) error { - ip4Count := 0 - ip6Count := 0 - for _, cidr := range cidrs { - ip := cidr.Ip - switch len(ip) { - case 4: - ip4Count++ - case 16: - ip6Count++ - default: - return fmt.Errorf("unexpect ip length: %d", len(ip)) - } - } - - cidrList := CIDRList(cidrs) - sort.Sort(&cidrList) - - m.ip4 = make([]uint32, 0, ip4Count) - m.prefix4 = make([]uint8, 0, ip4Count) - m.ip6 = make([]ipv6, 0, ip6Count) - m.prefix6 = make([]uint8, 0, ip6Count) - - for _, cidr := range cidrs { - ip := cidr.Ip - prefix := uint8(cidr.Prefix) - switch len(ip) { - case 4: - m.ip4 = append(m.ip4, normalize4(binary.BigEndian.Uint32(ip), prefix)) - m.prefix4 = append(m.prefix4, prefix) - case 16: - ip6 := ipv6{ - a: binary.BigEndian.Uint64(ip[0:8]), - b: binary.BigEndian.Uint64(ip[8:16]), - } - ip6 = normalize6(ip6, prefix) - - m.ip6 = append(m.ip6, ip6) - m.prefix6 = append(m.prefix6, prefix) + addr, ok := netip.AddrFromSlice(cidr.Ip) + if !ok { + return fmt.Errorf("error when loading GeoIP: invalid IP: %s", cidr.Ip) + } + err := m.cidrSet.AddIpCidr(netip.PrefixFrom(addr, int(cidr.Prefix))) + if err != nil { + return fmt.Errorf("error when loading GeoIP: %w", err) } } + m.cidrSet.Merge() return nil } @@ -238,91 +146,13 @@ func (m *GeoIPMatcher) SetReverseMatch(isReverseMatch bool) { m.reverseMatch = isReverseMatch } -func (m *GeoIPMatcher) match4(ip uint32) bool { - if len(m.ip4) == 0 { - return false - } - - if ip < m.ip4[0] { - return false - } - - size := uint32(len(m.ip4)) - l := uint32(0) - r := size - for l < r { - x := ((l + r) >> 1) - if ip < m.ip4[x] { - r = x - continue - } - - nip := normalize4(ip, m.prefix4[x]) - if nip == m.ip4[x] { - return true - } - - l = x + 1 - } - - return l > 0 && normalize4(ip, m.prefix4[l-1]) == m.ip4[l-1] -} - -func less6(a ipv6, b ipv6) bool { - return a.a < b.a || (a.a == b.a && a.b < b.b) -} - -func (m *GeoIPMatcher) match6(ip ipv6) bool { - if len(m.ip6) == 0 { - return false - } - - if less6(ip, m.ip6[0]) { - return false - } - - size := uint32(len(m.ip6)) - l := uint32(0) - r := size - for l < r { - x := (l + r) / 2 - if less6(ip, m.ip6[x]) { - r = x - continue - } - - if normalize6(ip, m.prefix6[x]) == m.ip6[x] { - return true - } - - l = x + 1 - } - - return l > 0 && normalize6(ip, m.prefix6[l-1]) == m.ip6[l-1] -} - // Match returns true if the given ip is included by the GeoIP. -func (m *GeoIPMatcher) Match(ip net.IP) bool { - switch len(ip) { - case 4: - if m.reverseMatch { - return !m.match4(binary.BigEndian.Uint32(ip)) - } - return m.match4(binary.BigEndian.Uint32(ip)) - case 16: - if m.reverseMatch { - return !m.match6(ipv6{ - a: binary.BigEndian.Uint64(ip[0:8]), - b: binary.BigEndian.Uint64(ip[8:16]), - }) - } - return m.match6(ipv6{ - a: binary.BigEndian.Uint64(ip[0:8]), - b: binary.BigEndian.Uint64(ip[8:16]), - }) - default: - return false +func (m *GeoIPMatcher) Match(ip netip.Addr) bool { + match := m.cidrSet.IsContain(ip) + if m.reverseMatch { + return !match } + return match } // GeoIPMatcherContainer is a container for GeoIPMatchers. It keeps unique copies of GeoIPMatcher by country code. @@ -344,6 +174,7 @@ func (c *GeoIPMatcherContainer) Add(geoip *GeoIP) (*GeoIPMatcher, error) { m := &GeoIPMatcher{ countryCode: geoip.CountryCode, reverseMatch: geoip.ReverseMatch, + cidrSet: cidr.NewIpCidrSet(), } if err := m.Init(geoip.Cidr); err != nil { return nil, err @@ -369,8 +200,7 @@ func NewGeoIPMatcher(geoip *GeoIP) (*GeoIPMatcher, error) { return matcher, nil } -func (m *MultiGeoIPMatcher) ApplyIp(ip net.IP) bool { - +func (m *MultiGeoIPMatcher) ApplyIp(ip netip.Addr) bool { for _, matcher := range m.matchers { if matcher.Match(ip) { return true diff --git a/dns/filters.go b/dns/filters.go index 46244c35e..d8633e8bb 100644 --- a/dns/filters.go +++ b/dns/filters.go @@ -41,7 +41,7 @@ func (gf *geoipFilter) Match(ip netip.Addr) bool { return false } } - return !geoIPMatcher.Match(ip.AsSlice()) + return !geoIPMatcher.Match(ip) } type ipnetFilter struct { diff --git a/rules/common/geoip.go b/rules/common/geoip.go index 3a29fae41..ebca1d162 100644 --- a/rules/common/geoip.go +++ b/rules/common/geoip.go @@ -48,7 +48,7 @@ func (g *GEOIP) Match(metadata *C.Metadata) (bool, string) { } return false, g.adapter } - return g.geoIPMatcher.Match(ip.AsSlice()), g.adapter + return g.geoIPMatcher.Match(ip), g.adapter } func (g *GEOIP) Adapter() string {