Feat: memory-efficient geodata decoder (#319)

.gitignore

@@ -19,3 +19,5 @@
 *.tar.gz
 *.crt
 *.key
+*.dat
+trojan-go

common/common.go

@@ -3,9 +3,10 @@ package common
 import (
 	"crypto/sha256"
 	"fmt"
-	"log"
 	"os"
 	"path/filepath"
+
+	"github.com/p4gefau1t/trojan-go/log"
 )
 
 type Runnable interface {
@@ -31,3 +32,11 @@ func GetProgramDir() string {
 	}
 	return dir
 }
+
+func GetAssetLocation(file string) string {
+	if loc := os.Getenv("TROJAN_GO_LOCATION_ASSET"); loc != "" {
+		log.Debugf("env set: TROJAN_GO_LOCATION_ASSET=%s", loc)
+		return filepath.Join(loc, file)
+	}
+	return filepath.Join(GetProgramDir(), file)
+}

common/geodata/cache.go

@@ -0,0 +1,144 @@
+package geodata
+
+import (
+	"io/ioutil"
+	"strings"
+
+	v2router "github.com/v2fly/v2ray-core/v4/app/router"
+	"google.golang.org/protobuf/proto"
+
+	"github.com/p4gefau1t/trojan-go/common"
+	"github.com/p4gefau1t/trojan-go/log"
+)
+
+type GeoIPCache map[string]*v2router.GeoIP
+
+func (g GeoIPCache) Has(key string) bool {
+	return !(g.Get(key) == nil)
+}
+
+func (g GeoIPCache) Get(key string) *v2router.GeoIP {
+	if g == nil {
+		return nil
+	}
+	return g[key]
+}
+
+func (g GeoIPCache) Set(key string, value *v2router.GeoIP) {
+	if g == nil {
+		g = make(map[string]*v2router.GeoIP)
+	}
+	g[key] = value
+}
+
+func (g GeoIPCache) Unmarshal(filename, code string) (*v2router.GeoIP, error) {
+	asset := common.GetAssetLocation(filename)
+	idx := strings.ToLower(asset + ":" + code)
+	if g.Has(idx) {
+		log.Debugf("geoip cache HIT: %s -> %s", code, idx)
+		return g.Get(idx), nil
+	}
+
+	geoipBytes, err := Decode(asset, code)
+	switch err {
+	case nil:
+		var geoip v2router.GeoIP
+		if err := proto.Unmarshal(geoipBytes, &geoip); err != nil {
+			return nil, err
+		}
+		g.Set(idx, &geoip)
+		return &geoip, nil
+
+	case ErrCodeNotFound:
+		return nil, common.NewError("country code " + code + " not found in " + filename)
+
+	case ErrFailedToReadBytes, ErrFailedToReadExpectedLenBytes,
+		ErrInvalidGeodataFile, ErrInvalidGeodataVarintLength:
+		log.Warnf("failed to decode geoip file: %s, fallback to the original ReadFile method", filename)
+		geoipBytes, err = ioutil.ReadFile(asset)
+		if err != nil {
+			return nil, err
+		}
+		var geoipList v2router.GeoIPList
+		if err := proto.Unmarshal(geoipBytes, &geoipList); err != nil {
+			return nil, err
+		}
+		for _, geoip := range geoipList.GetEntry() {
+			if strings.EqualFold(code, geoip.GetCountryCode()) {
+				g.Set(idx, geoip)
+				return geoip, nil
+			}
+		}
+
+	default:
+		return nil, err
+	}
+
+	return nil, common.NewError("country code " + code + " not found in " + filename)
+}
+
+type GeoSiteCache map[string]*v2router.GeoSite
+
+func (g GeoSiteCache) Has(key string) bool {
+	return !(g.Get(key) == nil)
+}
+
+func (g GeoSiteCache) Get(key string) *v2router.GeoSite {
+	if g == nil {
+		return nil
+	}
+	return g[key]
+}
+
+func (g GeoSiteCache) Set(key string, value *v2router.GeoSite) {
+	if g == nil {
+		g = make(map[string]*v2router.GeoSite)
+	}
+	g[key] = value
+}
+
+func (g GeoSiteCache) Unmarshal(filename, code string) (*v2router.GeoSite, error) {
+	asset := common.GetAssetLocation(filename)
+	idx := strings.ToLower(asset + ":" + code)
+	if g.Has(idx) {
+		log.Debugf("geosite cache HIT: %s -> %s", code, idx)
+		return g.Get(idx), nil
+	}
+
+	geositeBytes, err := Decode(asset, code)
+	switch err {
+	case nil:
+		var geosite v2router.GeoSite
+		if err := proto.Unmarshal(geositeBytes, &geosite); err != nil {
+			return nil, err
+		}
+		g.Set(idx, &geosite)
+		return &geosite, nil
+
+	case ErrCodeNotFound:
+		return nil, common.NewError("list " + code + " not found in " + filename)
+
+	case ErrFailedToReadBytes, ErrFailedToReadExpectedLenBytes,
+		ErrInvalidGeodataFile, ErrInvalidGeodataVarintLength:
+		log.Warnf("failed to decode geoip file: %s, fallback to the original ReadFile method", filename)
+		geositeBytes, err = ioutil.ReadFile(asset)
+		if err != nil {
+			return nil, err
+		}
+		var geositeList v2router.GeoSiteList
+		if err := proto.Unmarshal(geositeBytes, &geositeList); err != nil {
+			return nil, err
+		}
+		for _, geosite := range geositeList.GetEntry() {
+			if strings.EqualFold(code, geosite.GetCountryCode()) {
+				g.Set(idx, geosite)
+				return geosite, nil
+			}
+		}
+
+	default:
+		return nil, err
+	}
+
+	return nil, common.NewError("list " + code + " not found in " + filename)
+}

common/geodata/decode.go

@@ -0,0 +1,114 @@
+// Package geodata includes utilities to decode and parse the geoip & geosite dat files.
+//
+// It relies on the proto structure of GeoIP, GeoIPList, GeoSite and GeoSiteList in
+// github.com/v2fly/v2ray-core/v4/app/router/config.proto to comply with following rules:
+//
+// 1. GeoIPList and GeoSiteList cannot be changed
+// 2. The country_code in GeoIP and GeoSite must be
+//    a length-delimited `string`(wired type) and has field_number set to 1
+//
+package geodata
+
+import (
+	"errors"
+	"io"
+	"os"
+	"strings"
+
+	"google.golang.org/protobuf/encoding/protowire"
+)
+
+var (
+	ErrFailedToReadBytes            = errors.New("failed to read bytes")
+	ErrFailedToReadExpectedLenBytes = errors.New("failed to read expected length of bytes")
+	ErrInvalidGeodataFile           = errors.New("invalid geodata file")
+	ErrInvalidGeodataVarintLength   = errors.New("invalid geodata varint length")
+	ErrCodeNotFound                 = errors.New("code not found")
+)
+
+func EmitBytes(f io.ReadSeeker, code string) ([]byte, error) {
+	count := 1
+	isInner := false
+	tempContainer := make([]byte, 0, 5)
+
+	var result []byte
+	var advancedN uint64 = 1
+	var geoDataVarintLength, codeVarintLength, varintLenByteLen uint64 = 0, 0, 0
+
+Loop:
+	for {
+		container := make([]byte, advancedN)
+		bytesRead, err := f.Read(container)
+		if err == io.EOF {
+			return nil, ErrCodeNotFound
+		}
+		if err != nil {
+			return nil, ErrFailedToReadBytes
+		}
+		if bytesRead != len(container) {
+			return nil, ErrFailedToReadExpectedLenBytes
+		}
+
+		switch count {
+		case 1, 3: // data type ((field_number << 3) | wire_type)
+			if container[0] != 10 { // byte `0A` equals to `10` in decimal
+				return nil, ErrInvalidGeodataFile
+			}
+			advancedN = 1
+			count++
+		case 2, 4: // data length
+			tempContainer = append(tempContainer, container...)
+			if container[0] > 127 { // max one-byte-length byte `7F`(0FFF FFFF) equals to `127` in decimal
+				advancedN = 1
+				goto Loop
+			}
+			lenVarint, n := protowire.ConsumeVarint(tempContainer)
+			if n < 0 {
+				return nil, ErrInvalidGeodataVarintLength
+			}
+			tempContainer = nil
+			if !isInner {
+				isInner = true
+				geoDataVarintLength = lenVarint
+				advancedN = 1
+			} else {
+				isInner = false
+				codeVarintLength = lenVarint
+				varintLenByteLen = uint64(n)
+				advancedN = codeVarintLength
+			}
+			count++
+		case 5: // data value
+			if strings.EqualFold(string(container), code) {
+				count++
+				offset := -(1 + int64(varintLenByteLen) + int64(codeVarintLength))
+				f.Seek(offset, 1)               // back to the start of GeoIP or GeoSite varint
+				advancedN = geoDataVarintLength // the number of bytes to be read in next round
+			} else {
+				count = 1
+				offset := int64(geoDataVarintLength) - int64(codeVarintLength) - int64(varintLenByteLen) - 1
+				f.Seek(offset, 1) // skip the unmatched GeoIP or GeoSite varint
+				advancedN = 1     // the next round will be the start of another GeoIPList or GeoSiteList
+			}
+		case 6: // matched GeoIP or GeoSite varint
+			result = container
+			break Loop
+		}
+	}
+
+	return result, nil
+}
+
+func Decode(filename, code string) ([]byte, error) {
+	f, err := os.Open(filename)
+	if err != nil {
+		return nil, err
+	}
+	defer f.Close()
+
+	geoBytes, err := EmitBytes(f, code)
+	if err != nil {
+		return nil, err
+	}
+	return geoBytes, nil
+}

common/geodata/decode_test.go

@@ -0,0 +1,68 @@
+package geodata_test
+
+import (
+	"bytes"
+	"errors"
+	"io/fs"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/p4gefau1t/trojan-go/common"
+	"github.com/p4gefau1t/trojan-go/common/geodata"
+)
+
+const (
+	geoipURL   = "https://raw.githubusercontent.com/v2fly/geoip/release/geoip.dat"
+	geositeURL = "https://raw.githubusercontent.com/v2fly/domain-list-community/release/dlc.dat"
+)
+
+func init() {
+	wd, err := os.Getwd()
+	common.Must(err)
+
+	tempPath := filepath.Join(wd, "..", "..", "test", "temp")
+	os.Setenv("TROJAN_GO_LOCATION_ASSET", tempPath)
+
+	geoipPath := common.GetAssetLocation("geoip.dat")
+	geositePath := common.GetAssetLocation("geosite.dat")
+
+	if _, err := os.Stat(geoipPath); err != nil && errors.Is(err, fs.ErrNotExist) {
+		common.Must(os.MkdirAll(tempPath, 0755))
+		geoipBytes, err := common.FetchHTTPContent(geoipURL)
+		common.Must(err)
+		common.Must(common.WriteFile(geoipPath, geoipBytes))
+	}
+	if _, err := os.Stat(geositePath); err != nil && errors.Is(err, fs.ErrNotExist) {
+		common.Must(os.MkdirAll(tempPath, 0755))
+		geositeBytes, err := common.FetchHTTPContent(geositeURL)
+		common.Must(err)
+		common.Must(common.WriteFile(geositePath, geositeBytes))
+	}
+}
+
+func TestDecodeGeoIP(t *testing.T) {
+	filename := common.GetAssetLocation("geoip.dat")
+	result, err := geodata.Decode(filename, "test")
+	if err != nil {
+		t.Error(err)
+	}
+
+	expected := []byte{10, 4, 84, 69, 83, 84, 18, 8, 10, 4, 127, 0, 0, 0, 16, 8}
+	if !bytes.Equal(result, expected) {
+		t.Errorf("failed to load geoip:test, expected: %v, got: %v", expected, result)
+	}
+}
+
+func TestDecodeGeoSite(t *testing.T) {
+	filename := common.GetAssetLocation("geosite.dat")
+	result, err := geodata.Decode(filename, "test")
+	if err != nil {
+		t.Error(err)
+	}
+
+	expected := []byte{10, 4, 84, 69, 83, 84, 18, 20, 8, 3, 18, 16, 116, 101, 115, 116, 46, 101, 120, 97, 109, 112, 108, 101, 46, 99, 111, 109}
+	if !bytes.Equal(result, expected) {
+		t.Errorf("failed to load geosite:test, expected: %v, got: %v", expected, result)
+	}
+}

common/geodata/load.go

@@ -0,0 +1,36 @@
+package geodata
+
+import (
+	"runtime"
+
+	v2router "github.com/v2fly/v2ray-core/v4/app/router"
+)
+
+var geoipcache GeoIPCache = make(map[string]*v2router.GeoIP)
+var geositecache GeoSiteCache = make(map[string]*v2router.GeoSite)
+
+func LoadIP(filename, country string) ([]*v2router.CIDR, error) {
+	geoip, err := geoipcache.Unmarshal(filename, country)
+	if err != nil {
+		return nil, err
+	}
+	runtime.GC()
+	return geoip.Cidr, nil
+}
+
+func LoadGeoIP(country string) ([]*v2router.CIDR, error) {
+	return LoadIP("geoip.dat", country)
+}
+
+func LoadSite(filename, list string) ([]*v2router.Domain, error) {
+	geosite, err := geositecache.Unmarshal(filename, list)
+	if err != nil {
+		return nil, err
+	}
+	runtime.GC()
+	return geosite.Domain, nil
+}
+
+func LoadGeoSite(list string) ([]*v2router.Domain, error) {
+	return LoadSite("geosite.dat", list)
+}

common/net.go

@@ -2,8 +2,15 @@ package common
 
 import (
 	"fmt"
+	"io"
+	"io/ioutil"
 	"net"
+	"net/http"
+	"net/url"
+	"os"
 	"strconv"
+	"strings"
+	"time"
 )
 
 const (
@@ -58,3 +65,59 @@ func PickPort(network string, host string) int {
 	}
 	return 0
 }
+
+func WriteAllBytes(writer io.Writer, payload []byte) error {
+	for len(payload) > 0 {
+		n, err := writer.Write(payload)
+		if err != nil {
+			return err
+		}
+		payload = payload[n:]
+	}
+	return nil
+}
+
+func WriteFile(path string, payload []byte) error {
+	writer, err := os.Create(path)
+	if err != nil {
+		return err
+	}
+	defer writer.Close()
+
+	return WriteAllBytes(writer, payload)
+}
+
+func FetchHTTPContent(target string) ([]byte, error) {
+	parsedTarget, err := url.Parse(target)
+	if err != nil {
+		return nil, fmt.Errorf("invalid URL: %s", target)
+	}
+
+	if s := strings.ToLower(parsedTarget.Scheme); s != "http" && s != "https" {
+		return nil, fmt.Errorf("invalid scheme: %s", parsedTarget.Scheme)
+	}
+
+	client := &http.Client{
+		Timeout: 30 * time.Second,
+	}
+	resp, err := client.Do(&http.Request{
+		Method: "GET",
+		URL:    parsedTarget,
+		Close:  true,
+	})
+	if err != nil {
+		return nil, fmt.Errorf("failed to dial to %s", target)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("unexpected HTTP status code: %d", resp.StatusCode)
+	}
+
+	content, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read HTTP response")
+	}
+
+	return content, nil
+}

