pkg/driver: Add Driver.AdditionalSetupForSSH()

## pkg/driver: Add `Driver.AdditionalSetupForSSH()`
AdditionalSetupForSSH provides additional setup required for SSH connection.
It is called after the VM is started, before the first SSH connection.
It returns an error if the setup fails.

## pkg/driver/vz: Change waiting `waitSSHLocalPortAccessible` in `Driver.AdditionalSetupForSSH()` instead of `startVM`.
`waitSSHLocalPortAccessible` is being used for waiting configuration SSH over VSOCK.
It makes `LimaVzDriver.Start()` to return as soon as possible.

Signed-off-by: Norio Nomura <norio.nomura@gmail.com>

Author: norio-nomura

pkg/driver/driver.go

@@ -88,6 +88,11 @@ type Driver interface {
 	FillConfig(ctx context.Context, cfg *limatype.LimaYAML, filePath string) error
 
 	SSHAddress(ctx context.Context) (string, error)
+
+	// AdditionalSetupForSSH provides additional setup required for SSH connection.
+	// It is called after VM is started, before first SSH connection.
+	// It returns an error if the setup fails.
+	AdditionalSetupForSSH(ctx context.Context) error
 }
 
 type ConfiguredDriver struct {

pkg/driver/external/client/methods.go

@@ -323,3 +323,16 @@ func (d *DriverClient) BootScripts() (map[string][]byte, error) {
 	d.logger.Debugf("Boot scripts retrieved successfully: %d scripts", len(resp.Scripts))
 	return resp.Scripts, nil
 }
+
+func (d *DriverClient) AdditionalSetupForSSH(ctx context.Context) error {
+	d.logger.Debug("Performing additional setup for SSH connection")
+
+	_, err := d.DriverSvc.AdditionalSetupForSSH(ctx, &emptypb.Empty{})
+	if err != nil {
+		d.logger.Errorf("Failed to perform additional setup for SSH: %v", err)
+		return err
+	}
+
+	d.logger.Debug("Additional setup for SSH completed successfully")
+	return nil
+}

pkg/driver/external/driver.pb.desc

@@ -1,5 +1,5 @@
 
-�
+�
 driver.protogoogle/protobuf/empty.proto"�

 BootScriptsResponse;
 scripts
 (2!.BootScriptsResponse.ScriptsEntryRscripts:
@@ -30,7 +30,7 @@ connection")
 ListSnapshotsResponse
 	snapshots
 
(	R	snapshots"B
 ForwardGuestAgentResponse%
-should_forward
 
(RshouldForward2�	
+should_forward
 
(RshouldForward2�	
 Driver:
 Validate.google.protobuf.Empty.google.protobuf.Empty8
 Create.google.protobuf.Empty.google.protobuf.Empty<
@@ -52,4 +52,5 @@ CreateDisk.google.protobuf.Empty.google.protobuf.Empty1
 	Configure.SetConfigRequest.google.protobuf.Empty-
 Info.google.protobuf.Empty.InfoResponse9
 
-SSHAddress.google.protobuf.Empty.SSHAddressResponseB0Z.github.com/lima-vm/lima/v2/pkg/driver/externalbproto3
+SSHAddress.google.protobuf.Empty.SSHAddressResponseG
+AdditionalSetupForSSH.google.protobuf.Empty.google.protobuf.EmptyB0Z.github.com/lima-vm/lima/v2/pkg/driver/externalbproto3

pkg/driver/external/driver.pb.go

@@ -28,6 +28,10 @@ service Driver {
   rpc Configure(SetConfigRequest) returns (google.protobuf.Empty);
   rpc Info(google.protobuf.Empty) returns (InfoResponse);
   rpc SSHAddress(google.protobuf.Empty) returns (SSHAddressResponse);
+
+  // AdditionalSetupForSSH provides additional setup required for SSH connection.
+  // It is called after VM is started, before first SSH connection.
+  rpc AdditionalSetupForSSH(google.protobuf.Empty) returns (google.protobuf.Empty);
 }
 
 message BootScriptsResponse {

pkg/driver/external/driver.proto

@@ -284,3 +284,14 @@ func (s *DriverServer) ForwardGuestAgent(_ context.Context, _ *emptypb.Empty) (*
 	s.logger.Debug("Received ForwardGuestAgent request")
 	return &pb.ForwardGuestAgentResponse{ShouldForward: s.driver.ForwardGuestAgent()}, nil
 }
+
+func (s *DriverServer) AdditionalSetupForSSH(ctx context.Context, _ *emptypb.Empty) (*emptypb.Empty, error) {
+	s.logger.Debug("Received AdditionalSetupForSSH request")
+	err := s.driver.AdditionalSetupForSSH(ctx)
+	if err != nil {
+		s.logger.Errorf("AdditionalSetupForSSH failed: %v", err)
+		return &emptypb.Empty{}, err
+	}
+	s.logger.Debug("AdditionalSetupForSSH succeeded")
+	return &emptypb.Empty{}, nil
+}

pkg/driver/external/driver_grpc.pb.go

@@ -720,3 +720,7 @@ func (l *LimaQemuDriver) ForwardGuestAgent() bool {
 	// if driver is not providing, use host agent
 	return l.vSockPort == 0 && l.virtioPort == ""
 }
+
+func (l *LimaQemuDriver) AdditionalSetupForSSH(_ context.Context) error {
+	return nil
+}

pkg/driver/external/server/methods.go

@@ -52,28 +52,26 @@ type virtualMachineWrapper struct {
 // Hold all *os.File created via socketpair() so that they won't get garbage collected. f.FD() gets invalid if f gets garbage collected.
 var vmNetworkFiles = make([]*os.File, 1)
 
-func startVM(ctx context.Context, inst *limatype.Instance, sshLocalPort int) (*virtualMachineWrapper, chan error, error) {
+func startVM(ctx context.Context, inst *limatype.Instance, sshLocalPort int) (vm *virtualMachineWrapper, waitSSHLocalPortAccessible <-chan any, errCh chan error, err error) {
 	usernetClient, stopUsernet, err := startUsernet(ctx, inst)
 	if err != nil {
-		return nil, nil, err
+		return nil, nil, nil, err
 	}
 
 	machine, err := createVM(ctx, inst)
 	if err != nil {
-		return nil, nil, err
+		return nil, nil, nil, err
 	}
 
 	err = machine.Start()
 	if err != nil {
-		return nil, nil, err
+		return nil, nil, nil, err
 	}
 
 	wrapper := &virtualMachineWrapper{VirtualMachine: machine, stopped: false}
+	notifySSHLocalPortAccessible := make(chan any)
+	sendErrCh := make(chan error)
 
-	errCh := make(chan error)
-
-	waitSSHLocalPortAccessible := make(chan struct{})
-	defer close(waitSSHLocalPortAccessible)
 	go func() {
 		// Handle errors via errCh and handle stop vm during context close
 		defer func() {
@@ -95,42 +93,44 @@ func startVM(ctx context.Context, inst *limatype.Instance, sshLocalPort int) (*v
 					pidFile := filepath.Join(inst.Dir, filenames.PIDFile(*inst.Config.VMType))
 					if _, err := os.Stat(pidFile); !errors.Is(err, os.ErrNotExist) {
 						logrus.Errorf("pidfile %q already exists", pidFile)
-						errCh <- err
+						sendErrCh <- err
 					}
 					if err := os.WriteFile(pidFile, []byte(strconv.Itoa(os.Getpid())+"\n"), 0o644); err != nil {
 						logrus.Errorf("error writing to pid fil %q", pidFile)
-						errCh <- err
+						sendErrCh <- err
 					}
 					logrus.Info("[VZ] - vm state change: running")
 
-					usernetSSHLocalPort := sshLocalPort
-					useSSHOverVsock := true
-					if envVar := os.Getenv("LIMA_SSH_OVER_VSOCK"); envVar != "" {
-						b, err := strconv.ParseBool(envVar)
-						if err != nil {
-							logrus.WithError(err).Warnf("invalid LIMA_SSH_OVER_VSOCK value %q", envVar)
-						} else {
-							useSSHOverVsock = b
+					go func() {
+						defer close(notifySSHLocalPortAccessible)
+						usernetSSHLocalPort := sshLocalPort
+						useSSHOverVsock := true
+						if envVar := os.Getenv("LIMA_SSH_OVER_VSOCK"); envVar != "" {
+							b, err := strconv.ParseBool(envVar)
+							if err != nil {
+								logrus.WithError(err).Warnf("invalid LIMA_SSH_OVER_VSOCK value %q", envVar)
+							} else {
+								useSSHOverVsock = b
+							}
 						}
-					}
-					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
+						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 {
-							logrus.WithError(err).Warn("Failed to detect SSH server on vsock port, falling back to usernet forwarder")
+							logrus.WithError(err).Warn("Failed to wait for the guest SSH server to become available, 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")
-					}
-					err := usernetClient.ConfigureDriver(ctx, inst, usernetSSHLocalPort)
-					if err != nil {
-						errCh <- err
-					}
-					waitSSHLocalPortAccessible <- struct{}{}
+						err := usernetClient.ConfigureDriver(ctx, inst, usernetSSHLocalPort)
+						if err != nil {
+							sendErrCh <- err
+						}
+					}()
 				case vz.VirtualMachineStateStopped:
 					logrus.Info("[VZ] - vm state change: stopped")
 					wrapper.mu.Lock()
@@ -140,15 +140,14 @@ func startVM(ctx context.Context, inst *limatype.Instance, sshLocalPort int) (*v
 					if stopUsernet != nil {
 						stopUsernet()
 					}
-					errCh <- errors.New("vz driver state stopped")
+					sendErrCh <- errors.New("vz driver state stopped")
 				default:
 					logrus.Debugf("[VZ] - vm state change: %q", newState)
 				}
 			}
 		}
 	}()
-	<-waitSSHLocalPortAccessible
-	return wrapper, errCh, err
+	return wrapper, notifySSHLocalPortAccessible, sendErrCh, err
 }
 
 func startUsernet(ctx context.Context, inst *limatype.Instance) (*usernet.Client, context.CancelFunc, error) {