pkg/driver/qemu: Wait for SSH to be ready in AdditionalSetupForSSH()

- pkg/sshutil: Add `WaitSSHReady()`
`WaitSSHReady` waits until the SSH port is ready to accept connections.
The `dialContext` function is used to create a connection to the SSH server.
The `addressForLogging` parameter is used for logging purposes.
The `timeoutSeconds` parameter specifies the maximum number of seconds to wait.

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

Author: norio-nomura

pkg/driver/qemu/qemu_driver.go

@@ -37,6 +37,7 @@ import (
 	"github.com/lima-vm/lima/v2/pkg/osutil"
 	"github.com/lima-vm/lima/v2/pkg/ptr"
 	"github.com/lima-vm/lima/v2/pkg/reflectutil"
+	"github.com/lima-vm/lima/v2/pkg/sshutil"
 	"github.com/lima-vm/lima/v2/pkg/version/versionutil"
 )
 
@@ -721,6 +722,17 @@ func (l *LimaQemuDriver) ForwardGuestAgent() bool {
 	return l.vSockPort == 0 && l.virtioPort == ""
 }
 
-func (l *LimaQemuDriver) AdditionalSetupForSSH(_ context.Context) error {
+func (l *LimaQemuDriver) AdditionalSetupForSSH(ctx context.Context) error {
+	// Ensure that the QEMU instance is ready to accept SSH connections.
+	time.Sleep(10 * time.Second)
+	// Wait until the port is available.
+	addr := net.JoinHostPort("127.0.0.1", fmt.Sprintf("%d", l.SSHLocalPort))
+	dialContext := func(ctx context.Context) (net.Conn, error) {
+		dialer := net.Dialer{Timeout: 1 * time.Second}
+		return dialer.DialContext(ctx, "tcp", addr)
+	}
+	if err := sshutil.WaitSSHReady(ctx, dialContext, addr, 600); err != nil {
+		return err
+	}
 	return nil
 }

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 {
@@ -255,8 +257,7 @@ 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)))
 
 		timeoutSeconds := 10
 		if timeoutString := r.URL.Query().Get("timeout"); timeoutString != "" {
@@ -267,27 +268,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, timeoutSeconds); err != nil {
+			http.Error(w, err.Error(), http.StatusRequestTimeout)
+		} else {
+			w.WriteHeader(http.StatusOK)
 		}
 	})
 	return m

pkg/sshutil/sshutil.go

@@ -10,6 +10,7 @@ import (
 	"encoding/binary"
 	"errors"
 	"fmt"
+	"io"
 	"io/fs"
 	"net"
 	"os"
@@ -608,3 +609,57 @@ func findRegexpInSSHArgs(sshArgs []string, re *regexp.Regexp) string {
 	}
 	return ""
 }
+
+// WaitSSHReady waits until the SSH port is ready to accept connections.
+// The dialContext function is used to create a connection to the SSH server.
+// The addressForLogging parameter is used for logging purposes.
+// The timeoutSeconds parameter specifies the maximum number of seconds to wait.
+func WaitSSHReady(ctx context.Context, dialContext func(context.Context) (net.Conn, error), addressForLogging string, timeoutSeconds int) error {
+	ctx, cancel := context.WithTimeout(ctx, time.Duration(timeoutSeconds)*time.Second)
+	defer cancel()
+	// Wait until the SSH port is available.
+	for {
+		conn, err := dialContext(ctx)
+		if err == nil {
+			// Check if the SSH banner is received.
+			buf := make([]byte, 128)
+		RetryRead:
+			for {
+				if err := conn.SetReadDeadline(time.Now().Add(1 * time.Second)); err != nil {
+					logrus.WithError(err).Debugf("Failed to set read deadline on SSH connection to %s", addressForLogging)
+					break RetryRead
+				}
+				n, err := conn.Read(buf)
+				if err == nil {
+					if bytes.HasPrefix(buf[:n], []byte("SSH-")) {
+						conn.Close()
+						logrus.Debugf("SSH port is available on %s", addressForLogging)
+						return nil // SSH ready!
+					}
+					return fmt.Errorf("invalid SSH banner from %s: %q", addressForLogging, string(buf[:n]))
+				} else if ne, ok := err.(net.Error); ok && ne.Timeout() {
+					logrus.Debugf("Timeout reading SSH banner from %s, retrying...", addressForLogging)
+				} else if errors.Is(err, io.EOF) {
+					logrus.Debugf("EOF reading SSH banner from %s, retrying...", addressForLogging)
+					break RetryRead
+				} else {
+					return fmt.Errorf("failed to read SSH banner from %s: %w", addressForLogging, err)
+				}
+				select {
+				case <-ctx.Done():
+					break RetryRead
+				case <-time.After(1 * time.Second):
+					continue
+				}
+			}
+			conn.Close()
+		}
+		logrus.Debugf("Waiting for SSH port to accept connections on %s", addressForLogging)
+		select {
+		case <-ctx.Done():
+			return fmt.Errorf("failed to waiting for SSH port to become available on %s: %w", addressForLogging, ctx.Err())
+		case <-time.After(1 * time.Second):
+			continue
+		}
+	}
+}