go.mod

@@ -14,8 +14,8 @@ require (
 	github.com/txthinking/x v0.0.0-20210326105829-476fab902fbe // indirect
 	github.com/v2fly/v2ray-core/v4 v4.38.3
 	github.com/xtaci/smux v1.5.15
-	golang.org/x/crypto v0.0.0-20210421170649-83a5a9bb288b
-	golang.org/x/net v0.0.0-20210502030024-e5908800b52b
+	golang.org/x/crypto v0.0.0-20210503195802-e9a32991a82e
+	golang.org/x/net v0.0.0-20210503060351-7fd8e65b6420
 	golang.org/x/time v0.0.0-20210220033141-f8bda1e9f3ba
 	google.golang.org/grpc v1.37.0
 	google.golang.org/protobuf v1.26.0

go.sum

@@ -224,8 +224,8 @@ golang.org/x/crypto v0.0.0-20200221231518-2aa609cf4a9d/go.mod h1:LzIPMQfyMNhhGPh
 golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
 golang.org/x/crypto v0.0.0-20210220033148-5ea612d1eb83/go.mod h1:jdWPYTVW3xRLrWPugEBEK3UY2ZEsg3UU495nc5E+M+I=
 golang.org/x/crypto v0.0.0-20210415154028-4f45737414dc/go.mod h1:T9bdIzuCu7OtxOm1hfPfRQxPLYneinmdGuTeoZ9dtd4=
-golang.org/x/crypto v0.0.0-20210421170649-83a5a9bb288b h1:7mWr3k41Qtv8XlltBkDkl8LoP3mpSgBW8BUoxtEdbXg=
-golang.org/x/crypto v0.0.0-20210421170649-83a5a9bb288b/go.mod h1:T9bdIzuCu7OtxOm1hfPfRQxPLYneinmdGuTeoZ9dtd4=
+golang.org/x/crypto v0.0.0-20210503195802-e9a32991a82e h1:8foAy0aoO5GkqCvAEJ4VC4P3zksTg4X4aJCDpZzmgQI=
+golang.org/x/crypto v0.0.0-20210503195802-e9a32991a82e/go.mod h1:P+XmwS30IXTQdn5tA2iutPOUgjI07+tq3H3K9MVA1s8=
 golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
 golang.org/x/lint v0.0.0-20180702182130-06c8688daad7/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE=
 golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE=
@@ -251,8 +251,8 @@ golang.org/x/net v0.0.0-20200625001655-4c5254603344/go.mod h1:/O7V0waA8r7cgGh81R
 golang.org/x/net v0.0.0-20200707034311-ab3426394381/go.mod h1:/O7V0waA8r7cgGh81Ro3o1hOxt32SMVPicZroKQ2sZA=
 golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
 golang.org/x/net v0.0.0-20210415231046-e915ea6b2b7d/go.mod h1:9tjilg8BloeKEkVJvy7fQ90B1CfIiPueXVOjqfkSzI8=
-golang.org/x/net v0.0.0-20210502030024-e5908800b52b h1:jCRjgm6WJHzM8VQrm/es2wXYqqbq0NZ1yXFHHgzkiVQ=
-golang.org/x/net v0.0.0-20210502030024-e5908800b52b/go.mod h1:OJAsFXCWl8Ukc7SiCT/9KSuxbyM7479/AVlXFRxuMCk=
+golang.org/x/net v0.0.0-20210503060351-7fd8e65b6420 h1:a8jGStKg0XqKDlKqjLrXn0ioF5MH36pT7Z0BRTqLhbk=
+golang.org/x/net v0.0.0-20210503060351-7fd8e65b6420/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=

statistic/statistics.go

@@ -6,9 +6,8 @@ import (
 	"strings"
 	"sync"
 
-	"github.com/p4gefau1t/trojan-go/log"
-
 	"github.com/p4gefau1t/trojan-go/common"
+	"github.com/p4gefau1t/trojan-go/log"
 )
 
 type TrafficMeter interface {

tunnel/router/client.go

@@ -2,16 +2,15 @@ package router
 
 import (
 	"context"
-	"io/ioutil"
 	"net"
 	"regexp"
 	"strconv"
 	"strings"
 
 	v2router "github.com/v2fly/v2ray-core/v4/app/router"
-	"google.golang.org/protobuf/proto"
 
 	"github.com/p4gefau1t/trojan-go/common"
+	"github.com/p4gefau1t/trojan-go/common/geodata"
 	"github.com/p4gefau1t/trojan-go/config"
 	"github.com/p4gefau1t/trojan-go/log"
 	"github.com/p4gefau1t/trojan-go/tunnel"
@@ -39,7 +38,7 @@ func matchDomain(list []*v2router.Domain, target string) bool {
 		case v2router.Domain_Full:
 			domain := d.GetValue()
 			if domain == target {
-				log.Trace("domain:", target, "hit domain(full) rule:", domain)
+				log.Tracef("domain %s hit domain(full) rule: %s", target, domain)
 				return true
 			}
 		case v2router.Domain_Domain:
@@ -47,14 +46,14 @@ func matchDomain(list []*v2router.Domain, target string) bool {
 			if strings.HasSuffix(target, domain) {
 				idx := strings.Index(target, domain)
 				if idx == 0 || target[idx-1] == '.' {
-					log.Trace("domain:", target, "hit domain rule:", domain)
+					log.Tracef("domain %s hit domain rule: %s", target, domain)
 					return true
 				}
 			}
 		case v2router.Domain_Plain:
 			//keyword
 			if strings.Contains(target, d.GetValue()) {
-				log.Trace("domain:", target, "hit keyword rule:", d.GetValue())
+				log.Tracef("domain %s hit keyword rule: %s", target, d.GetValue())
 				return true
 			}
 		case v2router.Domain_Regex:
@@ -64,11 +63,11 @@ func matchDomain(list []*v2router.Domain, target string) bool {
 				return false
 			}
 			if matched {
-				log.Trace("domain:", target, "hit regex rule:", d.GetValue())
+				log.Tracef("domain %s hit regex rule: %s", target, d.GetValue())
 				return true
 			}
 		default:
-			log.Debug("unknown rule type:" + d.GetType().String())
+			log.Debug("unknown rule type:", d.GetType().String())
 		}
 	}
 	return false
@@ -233,7 +232,7 @@ func loadCode(cfg *Config, prefix string) []codeInfo {
 					strategy: Proxy,
 				})
 			} else {
-				log.Warn("invalid empty rule: ", s)
+				log.Warn("invalid empty rule:", s)
 			}
 		}
 	}
