[sshshim] Add and pass ctx parameter for sshshim execution

internal/boxcli/midcobra/midcobra.go

@@ -9,6 +9,7 @@ import (
 	"os/exec"
 
 	"github.com/spf13/cobra"
+
 	"go.jetpack.io/devbox/internal/boxcli/usererr"
 	"go.jetpack.io/devbox/internal/debug"
 	"go.jetpack.io/devbox/internal/ux"

internal/boxcli/root.go

@@ -11,6 +11,7 @@ import (
 	"strings"
 
 	"github.com/spf13/cobra"
+
 	"go.jetpack.io/devbox/internal/boxcli/midcobra"
 	"go.jetpack.io/devbox/internal/cloud/openssh/sshshim"
 	"go.jetpack.io/devbox/internal/debug"
@@ -97,17 +98,18 @@ func Execute(ctx context.Context, args []string) int {
 }
 
 func Main() {
+	ctx := context.Background()
 	if strings.HasSuffix(os.Args[0], "ssh") ||
 		strings.HasSuffix(os.Args[0], "scp") {
-		code := sshshim.Execute(os.Args)
-		os.Exit(code)
+		os.Exit(sshshim.Execute(ctx, os.Args))
 	}
+
 	if len(os.Args) > 1 && os.Args[1] == "bug" {
 		telemetry.ReportErrors()
 		return
 	}
-	code := Execute(context.Background(), os.Args[1:])
-	os.Exit(code)
+
+	os.Exit(Execute(ctx, os.Args[1:]))
 }
 
 func listAllCommands(cmd *cobra.Command, indent string) {

internal/cloud/openssh/sshshim/command.go

@@ -4,34 +4,37 @@
 package sshshim
 
 import (
+	"context"
 	"fmt"
 	"os"
 
 	"go.jetpack.io/devbox/internal/debug"
 	"go.jetpack.io/devbox/internal/telemetry"
 )
 
-func Execute(args []string) int {
+func Execute(ctx context.Context, args []string) int {
 	defer debug.Recover()
 	telemetry.Start(telemetry.AppSSHShim)
 	defer telemetry.Stop()
 
-	if err := execute(args); err != nil {
+	if err := execute(ctx, args); err != nil {
 		telemetry.Error(err, telemetry.Metadata{})
 		return 1
 	}
 	return 0
 }
 
-func execute(args []string) error {
+func execute(ctx context.Context, args []string) error {
 	EnableDebug() // Always enable for now.
 	debug.Log("os.Args: %v", args)
 
-	if alive, err := EnsureLiveVMOrTerminateMutagenSessions(args[1:]); err != nil {
+	alive, err := EnsureLiveVMOrTerminateMutagenSessions(ctx, args[1:])
+	if err != nil {
 		debug.Log("ensureLiveVMOrTerminateMutagenSessions error: %v", err)
 		fmt.Fprintf(os.Stderr, "%v", err)
 		return err
-	} else if !alive {
+	}
+	if !alive {
 		return nil
 	}
 
@@ -40,5 +43,6 @@ func execute(args []string) error {
 		fmt.Fprintf(os.Stderr, "%v", err)
 		return err
 	}
+
 	return nil
 }

internal/cloud/openssh/sshshim/mutagen.go

@@ -11,13 +11,14 @@ import (
 	"time"
 
 	"github.com/pkg/errors"
+
 	"go.jetpack.io/devbox/internal/cloud/mutagenbox"
 	"go.jetpack.io/devbox/internal/debug"
 )
 
 // EnsureLiveVMOrTerminateMutagenSessions returns true if a liveVM is found, OR sshArgs were connecting to a server that is not a devbox-VM.
 // EnsureLiveVMOrTerminateMutagenSessions returns false iff the sshArgs were connecting to a devbox VM AND a deadVM is found.
-func EnsureLiveVMOrTerminateMutagenSessions(sshArgs []string) (bool, error) {
+func EnsureLiveVMOrTerminateMutagenSessions(ctx context.Context, sshArgs []string) (bool, error) {
 	vmAddr := vmAddressIfAny(sshArgs)
 
 	debug.Log("Found vmAddr: %s", vmAddr)
@@ -28,9 +29,11 @@ func EnsureLiveVMOrTerminateMutagenSessions(sshArgs []string) (bool, error) {
 		return true, nil
 	}
 
-	if isActive, err := checkActiveVMWithRetries(vmAddr); err != nil {
+	isActive, err := checkActiveVMWithRetries(ctx, vmAddr)
+	if err != nil {
 		return false, errors.WithStack(err)
-	} else if !isActive {
+	}
+	if !isActive {
 		debug.Log("terminating mutagen session for vm: %s", vmAddr)
 		// If no vm is active, then we should terminate the running mutagen sessions
 		return false, terminateMutagenSessions(vmAddr)
@@ -57,25 +60,25 @@ func terminateMutagenSessions(vmAddr string) error {
 	return mutagenbox.ForwardTerminateByHost(hostname)
 }
 
-func checkActiveVMWithRetries(vmAddr string) (bool, error) {
-	var err error
+func checkActiveVMWithRetries(ctx context.Context, vmAddr string) (bool, error) {
+	var finalErr error
 
 	// Try 3 times:
 	for num := 0; num < 3; num++ {
-		var isActive bool
-		isActive, err = checkActiveVM(vmAddr)
+		isActive, err := checkActiveVM(ctx, vmAddr)
 		if err == nil && isActive {
 			// found an active VM
 			return true, nil
 		}
+		finalErr = err
 		time.Sleep(10 * time.Second)
 		debug.Log("Try %d failed to find activeVM for %s", num, vmAddr)
 	}
-	return false, err
+	return false, finalErr
 }
 
-func checkActiveVM(vmAddr string) (bool, error) {
-	ctx, cancel := context.WithTimeout(context.Background(), time.Minute*2)
+func checkActiveVM(ctx context.Context, vmAddr string) (bool, error) {
+	ctx, cancel := context.WithTimeout(ctx, time.Minute*2)
 	defer cancel()
 	cmd := exec.CommandContext(ctx, "ssh", vmAddr, "echo 'alive'")