改进DNS域名解析相关函数

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)
+			}
+			locker.Unlock()
 		}
-		break
+	})
+	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)
+				}
+				locker.Unlock()
+			}
 		}
-
-		break
+	})
+	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
+}