feat: use domain suffix tree instead of map to match

Mythologyli · Feb 6, 2024 · 761c21d · 761c21d
1 parent a3690b9
commit 761c21d
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 8 deletions.
diff --git 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
@@ -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
@@ -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,