Fix sshkeyscan to work with ipv6

- make use of standard net.SplitHostPort and net.JoinHostPort functions

Signed-off-by: Alex Couture-Beil <alex@earthly.dev>

util/sshutil/keyscan.go

@@ -11,28 +11,26 @@ import (
 
 const defaultPort = 22
 
-var ErrMalformedServer = fmt.Errorf("invalid server, must be of the form hostname, or hostname:port")
-
 var errCallbackDone = fmt.Errorf("callback failed on purpose")
 
+// addDefaultPort appends a default port if hostport doesn't contain one
+func addDefaultPort(hostport string, defaultPort int) string {
+	_, _, err := net.SplitHostPort(hostport)
+	if err == nil {
+		return hostport
+	}
+	hostport = net.JoinHostPort(hostport, strconv.Itoa(defaultPort))
+	return hostport
+}
+
 // SshKeyScan scans a ssh server for the hostkey; server should be in the form hostname, or hostname:port
 func SSHKeyScan(server string) (string, error) {
-	port := defaultPort
-	parts := strings.Split(server, ":")
-	if len(parts) == 2 {
-		var err error
-		server = parts[0]
-		port, err = strconv.Atoi(parts[1])
-		if err != nil {
-			return "", ErrMalformedServer
-		}
-	} else if len(parts) > 2 {
-		return "", ErrMalformedServer
-	}
-
 	var key string
-	KeyScanCallback := func(hostname string, remote net.Addr, pubKey ssh.PublicKey) error {
-		hostname = strings.TrimSuffix(hostname, fmt.Sprintf(":%d", port))
+	KeyScanCallback := func(hostport string, remote net.Addr, pubKey ssh.PublicKey) error {
+		hostname, _, err := net.SplitHostPort(hostport)
+		if err != nil {
+			return err
+		}
 		key = strings.TrimSpace(fmt.Sprintf("%s %s", hostname, string(ssh.MarshalAuthorizedKey(pubKey))))
 		return errCallbackDone
 	}
@@ -40,7 +38,8 @@ func SSHKeyScan(server string) (string, error) {
 		HostKeyCallback: KeyScanCallback,
 	}
 
-	conn, err := ssh.Dial("tcp", fmt.Sprintf("%s:%d", server, port), config)
+	server = addDefaultPort(server, defaultPort)
+	conn, err := ssh.Dial("tcp", server, config)
 	if key != "" {
 		// as long as we get the key, the function worked
 		err = nil