diff --git a/.gitignore b/.gitignore index d072238..2551e35 100644 --- a/.gitignore +++ b/.gitignore @@ -19,3 +19,5 @@ *.tar.gz *.crt *.key +*.dat +trojan-go diff --git a/common/common.go b/common/common.go index eaabb54..046f0b6 100644 --- a/common/common.go +++ b/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) +} diff --git a/common/geodata/cache.go b/common/geodata/cache.go new file mode 100644 index 0000000..65ffd06 --- /dev/null +++ b/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) +} diff --git a/common/geodata/decode.go b/common/geodata/decode.go new file mode 100644 index 0000000..677fd55 --- /dev/null +++ b/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 +} diff --git a/common/geodata/decode_test.go b/common/geodata/decode_test.go new file mode 100644 index 0000000..2697e19 --- /dev/null +++ b/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) + } +} diff --git a/common/geodata/load.go b/common/geodata/load.go new file mode 100644 index 0000000..16ca509 --- /dev/null +++ b/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) +} diff --git a/common/net.go b/common/net.go index 4dd5031..0533ce1 100644 --- a/common/net.go +++ b/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 +} diff --git a/go.mod b/go.mod index 90a5694..2f94c12 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index 61344c7..4a56227 100644 --- a/go.sum +++ b/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= diff --git a/statistic/statistics.go b/statistic/statistics.go index 9551ad2..d796523 100644 --- a/statistic/statistics.go +++ b/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 { diff --git a/tunnel/router/client.go b/tunnel/router/client.go index 5130bd1..01876f2 100644 --- a/tunnel/router/client.go +++ b/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) } } } diff --git a/tunnel/router/config.go b/tunnel/router/config.go index c975a14..d96bf41 100644 --- a/tunnel/router/config.go +++ b/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 }) }