Geodata: Support automatically updating .dat files and hot reloading (#5992)

https://github.com/XTLS/Xray-core/pull/5992#issuecomment-4320551920

Usage: https://github.com/XTLS/Xray-core/pull/5992#issuecomment-4291168039

common/geodata/domain_registry.go

@@ -1,11 +1,59 @@
 package geodata
 
+import (
+	"context"
+	"sync"
+	"sync/atomic"
+
+	"github.com/xtls/xray-core/common/errors"
+)
+
 type DomainRegistry struct {
-	factory DomainMatcherFactory
+	mu       sync.Mutex
+	factory  DomainMatcherFactory
+	matchers []*DynamicDomainMatcher
 }
 
 func (r *DomainRegistry) BuildDomainMatcher(rules []*DomainRule) (DomainMatcher, error) {
-	return r.factory.BuildMatcher(rules)
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	m, err := r.factory.BuildMatcher(rules)
+	if err != nil {
+		return nil, err
+	}
+
+	d := NewDynamicDomainMatcher(rules, m)
+	r.matchers = append(r.matchers, d)
+	return d, nil
+}
+
+func (r *DomainRegistry) Reload() error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	errors.LogInfo(context.Background(), "reloading GeoSite data for ", len(r.matchers), " domain matcher(s)")
+
+	factory := newDomainMatcherFactory()
+	type reloadEntry struct {
+		dynamic *DynamicDomainMatcher
+		matcher DomainMatcher
+	}
+	reloaded := make([]reloadEntry, len(r.matchers))
+	for i, d := range r.matchers {
+		m, err := factory.BuildMatcher(d.rules)
+		if err != nil {
+			errors.LogErrorInner(context.Background(), err, "failed to reload GeoSite data for domain matcher ", i)
+			return err
+		}
+		reloaded[i] = reloadEntry{dynamic: d, matcher: m}
+	}
+	for _, entry := range reloaded {
+		entry.dynamic.Reload(entry.matcher)
+	}
+	r.factory = factory
+	errors.LogInfo(context.Background(), "reloaded GeoSite data for ", len(r.matchers), " domain matcher(s)")
+	return nil
 }
 
 func newDomainRegistry() *DomainRegistry {
@@ -15,3 +63,32 @@ func newDomainRegistry() *DomainRegistry {
 }
 
 var DomainReg = newDomainRegistry()
+
+type domainMatcherState struct {
+	matcher DomainMatcher
+}
+
+type DynamicDomainMatcher struct {
+	rules []*DomainRule
+	state atomic.Pointer[domainMatcherState]
+}
+
+// Match implements DomainMatcher.
+func (d *DynamicDomainMatcher) Match(input string) []uint32 {
+	return d.state.Load().matcher.Match(input)
+}
+
+// MatchAny implements DomainMatcher.
+func (d *DynamicDomainMatcher) MatchAny(input string) bool {
+	return d.state.Load().matcher.MatchAny(input)
+}
+
+func (d *DynamicDomainMatcher) Reload(newMatcher DomainMatcher) {
+	d.state.Store(&domainMatcherState{matcher: newMatcher})
+}
+
+func NewDynamicDomainMatcher(rules []*DomainRule, matcher DomainMatcher) *DynamicDomainMatcher {
+	d := &DynamicDomainMatcher{rules: rules}
+	d.Reload(matcher)
+	return d
+}

common/geodata/ip_matcher.go

@@ -1016,3 +1016,7 @@ func buildOptimizedIPMatcher(f *IPSetFactory, rules []*IPRule) (IPMatcher, error
 		return &HeuristicMultiIPMatcher{matchers: subs}, nil
 	}
 }
+
+func newIPSetFactory() *IPSetFactory {
+	return &IPSetFactory{shared: make(map[string]*IPSet)}
+}

common/geodata/ip_registry.go

@@ -1,17 +1,135 @@
 package geodata
 
+import (
+	"context"
+	"sync"
+	"sync/atomic"
+
+	"github.com/xtls/xray-core/common/errors"
+	"github.com/xtls/xray-core/common/net"
+)
+
 type IPRegistry struct {
+	mu           sync.Mutex
 	ipsetFactory *IPSetFactory
+	matchers     []*DynamicIPMatcher
 }
 
 func (r *IPRegistry) BuildIPMatcher(rules []*IPRule) (IPMatcher, error) {
-	return buildOptimizedIPMatcher(r.ipsetFactory, rules)
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	m, err := buildOptimizedIPMatcher(r.ipsetFactory, rules)
+	if err != nil {
+		return nil, err
+	}
+
+	d := NewDynamicIPMatcher(rules, m)
+	r.matchers = append(r.matchers, d)
+	return d, nil
+}
+
+func (r *IPRegistry) Reload() error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	errors.LogInfo(context.Background(), "reloading GeoIP data for ", len(r.matchers), " IP matcher(s)")
+
+	factory := newIPSetFactory()
+	type reloadEntry struct {
+		dynamic *DynamicIPMatcher
+		matcher IPMatcher
+	}
+	reloaded := make([]reloadEntry, len(r.matchers))
+	for i, d := range r.matchers {
+		m, err := buildOptimizedIPMatcher(factory, d.rules)
+		if err != nil {
+			errors.LogErrorInner(context.Background(), err, "failed to reload GeoIP data for IP matcher ", i)
+			return err
+		}
+		reloaded[i] = reloadEntry{dynamic: d, matcher: m}
+	}
+	for _, entry := range reloaded {
+		entry.dynamic.Reload(entry.matcher)
+	}
+	r.ipsetFactory = factory
+	errors.LogInfo(context.Background(), "reloaded GeoIP data for ", len(r.matchers), " IP matcher(s)")
+	return nil
 }
 
 func newIPRegistry() *IPRegistry {
 	return &IPRegistry{
-		ipsetFactory: &IPSetFactory{shared: make(map[string]*IPSet)},
+		ipsetFactory: newIPSetFactory(),
 	}
 }
 
 var IPReg = newIPRegistry()
