pkg/sshutil: Update ssh_known_hosts if the SSH server does not accept host keys via cloud-init

Prepare HostKeyCallback that corrects known_hosts when the host key is not found.
This is required when the SSH server generates a new host key on the first boot,
instead of accepting the provided host keys via cloud-init.
e.g. https://github.com/lima-vm/alpine-lima

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

Author: norio-nomura

pkg/driver/krunkit/krunkit_driver_darwin_arm64.go

@@ -27,6 +27,7 @@ import (
 	"github.com/lima-vm/lima/v2/pkg/networks/usernet"
 	"github.com/lima-vm/lima/v2/pkg/osutil"
 	"github.com/lima-vm/lima/v2/pkg/ptr"
+	"github.com/lima-vm/lima/v2/pkg/sshutil"
 )
 
 type LimaKrunkitDriver struct {
@@ -315,6 +316,17 @@ func (l *LimaKrunkitDriver) GuestAgentConn(_ context.Context) (net.Conn, string,
 	return nil, "unix", nil
 }
 
-func (l *LimaKrunkitDriver) AdditionalSetupForSSH(_ context.Context) error {
+func (l *LimaKrunkitDriver) AdditionalSetupForSSH(ctx context.Context) error {
+	// 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)
+	}
+	user := *l.Instance.Config.User.Name
+	instanceName := l.Instance.Name
+	if err := sshutil.WaitSSHReady(ctx, dialContext, addr, user, instanceName, 600); err != nil {
+		return err
+	}
 	return nil
 }

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 {
+	// 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)
+	}
+	user := *l.Instance.Config.User.Name
+	instanceName := l.Instance.Name
+	if err := sshutil.WaitSSHReady(ctx, dialContext, addr, user, instanceName, 600); err != nil {
+		return err
+	}
 	return nil
 }

pkg/driver/wsl2/wsl_driver_windows.go

