🚸 优化 Agent 取 IP

pkg/utils/http.go

@@ -6,10 +6,7 @@ import (
 	"net"
 	"net/http"
 	"strings"
-	"sync"
 	"time"
-
-	"github.com/miekg/dns"
 )
 
 func NewSingleStackHTTPClient(httpTimeout, dialTimeout, keepAliveTimeout time.Duration, ipv6 bool) *http.Client {
@@ -39,60 +36,52 @@ func NewSingleStackHTTPClient(httpTimeout, dialTimeout, keepAliveTimeout time.Du
 func resolveIP(addr string, ipv6 bool) (string, error) {
 	url := strings.Split(addr, ":")
 
-	m := new(dns.Msg)
-	if ipv6 {
-		m.SetQuestion(dns.Fqdn(url[0]), dns.TypeAAAA)
-	} else {
-		m.SetQuestion(dns.Fqdn(url[0]), dns.TypeA)
-	}
-	m.RecursionDesired = true
-
 	dnsServers := []string{"2606:4700:4700::1001", "2001:4860:4860::8844", "2400:3200::1", "2400:3200:baba::1"}
 	if !ipv6 {
 		dnsServers = []string{"1.0.0.1", "8.8.4.4", "223.5.5.5", "223.6.6.6"}
 	}
 
-	var wg sync.WaitGroup
-	var resolveLock sync.RWMutex
+	res, err := net.LookupIP(url[0])
+	if err != nil {
+		for i := 0; i < len(dnsServers); i++ {
+			dnsServer := dnsServers[i]
+			if ipv6 {
+				dnsServer = "[" + dnsServer + "]"
+			}
+			r := &net.Resolver{
+				PreferGo: true,
+				Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
+					d := net.Dialer{
+						Timeout: time.Second * 10,
+					}
+					return d.DialContext(ctx, "udp", dnsServer+":53")
+				},
+			}
+			res, err = r.LookupIP(context.Background(), "ip", url[0])
+			if err == nil {
+				break
+			}
+		}
+	}
+
+	if err != nil {
+		return "", err
+	}
+
 	var ipv4Resolved, ipv6Resolved bool
 
-	wg.Add(len(dnsServers))
-	go func() {
-
-	}()
-	for i := 0; i < len(dnsServers); i++ {
-		go func(i int) {
-			defer wg.Done()
-			c := new(dns.Client)
-			c.Timeout = time.Second * 3
-			r, _, err := c.Exchange(m, net.JoinHostPort(dnsServers[i], "53"))
-			if err != nil {
-				return
-			}
-			resolveLock.Lock()
-			defer resolveLock.Unlock()
-			if ipv6 && ipv6Resolved {
-				return
-			}
-			if !ipv6 && ipv4Resolved {
-				return
-			}
-			for _, ans := range r.Answer {
-				if ipv6 {
-					if aaaa, ok := ans.(*dns.AAAA); ok {
-						url[0] = "[" + aaaa.AAAA.String() + "]"
-						ipv6Resolved = true
-					}
-				} else {
-					if a, ok := ans.(*dns.A); ok {
-						url[0] = a.A.String()
-						ipv4Resolved = true
-					}
-				}
-			}
-		}(i)
+	for i := 0; i < len(res); i++ {
+		if ip4 := res[i].To4(); ip4 != nil && !ipv6 {
+			ipv4Resolved = true
+			url[0] = ip4.String()
+			break
+		}
+		if ip6 := res[i].To16(); ip6 != nil && ipv6 {
+			ipv6Resolved = true
+			url[0] = "[" + ip6.String() + "]"
+			break
+		}
 	}
-	wg.Wait()
 
 	if ipv6 && !ipv6Resolved {
 		return "", errors.New("the AAAA record not resolved")