fix: properly format IPv6 DNS server addresses when appending port (#51)

chore: add test for EnsureDNSAddress helper function

Signed-off-by: Rui Chen <[email protected]>
Co-authored-by: Bruno Schaatsbergen <[email protected]>

Author: chenrui333

cmd/root.go

@@ -2,6 +2,7 @@ package cmd
 
 import (
 	"fmt"
+	"net"
 	"os"
 	"runtime"
 	"sort"
@@ -29,6 +30,20 @@ var (
 	qtype   string
 )
 
+// EnsureDNSAddress formats the DNS server address properly.
+func EnsureDNSAddress(server string) string {
+    if strings.Contains(server, "]") || strings.Contains(server, ":") && net.ParseIP(server) == nil {
+        return server
+    }
+
+    ip := net.ParseIP(server)
+    if ip != nil && ip.To4() == nil { // It's IPv6 (and not IPv4)
+        return "[" + server + "]:53"
+    }
+    // Otherwise, assume IPv4 or hostname, so append port normally.
+    return server + ":53"
+}
+
 func NewRootCommand() *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "zns",
@@ -138,11 +153,7 @@ func NewRootCommand() *cobra.Command {
 				}
 			}
 
-			// If the server address does not already include a port,
-			// append the default DNS port (53) to it.
-			if !strings.Contains(server, ":") {
-				server = fmt.Sprintf("%s:53", server)
-			}
+			server = EnsureDNSAddress(server)
 
 			querier := query.NewQueryClient(server, new(dns.Client), logger)
 

cmd/root_test.go

@@ -187,3 +187,22 @@ func Test_Cmd_LogFile_Debug(t *testing.T) {
 	assert.Contains(t, string(logFile), "A       |example.com.   |01m00s   |93.184.216.34")
 	assert.Contains(t, string(logFile), "CNAME   |example.com.   |01m00s   |example.org.")
 }
+
+func TestEnsureDNSAddress(t *testing.T) {
+    testCases := []struct {
+        input    string
+        expected string
+    }{
+        {"127.0.0.1", "127.0.0.1:53"},
+        {"2001:558:feed::1", "[2001:558:feed::1]:53"},
+        {"[2001:558:feed::1]:53", "[2001:558:feed::1]:53"},
+        {"example.com", "example.com:53"},
+    }
+
+    for _, tc := range testCases {
+        t.Run(fmt.Sprintf("input=%s", tc.input), func(t *testing.T) {
+            result := EnsureDNSAddress(tc.input)
+            assert.Equal(t, tc.expected, result)
+        })
+    }
+}