@@ -10,6 +10,7 @@ import (
 	"fmt"
 	"net"
 	"regexp"
+	"time"
 
 	"github.com/Microsoft/go-winio"
 	"github.com/Microsoft/go-winio/pkg/guid"
@@ -21,6 +22,7 @@ import (
 	"github.com/lima-vm/lima/v2/pkg/limayaml"
 	"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/windows"
 )
 
@@ -358,6 +360,17 @@ func (l *LimaWslDriver) ForwardGuestAgent() bool {
 	return l.vSockPort == 0 && l.virtioPort == ""
 }
 
-func (l *LimaWslDriver) AdditionalSetupForSSH(_ context.Context) error {
+func (l *LimaWslDriver) AdditionalSetupForSSH(ctx context.Context) error {
+	// 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)
+	}
+	user := *l.Instance.Config.User.Name
+	instanceName := l.Instance.Name
+	if err := sshutil.WaitSSHReady(ctx, dialContext, addr, user, instanceName, 600); err != nil {
+		return err
+	}
 	return nil
 }

pkg/sshutil/sshutil.go

@@ -533,8 +533,9 @@ func detectAESAcceleration() bool {
 
 // 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, instanceName parameter is used for ssh.ClientConn creation.
+// The addr, user, and instanceName parameters are used for ssh.ClientConn creation.
 // The timeoutSeconds parameter specifies the maximum number of seconds to wait.
+// If the SSH server uses a host key not in the instance's ssh_known_hosts, the host key will be updated.
 func WaitSSHReady(ctx context.Context, dialContext func(context.Context) (net.Conn, error), addr, user, instanceName string, timeoutSeconds int) error {
 	ctx, cancel := context.WithTimeout(ctx, time.Duration(timeoutSeconds)*time.Second)
 	defer cancel()
@@ -544,16 +545,22 @@ func WaitSSHReady(ctx context.Context, dialContext func(context.Context) (net.Co
 	if err != nil {
 		return err
 	}
-	// Prepare HostKeyCallback
+	// Prepare HostKeyCallback that corrects known_hosts when the host key is not found.
+	// This is required when the SSH server generates a new host key on the first boot,
+	// instead of accepting the provided host keys via cloud-init.
+	// e.g. https://github.com/lima-vm/alpine-lima
 	hostKeyChecker, err := HostKeyCheckerWithKeysInKnownHosts(instanceName)
 	if err != nil {
 		return err
 	}
+	// Ensure known_hosts is updated at the end
+	defer hostKeyChecker.updateKnownHosts()
+
 	// Prepare ssh client config
 	sshConfig := &ssh.ClientConfig{
 		User:            user,
 		Auth:            []ssh.AuthMethod{ssh.PublicKeys(signer)},
-		HostKeyCallback: hostKeyChecker,
+		HostKeyCallback: hostKeyChecker.check,
 		Timeout:         10 * time.Second,
 	}
 	// Wait until the SSH server is available.
@@ -580,11 +587,16 @@ func WaitSSHReady(ctx context.Context, dialContext func(context.Context) (net.Co
 	}
 }
 
+// errHostKeyMismatch is returned when the SSH host key does not match known hosts.
+var errHostKeyMismatch = errors.New("ssh: host key mismatch")
+
 func isRetryableError(err error) bool {
 	// Port forwarder accepted the connection, but the destination is not ready yet.
 	return osutil.IsConnectionResetError(err) ||
 		// SSH server not ready yet (e.g. host key not generated on initial boot).
-		strings.HasSuffix(err.Error(), "no supported methods remain")
+		strings.HasSuffix(err.Error(), "no supported methods remain") ||
+		// Host key is not yet in known_hosts, but will be corrected, so we can retry.
+		errors.Is(err, errHostKeyMismatch)
 }
 
 // userPrivateKeySigner returns the user's private key signer.
@@ -606,19 +618,71 @@ func userPrivateKeySigner() (ssh.Signer, error) {
 	return signer, nil
 }
 
-func HostKeyCheckerWithKeysInKnownHosts(instanceName string) (ssh.HostKeyCallback, error) {
+type hostKeyCheckerWithCorrectingUnknownKeys struct {
+	instanceName string
+	publicKeys   map[string]ssh.PublicKey
+	unknownKeys  map[string]ssh.PublicKey
+}
+
+// check checks whether the given host key is in the known hosts.
+// If the host key is not found, it is recorded in unknownKeys for later correction.
+func (h *hostKeyCheckerWithCorrectingUnknownKeys) check(hostname string, remote net.Addr, key ssh.PublicKey) error {
+	marshaledKey := string(key.Marshal())
+	if _, ok := h.publicKeys[marshaledKey]; ok {
+		return nil
+	}
+	if _, ok := h.unknownKeys[marshaledKey]; ok {
+		return nil
+	}
+	logrus.Warnf("SSH host key for instance %q not found in %s; adding it", h.instanceName, filenames.SSHKnownHosts)
+	h.unknownKeys[marshaledKey] = key
+	// If always returning nil here, GitHub Advanced Security may report "Use of insecure HostKeyCallback implementation".
+	// So, we return an error here to make the SSH client report the host key mismatch.
+	return errHostKeyMismatch
+}
+
+// updateKnownHosts updates the known_hosts file with any unknown keys recorded during checks.
+// It is required to call this method after using the hostKeyCheckerWithCorrectingUnknownKeys
+// to ensure that the known_hosts file is updated appropriately.
+func (h *hostKeyCheckerWithCorrectingUnknownKeys) updateKnownHosts() error {
+	if len(h.unknownKeys) == 0 {
+		return nil
+	}
+	// If there are unknown keys, our provided host keys via cloud-init were not accepted by the Guest SSH server.
+	// We need to replace known_hosts file with the unknown keys that guest SSH server is actually using.
+	logrus.Infof("Updating %s file for instance %q with %d new host key(s)", filenames.SSHKnownHosts, h.instanceName, len(h.unknownKeys))
+	instanceDir, err := dirnames.InstanceDir(h.instanceName)
+	if err != nil {
+		return fmt.Errorf("failed to get instance dir for instance %q: %w", h.instanceName, err)
+	}
+	knownHostsPath := filepath.Join(instanceDir, filenames.SSHKnownHosts)
+	hostname := hostname.FromInstName(h.instanceName)
+	var sshKnownHosts []byte
+	for _, key := range h.unknownKeys {
+		sshKnownHosts = append(sshKnownHosts, knownhosts.Line([]string{hostname}, key)...)
+		sshKnownHosts = append(sshKnownHosts, '\n')
+	}
+	if err := os.WriteFile(knownHostsPath, sshKnownHosts, 0o644); err != nil {
+		return fmt.Errorf("failed to write known_hosts file at %s: %w", knownHostsPath, err)
+	}
+	return nil
+}
+
+// HostKeyCheckerWithKeysInKnownHosts creates a host key checker using the known hosts file
+// located in the instance directory.
+func HostKeyCheckerWithKeysInKnownHosts(instanceName string) (hostKeyCheckerWithCorrectingUnknownKeys, error) {
 	publicKeys, err := PublicKeysFromKnownHosts(instanceName)
 	if err != nil {
-		return nil, err
+		return hostKeyCheckerWithCorrectingUnknownKeys{}, err
 	}
-	return func(_ string, _ net.Addr, key ssh.PublicKey) error {
-		keyBytes := key.Marshal()
-		for _, pk := range publicKeys {
-			if bytes.Equal(keyBytes, pk.Marshal()) {
-				return nil
-			}
-		}
-		return errors.New("ssh: host key mismatch")
+	publicKeysMap := make(map[string]ssh.PublicKey, len(publicKeys))
+	for _, pk := range publicKeys {
+		publicKeysMap[string(pk.Marshal())] = pk
+	}
+	return hostKeyCheckerWithCorrectingUnknownKeys{
+		instanceName: instanceName,
+		publicKeys:   publicKeysMap,
+		unknownKeys:  make(map[string]ssh.PublicKey),
 	}, nil
 }