Add --reuse-sock flag so browser IDEs can reuse another SSH connectio…

…ns SSH_AUTH_SOCK (#1471) * Add --reuse-sock flag so browser IDEs can reuse another SSH connections SSH_AUTH_SOCK * Fix comment * Fix lint error * Fix lint error * Update UML * Use random string to identify auth sock * Detect default IDE as well as explicit * Add some docs * Use official ssh dep * Add doc blocks] * Pass auth sock id to other IDEs * Fix marimo and jlab * Move SSHAuthSockID to CLi options * PR tidy up * Wrap backhaul connection in check if reuse sock is set (not in proxy mode) * Make info logs debug * Remove log
loft-sh · Jan 8, 2025 · a4e5a36 · a4e5a36
1 parent 9624d9c
commit a4e5a36
Show file tree

Hide file tree

Showing 19 changed files with 222 additions and 45 deletions.
diff --git a/cmd/helper/ssh_server.go b/cmd/helper/ssh_server.go
@@ -22,11 +22,12 @@ import (
 type SSHServerCmd struct {
 	*flags.GlobalFlags
 
-	Token         string
-	Address       string
-	Stdio         bool
-	TrackActivity bool
-	Workdir       string
+	Token            string
+	Address          string
+	Stdio            bool
+	TrackActivity    bool
+	ReuseSSHAuthSock string
+	Workdir          string
 }
 
 // NewSSHServerCmd creates a new ssh command
@@ -44,6 +45,8 @@ func NewSSHServerCmd(flags *flags.GlobalFlags) *cobra.Command {
 	sshCmd.Flags().StringVar(&cmd.Address, "address", fmt.Sprintf("0.0.0.0:%d", helperssh.DefaultPort), "Address to listen to")
 	sshCmd.Flags().BoolVar(&cmd.Stdio, "stdio", false, "Will listen on stdout and stdin instead of an address")
 	sshCmd.Flags().BoolVar(&cmd.TrackActivity, "track-activity", false, "If enabled will write the last activity time to a file")
+	sshCmd.Flags().StringVar(&cmd.ReuseSSHAuthSock, "reuse-ssh-auth-sock", "", "If set, the SSH_AUTH_SOCK is expected to already be available in the workspace (under /tmp using the key provided) and the connection reuses this instead of creating a new one")
+	_ = sshCmd.Flags().MarkHidden("reuse-ssh-auth-sock")
 	sshCmd.Flags().StringVar(&cmd.Token, "token", "", "Base64 encoded token to use")
 	sshCmd.Flags().StringVar(&cmd.Workdir, "workdir", "", "Directory where commands will run on the host")
 	return sshCmd
@@ -89,7 +92,7 @@ func (cmd *SSHServerCmd) Run(_ *cobra.Command, _ []string) error {
 	}
 
 	// start the server
-	server, err := helperssh.NewServer(cmd.Address, hostKey, keys, cmd.Workdir, log.Default.ErrorStreamOnly())
+	server, err := helperssh.NewServer(cmd.Address, hostKey, keys, cmd.Workdir, cmd.ReuseSSHAuthSock, log.Default.ErrorStreamOnly())
 	if err != nil {
 		return err
 	}

diff --git a/cmd/ssh.go b/cmd/ssh.go
@@ -56,6 +56,7 @@ type SSHCmd struct {
 
 	Stdio                     bool
 	JumpContainer             bool
+	ReuseSSHAuthSock          string
 	AgentForwarding           bool
 	GPGAgentForwarding        bool
 	GitSSHSignatureForwarding bool
@@ -114,6 +115,8 @@ func NewSSHCmd(f *flags.GlobalFlags) *cobra.Command {
 	sshCmd.Flags().StringVar(&cmd.WorkDir, "workdir", "", "The working directory in the container")
 	sshCmd.Flags().BoolVar(&cmd.Proxy, "proxy", false, "If true will act as intermediate proxy for a proxy provider")
 	sshCmd.Flags().BoolVar(&cmd.AgentForwarding, "agent-forwarding", true, "If true forward the local ssh keys to the remote machine")
+	sshCmd.Flags().StringVar(&cmd.ReuseSSHAuthSock, "reuse-ssh-auth-sock", "", "If set, the SSH_AUTH_SOCK is expected to already be available in the workspace (under /tmp using the key provided) and the connection reuses this instead of creating a new one")
+	_ = sshCmd.Flags().MarkHidden("reuse-ssh-auth-sock")
 	sshCmd.Flags().BoolVar(&cmd.GPGAgentForwarding, "gpg-agent-forwarding", false, "If true forward the local gpg-agent to the remote machine")
 	sshCmd.Flags().BoolVar(&cmd.Stdio, "stdio", false, "If true will tunnel connection through stdout and stdin")
 	sshCmd.Flags().BoolVar(&cmd.StartServices, "start-services", true, "If false will not start any port-forwarding or git / docker credentials helper")
@@ -446,6 +449,10 @@ func (cmd *SSHCmd) startTunnel(ctx context.Context, devPodConfig *config.Config,
 
 	log.Debugf("Run outer container tunnel")
 	command := fmt.Sprintf("'%s' helper ssh-server --track-activity --stdio --workdir '%s'", agent.ContainerDevPodHelperLocation, workdir)
+	if cmd.ReuseSSHAuthSock != "" {
+		log.Debug("Reusing SSH_AUTH_SOCK")
+		command += fmt.Sprintf(" --reuse-ssh-auth-sock=%s", cmd.ReuseSSHAuthSock)
+	}
 	if cmd.Debug {
 		command += " --debug"
 	}
@@ -513,7 +520,7 @@ func (cmd *SSHCmd) startServices(
 	log log.Logger,
 ) {
 	if cmd.User != "" {
-		err := tunnel.RunInContainer(
+		err := tunnel.RunServices(
 			ctx,
 			devPodConfig,
 			containerClient,

diff --git a/cmd/up.go b/cmd/up.go
@@ -23,6 +23,7 @@ import (
 	config2 "github.com/loft-sh/devpod/pkg/devcontainer/config"
 	"github.com/loft-sh/devpod/pkg/devcontainer/sshtunnel"
 	dpFlags "github.com/loft-sh/devpod/pkg/flags"
+	"github.com/loft-sh/devpod/pkg/ide"
 	"github.com/loft-sh/devpod/pkg/ide/fleet"
 	"github.com/loft-sh/devpod/pkg/ide/jetbrains"
 	"github.com/loft-sh/devpod/pkg/ide/jupyter"
@@ -37,6 +38,7 @@ import (
 	devssh "github.com/loft-sh/devpod/pkg/ssh"
 	"github.com/loft-sh/devpod/pkg/telemetry"
 	"github.com/loft-sh/devpod/pkg/tunnel"
+	"github.com/loft-sh/devpod/pkg/util"
 	"github.com/loft-sh/devpod/pkg/version"
 	workspace2 "github.com/loft-sh/devpod/pkg/workspace"
 	"github.com/loft-sh/log"
@@ -153,6 +155,19 @@ func (cmd *UpCmd) Run(
 		cmd.Recreate = true
 	}
 
+	// check if we are a browser IDE and need to reuse the SSH_AUTH_SOCK
+	targetIDE := client.WorkspaceConfig().IDE.Name
+	// Check override
+	if cmd.IDE != "" {
+		targetIDE = cmd.IDE
+	}
+	if !cmd.Proxy && ide.ReusesAuthSock(targetIDE) {
+		cmd.SSHAuthSockID = util.RandStringBytes(10)
+		log.Debug("Reusing SSH_AUTH_SOCK", cmd.SSHAuthSockID)
+	} else if cmd.Proxy && ide.ReusesAuthSock(targetIDE) {
+		log.Debug("Reusing SSH_AUTH_SOCK is not supported with proxy mode, consider launching the IDE from the platform UI")
+	}
+
 	// run devpod agent up
 	result, err := cmd.devPodUp(ctx, devPodConfig, client, log)
 	if err != nil {
@@ -267,6 +282,7 @@ func (cmd *UpCmd) Run(
 				ideConfig.Options,
 				cmd.GitUsername,
 				cmd.GitToken,
+				cmd.SSHAuthSockID,
 				log,
 			)
 		case string(config.IDERustRover):
@@ -303,6 +319,7 @@ func (cmd *UpCmd) Run(
 				ideConfig.Options,
 				cmd.GitUsername,
 				cmd.GitToken,
+				cmd.SSHAuthSockID,
 				log,
 			)
 		case string(config.IDEJupyterDesktop):
@@ -315,6 +332,7 @@ func (cmd *UpCmd) Run(
 				ideConfig.Options,
 				cmd.GitUsername,
 				cmd.GitToken,
+				cmd.SSHAuthSockID,
 				log)
 		case string(config.IDEMarimo):
 			return startMarimoInBrowser(
@@ -326,6 +344,7 @@ func (cmd *UpCmd) Run(
 				ideConfig.Options,
 				cmd.GitUsername,
 				cmd.GitToken,
+				cmd.SSHAuthSockID,
 				log)
 		}
 	}
@@ -552,7 +571,7 @@ func startMarimoInBrowser(
 	client client2.BaseWorkspaceClient,
 	user string,
 	ideOptions map[string]config.OptionValue,
-	gitUsername, gitToken string,
+	gitUsername, gitToken, authSockID string,
 	logger log.Logger,
 ) error {
 	if forwardGpg {
@@ -599,6 +618,7 @@ func startMarimoInBrowser(
 		extraPorts,
 		gitUsername,
 		gitToken,
+		authSockID,
 		logger,
 	)
 }
@@ -610,7 +630,7 @@ func startJupyterNotebookInBrowser(
 	client client2.BaseWorkspaceClient,
 	user string,
 	ideOptions map[string]config.OptionValue,
-	gitUsername, gitToken string,
+	gitUsername, gitToken, authSockID string,
 	logger log.Logger,
 ) error {
 	if forwardGpg {
@@ -657,6 +677,7 @@ func startJupyterNotebookInBrowser(
 		extraPorts,
 		gitUsername,
 		gitToken,
+		authSockID,
 		logger,
 	)
 }
@@ -668,7 +689,7 @@ func startJupyterDesktop(
 	client client2.BaseWorkspaceClient,
 	user string,
 	ideOptions map[string]config.OptionValue,
-	gitUsername, gitToken string,
+	gitUsername, gitToken, authSockID string,
 	logger log.Logger,
 ) error {
 	if forwardGpg {
@@ -712,6 +733,7 @@ func startJupyterDesktop(
 		extraPorts,
 		gitUsername,
 		gitToken,
+		authSockID,
 		logger,
 	)
 }
@@ -758,7 +780,7 @@ func startVSCodeInBrowser(
 	client client2.BaseWorkspaceClient,
 	workspaceFolder, user string,
 	ideOptions map[string]config.OptionValue,
-	gitUsername, gitToken string,
+	gitUsername, gitToken, authSockID string,
 	logger log.Logger,
 ) error {
 	if forwardGpg {
@@ -806,6 +828,7 @@ func startVSCodeInBrowser(
 		extraPorts,
 		gitUsername,
 		gitToken,
+		authSockID,
 		logger,
 	)
 }
@@ -841,16 +864,75 @@ func parseAddressAndPort(bindAddressOption string, defaultPort int) (string, int
 	return address, portName, nil
 }
 
+// setupBackhaul sets up a long running command in the container to ensure an SSH connection is kept alive
+func setupBackhaul(client client2.BaseWorkspaceClient, authSockId string, log log.Logger) error {
+	execPath, err := os.Executable()
+	if err != nil {
+		return err
+	}
+
+	remoteUser, err := devssh.GetUser(client.WorkspaceConfig().ID, client.WorkspaceConfig().SSHConfigPath)
+	if err != nil {
+		remoteUser = "root"
+	}
+
+	dotCmd := exec.Command(
+		execPath,
+		"ssh",
+		"--agent-forwarding=true",
+		fmt.Sprintf("--reuse-ssh-auth-sock=%s", authSockId),
+		"--start-services=false",
+		"--user",
+		remoteUser,
+		"--context",
+		client.Context(),
+		client.Workspace(),
+		"--log-output=raw",
+		"--command",
+		"while true; do sleep 6000000; done", // sleep infinity is not available on all systems
+	)
+
+	if log.GetLevel() == logrus.DebugLevel {
+		dotCmd.Args = append(dotCmd.Args, "--debug")
+	}
+
+	log.Info("Setting up backhaul SSH connection")
+
+	writer := log.Writer(logrus.InfoLevel, false)
+
+	dotCmd.Stdout = writer
+	dotCmd.Stderr = writer
+
+	err = dotCmd.Run()
+	if err != nil {
+		return err
+	}
+
+	log.Infof("Done setting up backhaul")
+
+	return nil
+}
+
 func startBrowserTunnel(
 	ctx context.Context,
 	devPodConfig *config.Config,
 	client client2.BaseWorkspaceClient,
 	user, targetURL string,
 	forwardPorts bool,
 	extraPorts []string,
-	gitUsername, gitToken string,
+	gitUsername, gitToken, authSockID string,
 	logger log.Logger,
 ) error {
+	// Setup a backhaul SSH connection using the remote user so there is an AUTH SOCK to use
+	// With normal IDEs this would be the SSH connection made by the IDE
+	// authSockID is not set when in proxy mode since we cannot use the proxies ssh-agent
+	if authSockID != "" {
+		go func() {
+			if err := setupBackhaul(client, authSockID, logger); err != nil {
+				logger.Error("Failed to setup backhaul SSH connection: ", err)
+			}
+		}()
+	}
 	err := tunnel.NewTunnel(
 		ctx,
 		func(ctx context.Context, stdin io.Reader, stdout io.Writer) error {
@@ -859,6 +941,7 @@ func startBrowserTunnel(
 
 			cmd, err := createSSHCommand(ctx, client, logger, []string{
 				"--log-output=raw",
+				fmt.Sprintf("--reuse-ssh-auth-sock=%s", authSockID),
 				"--stdio",
 			})
 			if err != nil {
@@ -880,7 +963,7 @@ func startBrowserTunnel(
 			}
 
 			// run in container
-			err := tunnel.RunInContainer(
+			err := tunnel.RunServices(
 				ctx,
 				devPodConfig,
 				containerClient,

diff --git a/docs/uml/up_sequence.puml b/docs/uml/up_sequence.puml
@@ -63,5 +63,20 @@ deactivate ContainerAgent
 Agent --> DevPod: 
 deactivate Agent
 
+alt if using browser based IDE (openvscode, marimo, jupyter)
+DevPod -> ContainerAgent: devpod ssh --reuse-ssh-auth-sock
+end
+
 DevPod -> IDE: Start
+
+alt if using normal IDE (vscode, intilliJ)
+IDE -> ContainerAgent: devpod ssh
+ContainerAgent --> IDE: ssh close
+end
+
+alt if using browser based IDE (openvscode, marimo, jupyter)
+ContainerAgent -> DevPod: ssh close
+end
+
+
 @enduml
diff --git a/go.mod b/go.mod
@@ -26,7 +26,7 @@ require (
 	github.com/loft-sh/api/v4 v4.0.0-alpha.6.0.20241129074910-a24d4104d586
 	github.com/loft-sh/log v0.0.0-20240219160058-26d83ffb46ac
 	github.com/loft-sh/programming-language-detection v0.0.5
-	github.com/loft-sh/ssh v0.0.4
+	github.com/loft-sh/ssh v0.0.5
 	github.com/mattn/go-isatty v0.0.20
 	github.com/mgutz/ansi v0.0.0-20200706080929-d51e80ef957d
 	github.com/mitchellh/go-homedir v1.1.0

diff --git a/go.sum b/go.sum
@@ -362,8 +362,8 @@ github.com/loft-sh/log v0.0.0-20240219160058-26d83ffb46ac h1:Gz/7Lb7WgdgIv+KJz87
 github.com/loft-sh/log v0.0.0-20240219160058-26d83ffb46ac/go.mod h1:YImeRjXH34Yf5E79T7UHBQpDZl9fIaaFRgyZ/bkY+UQ=
 github.com/loft-sh/programming-language-detection v0.0.5 h1:XiWlxtrf4t6Z7SQiob0JMKaCeMHCP3kWhB80wLt+EMY=
 github.com/loft-sh/programming-language-detection v0.0.5/go.mod h1:QGPQGKr9q1+rQS4OyisS5CPGY1a76SdNaZuk9oy+2cE=
-github.com/loft-sh/ssh v0.0.4 h1:Ybopo9SQpkZjMQ1hbnD71ZcN1fwe5n3dS1qiFfJRIAA=
-github.com/loft-sh/ssh v0.0.4/go.mod h1:jgAfPSNioyL2wdFesXY5Wi4pYpdNo4u7AzworofHeyU=
+github.com/loft-sh/ssh v0.0.5 h1:CmLfBrbekAZmYhpS+urhqmUZW1XU9kUo2bi4lJiUFH8=
+github.com/loft-sh/ssh v0.0.5/go.mod h1:jgAfPSNioyL2wdFesXY5Wi4pYpdNo4u7AzworofHeyU=
 github.com/lucasb-eyer/go-colorful v1.2.0 h1:1nnpGOrhyZZuNyfu1QjKiUICQ74+3FNCN69Aj6K7nkY=
 github.com/lucasb-eyer/go-colorful v1.2.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0=
 github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0=

diff --git a/pkg/devcontainer/setup.go b/pkg/devcontainer/setup.go
@@ -18,6 +18,7 @@ import (
 	"github.com/loft-sh/devpod/pkg/devcontainer/crane"
 	"github.com/loft-sh/devpod/pkg/devcontainer/sshtunnel"
 	"github.com/loft-sh/devpod/pkg/driver"
+	"github.com/loft-sh/devpod/pkg/ide"
 	provider2 "github.com/loft-sh/devpod/pkg/provider"
 	"github.com/loft-sh/log"
 	"github.com/pkg/errors"
@@ -97,6 +98,9 @@ func (r *runner) setupContainer(
 
 	// ssh tunnel
 	sshTunnelCmd := fmt.Sprintf("'%s' helper ssh-server --stdio", agent.ContainerDevPodHelperLocation)
+	if ide.ReusesAuthSock(r.WorkspaceConfig.Workspace.IDE.Name) {
+		sshTunnelCmd += fmt.Sprintf(" --reuse-ssh-auth-sock=%s", r.WorkspaceConfig.CLIOptions.SSHAuthSockID)
+	}
 	if r.Log.GetLevel() == logrus.DebugLevel {
 		sshTunnelCmd += " --debug"
 	}

diff --git a/pkg/ide/jupyter/jupyter.go b/pkg/ide/jupyter/jupyter.go
@@ -102,7 +102,7 @@ func (o *JupyterNotbookServer) Start() error {
 		runCommand := fmt.Sprintf("jupyter notebook --ip='*' --NotebookApp.notebook_dir='%s' --NotebookApp.token='' --NotebookApp.password='' --no-browser --port '%s' --allow-root", o.workspaceFolder, strconv.Itoa(DefaultServerPort))
 		args := []string{}
 		if o.userName != "" {
-			args = append(args, "su", o.userName, "-l", "-c", runCommand)
+			args = append(args, "su", o.userName, "-w", "SSH_AUTH_SOCK", "-l", "-c", runCommand)
 		} else {
 			args = append(args, "sh", "-l", "-c", runCommand)
 		}

diff --git a/pkg/ide/marimo/marimo.go b/pkg/ide/marimo/marimo.go
@@ -98,7 +98,7 @@ func (s *Server) start() error {
 		runCommand := fmt.Sprintf("marimo edit --headless --host 0.0.0.0 --port %s --token-password %s", strconv.Itoa(DefaultServerPort), token)
 		args := []string{}
 		if s.userName != "" {
-			args = append(args, "su", s.userName, "-l", "-c", runCommand)
+			args = append(args, "su", s.userName, "-w", "SSH_AUTH_SOCK", "-l", "-c", runCommand)
 		} else {
 			args = append(args, "sh", "-l", "-c", runCommand)
 		}

diff --git a/pkg/ide/types.go b/pkg/ide/types.go
@@ -37,3 +37,9 @@ func (o Options) GetValue(values map[string]config.OptionValue, key string) stri
 
 	return ""
 }
+
+// ReusesAuthSock determines if the --reuse-ssh-auth-sock flag should be passed to the ssh server helper based on the IDE.
+// Browser based IDEs use a browser tunnel to communicate with the remote server instead of an independent ssh connection
+func ReusesAuthSock(ide string) bool {
+	return ide == "openvscode" || ide == "marimo" || ide == "jupyternotebook" || ide == "jlab"
+}
diff --git a/pkg/provider/workspace.go b/pkg/provider/workspace.go
@@ -219,6 +219,7 @@ type CLIOptions struct {
 	GitCloneRecursiveSubmodules bool              `json:"gitCloneRecursive,omitempty"`
 	FallbackImage               string            `json:"fallbackImage,omitempty"`
 	GitSSHSigningKey            string            `json:"gitSshSigningKey,omitempty"`
+	SSHAuthSockID               string            `json:"sshAuthSockID,omitempty"` // ID to use when looking for SSH_AUTH_SOCK, defaults to a new random ID if not set (only used for browser IDEs)
 
 	// build options
 	Repository string   `json:"repository,omitempty"`