feat: add config dir as extern geosite search dir. fix #41

cmd/run.go

@@ -74,17 +74,17 @@ var (
 			logrus.SetLevel(log.Level)
 
 			log.Infof("Include config files: [%v]", strings.Join(includes, ", "))
-			if err := Run(log, conf); err != nil {
+			if err := Run(log, conf, []string{filepath.Dir(cfgFile)}); err != nil {
 				logrus.Fatalln(err)
 			}
 		},
 	}
 )
 
-func Run(log *logrus.Logger, conf *config.Config) (err error) {
+func Run(log *logrus.Logger, conf *config.Config, externGeoDataDirs []string) (err error) {
 
 	// New ControlPlane.
-	c, err := newControlPlane(log, nil, nil, conf)
+	c, err := newControlPlane(log, nil, nil, conf, externGeoDataDirs)
 	if err != nil {
 		return err
 	}
@@ -172,13 +172,13 @@ loop:
 			obj := c.EjectBpf()
 			dnsCache := c.CloneDnsCache()
 			log.Warnln("[Reload] Load new control plane")
-			newC, err := newControlPlane(log, obj, dnsCache, newConf)
+			newC, err := newControlPlane(log, obj, dnsCache, newConf, externGeoDataDirs)
 			if err != nil {
 				log.WithFields(logrus.Fields{
 					"err": err,
 				}).Errorln("[Reload] Failed to reload; try to roll back configuration")
 				// Load last config back.
-				newC, err = newControlPlane(log, obj, dnsCache, conf)
+				newC, err = newControlPlane(log, obj, dnsCache, conf, externGeoDataDirs)
 				if err != nil {
 					sdnotify.Stopping()
 					obj.Close()
@@ -216,7 +216,7 @@ loop:
 	return nil
 }
 
-func newControlPlane(log *logrus.Logger, bpf interface{}, dnsCache map[string]*control.DnsCache, conf *config.Config) (c *control.ControlPlane, err error) {
+func newControlPlane(log *logrus.Logger, bpf interface{}, dnsCache map[string]*control.DnsCache, conf *config.Config, externGeoDataDirs []string) (c *control.ControlPlane, err error) {
 	// Deep copy to prevent modification.
 	conf = deepcopy.Copy(conf).(*config.Config)
 
@@ -277,6 +277,7 @@ func newControlPlane(log *logrus.Logger, bpf interface{}, dnsCache map[string]*c
 		&conf.Routing,
 		&conf.Global,
 		&conf.Dns,
+		externGeoDataDirs,
 	)
 	if err != nil {
 		return nil, err

common/assets/assets.go

@@ -31,18 +31,21 @@ type CacheItem struct {
 type LocationFinder struct {
 	mu         sync.Mutex
 	m          map[string]CacheItem
+	externDirs []string
 }
 
-func NewLocationFinder() *LocationFinder {
+func NewLocationFinder(externDirPath []string) *LocationFinder {
 	return &LocationFinder{
 		mu:         sync.Mutex{},
 		m:          map[string]CacheItem{},
+		externDirs: externDirPath,
 	}
 }
 
 func (c *LocationFinder) GetLocationAsset(log *logrus.Logger, filename string) (path string, err error) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
+	// Search cache.
 	if item, ok := c.m[filename]; ok && time.Now().Before(item.CacheDeadline) {
 		return item.Path, nil
 	}
@@ -63,63 +66,55 @@ func (c *LocationFinder) GetLocationAsset(log *logrus.Logger, filename string) (
 		}
 	}()
 
+	// Search dirs.
+	var searchDirs []string
 	folder := "dae"
 	location := os.Getenv("DAE_LOCATION_ASSET")
 	// check if DAE_LOCATION_ASSET is set
 	if location != "" {
 		// add DAE_LOCATION_ASSET to search path
-		searchPaths := []string{
+		searchDirs = []string{
-			filepath.Join(location, filename),
+			location,
 		}
 		// additional paths for non windows platforms
 		if runtime.GOOS != "windows" {
-			searchPaths = append(
+			searchDirs = append(
-				searchPaths,
+				searchDirs,
-				filepath.Join("/usr/local/share", folder, filename),
+				filepath.Join("/usr/local/share", folder),
-				filepath.Join("/usr/share", folder, filename),
+				filepath.Join("/usr/share", folder),
 			)
 		}
-		searchDirs := make([]string, len(searchPaths))
+		searchDirs = append(searchDirs, c.externDirs...)
+	} else {
+		if runtime.GOOS != "windows" {
+			// Search XDG data directories on non windows platform
+			searchDirs = append([]string{xdg.DataHome}, xdg.DataDirs...)
 			for i := range searchDirs {
-			searchDirs[i] = filepath.Dir(searchPaths[i])
+				searchDirs[i] = filepath.Join(searchDirs[i], folder)
+			}
+			searchDirs = append(searchDirs, c.externDirs...)
+		} else {
+			searchDirs = append([]string{}, c.externDirs...)
+			// fallback to the old behavior of using only current dir on Windows
+			pwd := "./"
+			if absPath, e := filepath.Abs(pwd); e == nil {
+				pwd = absPath
+			}
+			searchDirs = append(searchDirs, pwd)
+		}
 	}
 	log.Debugf(`Search "%v" in [%v]`, filename, strings.Join(searchDirs, ", "))
-		for _, searchPath := range searchPaths {
+	for _, searchDir := range searchDirs {
-			if _, err = os.Stat(searchPath); err != nil && errors.Is(err, fs.ErrNotExist) {
+		searchPath := filepath.Join(searchDir, filename)
+		if _, err = os.Stat(searchPath); err != nil {
+			if errors.Is(err, fs.ErrNotExist) {
 				continue
 			}
+			return "", err
+		}
 		log.Debugf(`Found "%v" at %v`, filename, searchPath)
 		// return the first path that exists
 		return searchPath, nil
 	}
 	return "", fmt.Errorf("%v: %w in [%v]", filename, os.ErrNotExist, strings.Join(searchDirs, ", "))
-	} else {
-		if runtime.GOOS != "windows" {
-			// search XDG data directories on non windows platform
-			searchDirs := append([]string{xdg.DataHome}, xdg.DataDirs...)
-			for i := range searchDirs {
-				searchDirs[i] = filepath.Join(searchDirs[i], folder)
 }
-			log.Debugf(`Search "%v" in [%v]`, filename, strings.Join(searchDirs, ", "))
-			relpath := filepath.Join(folder, filename)
-			fullpath, err := xdg.SearchDataFile(relpath)
-			if err != nil {
-				return "", fmt.Errorf("%v: %w in [%v]", filename, os.ErrNotExist, strings.Join(searchDirs, ", "))
-			}
-			log.Debugf(`Found "%v" at %v`, filename, fullpath)
-			return fullpath, nil
-		} else {
-			// fallback to the old behavior of using only current dir on Windows
-			path := filepath.Join("./", filename)
-			if absPath, e := filepath.Abs(path); e == nil {
-				path = absPath
-			}
-			if _, err := os.Stat(path); errors.Is(err, os.ErrNotExist) {
-				return "", fmt.Errorf("%v: %w in %v", filename, os.ErrNotExist, path)
-			}
-			return path, nil
-		}
-	}
-}
-
-var DefaultLocationFinder = NewLocationFinder()

component/dns/dns.go

@@ -8,6 +8,7 @@ package dns
 import (
 	"fmt"
 	"github.com/daeuniverse/dae/common"
+	"github.com/daeuniverse/dae/common/assets"
 	"github.com/daeuniverse/dae/common/consts"
 	"github.com/daeuniverse/dae/component/routing"
 	"github.com/daeuniverse/dae/config"
@@ -30,12 +31,14 @@ type Dns struct {
 }
 
 type NewOption struct {
+	Logger                *logrus.Logger
+	LocationFinder        *assets.LocationFinder
 	UpstreamReadyCallback func(dnsUpstream *Upstream) (err error)
 }
 
-func New(log *logrus.Logger, dns *config.Dns, opt *NewOption) (s *Dns, err error) {
+func New(dns *config.Dns, opt *NewOption) (s *Dns, err error) {
 	s = &Dns{
-		log: log,
+		log: opt.Logger,
 		upstream2Index: map[*Upstream]int{
 			nil: int(consts.DnsRequestOutboundIndex_AsIs),
 		},
@@ -79,21 +82,21 @@ func New(log *logrus.Logger, dns *config.Dns, opt *NewOption) (s *Dns, err error
 	}
 	// Optimize routings.
 	if dns.Routing.Request.Rules, err = routing.ApplyRulesOptimizers(dns.Routing.Request.Rules,
-		&routing.DatReaderOptimizer{Logger: log},
+		&routing.DatReaderOptimizer{Logger: opt.Logger, LocationFinder: opt.LocationFinder},
 		&routing.MergeAndSortRulesOptimizer{},
 		&routing.DeduplicateParamsOptimizer{},
 	); err != nil {
 		return nil, err
 	}
 	if dns.Routing.Response.Rules, err = routing.ApplyRulesOptimizers(dns.Routing.Response.Rules,
-		&routing.DatReaderOptimizer{Logger: log},
+		&routing.DatReaderOptimizer{Logger: opt.Logger, LocationFinder: opt.LocationFinder},
 		&routing.MergeAndSortRulesOptimizer{},
 		&routing.DeduplicateParamsOptimizer{},
 	); err != nil {
 		return nil, err
 	}
 	// Parse request routing.
-	reqMatcherBuilder, err := NewRequestMatcherBuilder(log, dns.Routing.Request.Rules, upstreamName2Id, dns.Routing.Request.Fallback)
+	reqMatcherBuilder, err := NewRequestMatcherBuilder(opt.Logger, dns.Routing.Request.Rules, upstreamName2Id, dns.Routing.Request.Fallback)
 	if err != nil {
 		return nil, fmt.Errorf("failed to build DNS request routing: %w", err)
 	}
@@ -102,7 +105,7 @@ func New(log *logrus.Logger, dns *config.Dns, opt *NewOption) (s *Dns, err error
 		return nil, fmt.Errorf("failed to build DNS request routing: %w", err)
 	}
 	// Parse response routing.
-	respMatcherBuilder, err := NewResponseMatcherBuilder(log, dns.Routing.Response.Rules, upstreamName2Id, dns.Routing.Response.Fallback)
+	respMatcherBuilder, err := NewResponseMatcherBuilder(opt.Logger, dns.Routing.Response.Rules, upstreamName2Id, dns.Routing.Response.Fallback)
 	if err != nil {
 		return nil, fmt.Errorf("failed to build DNS response routing: %w", err)
 	}

component/routing/optimizer.go

@@ -7,12 +7,12 @@ package routing
 
 import (
 	"fmt"
-	"github.com/mohae/deepcopy"
-	"github.com/sirupsen/logrus"
 	"github.com/daeuniverse/dae/common/assets"
 	"github.com/daeuniverse/dae/common/consts"
 	"github.com/daeuniverse/dae/pkg/config_parser"
 	"github.com/daeuniverse/dae/pkg/geodata"
+	"github.com/mohae/deepcopy"
+	"github.com/sirupsen/logrus"
 	"net/netip"
 	"sort"
 	"strings"
@@ -153,6 +153,7 @@ func (o *DeduplicateParamsOptimizer) Optimize(rules []*config_parser.RoutingRule
 }
 
 type DatReaderOptimizer struct {
+	LocationFinder *assets.LocationFinder
 	Logger         *logrus.Logger
 }
 
@@ -160,7 +161,7 @@ func (o *DatReaderOptimizer) loadGeoSite(filename string, code string) (params [
 	if !strings.HasSuffix(filename, ".dat") {
 		filename += ".dat"
 	}
-	filePath, err := assets.DefaultLocationFinder.GetLocationAsset(o.Logger, filename)
+	filePath, err := o.LocationFinder.GetLocationAsset(o.Logger, filename)
 	if err != nil {
 		o.Logger.Debugf("Failed to read geosite \"%v:%v\": %v", filename, code, err)
 		return nil, err
@@ -205,7 +206,7 @@ func (o *DatReaderOptimizer) loadGeoIp(filename string, code string) (params []*
 	if !strings.HasSuffix(filename, ".dat") {
 		filename += ".dat"
 	}
-	filePath, err := assets.DefaultLocationFinder.GetLocationAsset(o.Logger, filename)
+	filePath, err := o.LocationFinder.GetLocationAsset(o.Logger, filename)
 	if err != nil {
 		o.Logger.Debugf("Failed to read geoip \"%v:%v\": %v", filename, code, err)
 		return nil, err

control/control_plane.go

@@ -12,6 +12,7 @@ import (
 	"github.com/cilium/ebpf"
 	"github.com/cilium/ebpf/rlimit"
 	"github.com/daeuniverse/dae/common"
+	"github.com/daeuniverse/dae/common/assets"
 	"github.com/daeuniverse/dae/common/consts"
 	"github.com/daeuniverse/dae/common/netutils"
 	"github.com/daeuniverse/dae/component/dns"
@@ -71,6 +72,7 @@ func NewControlPlane(
 	routingA *config.Routing,
 	global *config.Global,
 	dnsConfig *config.Dns,
+	externGeoDataDirs []string,
 ) (*ControlPlane, error) {
 	var err error
 
@@ -290,10 +292,11 @@ func NewControlPlane(
 		outboundId2Name[uint8(i)] = o.Name
 	}
 	// Apply rules optimizers.
+	locationFinder := assets.NewLocationFinder(externGeoDataDirs)
 	var rules []*config_parser.RoutingRule
 	if rules, err = routing.ApplyRulesOptimizers(routingA.Rules,
 		&routing.AliasOptimizer{},
-		&routing.DatReaderOptimizer{Logger: log},
+		&routing.DatReaderOptimizer{Logger: log, LocationFinder: locationFinder},
 		&routing.MergeAndSortRulesOptimizer{},
 		&routing.DeduplicateParamsOptimizer{},
 	); err != nil {
@@ -348,7 +351,9 @@ func NewControlPlane(
 	}()
 
 	/// DNS upstream.
-	dnsUpstream, err := dns.New(log, dnsConfig, &dns.NewOption{
+	dnsUpstream, err := dns.New(dnsConfig, &dns.NewOption{
+		Logger:                log,
+		LocationFinder:        locationFinder,
 		UpstreamReadyCallback: plane.dnsUpstreamReadyCallback,
 	})
 	if err != nil {