pkg/driver/vz: Try SSH handshake to check if SSH port is available.

Check the SSH server in a way that complies with the SSH protocol using x/crypto/ssh.
This change fixes #4334 by falling back to usernet port forwarder on failing SSH connections over VSOCK.

- pkg/networks/usernet: Rename entry point from `/extension/wait_port` to `/extension/wait_ssh_server`
  Because it changed to an SSH server-specific entry point.
  When a client accesses the old entry point, it fails and continues with falling back to the usernet forwarder.

- pkg/sshutil: Add `WaitSSHReady()`
  WaitSSHReady waits until the SSH server is ready to accept connections.
  The dialContext function is used to create a connection to the SSH server.
  The addr, user, privateKeyPath parameter is used for ssh.ClientConn creation.
  The timeoutSeconds parameter specifies the maximum number of seconds to wait.

Signed-off-by: Norio Nomura <[email protected]>

Author: norio-nomura

go.mod

@@ -117,7 +117,7 @@ require (
 	github.com/x448/float16 v0.8.4 // indirect
 	github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
 	github.com/yuin/gopher-lua v1.1.1 // indirect
-	golang.org/x/crypto v0.43.0 // indirect
+	golang.org/x/crypto v0.43.0
 	golang.org/x/mod v0.29.0 // indirect
 	golang.org/x/oauth2 v0.30.0 // indirect
 	golang.org/x/term v0.36.0 // indirect

pkg/driver/vz/vm_darwin.go

@@ -113,18 +113,18 @@ func startVM(ctx context.Context, inst *limatype.Instance, sshLocalPort int) (vm
 								useSSHOverVsock = b
 							}
 						}
