diff --git a/go.mod b/go.mod index 0a4efb3..1be9930 100644 --- a/go.mod +++ b/go.mod @@ -11,6 +11,7 @@ require ( github.com/beevik/etree v1.2.0 github.com/cloverstd/tcping v0.1.1 github.com/cxz66666/sing-tun v0.0.0-20231028191617-2867d9374292 + github.com/golang-infrastructure/go-domain-suffix-trie v0.0.2 github.com/miekg/dns v1.1.57 github.com/patrickmn/go-cache v2.1.0+incompatible github.com/shadowsocks/go-shadowsocks2 v0.1.5 diff --git a/go.sum b/go.sum index 78e574f..3f30e4c 100644 --- a/go.sum +++ b/go.sum @@ -23,6 +23,8 @@ github.com/go-ole/go-ole v1.3.0 h1:Dt6ye7+vXGIKZ7Xtk4s6/xVdGDQynvom7xCFEdWr6uE= github.com/go-ole/go-ole v1.3.0/go.mod h1:5LS6F96DhAwUc7C+1HLexzMXY1xGRSryjyPPKW6zv78= github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572 h1:tfuBGBXKqDEevZMzYi5KSi8KkcZtzBcTgAUUtapy0OI= github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572/go.mod h1:9Pwr4B2jHnOSGXyyzV8ROjYa2ojvAY6HCGYYfMoC3Ls= +github.com/golang-infrastructure/go-domain-suffix-trie v0.0.2 h1:8AyA1jDDahI6moCH9xtmOlMS07OncQTK8INinJxB0tU= +github.com/golang-infrastructure/go-domain-suffix-trie v0.0.2/go.mod h1:PA7neNAxdqgeAh3Ggyssa8uhJ2VUuKxft25tGdMSfjE= github.com/google/btree v1.1.2 h1:xf4v41cLI2Z6FxbKm+8Bu+m8ifhj15JuZ9sa0jZCMUU= github.com/google/btree v1.1.2/go.mod h1:qOPhT0dTNdNzV6Z/lhRX0YXUafgPLFUh+gZMl761Gm4= github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38= diff --git a/resolve/resolver.go b/resolve/resolver.go index a6ed994..39e5f7f 100644 --- a/resolve/resolver.go +++ b/resolve/resolver.go @@ -3,11 +3,11 @@ package resolve import ( "context" "errors" + domainsuffixtrie "github.com/golang-infrastructure/go-domain-suffix-trie" "github.com/mythologyli/zju-connect/log" "github.com/mythologyli/zju-connect/stack" "github.com/patrickmn/go-cache" "net" - "strings" "sync" "time" ) @@ -17,7 +17,7 @@ type Resolver struct { remoteTCPResolver *net.Resolver secondaryResolver *net.Resolver ttl uint64 - domainResource map[string]bool + domainResource *domainsuffixtrie.DomainSuffixTrieNode[bool] dnsResource map[string]net.IP useRemoteDNS bool @@ -37,11 +37,8 @@ type Resolver struct { func (r *Resolver) Resolve(ctx context.Context, host string) (context.Context, net.IP, error) { var useVPN = false if r.domainResource != nil { - for domain := range r.domainResource { - if strings.Contains(host, domain) { - useVPN = true - break - } + if r.domainResource.FindMatchDomainSuffixPayload(host) { + useVPN = true } } @@ -161,6 +158,11 @@ func (r *Resolver) CleanCache(duration time.Duration) { } func NewResolver(stack stack.Stack, remoteDNSServer, secondaryDNSServer string, ttl uint64, domainResource map[string]bool, dnsResource map[string]net.IP, useRemoteDNS bool) *Resolver { + domainSuffixTree := domainsuffixtrie.NewDomainSuffixTrie[bool]() + for domain := range domainResource { + _ = domainSuffixTree.AddDomainSuffix(domain, true) + } + resolver := &Resolver{ remoteUDPResolver: &net.Resolver{ PreferGo: true, @@ -181,7 +183,7 @@ func NewResolver(stack stack.Stack, remoteDNSServer, secondaryDNSServer string, }, }, ttl: ttl, - domainResource: domainResource, + domainResource: domainSuffixTree, dnsResource: dnsResource, dnsCache: cache.New(time.Duration(ttl)*time.Second, time.Duration(ttl)*2*time.Second), useRemoteDNS: useRemoteDNS,