+
+type ipMatcherState struct {
+	matcher IPMatcher
+}
+
+type DynamicIPMatcher struct {
+	rules      []*IPRule
+	state      atomic.Pointer[ipMatcherState]
+	mu         sync.Mutex
+	reverse    bool
+	reverseSet bool
+}
+
+// Match implements IPMatcher.
+func (d *DynamicIPMatcher) Match(ip net.IP) bool {
+	return d.state.Load().matcher.Match(ip)
+}
+
+// AnyMatch implements IPMatcher.
+func (d *DynamicIPMatcher) AnyMatch(ips []net.IP) bool {
+	return d.state.Load().matcher.AnyMatch(ips)
+}
+
+// Matches implements IPMatcher.
+func (d *DynamicIPMatcher) Matches(ips []net.IP) bool {
+	return d.state.Load().matcher.Matches(ips)
+}
+
+// FilterIPs implements IPMatcher.
+func (d *DynamicIPMatcher) FilterIPs(ips []net.IP) (matched []net.IP, unmatched []net.IP) {
+	return d.state.Load().matcher.FilterIPs(ips)
+}
+
+// ToggleReverse implements IPMatcher.
+func (d *DynamicIPMatcher) ToggleReverse() {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+
+	d.reverse = !d.reverse
+	d.state.Load().matcher.ToggleReverse()
+}
+
+// SetReverse implements IPMatcher.
+func (d *DynamicIPMatcher) SetReverse(reverse bool) {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+
+	d.reverse = reverse
+	d.reverseSet = true
+	d.state.Load().matcher.SetReverse(reverse)
+}
+
+func (d *DynamicIPMatcher) Reload(newMatcher IPMatcher) {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+
+	if d.reverseSet {
+		newMatcher.SetReverse(d.reverse)
+	} else if d.reverse {
+		newMatcher.ToggleReverse()
+	}
+	d.state.Store(&ipMatcherState{matcher: newMatcher})
+}
+
+func NewDynamicIPMatcher(rules []*IPRule, matcher IPMatcher) *DynamicIPMatcher {
+	d := &DynamicIPMatcher{rules: rules}
+	d.Reload(matcher)
+	return d
+}

common/platform/filesystem/file.go

@@ -1,6 +1,7 @@
 package filesystem
 
 import (
+	"errors"
 	"io"
 	"os"
 	"path/filepath"
@@ -26,11 +27,48 @@ func ReadFile(path string) ([]byte, error) {
 }
 
 func ReadAsset(file string) ([]byte, error) {
-	return ReadFile(platform.GetAssetLocation(file))
+	path, _, err := getAssetFileLocation(file)
+	if err != nil {
+		return nil, err
+	}
+	return ReadFile(path)
 }
 
 func OpenAsset(file string) (io.ReadCloser, error) {
-	return NewFileReader(platform.GetAssetLocation(file))
+	path, _, err := getAssetFileLocation(file)
+	if err != nil {
+		return nil, err
+	}
+	return NewFileReader(path)
+}
+
+func StatAsset(file string) (os.FileInfo, error) {
+	_, info, err := getAssetFileLocation(file)
+	return info, err
+}
+
+func ResolveAsset(file string) (string, error) {
+	path, _, err := getAssetFileLocation(file)
+	return path, err
+}
+
+func getAssetFileLocation(file string) (string, os.FileInfo, error) {
+	if !filepath.IsLocal(file) || file == "." {
+		return "", nil, errors.New("asset path must stay in asset directory")
+	}
+	local, err := filepath.Localize(file)
+	if err != nil {
+		return "", nil, err
+	}
+	path := platform.GetAssetLocation(local)
+	info, err := os.Stat(path)
+	if err != nil {
+		return "", nil, err
+	}
+	if !info.Mode().IsRegular() {
+		return "", nil, errors.New("asset is not a regular file")
+	}
+	return path, info, nil
 }
 
 func ReadCert(file string) ([]byte, error) {

common/platform/filesystem/file_test.go

@@ -0,0 +1,32 @@
+package filesystem_test
+
+import (
+	"path/filepath"
+	"testing"
+
+	. "github.com/xtls/xray-core/common/platform/filesystem"
+)
+
+func TestStatAssetRejectsInvalidPath(t *testing.T) {
+	for _, file := range []string{
+		"",
+		".",
+		"..",
+		"../geoip.dat",
+		"nested/..",
+		"nested/../geoip.dat",
+		"nested//geoip.dat",
+		"/geoip.dat",
+		"/tmp/geoip.dat",
+		`C:\geoip.dat`,
+		`C:geoip.dat`,
+		`\\server\share\geoip.dat`,
+		`nested\geoip.dat`,
+		`nested\..\geoip.dat`,
+		filepath.Join(t.TempDir(), "geoip.dat"),
+	} {
+		if _, err := StatAsset(file); err == nil {
+			t.Fatalf("expected error for %q", file)
+		}
+	}
+}