Merge pull request #1725 from afbjorklund/reverse-remove

Don't delete local socket when reverse
lima-vm · Aug 12, 2023 · 27b4e78 · 27b4e78
2 parents ca9f32f + 05cff79
commit 27b4e78
Showing 1 changed file with 39 additions and 11 deletions.
diff --git a/pkg/hostagent/hostagent.go b/pkg/hostagent/hostagent.go
@@ -567,6 +567,21 @@ const (
 	verbCancel  = "cancel"
 )
 
+func executeSSH(ctx context.Context, sshConfig *ssh.SSHConfig, port int, command ...string) error {
+	args := sshConfig.Args()
+	args = append(args,
+		"-p", strconv.Itoa(port),
+		"127.0.0.1",
+		"--",
+	)
+	args = append(args, command...)
+	cmd := exec.CommandContext(ctx, sshConfig.Binary(), args...)
+	if out, err := cmd.Output(); err != nil {
+		return fmt.Errorf("failed to run %v: %q: %w", cmd.Args, string(out), err)
+	}
+	return nil
+}
+
 func forwardSSH(ctx context.Context, sshConfig *ssh.SSHConfig, port int, local, remote string, verb string, reverse bool) error {
 	args := sshConfig.Args()
 	args = append(args,
@@ -594,36 +609,49 @@ func forwardSSH(ctx context.Context, sshConfig *ssh.SSHConfig, port int, local,
 		case verbForward:
 			if reverse {
 				logrus.Infof("Forwarding %q (host) to %q (guest)", local, remote)
+				if err := executeSSH(ctx, sshConfig, port, "rm", "-f", remote); err != nil {
+					logrus.WithError(err).Warnf("Failed to clean up %q (guest) before setting up forwarding", remote)
+				}
 			} else {
 				logrus.Infof("Forwarding %q (guest) to %q (host)", remote, local)
-			}
-			if err := os.RemoveAll(local); err != nil {
-				logrus.WithError(err).Warnf("Failed to clean up %q (host) before setting up forwarding", local)
+				if err := os.RemoveAll(local); err != nil {
+					logrus.WithError(err).Warnf("Failed to clean up %q (host) before setting up forwarding", local)
+				}
 			}
 			if err := os.MkdirAll(filepath.Dir(local), 0750); err != nil {
 				return fmt.Errorf("can't create directory for local socket %q: %w", local, err)
 			}
 		case verbCancel:
 			if reverse {
 				logrus.Infof("Stopping forwarding %q (host) to %q (guest)", local, remote)
+				if err := executeSSH(ctx, sshConfig, port, "rm", "-f", remote); err != nil {
+					logrus.WithError(err).Warnf("Failed to clean up %q (guest) after stopping forwarding", remote)
+				}
 			} else {
 				logrus.Infof("Stopping forwarding %q (guest) to %q (host)", remote, local)
+				defer func() {
+					if err := os.RemoveAll(local); err != nil {
+						logrus.WithError(err).Warnf("Failed to clean up %q (host) after stopping forwarding", local)
+					}
+				}()
 			}
-			defer func() {
-				if err := os.RemoveAll(local); err != nil {
-					logrus.WithError(err).Warnf("Failed to clean up %q (host) after stopping forwarding", local)
-				}
-			}()
 		default:
 			panic(fmt.Errorf("invalid verb %q", verb))
 		}
 	}
 	cmd := exec.CommandContext(ctx, sshConfig.Binary(), args...)
 	if out, err := cmd.Output(); err != nil {
 		if verb == verbForward && strings.HasPrefix(local, "/") {
-			logrus.WithError(err).Warnf("Failed to set up forward from %q (guest) to %q (host)", remote, local)
-			if removeErr := os.RemoveAll(local); err != nil {
-				logrus.WithError(removeErr).Warnf("Failed to clean up %q (host) after forwarding failed", local)
+			if reverse {
+				logrus.WithError(err).Warnf("Failed to set up forward from %q (host) to %q (guest)", local, remote)
+				if err := executeSSH(ctx, sshConfig, port, "rm", "-f", remote); err != nil {
+					logrus.WithError(err).Warnf("Failed to clean up %q (guest) after forwarding failed", remote)
+				}
+			} else {
+				logrus.WithError(err).Warnf("Failed to set up forward from %q (guest) to %q (host)", remote, local)
+				if removeErr := os.RemoveAll(local); err != nil {
+					logrus.WithError(removeErr).Warnf("Failed to clean up %q (host) after forwarding failed", local)
+				}
 			}
 		}
 		return fmt.Errorf("failed to run %v: %q: %w", cmd.Args, string(out), err)