fix(app/router): fixed a bug in geoip matching with refactoring (#2489)

* Refactor the IP address matching with netipx library
* Add a regression test for previous bug

Fixes https://github.com/XTLS/Xray-core/issues/1933

---------

Co-authored-by: Loyalsoldier <[email protected]>

app/router/condition_geoip.go

@@ -1,81 +1,49 @@
 package router
 
 import (
-	"encoding/binary"
-	"sort"
+	"net/netip"
+	"strconv"
 
 	"github.com/xtls/xray-core/common/net"
+	"go4.org/netipx"
 )
 
-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          *netipx.IPSet
+	ip6          *netipx.IPSet
 }
 
 func (m *GeoIPMatcher) Init(cidrs []*CIDR) error {
-	ip4Count := 0
-	ip6Count := 0
+	var builder4, builder6 netipx.IPSetBuilder
 
 	for _, cidr := range cidrs {
-		ip := cidr.Ip
+		ip := net.IP(cidr.GetIp())
+		ipPrefixString := ip.String() + "/" + strconv.Itoa(int(cidr.GetPrefix()))
+		ipPrefix, err := netip.ParsePrefix(ipPrefixString)
+		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)
+	if ip4, err := builder4.IPSet(); err != nil {
+		return err
+	} else {
+		m.ip4 = ip4
+	}
 
-			m.ip6 = append(m.ip6, ip6)
-			m.prefix6 = append(m.prefix6, prefix)
-		}
+	if ip6, err := builder6.IPSet(); err != nil {
+		return err
+	} else {
+		m.ip6 = ip6
 	}
 
 	return nil
@@ -85,91 +53,37 @@ 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] {
+func (m *GeoIPMatcher) match4(ip net.IP) bool {
+	nip, ok := netipx.FromStdIP(ip)
+	if !ok {
 		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 {
-		return false
-	}
-
-	if less6(ip, m.ip6[0]) {
+func (m *GeoIPMatcher) match6(ip net.IP) bool {
+	nip, ok := netipx.FromStdIP(ip)
+	if !ok {
 		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.

app/router/condition_geoip_test.go

@@ -53,7 +53,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},
@@ -124,8 +124,40 @@ func TestGeoIPMatcher(t *testing.T) {
 	}
 }
 
+func TestGeoIPMatcherRegression(t *testing.T) {
+	cidrList := []*router.CIDR{
+		{Ip: []byte{98, 108, 20, 0}, Prefix: 22},
+		{Ip: []byte{98, 108, 20, 0}, Prefix: 23},
+	}
+
+	matcher := &router.GeoIPMatcher{}
+	common.Must(matcher.Init(cidrList))
+
+	testCases := []struct {
+		Input  string
+		Output bool
+	}{
+		{
+			Input:  "98.108.22.11",
+			Output: true,
+		},
+		{
+			Input:  "98.108.25.0",
+			Output: false,
+		},
+	}
+
+	for _, testCase := range testCases {
+		ip := net.ParseAddress(testCase.Input).IP()
+		actual := matcher.Match(ip)
+		if actual != testCase.Output {
+			t.Error("expect input", testCase.Input, "to be", testCase.Output, ", but actually", actual)
+		}
+	}
+}
+
 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/xtls/xray-core/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

@@ -47,6 +47,7 @@ require (
 	github.com/quic-go/qtls-go1-20 v0.3.3 // indirect
 	github.com/riobard/go-bloom v0.0.0-20200614022211-cdc8013cb5b3 // indirect
 	go.uber.org/atomic v1.11.0 // indirect
+	go4.org/netipx v0.0.0-20230824141953-6213f710f925 // indirect
 	golang.org/x/exp v0.0.0-20230725093048-515e97ebf090 // indirect
 	golang.org/x/mod v0.12.0 // indirect
 	golang.org/x/text v0.12.0 // indirect

go.sum

@@ -173,6 +173,8 @@ go.opencensus.io v0.18.0/go.mod h1:vKdFvxhtzZ9onBp9VKHK8z/sRpBMnKAsufL7wlDrCOA=
 go.uber.org/atomic v1.11.0 h1:ZvwS0R+56ePWxUNi+Atn9dWONBPp/AUETXlHW0DxSjE=
 go.uber.org/atomic v1.11.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0=
 go4.org v0.0.0-20180809161055-417644f6feb5/go.mod h1:MkTOUMDaeVYJUOUsaDXIhWPZYa1yOyC1qaOBpL57BhE=
+go4.org/netipx v0.0.0-20230824141953-6213f710f925 h1:eeQDDVKFkx0g4Hyy8pHgmZaK0EqB4SD6rvKbUdN3ziQ=
+go4.org/netipx v0.0.0-20230824141953-6213f710f925/go.mod h1:PLyyIXexvUFg3Owu6p/WfdlivPbZJsZdgWZlrGope/Y=
 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-20181030102418-4d3f4d9ffa16/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=