@@ -245,7 +244,7 @@ func loadCode(cfg *Config, prefix string) []codeInfo {
 					strategy: Bypass,
 				})
 			} else {
-				log.Warn("invalid empty rule: ", s)
+				log.Warn("invalid empty rule:", s)
 			}
 		}
 	}
@@ -257,7 +256,7 @@ func loadCode(cfg *Config, prefix string) []codeInfo {
 					strategy: Block,
 				})
 			} else {
-				log.Warn("invalid empty rule: ", s)
+				log.Warn("invalid empty rule:", s)
 			}
 		}
 	}
@@ -305,86 +304,58 @@ func NewClient(ctx context.Context, underlay tunnel.Client) (*Client, error) {
 		return nil, common.NewError("unknown strategy: " + cfg.Router.DomainStrategy)
 	}
 
-	geoipData, err := ioutil.ReadFile(cfg.Router.GeoIPFilename)
-	if err != nil {
-		log.Warn("failed to read geoip.dat file: ", err)
-	} else {
-		geoip := new(v2router.GeoIPList)
-		if err := proto.Unmarshal(geoipData, geoip); err != nil {
-			return nil, err
-		}
-		ipCode := loadCode(cfg, "geoip:")
-		for _, c := range ipCode {
-			c.code = strings.ToUpper(c.code)
-			found := false
-			for _, e := range geoip.GetEntry() {
-				code := e.GetCountryCode()
-				if strings.EqualFold(c.code, code) {
-					client.cidrs[c.strategy] = append(client.cidrs[c.strategy], e.GetCidr()...)
-					found = true
-					break
-				}
-			}
-			if found {
-				log.Info("geoip info", c, "loaded")
-			} else {
-				log.Warn("geoip info", c, "not found")
-			}
+	ipCode := loadCode(cfg, "geoip:")
+	for _, c := range ipCode {
+		code := c.code
+		cidrs, err := geodata.LoadGeoIP(code)
+		if err != nil {
+			log.Error(err)
+		} else {
+			log.Infof("geoip:%s loaded", code)
+			client.cidrs[c.strategy] = append(client.cidrs[c.strategy], cidrs...)
 		}
 	}
 
-	geositeData, err := ioutil.ReadFile(cfg.Router.GeoSiteFilename)
-	if err != nil {
-		log.Warn("failed to read geosite.dat file: ", err)
-	} else {
-		geosite := new(v2router.GeoSiteList)
-		if err := proto.Unmarshal(geositeData, geosite); err != nil {
-			return nil, err
-		}
-		siteCode := loadCode(cfg, "geosite:")
-		for _, c := range siteCode {
-			attrWanted := ""
-			// Test if user wants domains that have an attribute
-			if attrIdx := strings.Index(c.code, "@"); attrIdx > 0 {
-				if attrIdx+1 < len(c.code) {
-					c.code = strings.ToUpper(c.code[:attrIdx])
-					attrWanted = c.code[attrIdx+1:]
-				} else { // "geosite:google@" is invalid
-					log.Warn("geosite info", c.code, "invalid")
-					continue
-				}
-			} else if attrIdx == 0 { // "geosite:@cn" is invalid
-				log.Warn("geosite info", c.code, "invalid")
+	siteCode := loadCode(cfg, "geosite:")
+	for _, c := range siteCode {
+		code := c.code
+		attrWanted := ""
+		// Test if user wants domains that have an attribute
+		if attrIdx := strings.Index(code, "@"); attrIdx > 0 {
+			if !strings.HasSuffix(code, "@") {
+				code = c.code[:attrIdx]
+				attrWanted = c.code[attrIdx+1:]
+			} else { // "geosite:google@" is invalid
+				log.Warnf("geosite:%s invalid", code)
 				continue
-			} else {
-				c.code = strings.ToUpper(c.code)
 			}
+		} else if attrIdx == 0 { // "geosite:@cn" is invalid
+			log.Warnf("geosite:%s invalid", code)
+			continue
+		}
+
+		domainList, err := geodata.LoadGeoSite(code)
+		if err != nil {
+			log.Error(err)
+		} else {
 			found := false
-			for _, e := range geosite.GetEntry() {
-				code := e.GetCountryCode()
-				if strings.EqualFold(c.code, code) {
-					domainList := e.GetDomain()
-					if attrWanted != "" {
-						for _, domain := range domainList {
-							for _, attr := range domain.GetAttribute() {
-								if strings.EqualFold(attrWanted, attr.GetKey()) {
-									client.domains[c.strategy] = append(client.domains[c.strategy], domain)
-									found = true
-								}
-							}
+			if attrWanted != "" {
+				for _, domain := range domainList {
+					for _, attr := range domain.GetAttribute() {
+						if strings.EqualFold(attrWanted, attr.GetKey()) {
+							client.domains[c.strategy] = append(client.domains[c.strategy], domain)
+							found = true
 						}
-						break
-					} else {
-						client.domains[c.strategy] = append(client.domains[c.strategy], domainList...)
-						found = true
-						break
 					}
 				}
+			} else {
+				client.domains[c.strategy] = append(client.domains[c.strategy], domainList...)
+				found = true
 			}
 			if found {
-				log.Info("geosite info", c, "loaded")
+				log.Infof("geosite:%s loaded", c.code)
 			} else {
-				log.Warn("geosite info", c, "not found")
+				log.Errorf("geosite:%s not found", c.code)
 			}
 		}
 	}

tunnel/router/config.go

@@ -1,12 +1,8 @@
 package router
 
 import (
-	"os"
-	"path/filepath"
-
 	"github.com/p4gefau1t/trojan-go/common"
 	"github.com/p4gefau1t/trojan-go/config"
-	"github.com/p4gefau1t/trojan-go/log"
 )
 
 type Config struct {
@@ -30,15 +26,10 @@ func init() {
 			Router: RouterConfig{
 				DefaultPolicy:   "proxy",
 				DomainStrategy:  "as_is",
-				GeoIPFilename:   filepath.Join(common.GetProgramDir(), "geoip.dat"),
-				GeoSiteFilename: filepath.Join(common.GetProgramDir(), "geosite.dat"),
+				GeoIPFilename:   common.GetAssetLocation("geoip.dat"),
+				GeoSiteFilename: common.GetAssetLocation("geosite.dat"),
 			},
 		}
-		if path := os.Getenv("TROJAN_GO_LOCATION_ASSET"); path != "" {
-			cfg.Router.GeoIPFilename = filepath.Join(path, "geoip.dat")
-			cfg.Router.GeoSiteFilename = filepath.Join(path, "geosite.dat")
-			log.Debug("env set:", path)
-		}
 		return cfg
 	})
 }