refactor: more clean handling of dns server ip

bschaatsbergen · Jan 9, 2024 · 461962e · 461962e
1 parent b96372a
commit 461962e
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 8 deletions.
diff --git a/cmd/root.go b/cmd/root.go
@@ -51,7 +51,6 @@ var (
 			for _, queryType := range queryTypes {
 				msg := core.PrepareDNSQuery(domainName, queryType.Type)
 
-				response, _, err := core.SendDNSQuery(&client, msg, flagStore.DNSServerIP)
 				response, _, err := core.SendDNSQuery(&client, msg, flagStore.DNSServerIP, flagStore.DNSServerPort)
 				if err != nil {
 					log.Fatal(err)

diff --git a/pkg/core/core.go b/pkg/core/core.go
@@ -49,7 +49,6 @@ func PrepareDNSQuery(domainName string, queryType uint16) dns.Msg {
 }
 
 // SendDNSQuery sends a DNS query to a given DNS server.
-func SendDNSQuery(client *dns.Client, msg dns.Msg, dnsServerIP string) (*dns.Msg, time.Duration, error) {
 func SendDNSQuery(client *dns.Client, msg dns.Msg, dnsServerIP, dnsServerPort string) (*dns.Msg, time.Duration, error) {
 	if dnsServerIP == "" {
 		goOS := runtime.GOOS
@@ -64,13 +63,10 @@ func SendDNSQuery(client *dns.Client, msg dns.Msg, dnsServerIP, dnsServerPort st
 		dnsServerIP = conf.Servers[0]
 
 	}
-	// If the server IP is IPv6, wrap it in square brackets to remove ambiguity from port number and address.
-	if net.ParseIP(dnsServerIP).To4() == nil {
-		dnsServerIP = "[" + dnsServerIP + "]"
-	}
 
-	logrus.Debugf("Sending DNS query to %s", dnsServerIP)
-	response, timeDuration, err := client.Exchange(&msg, dnsServerIP+":53")
+	addr := net.JoinHostPort(dnsServerIP, dnsServerPort)
+
+	response, timeDuration, err := client.Exchange(&msg, addr)
 
 	if err != nil {
 		logrus.Debug("Failed to receive DNS response.")