:recycle: refactor geoip match logic (#1157)

README.md

@@ -37,6 +37,8 @@ This repo relies on the following third-party projects:
   - [seiflotfy/cuckoofilter](https://github.com/seiflotfy/cuckoofilter)
   - [google/starlark-go](https://github.com/google/starlark-go)
   - [jhump/protoreflect](https://github.com/jhump/protoreflect)
+  - [inetaf/netaddr](https://github.com/inetaf/netaddr)
+
 - For testing only:
   - [miekg/dns](https://github.com/miekg/dns)
   - [h12w/socks](https://github.com/h12w/socks)

app/router/condition_geoip.go

@@ -4,81 +4,45 @@
 package router
 
 import (
-	"encoding/binary"
-	"sort"
+	"strconv"
+
+	"inet.af/netaddr"
 
 	"github.com/v2fly/v2ray-core/v4/common/net"
 )
 
-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
+	ip4          *netaddr.IPSet
+	ip6          *netaddr.IPSet
 }
 
 func (m *GeoIPMatcher) Init(cidrs []*CIDR) error {
-	ip4Count := 0
-	ip6Count := 0
-
+	var builder4, builder6 netaddr.IPSetBuilder
 	for _, cidr := range cidrs {
-		ip := cidr.Ip
+		ip := net.IP(cidr.GetIp())
+		ipStr := ip.String() + "/" + strconv.Itoa(int(cidr.GetPrefix()))
+		ipPrefix, err := netaddr.ParseIPPrefix(ipStr)
+		if err != nil {
+			return err
+		}
 		switch len(ip) {
-		case 4:
-			ip4Count++
-		case 16:
-			ip6Count++
-		default:
-			return newError("unexpect ip length: ", len(ip))
+		case net.IPv4len:
+			builder4.AddPrefix(ipPrefix)
+		case net.IPv6len:
+			builder6.AddPrefix(ipPrefix)
 		}
 	}
 
-	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 cidrList {
-		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)
-		}
+	var err error
+	m.ip4, err = builder4.IPSet()
+	if err != nil {
+		return err
+	}
+	m.ip6, err = builder6.IPSet()
+	if err != nil {
+		return err
 	}
 
 	return nil
@@ -88,91 +52,35 @@ func (m *GeoIPMatcher) SetReverseMatch(isReverseMatch bool) {
 	m.reverseMatch = isReverseMatch
 }
 
-func (m *GeoIPMatcher) match4(ip uint32) bool {
-	if len(m.ip4) == 0 {
+func (m *GeoIPMatcher) match4(ip net.IP) bool {
+	nip, ok := netaddr.FromStdIP(ip)
+	if !ok {
 		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]
+	return m.ip4.Contains(nip)
 }
 
-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 {
+func (m *GeoIPMatcher) match6(ip net.IP) bool {
+	nip, ok := netaddr.FromStdIP(ip)
+	if !ok {
 		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]
+	return m.ip6.Contains(nip)
 }
 
 // Match returns true if the given ip is included by the GeoIP.
 func (m *GeoIPMatcher) Match(ip net.IP) bool {
+	isMatched := false
 	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
+	case net.IPv4len:
+		isMatched = m.match4(ip)
+	case net.IPv6len:
+		isMatched = m.match6(ip)
+	}
+	if m.reverseMatch {
+		return !isMatched
 	}
+	return isMatched
 }
 
 // GeoIPMatcherContainer is a container for GeoIPMatchers. It keeps unique copies of GeoIPMatcher by country code.
@@ -183,7 +91,7 @@ type GeoIPMatcherContainer struct {
 // Add adds a new GeoIP set into the container.
 // If the country code of GeoIP is not empty, GeoIPMatcherContainer will try to find an existing one, instead of adding a new one.
 func (c *GeoIPMatcherContainer) Add(geoip *GeoIP) (*GeoIPMatcher, error) {
-	if len(geoip.CountryCode) > 0 {
+	if geoip.CountryCode != "" {
 		for _, m := range c.matchers {
 			if m.countryCode == geoip.CountryCode && m.reverseMatch == geoip.ReverseMatch {
 				return m, nil
@@ -198,7 +106,7 @@ func (c *GeoIPMatcherContainer) Add(geoip *GeoIP) (*GeoIPMatcher, error) {
 	if err := m.Init(geoip.Cidr); err != nil {
 		return nil, err
 	}
-	if len(geoip.CountryCode) > 0 {
+	if geoip.CountryCode != "" {
 		c.matchers = append(c.matchers, m)
 	}
 	return m, nil

app/router/condition_geoip_test.go

@@ -63,7 +63,7 @@ func TestGeoIPMatcherContainer(t *testing.T) {
 }
 
 func TestGeoIPMatcher(t *testing.T) {
-	cidrList := router.CIDRList{
+	cidrList := []*router.CIDR{
 		{Ip: []byte{0, 0, 0, 0}, Prefix: 8},
 		{Ip: []byte{10, 0, 0, 0}, Prefix: 8},
 		{Ip: []byte{100, 64, 0, 0}, Prefix: 10},
@@ -135,7 +135,7 @@ func TestGeoIPMatcher(t *testing.T) {
 }
 
 func TestGeoIPReverseMatcher(t *testing.T) {
-	cidrList := router.CIDRList{
+	cidrList := []*router.CIDR{
 		{Ip: []byte{8, 8, 8, 8}, Prefix: 32},
 		{Ip: []byte{91, 108, 4, 0}, Prefix: 16},
 	}

app/router/config.go

@@ -9,45 +9,6 @@ import (
 	"github.com/v2fly/v2ray-core/v4/features/routing"
 )
 
-// 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 Rule struct {
 	Tag       string
 	Balancer  *Balancer

go.mod

@@ -16,14 +16,15 @@ require (
 	github.com/v2fly/BrowserBridge v0.0.0-20210430233438-0570fc1d7d08
 	github.com/v2fly/VSign v0.0.0-20201108000810-e2adc24bf848
 	github.com/v2fly/ss-bloomring v0.0.0-20210312155135-28617310f63e
-	go.starlark.net v0.0.0-20210602144842-1cdb82c9e17a
+	go.starlark.net v0.0.0-20210901212718-87f333178d59
 	golang.org/x/crypto v0.0.0-20210817164053-32db794688a5
-	golang.org/x/net v0.0.0-20210813160813-60bc85c4be6d
+	golang.org/x/net v0.0.0-20210903162142-ad29c8ab022f
 	golang.org/x/sync v0.0.0-20210220032951-036812b2e83c
-	golang.org/x/sys v0.0.0-20210820121016-41cdb8703e55
+	golang.org/x/sys v0.0.0-20210903071746-97244b99971b
 	google.golang.org/grpc v1.40.0
 	google.golang.org/protobuf v1.27.1
 	h12.io/socks v1.0.3
+	inet.af/netaddr v0.0.0-20210903134321-85fa6c94624e
 )
 
 require (
@@ -41,6 +42,8 @@ require (
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/riobard/go-bloom v0.0.0-20200614022211-cdc8013cb5b3 // indirect
 	github.com/xtaci/smux v1.5.15 // indirect
+	go4.org/intern v0.0.0-20210108033219-3eb7198706b2 // indirect
+	go4.org/unsafe/assume-no-moving-gc v0.0.0-20201222180813-1025295fd063 // indirect
 	golang.org/x/mod v0.4.2 // indirect
 	golang.org/x/text v0.3.6 // indirect
 	golang.org/x/tools v0.1.1 // indirect

go.sum

@@ -60,6 +60,7 @@ github.com/dgryski/go-metro v0.0.0-20200812162917-85c65e2d0165 h1:BS21ZUJ/B5X2UV
 github.com/dgryski/go-metro v0.0.0-20200812162917-85c65e2d0165/go.mod h1:c9O8+fpSOX1DM8cPNSkX/qsBWdkD4yd2dpciOWQjpBw=
 github.com/dgryski/go-sip13 v0.0.0-20181026042036-e10d5fee7954/go.mod h1:vAd38F8PWV+bWy6jNmig1y/TA+kYO4g3RSRF0IAv0no=
 github.com/dustin/go-humanize v1.0.0/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
+github.com/dvyukov/go-fuzz v0.0.0-20210103155950-6a8e9d1f2415/go.mod h1:11Gm+ccJnvAhCNLlf5+cS9KjtbaD5I5zaZpFMsTHWTw=
 github.com/ebfe/bcrypt_pbkdf v0.0.0-20140212075826-3c8d2dcb253a h1:YtdtTUN1iH97s+6PUjLnaiKSQj4oG1/EZ3N9bx6g4kU=
 github.com/ebfe/bcrypt_pbkdf v0.0.0-20140212075826-3c8d2dcb253a/go.mod h1:/CZpbhAusDOobpcb9yubw46kdYjq0zRC0Wpg9a9zFQM=
 github.com/envoyproxy/go-control-plane v0.9.0/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4=
@@ -342,12 +343,18 @@ go.opencensus.io v0.18.0/go.mod h1:vKdFvxhtzZ9onBp9VKHK8z/sRpBMnKAsufL7wlDrCOA=
 go.opencensus.io v0.21.0/go.mod h1:mSImk1erAIZhrmZN+AvHh14ztQfjbGwt4TtuofqLduU=
 go.opencensus.io v0.22.0/go.mod h1:+kGneAE2xo2IficOXnaByMWTGM9T73dGwxeWcUqIpI8=
 go.opentelemetry.io/proto/otlp v0.7.0/go.mod h1:PqfVotwruBrMGOCsRd/89rSnXhoiJIqeYNgFYFoEGnI=
-go.starlark.net v0.0.0-20210602144842-1cdb82c9e17a h1:wDtSCWGrX9tusypq2Qq9xzaA3Tf/+4D2KaWO+HQvGZE=
-go.starlark.net v0.0.0-20210602144842-1cdb82c9e17a/go.mod h1:t3mmBBPzAVvK0L0n1drDmrQsJ8FoIx4INCqVMTr/Zo0=
+go.starlark.net v0.0.0-20210901212718-87f333178d59 h1:F8ArBy9n1l7HE1JjzOIYqweEqoUlywy5+L3bR0tIa9g=
+go.starlark.net v0.0.0-20210901212718-87f333178d59/go.mod h1:t3mmBBPzAVvK0L0n1drDmrQsJ8FoIx4INCqVMTr/Zo0=
 go.uber.org/atomic v1.4.0/go.mod h1:gD2HeocX3+yG+ygLZcrzQJaqmWj9AIm7n08wl/qW/PE=
 go.uber.org/multierr v1.1.0/go.mod h1:wR5kodmAFQ0UK8QlbwjlSNy0Z68gJhDJUG5sjR94q/0=
 go.uber.org/zap v1.10.0/go.mod h1:vwi/ZaCAaUcBkycHslxD9B2zi4UTXhF60s6SWpuDF0Q=
+go4.org v0.0.0-20180809161055-417644f6feb5 h1:+hE86LblG4AyDgwMCLTE6FOlM9+qjHSYS+rKqxUVdsM=
 go4.org v0.0.0-20180809161055-417644f6feb5/go.mod h1:MkTOUMDaeVYJUOUsaDXIhWPZYa1yOyC1qaOBpL57BhE=
+go4.org/intern v0.0.0-20210108033219-3eb7198706b2 h1:VFTf+jjIgsldaz/Mr00VaCSswHJrI2hIjQygE/W4IMg=
+go4.org/intern v0.0.0-20210108033219-3eb7198706b2/go.mod h1:vLqJ+12kCw61iCWsPto0EOHhBS+o4rO5VIucbc9g2Cc=
+go4.org/unsafe/assume-no-moving-gc v0.0.0-20201222175341-b30ae309168e/go.mod h1:FftLjUGFEDu5k8lt0ddY+HcrH/qU/0qk+H8j9/nTl3E=
+go4.org/unsafe/assume-no-moving-gc v0.0.0-20201222180813-1025295fd063 h1:1tk03FUNpulq2cuWpXZWj649rwJpk0d20rxWiopKRmc=
+go4.org/unsafe/assume-no-moving-gc v0.0.0-20201222180813-1025295fd063/go.mod h1:FftLjUGFEDu5k8lt0ddY+HcrH/qU/0qk+H8j9/nTl3E=
 golang.org/x/build v0.0.0-20190111050920-041ab4dc3f9d/go.mod h1:OWs+y06UdEOHN4y+MfF/py+xQ/tYqIWW03b70/CG9Rw=
 golang.org/x/crypto v0.0.0-20180904163835-0709b304e793/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4=
 golang.org/x/crypto v0.0.0-20181029021203-45a5f77698d3/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4=
@@ -412,8 +419,8 @@ golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwY
 golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
 golang.org/x/net v0.0.0-20210405180319-a5a99cb37ef4/go.mod h1:p54w0d4576C0XHj96bSt6lcn1PtDYWL6XObtHCRCNQM=
 golang.org/x/net v0.0.0-20210428140749-89ef3d95e781/go.mod h1:OJAsFXCWl8Ukc7SiCT/9KSuxbyM7479/AVlXFRxuMCk=
-golang.org/x/net v0.0.0-20210813160813-60bc85c4be6d h1:LO7XpTYMwTqxjLcGWPijK3vRXg1aWdlNOVOHRq45d7c=
-golang.org/x/net v0.0.0-20210813160813-60bc85c4be6d/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=
+golang.org/x/net v0.0.0-20210903162142-ad29c8ab022f h1:w6wWR0H+nyVpbSAQbzVEIACVyr/h8l/BEkY6Sokc7Eg=
+golang.org/x/net v0.0.0-20210903162142-ad29c8ab022f/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=
 golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U=
 golang.org/x/oauth2 v0.0.0-20181017192945-9dcd33a902f4/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U=
 golang.org/x/oauth2 v0.0.0-20181203162652-d668ce993890/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U=
@@ -463,8 +470,8 @@ golang.org/x/sys v0.0.0-20210403161142-5e06dd20ab57/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/sys v0.0.0-20210423082822-04245dca01da/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210510120138-977fb7262007/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.0.0-20210820121016-41cdb8703e55 h1:rw6UNGRMfarCepjI8qOepea/SXwIBVfTKjztZ5gBbq4=
-golang.org/x/sys v0.0.0-20210820121016-41cdb8703e55/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20210903071746-97244b99971b h1:3Dq0eVHn0uaQJmPO+/aYPI/fRMqdrVDbu7MQcku54gg=
+golang.org/x/sys v0.0.0-20210903071746-97244b99971b/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.1-0.20180807135948-17ff2d5776d2/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
@@ -597,6 +604,8 @@ honnef.co/go/tools v0.0.0-20190418001031-e561f6794a2a/go.mod h1:rf3lG4BRIbNafJWh
 honnef.co/go/tools v0.0.0-20190523083050-ea95bdfd59fc/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
 honnef.co/go/tools v0.0.1-2019.2.3/go.mod h1:a3bituU0lyd329TUQxRnasdCoJDkEUEAqEt0JzvZhAg=
 honnef.co/go/tools v0.0.1-2020.1.4/go.mod h1:X/FiERA/W4tHapMX5mGpAtMSVEeEUOyHaw9vFzvIQ3k=
+inet.af/netaddr v0.0.0-20210903134321-85fa6c94624e h1:tvgqez5ZQoBBiBAGNU/fmJy247yB/7++kcLOEoMYup0=
+inet.af/netaddr v0.0.0-20210903134321-85fa6c94624e/go.mod h1:z0nx+Dh+7N7CC8V5ayHtHGpZpxLQZZxkIaaz6HN65Ls=
 rsc.io/binaryregexp v0.2.0/go.mod h1:qTv7/COck+e2FymRvadv62gMdZztPaShugOCi3I+8D8=
 sourcegraph.com/sourcegraph/go-diff v0.5.0/go.mod h1:kuch7UrkMzY0X+p9CRK03kfuPQ2zzQcaEFbx8wA8rck=
 sourcegraph.com/sqs/pbtypes v0.0.0-20180604144634-d3ebe8f20ae4/go.mod h1:ketZ/q3QxT9HOBeFhu6RdvsftgpsbFHBF5Cas6cDKZ0=