merged in commit ccdd03 into refactor, improved nameserver parsing

cmd/worker_manager.go

@@ -138,7 +138,11 @@ func populateCLIConfig(gc *CLIConf, flags *pflag.FlagSet) *CLIConf {
 			ns = strings.Split(gc.NameServersString, ",")
 		}
 		for i, s := range ns {
-			ns[i] = util.AddDefaultPortToDNSServerName(s)
+			nsWithPort, err := util.AddDefaultPortToDNSServerName(s)
+			if err != nil {
+				log.Fatalf("unable to parse name server: %s", s)
+			}
+			ns[i] = nsWithPort
 		}
 		gc.NameServers = ns
 	}
@@ -437,7 +441,10 @@ func doLookupWorker(gc *CLIConf, lookup LookupModule, rc *zdns.ResolverConfig, i
 			rawName, entryMetadata = parseMetadataInputLine(line)
 			res.Metadata = entryMetadata
 		} else if gc.NameServerMode {
-			nameServer = util.AddDefaultPortToDNSServerName(line)
+			nameServer, err = util.AddDefaultPortToDNSServerName(line)
+			if err != nil {
+				log.Fatal("unable to parse name server: ", line)
+			}
 		} else {
 			rawName, nameServer = parseNormalInputLine(line)
 		}
@@ -519,7 +526,10 @@ func parseNormalInputLine(line string) (string, string) {
 	if len(s) == 1 {
 		return s[0], ""
 	} else {
-		ns := util.AddDefaultPortToDNSServerName(s[1])
+		ns, err := util.AddDefaultPortToDNSServerName(s[1])
+		if err != nil {
+			log.Fatal("unable to parse name server: ", s[1])
+		}
 		return s[0], ns
 	}
 }

internal/util/util.go

@@ -16,30 +16,39 @@ package util
 
 import (
 	"fmt"
-	"regexp"
+	"net"
 	"strings"
 
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
 	"github.com/spf13/viper"
 )
 
-var rePort *regexp.Regexp
-var reV6 *regexp.Regexp
-
 const (
 	EnvPrefix              = "ZDNS"
 	DefaultFilePermissions = 0644 // rw-r--r--
 )
 
-func AddDefaultPortToDNSServerName(s string) string {
-	if !rePort.MatchString(s) {
-		return s + ":53"
-	} else if reV6.MatchString(s) {
-		return "[" + s + "]:53"
-	} else {
-		return s
+func AddDefaultPortToDNSServerName(inAddr string) (string, error) {
+	// Try to split host and port to see if the port is already specified.
+	host, port, err := net.SplitHostPort(inAddr)
+	if err != nil {
+		// might mean there's no port specified
+		host = inAddr
+	}
+
+	// Validate the host part as an IP address.
+	ip := net.ParseIP(host)
+	if ip == nil {
+		return "", fmt.Errorf("invalid IP address")
+	}
+
+	// If the original input does not have a port, specify port 53
+	if port == "" {
+		port = "53"
 	}
+
+	return net.JoinHostPort(ip.String(), port), nil
 }
 
 // Reference: https://github.com/carolynvs/stingoftheviper/blob/main/main.go
@@ -66,8 +75,3 @@ func BindFlags(cmd *cobra.Command, v *viper.Viper, envPrefix string) {
 func GetDefaultResolvers() []string {
 	return []string{"8.8.8.8:53", "8.8.4.4:53", "1.1.1.1:53", "1.0.0.1:53"}
 }
-
-func init() {
-	rePort = regexp.MustCompile(":\\d+$")      // string ends with potential port number
-	reV6 = regexp.MustCompile("^([0-9a-f]*:)") // string starts like valid IPv6 address
-}