+						hostAddress := net.JoinHostPort(inst.SSHAddress, strconv.Itoa(usernetSSHLocalPort))
 						if !useSSHOverVsock {
 							logrus.Info("LIMA_SSH_OVER_VSOCK is false, skipping detection of SSH server on vsock port")
-						} else if err := usernetClient.WaitOpeningSSHPort(ctx, inst); err == nil {
-							hostAddress := net.JoinHostPort(inst.SSHAddress, strconv.Itoa(usernetSSHLocalPort))
-							if err := wrapper.startVsockForwarder(ctx, 22, hostAddress); err == nil {
-								logrus.Infof("Detected SSH server is listening on the vsock port; changed %s to proxy for the vsock port", hostAddress)
-								usernetSSHLocalPort = 0 // disable gvisor ssh port forwarding
-							} else {
-								logrus.WithError(err).Warn("Failed to detect SSH server on vsock port, falling back to usernet forwarder")
-							}
+						} else if err := usernetClient.WaitOpeningSSHPort(ctx, inst); err != nil {
+							logrus.WithError(err).Info("Failed to wait for the guest SSH server to become available, falling back to usernet forwarder")
+						} else if err := wrapper.checkSSHOverVsockAvailable(ctx, inst); err != nil {
+							logrus.WithError(err).Info("Failed to detect SSH server on vsock port, falling back to usernet forwarder")
+						} else if err := wrapper.startVsockForwarder(ctx, 22, hostAddress); err != nil {
+							logrus.WithError(err).Info("Failed to start SSH server forwarder on vsock port, falling back to usernet forwarder")
 						} else {
-							logrus.WithError(err).Warn("Failed to wait for the guest SSH server to become available, falling back to usernet forwarder")
+							logrus.Infof("Detected SSH server is listening on the vsock port; changed %s to proxy for the vsock port", hostAddress)
+							usernetSSHLocalPort = 0 // disable gvisor ssh port forwarding
 						}
 						err := usernetClient.ConfigureDriver(ctx, inst, usernetSSHLocalPort)
 						if err != nil {

pkg/driver/vz/vsock_forwarder.go

@@ -9,20 +9,20 @@ import (
 	"context"
 	"errors"
 	"net"
+	"path/filepath"
 
 	"github.com/containers/gvisor-tap-vsock/pkg/tcpproxy"
 	"github.com/sirupsen/logrus"
+
+	"github.com/lima-vm/lima/v2/pkg/limatype"
+	"github.com/lima-vm/lima/v2/pkg/limatype/dirnames"
+	"github.com/lima-vm/lima/v2/pkg/limatype/filenames"
+	"github.com/lima-vm/lima/v2/pkg/sshutil"
 )
 
 func (m *virtualMachineWrapper) startVsockForwarder(ctx context.Context, vsockPort uint32, hostAddress string) error {
-	// Test if the vsock port is open
-	conn, err := m.dialVsock(ctx, vsockPort)
-	if err != nil {
-		return err
-	}
-	conn.Close()
 	// Start listening on localhost:hostPort and forward to vsock:vsockPort
-	_, _, err = net.SplitHostPort(hostAddress)
+	_, _, err := net.SplitHostPort(hostAddress)
 	if err != nil {
 		return err
 	}
@@ -73,3 +73,17 @@ func (m *virtualMachineWrapper) dialVsock(_ context.Context, port uint32) (conn
 	}
 	return nil, err
 }
+
+func (m *virtualMachineWrapper) checkSSHOverVsockAvailable(ctx context.Context, inst *limatype.Instance) error {
+	user := *inst.Config.User.Name
+	configDir, err := dirnames.LimaConfigDir()
+	if err != nil {
+		return err
+	}
+	privateKeyPath := filepath.Join(configDir, filenames.UserPrivateKey)
+	vsockPort := uint32(22)
+	addr := "vsock:22"
+	return sshutil.WaitSSHReady(ctx, func(ctx context.Context) (net.Conn, error) {
+		return m.dialVsock(ctx, vsockPort)
+	}, addr, user, privateKeyPath, 1)
+}

pkg/networks/usernet/client.go

@@ -11,6 +11,7 @@ import (
 	"net"
 	"net/http"
 	"os"
+	"path/filepath"
 	"strconv"
 	"time"
 
@@ -19,6 +20,8 @@ import (
 
 	"github.com/lima-vm/lima/v2/pkg/httpclientutil"
 	"github.com/lima-vm/lima/v2/pkg/limatype"
+	"github.com/lima-vm/lima/v2/pkg/limatype/dirnames"
+	"github.com/lima-vm/lima/v2/pkg/limatype/filenames"
 	"github.com/lima-vm/lima/v2/pkg/limayaml"
 	"github.com/lima-vm/lima/v2/pkg/networks/usernet/dnshosts"
 )
@@ -140,8 +143,14 @@ func (c *Client) WaitOpeningSSHPort(ctx context.Context, inst *limatype.Instance
 	if err != nil {
 		return err
 	}
+	user := *inst.Config.User.Name
+	configDir, err := dirnames.LimaConfigDir()
+	if err != nil {
+		return err
+	}
+	privateKeyPath := filepath.Join(configDir, filenames.UserPrivateKey)
 	// -1 avoids both sides timing out simultaneously.
-	u := fmt.Sprintf("%s/extension/wait_port?ip=%s&port=22&timeout=%d", c.base, ipAddr, timeoutSeconds-1)
+	u := fmt.Sprintf("%s/extension/wait_ssh_server?ip=%s&port=22&timeout=%d&user=%s&privateKeyPath=%s", c.base, ipAddr, timeoutSeconds-1, user, privateKeyPath)
 	res, err := httpclientutil.Get(ctx, c.client, u)
 	if err != nil {
 		return err

pkg/networks/usernet/gvproxy.go

@@ -22,6 +22,8 @@ import (
 	"github.com/containers/gvisor-tap-vsock/pkg/virtualnetwork"
 	"github.com/sirupsen/logrus"
 	"golang.org/x/sync/errgroup"
+
+	"github.com/lima-vm/lima/v2/pkg/sshutil"
 )
 
 type GVisorNetstackOpts struct {
@@ -243,7 +245,7 @@ func httpServe(ctx context.Context, g *errgroup.Group, ln net.Listener, mux http
 
 func muxWithExtension(n *virtualnetwork.VirtualNetwork) *http.ServeMux {
 	m := n.Mux()
-	m.HandleFunc("/extension/wait_port", func(w http.ResponseWriter, r *http.Request) {
+	m.HandleFunc("/extension/wait_ssh_server", func(w http.ResponseWriter, r *http.Request) {
 		ip := r.URL.Query().Get("ip")
 		if net.ParseIP(ip) == nil {
 			msg := fmt.Sprintf("invalid ip address: %s", ip)
@@ -255,8 +257,15 @@ func muxWithExtension(n *virtualnetwork.VirtualNetwork) *http.ServeMux {
 			http.Error(w, err.Error(), http.StatusBadRequest)
 			return
 		}
-		port := uint16(port16)
-		addr := fmt.Sprintf("%s:%d", ip, port)
+		addr := net.JoinHostPort(ip, fmt.Sprintf("%d", uint16(port16)))
+
+		user := r.URL.Query().Get("user")
+		privateKeyPath := r.URL.Query().Get("privateKeyPath")
+		if user == "" || privateKeyPath == "" {
+			msg := "user and privateKeyPath query parameters are required"
+			http.Error(w, msg, http.StatusBadRequest)
+			return
+		}
 
 		timeoutSeconds := 10
 		if timeoutString := r.URL.Query().Get("timeout"); timeoutString != "" {
@@ -267,27 +276,14 @@ func muxWithExtension(n *virtualnetwork.VirtualNetwork) *http.ServeMux {
 			}
 			timeoutSeconds = int(timeout16)
 		}
-		ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeoutSeconds)*time.Second)
-		defer cancel()
+		dialContext := func(ctx context.Context) (net.Conn, error) {
+			return n.DialContextTCP(ctx, addr)
+		}
 		// Wait until the port is available.
-		for {
-			conn, err := n.DialContextTCP(ctx, addr)
-			if err == nil {
-				conn.Close()
-				logrus.Debugf("Port is available on %s", addr)
-				w.WriteHeader(http.StatusOK)
-				break
-			}
-			select {
-			case <-ctx.Done():
-				msg := fmt.Sprintf("timed out waiting for port to become available on %s", addr)
-				logrus.Warn(msg)
-				http.Error(w, msg, http.StatusRequestTimeout)
-				return
-			default:
-			}
-			logrus.Debugf("Waiting for port to become available on %s", addr)
-			time.Sleep(1 * time.Second)
+		if err = sshutil.WaitSSHReady(r.Context(), dialContext, addr, user, privateKeyPath, timeoutSeconds); err != nil {
+			http.Error(w, err.Error(), http.StatusRequestTimeout)
+		} else {
+			w.WriteHeader(http.StatusOK)
 		}
 	})
 	return m

pkg/sshutil/sshutil.go

@@ -11,6 +11,7 @@ import (
 	"errors"
 	"fmt"
 	"io/fs"
+	"net"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -19,11 +20,13 @@ import (
 	"slices"
 	"strings"
 	"sync"
+	"syscall"
 	"time"
 
 	"github.com/coreos/go-semver/semver"
 	"github.com/mattn/go-shellwords"
 	"github.com/sirupsen/logrus"
+	"golang.org/x/crypto/ssh"
 	"golang.org/x/sys/cpu"
 
 	"github.com/lima-vm/lima/v2/pkg/ioutilx"
@@ -509,3 +512,59 @@ func detectAESAcceleration() bool {
 	}
 	return cpu.ARM.HasAES || cpu.ARM64.HasAES || cpu.PPC64.IsPOWER8 || cpu.S390X.HasAES || cpu.X86.HasAES
 }
+
+// WaitSSHReady waits until the SSH server is ready to accept connections.
+// The dialContext function is used to create a connection to the SSH server.
+// The addr, user, privateKeyPath parameter is used for ssh.ClientConn creation.
+// The timeoutSeconds parameter specifies the maximum number of seconds to wait.
+func WaitSSHReady(ctx context.Context, dialContext func(context.Context) (net.Conn, error), addr, user, privateKeyPath string, timeoutSeconds int) error {
+	ctx, cancel := context.WithTimeout(ctx, time.Duration(timeoutSeconds)*time.Second)
+	defer cancel()
+
+	// Prepare signer
+	key, err := os.ReadFile(privateKeyPath)
+	if err != nil {
+		return fmt.Errorf("failed to read private key %q: %w", privateKeyPath, err)
+	}
+	signer, err := ssh.ParsePrivateKey(key)
+	if err != nil {
+		return fmt.Errorf("failed to parse private key %q: %w", privateKeyPath, err)
+	}
+	// Prepare ssh client config
+	sshConfig := &ssh.ClientConfig{
+		User:            user,
+		Auth:            []ssh.AuthMethod{ssh.PublicKeys(signer)},
+		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
+		Timeout:         10 * time.Second,
+	}
+
+	// Wait until the SSH server is available.
+	for {
+		conn, err := dialContext(ctx)
+		if err == nil {
+			sshConn, chans, reqs, err := ssh.NewClientConn(conn, addr, sshConfig)
+			if err == nil {
+				sshClient := ssh.NewClient(sshConn, chans, reqs)
+				return sshClient.Close()
+			}
+			conn.Close()
+			if !isRetryableError(err) {
+				return fmt.Errorf("failed to create ssh.Conn to %q: %w", addr, err)
+			}
+		}
+		logrus.Debugf("Waiting for SSH port to accept connections on %s", addr)
+		select {
+		case <-ctx.Done():
+			return fmt.Errorf("failed to waiting for SSH port to become available on %s: %w", addr, ctx.Err())
+		case <-time.After(1 * time.Second):
+			continue
+		}
+	}
+}
+
+func isRetryableError(err error) bool {
+	// Port forwarder accepted the connection, but the destination is not ready yet.
+	return errors.Is(err, syscall.ECONNRESET) ||
+		// SSH server not ready yet (e.g. host key not generated on initial boot).
+		strings.HasSuffix(err.Error(), "no supported methods remain")
+}