改进DNS域名解析相关函数

2025-11-04 16:00:24 +08:00 · 2023-06-05 12:36:29 +08:00
parent fdfb6c70dd
commit 91c9b5a496
6 changed files with 199 additions and 30 deletions
--- a/go.mod
+++ b/go.mod
@@ -10,7 +10,7 @@ require (
 	github.com/andybalholm/brotli v1.0.4
 	github.com/cespare/xxhash v1.1.0
 	github.com/cespare/xxhash/v2 v2.1.1
-	github.com/go-acme/lego/v4 v4.9.0
+	github.com/go-acme/lego/v4 v4.10.2
 	github.com/go-sql-driver/mysql v1.7.0
 	github.com/go-telegram-bot-api/telegram-bot-api v4.6.4+incompatible
 	github.com/iwind/TeaGo v0.0.0-20230304012706-c1f4a4e27470
@@ -20,15 +20,16 @@ require (
 	github.com/pkg/sftp v1.12.0
 	github.com/shirou/gopsutil/v3 v3.22.2
 	github.com/smartwalle/alipay/v3 v3.1.7
-	golang.org/x/crypto v0.1.0
+	golang.org/x/crypto v0.5.0
 	golang.org/x/net v0.8.0
-	golang.org/x/sys v0.6.0
+	golang.org/x/sys v0.8.0
 	google.golang.org/grpc v1.45.0
 	gopkg.in/yaml.v3 v3.0.1
 )

 require (
-	github.com/cenkalti/backoff/v4 v4.1.3 // indirect
+	github.com/cenkalti/backoff/v4 v4.2.0 // indirect
+	github.com/fsnotify/fsnotify v1.6.0 // indirect
 	github.com/go-ole/go-ole v1.2.6 // indirect
 	github.com/golang/protobuf v1.5.2 // indirect
 	github.com/jmespath/go-jmespath v0.4.0 // indirect
--- a/go.sum
+++ b/go.sum
@@ -12,8 +12,8 @@ github.com/andybalholm/brotli v1.0.4 h1:V7DdXeJtZscaqfNuAdSRuRFzuiKlHSC/Zh3zl9qY
 github.com/andybalholm/brotli v1.0.4/go.mod h1:fO7iG3H7G2nSZ7m0zPUDn85XEX2GTukHGRSepvi9Eig=
 github.com/antihax/optional v1.0.0/go.mod h1:uupD/76wgC+ih3iEmQUL+0Ugr19nfwCT1kdvxnR2qWY=
 github.com/benbjohnson/clock v1.0.3/go.mod h1:bGMdMPoPVvcYyt1gHDf4J2KE153Yf9BuiUKYMaxlTDM=
-github.com/cenkalti/backoff/v4 v4.1.3 h1:cFAlzYUlVYDysBEH2T5hyJZMh3+5+WCBvSnK6Q8UtC4=
-github.com/cenkalti/backoff/v4 v4.1.3/go.mod h1:scbssz8iZGpm3xbr14ovlUdkxfGXNInqkPWOWmG2CLw=
+github.com/cenkalti/backoff/v4 v4.2.0 h1:HN5dHm3WBOgndBH6E8V0q2jIYIR3s9yglV8k/+MN3u4=
+github.com/cenkalti/backoff/v4 v4.2.0/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE=
 github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
 github.com/cespare/xxhash v1.1.0 h1:a6HrQnmkObjyL+Gs60czilIUGqrzKutQD6XZog3p+ko=
 github.com/cespare/xxhash v1.1.0/go.mod h1:XrSqR1VqqWfGrhpAt58auRo0WTKS1nRRg3ghfAqPWnc=
@@ -38,6 +38,8 @@ github.com/envoyproxy/go-control-plane v0.9.10-0.20210907150352-cf90f659a021/go.
 github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c=
 github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo=
 github.com/fsnotify/fsnotify v1.4.9/go.mod h1:znqG4EE+3YCdAaPaxE2ZRY/06pZUdp0tY4IgpuI1SZQ=
+github.com/fsnotify/fsnotify v1.6.0 h1:n+5WquG0fcWoWp6xPWfHdbskMCQaFnG6PfBrh1Ky4HY=
+github.com/fsnotify/fsnotify v1.6.0/go.mod h1:sl3t1tCWJFWoRz9R8WJCbQihKKwmorjAbSClcnxKAGw=
 github.com/ghodss/yaml v1.0.0/go.mod h1:4dBDuWmgqj2HViK6kFavaiC9ZROes6MMH2rRYeMEF04=
 github.com/go-gl/glfw/v3.3/glfw v0.0.0-20200222043503-6f7a984d4dc4/go.mod h1:tQ2UAYgL5IevRw8kRxooKSPJfGvJ9fJQFa0TUsXzTg8=
 github.com/go-ole/go-ole v1.2.6 h1:/Fpf6oFPoeFik9ty7siob0G6Ke8QvQEuVcuChpwXzpY=
@@ -162,8 +164,8 @@ golang.org/x/crypto v0.0.0-20190506204251-e1dfcc566284/go.mod h1:yigFU9vqHzYiE8U
 golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
 golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
 golang.org/x/crypto v0.0.0-20200820211705-5c72a883971a/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
-golang.org/x/crypto v0.1.0 h1:MDRAIl0xIo9Io2xV565hzXHw3zVseKrJKodhohM5CjU=
-golang.org/x/crypto v0.1.0/go.mod h1:RecgLatLF4+eUMCP1PoPZQb+cVrJcOPbHkTkbkB9sbw=
+golang.org/x/crypto v0.5.0 h1:U/0M97KRkSFvyD/3FSmdP5W5swImpNgle/EHFhOsQPE=
+golang.org/x/crypto v0.5.0/go.mod h1:NK/OQwhpMQP3MwtdjgLlYHnH9ebylxKWv3e0fK+mkQU=
 golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
 golang.org/x/exp v0.0.0-20190306152737-a1d7652674e8/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
 golang.org/x/exp v0.0.0-20200513190911-00229845015e/go.mod h1:4M0jN8W1tt0AVLNr8HDosyJCDCDuyL9N9+3m7wDWgKw=
@@ -222,8 +224,11 @@ golang.org/x/sys v0.0.0-20210510120138-977fb7262007/go.mod h1:oPkhp1MJrh7nUepCBc
 golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20210816074244-15123e1e1f71/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220111092808-5a964db01320/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20220908164124-27713097b956/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.6.0 h1:MVltZSvRTcU2ljQOhs94SXPftV6DCNnZViHeQps87pQ=
 golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.8.0 h1:EBmGv8NaZBZTWvrbjNoL6HVt+IVy3QDQpJs7VRIw3tU=
+golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/term v0.6.0 h1:clScbb1cHjoCkyRbWwBEUZ5H/tIFu5TAXIqaZD0Gcjw=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
--- a/internal/utils/lookup.go
+++ b/internal/utils/lookup.go
@@ -3,20 +3,27 @@ package utils
 import (
 	"errors"
 	teaconst "github.com/TeaOSLab/EdgeAPI/internal/const"
+	"github.com/TeaOSLab/EdgeAPI/internal/utils/taskutils"
 	"github.com/TeaOSLab/EdgeCommon/pkg/configutils"
+	"github.com/TeaOSLab/EdgeCommon/pkg/dnsconfigs"
+	"github.com/fsnotify/fsnotify"
+	"github.com/iwind/TeaGo/lists"
 	"github.com/iwind/TeaGo/logs"
 	"github.com/miekg/dns"
+	"sync"
 )

 var sharedDNSClient *dns.Client
 var sharedDNSConfig *dns.ClientConfig
+var sharedDNSLocker = &sync.RWMutex{}

 func init() {
 	if !teaconst.IsMain {
 		return
 	}

-	config, err := dns.ClientConfigFromFile("/etc/resolv.conf")
+	var resolvConfFile = "/etc/resolv.conf"
+	config, err := dns.ClientConfigFromFile(resolvConfFile)
 	if err != nil {
 		logs.Println("ERROR: configure dns client failed: " + err.Error())
 		return
@@ -25,6 +32,21 @@ func init() {
 	sharedDNSConfig = config
 	sharedDNSClient = &dns.Client{}

+	// 监视文件变化，以便及时更新配置
+	go func() {
+		watcher, watcherErr := fsnotify.NewWatcher()
+		if watcherErr == nil {
+			err = watcher.Add(resolvConfFile)
+			for range watcher.Events {
+				newConfig, err := dns.ClientConfigFromFile(resolvConfFile)
+				if err == nil && newConfig != nil {
+					sharedDNSLocker.Lock()
+					sharedDNSConfig = newConfig
+					sharedDNSLocker.Unlock()
+				}
+			}
+		}
+	}()
 }

 // LookupCNAME 查询CNAME记录
@@ -40,8 +62,10 @@ func LookupCNAME(host string) (string, error) {
 	m.RecursionDesired = true

 	var lastErr error
-	for _, serverAddr := range sharedDNSConfig.Servers {
-		r, _, err := sharedDNSClient.Exchange(m, configutils.QuoteIP(serverAddr)+":"+sharedDNSConfig.Port)
+	var serverAddrs = composeDNSResolverAddrs(nil)
+
+	for _, serverAddr := range serverAddrs {
+		r, _, err := sharedDNSClient.Exchange(m, serverAddr)
 		if err != nil {
 			lastErr = err
 			continue
@@ -56,8 +80,7 @@ func LookupCNAME(host string) (string, error) {
 }

 // LookupNS 查询NS记录
-// TODO 可以设置使用的DNS主机地址
-func LookupNS(host string) ([]string, error) {
+func LookupNS(host string, extraResolvers []*dnsconfigs.DNSResolver) ([]string, error) {
 	var m = new(dns.Msg)

 	m.SetQuestion(host+".", dns.TypeNS)
@@ -67,23 +90,36 @@ func LookupNS(host string) ([]string, error) {

 	var lastErr error
 	var hasValidServer = false
-	for _, serverAddr := range sharedDNSConfig.Servers {
-		r, _, err := sharedDNSClient.Exchange(m, configutils.QuoteIP(serverAddr)+":"+sharedDNSConfig.Port)
+	var serverAddrs = composeDNSResolverAddrs(extraResolvers)
+	if len(serverAddrs) == 0 {
+		return nil, nil
+	}
+
+	taskErr := taskutils.RunConcurrent(serverAddrs, taskutils.DefaultConcurrent, func(task any, locker *sync.RWMutex) {
+		var serverAddr = task.(string)
+		r, _, err := sharedDNSClient.Exchange(m, serverAddr)
 		if err != nil {
 			lastErr = err
-			continue
+			return
 		}

 		hasValidServer = true

 		if len(r.Answer) == 0 {
-			continue
+			return
 		}

 		for _, answer := range r.Answer {
-			result = append(result, answer.(*dns.NS).Ns)
+			var value = answer.(*dns.NS).Ns
+			locker.Lock()
+			if len(value) > 0 && !lists.ContainsString(result, value) {
+				result = append(result, value)
 			}
-		break
+			locker.Unlock()
+		}
+	})
+	if taskErr != nil {
+		return result, taskErr
 	}

 	if hasValidServer {
@@ -94,8 +130,7 @@ func LookupNS(host string) ([]string, error) {
 }

 // LookupTXT 获取CNAME
-// TODO 可以设置使用的DNS主机地址
-func LookupTXT(host string) ([]string, error) {
+func LookupTXT(host string, extraResolvers []*dnsconfigs.DNSResolver) ([]string, error) {
 	var m = new(dns.Msg)

 	m.SetQuestion(host+".", dns.TypeTXT)
@@ -104,23 +139,36 @@ func LookupTXT(host string) ([]string, error) {
 	var lastErr error
 	var result = []string{}
 	var hasValidServer = false
-	for _, serverAddr := range sharedDNSConfig.Servers {
-		r, _, err := sharedDNSClient.Exchange(m, configutils.QuoteIP(serverAddr)+":"+sharedDNSConfig.Port)
+	var serverAddrs = composeDNSResolverAddrs(extraResolvers)
+	if len(serverAddrs) == 0 {
+		return nil, nil
+	}
+
+	taskErr := taskutils.RunConcurrent(serverAddrs, taskutils.DefaultConcurrent, func(task any, locker *sync.RWMutex) {
+		var serverAddr = task.(string)
+		r, _, err := sharedDNSClient.Exchange(m, serverAddr)
 		if err != nil {
 			lastErr = err
-			continue
+			return
 		}
 		hasValidServer = true

 		if len(r.Answer) == 0 {
-			continue
+			return
 		}

 		for _, answer := range r.Answer {
-			result = append(result, answer.(*dns.TXT).Txt...)
+			for _, txt := range answer.(*dns.TXT).Txt {
+				locker.Lock()
+				if len(txt) > 0 && !lists.ContainsString(result, txt) {
+					result = append(result, txt)
 				}
-
-		break
+				locker.Unlock()
+			}
+		}
+	})
+	if taskErr != nil {
+		return result, taskErr
 	}

 	if hasValidServer {
@@ -129,3 +177,22 @@ func LookupTXT(host string) ([]string, error) {

 	return nil, lastErr
 }
+
+// 组合DNS解析服务器地址
+func composeDNSResolverAddrs(extraResolvers []*dnsconfigs.DNSResolver) []string {
+	sharedDNSLocker.RLock()
+	defer sharedDNSLocker.RUnlock()
+
+	// 这里不处理重复，方便我们可以多次重试
+	var servers = sharedDNSConfig.Servers
+	var port = sharedDNSConfig.Port
+
+	var serverAddrs = []string{}
+	for _, serverAddr := range servers {
+		serverAddrs = append(serverAddrs, configutils.QuoteIP(serverAddr)+":"+port)
+	}
+	for _, resolver := range extraResolvers {
+		serverAddrs = append(serverAddrs, resolver.Addr())
+	}
+	return serverAddrs
+}
--- a/internal/utils/lookup_test.go
+++ b/internal/utils/lookup_test.go
@@ -4,6 +4,7 @@ package utils_test

 import (
 	"github.com/TeaOSLab/EdgeAPI/internal/utils"
+	"github.com/TeaOSLab/EdgeCommon/pkg/dnsconfigs"
 	"testing"
 )

@@ -12,9 +13,25 @@ func TestLookupCNAME(t *testing.T) {
 }

 func TestLookupNS(t *testing.T) {
-	t.Log(utils.LookupNS("goedge.cn"))
+	t.Log(utils.LookupNS("goedge.cn", nil))
+}
+
+func TestLookupNSExtra(t *testing.T) {
+	t.Log(utils.LookupNS("goedge.cn", []*dnsconfigs.DNSResolver{
+		{
+			Host: "192.168.2.2",
+		},
+		{
+			Host: "192.168.2.2",
+			Port: 58,
+		},
+		{
+			Host: "8.8.8.8",
+			Port: 53,
+		},
+	}))
 }

 func TestLookupTXT(t *testing.T) {
-	t.Log(utils.LookupTXT("yanzheng.goedge.cn"))
+	t.Log(utils.LookupTXT("yanzheng.goedge.cn", nil))
 }
--- a/internal/utils/taskutils/concurrent.go
+++ b/internal/utils/taskutils/concurrent.go
@@ -0,0 +1,61 @@
+// Copyright 2023 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package taskutils
+
+import (
+	"errors"
+	"reflect"
+	"sync"
+)
+
+const DefaultConcurrent = 16
+
+func RunConcurrent(tasks any, concurrent int, f func(task any, locker *sync.RWMutex)) error {
+	if tasks == nil {
+		return nil
+	}
+	var tasksValue = reflect.ValueOf(tasks)
+	if tasksValue.Type().Kind() != reflect.Slice {
+		return errors.New("ony works for slice")
+	}
+
+	var countTasks = tasksValue.Len()
+	if countTasks == 0 {
+		return nil
+	}
+
+	if concurrent <= 0 {
+		concurrent = 8
+	}
+	if concurrent > countTasks {
+		concurrent = countTasks
+	}
+
+	var taskChan = make(chan any, countTasks)
+	for i := 0; i < countTasks; i++ {
+		taskChan <- tasksValue.Index(i).Interface()
+	}
+
+	var wg = &sync.WaitGroup{}
+	wg.Add(concurrent)
+
+	var locker = &sync.RWMutex{}
+
+	for i := 0; i < concurrent; i++ {
+		go func() {
+			defer wg.Done()
+
+			for {
+				select {
+				case task := <-taskChan:
+					f(task, locker)
+				default:
+					return
+				}
+			}
+		}()
+	}
+	wg.Wait()
+
+	return nil
+}
--- a/internal/utils/taskutils/concurrent_test.go
+++ b/internal/utils/taskutils/concurrent_test.go
@@ -0,0 +1,18 @@
+// Copyright 2023 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package taskutils_test
+
+import (
+	"github.com/TeaOSLab/EdgeAPI/internal/utils/taskutils"
+	"sync"
+	"testing"
+)
+
+func TestRunConcurrent(t *testing.T) {
+	err := taskutils.RunConcurrent([]string{"a", "b", "c", "d", "e"}, 3, func(task any, locker *sync.RWMutex) {
+		t.Log("run", task)
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+}