Merge pull request loft-sh#734 from neogopher/refactor-allow-ssh-agen…

…t-forwarding Refactor DevPod up command to allow SSH agent forwarding
BandhiyaHardik · Oct 27, 2023 · 4c61f0c · 4c61f0c
2 parents 6c1dc3e + 72b8394
commit 4c61f0c
Show file tree

Hide file tree

Showing 4 changed files with 221 additions and 123 deletions.
diff --git a/cmd/up.go b/cmd/up.go
@@ -19,6 +19,7 @@ import (
 	"github.com/loft-sh/devpod/pkg/command"
 	"github.com/loft-sh/devpod/pkg/config"
 	config2 "github.com/loft-sh/devpod/pkg/devcontainer/config"
+	"github.com/loft-sh/devpod/pkg/devcontainer/sshtunnel"
 	"github.com/loft-sh/devpod/pkg/ide/fleet"
 	"github.com/loft-sh/devpod/pkg/ide/jetbrains"
 	"github.com/loft-sh/devpod/pkg/ide/jupyter"
@@ -353,41 +354,25 @@ func (cmd *UpCmd) devPodUpMachine(
 	// create container etc.
 	log.Infof("Creating devcontainer...")
 	defer log.Debugf("Done creating devcontainer")
-	command := fmt.Sprintf(
+
+	// ssh tunnel command
+	sshTunnelCmd := fmt.Sprintf("'%s' helper ssh-server --stdio", client.AgentPath())
+	if log.GetLevel() == logrus.DebugLevel {
+		sshTunnelCmd += " --debug"
+	}
+
+	// create agent command
+	agentCommand := fmt.Sprintf(
 		"'%s' agent workspace up --workspace-info '%s'",
 		client.AgentPath(),
 		workspaceInfo,
 	)
 	if log.GetLevel() == logrus.DebugLevel {
-		command += " --debug"
-	}
-
-	// create pipes
-	stdoutReader, stdoutWriter, err := os.Pipe()
-	if err != nil {
-		return nil, err
-	}
-	stdinReader, stdinWriter, err := os.Pipe()
-	if err != nil {
-		return nil, err
+		agentCommand += " --debug"
 	}
-	defer stdoutWriter.Close()
-	defer stdinWriter.Close()
-
-	// start machine on stdio
-	cancelCtx, cancel := context.WithCancel(ctx)
-	defer cancel()
 
-	errChan := make(chan error, 1)
-	go func() {
-		defer log.Debugf("Done executing up command")
-		defer cancel()
-
-		writer := log.Writer(logrus.InfoLevel, false)
-		defer writer.Close()
-
-		log.Debugf("Inject and run command: %s", command)
-		err := agent.InjectAgentAndExecute(
+	agentInjectFunc := func(cancelCtx context.Context, sshCmd string, sshTunnelStdinReader, sshTunnelStdoutWriter *os.File, writer io.WriteCloser) error {
+		return agent.InjectAgentAndExecute(
 			cancelCtx,
 			func(ctx context.Context, command string, stdin io.Reader, stdout io.Writer, stderr io.Writer) error {
 				return client.Command(ctx, client2.CommandOptions{
@@ -401,56 +386,15 @@ func (cmd *UpCmd) devPodUpMachine(
 			client.AgentPath(),
 			client.AgentURL(),
 			true,
-			command,
-			stdinReader,
-			stdoutWriter,
+			sshCmd,
+			sshTunnelStdinReader,
+			sshTunnelStdoutWriter,
 			writer,
 			log.ErrorStreamOnly(),
 		)
-		if err != nil {
-			errChan <- fmt.Errorf("executing agent command: %w", err)
-		} else {
-			errChan <- nil
-		}
-	}()
-
-	// create container etc.
-	var result *config2.Result
-	if cmd.Proxy {
-		// create client on stdin & stdout
-		tunnelClient, err := tunnelserver.NewTunnelClient(os.Stdin, os.Stdout, true)
-		if err != nil {
-			return nil, errors.Wrap(err, "create tunnel client")
-		}
-
-		// create proxy server
-		result, err = tunnelserver.RunProxyServer(
-			cancelCtx,
-			tunnelClient,
-			stdoutReader,
-			stdinWriter,
-			log,
-		)
-		if err != nil {
-			return nil, errors.Wrap(err, "run proxy tunnel")
-		}
-	} else {
-		result, err = tunnelserver.RunUpServer(
-			cancelCtx,
-			stdoutReader,
-			stdinWriter,
-			client.AgentInjectGitCredentials(),
-			client.AgentInjectDockerCredentials(),
-			client.WorkspaceConfig(),
-			log,
-		)
-		if err != nil {
-			return nil, errors.Wrap(err, "run tunnel machine")
-		}
 	}
 
-	// wait until command finished
-	return result, <-errChan
+	return sshtunnel.ExecuteCommand(ctx, client, agentInjectFunc, sshTunnelCmd, agentCommand, cmd.Proxy, false, false, nil, log)
 }
 
 func startJupyterNotebookInBrowser(

diff --git a/pkg/devcontainer/setup.go b/pkg/devcontainer/setup.go
@@ -9,9 +9,9 @@ import (
 	"runtime"
 
 	"github.com/loft-sh/devpod/pkg/agent"
-	"github.com/loft-sh/devpod/pkg/agent/tunnelserver"
 	"github.com/loft-sh/devpod/pkg/compress"
 	"github.com/loft-sh/devpod/pkg/devcontainer/config"
+	"github.com/loft-sh/devpod/pkg/devcontainer/sshtunnel"
 	"github.com/loft-sh/devpod/pkg/driver"
 	provider2 "github.com/loft-sh/devpod/pkg/provider"
 	"github.com/pkg/errors"
@@ -66,64 +66,28 @@ func (r *runner) setupContainer(
 	// check if docker driver
 	_, isDockerDriver := r.Driver.(driver.DockerDriver)
 
+	// ssh tunnel
+	sshTunnelCmd := fmt.Sprintf("'%s' helper ssh-server --stdio", agent.ContainerDevPodHelperLocation)
+	if r.Log.GetLevel() == logrus.DebugLevel {
+		sshTunnelCmd += " --debug"
+	}
+
 	// setup container
 	r.Log.Infof("Setup container...")
-	command := fmt.Sprintf("'%s' agent container setup --setup-info '%s' --container-workspace-info '%s'", agent.ContainerDevPodHelperLocation, compressed, workspaceConfigCompressed)
+	setupCommand := fmt.Sprintf("'%s' agent container setup --setup-info '%s' --container-workspace-info '%s'", agent.ContainerDevPodHelperLocation, compressed, workspaceConfigCompressed)
 	if runtime.GOOS == "linux" || !isDockerDriver {
-		command += " --chown-workspace"
+		setupCommand += " --chown-workspace"
 	}
 	if !isDockerDriver {
-		command += " --stream-mounts"
+		setupCommand += " --stream-mounts"
 	}
 	if r.Log.GetLevel() == logrus.DebugLevel {
-		command += " --debug"
+		setupCommand += " --debug"
 	}
 
-	// create pipes
-	stdoutReader, stdoutWriter, err := os.Pipe()
-	if err != nil {
-		return nil, err
-	}
-	stdinReader, stdinWriter, err := os.Pipe()
-	if err != nil {
-		return nil, err
-	}
-	defer stdoutWriter.Close()
-	defer stdinWriter.Close()
-
-	// start machine on stdio
-	cancelCtx, cancel := context.WithCancel(ctx)
-	defer cancel()
-
-	errChan := make(chan error, 1)
-	go func() {
-		defer r.Log.Debugf("Done executing up command")
-		defer cancel()
-
-		writer := r.Log.Writer(logrus.InfoLevel, false)
-		defer writer.Close()
-
-		r.Log.Debugf("Run command in container: %s", command)
-		err = r.Driver.CommandDevContainer(cancelCtx, r.ID, "root", command, stdinReader, stdoutWriter, writer)
-		if err != nil {
-			errChan <- fmt.Errorf("executing container command: %w", err)
-		} else {
-			errChan <- nil
-		}
-	}()
-
-	// start server
-	result, err = tunnelserver.RunSetupServer(
-		cancelCtx,
-		stdoutReader,
-		stdinWriter,
-		r.WorkspaceConfig.Agent.InjectDockerCredentials != "false",
-		config.GetMounts(result),
-		r.Log,
-	)
-	if err != nil {
-		return nil, errors.Wrap(err, "run tunnel machine")
+	agentInjectFunc := func(cancelCtx context.Context, sshCmd string, sshTunnelStdinReader, sshTunnelStdoutWriter *os.File, writer io.WriteCloser) error {
+		return r.Driver.CommandDevContainer(cancelCtx, r.ID, "root", sshCmd, sshTunnelStdinReader, sshTunnelStdoutWriter, writer)
 	}
 
-	return result, <-errChan
+	return sshtunnel.ExecuteCommand(ctx, nil, agentInjectFunc, sshTunnelCmd, setupCommand, false, true, r.WorkspaceConfig.Agent.InjectDockerCredentials != "false", result, r.Log)
 }
diff --git a/pkg/devcontainer/setup/setup.go b/pkg/devcontainer/setup/setup.go
@@ -57,6 +57,12 @@ func SetupContainer(setupInfo *config.Result, extraWorkspaceEnv []string, chownW
 		log.Errorf("Error linking /home/root: %v", err)
 	}
 
+	// chown agent sock file
+	err = ChownAgentSock(setupInfo, log)
+	if err != nil {
+		return errors.Wrap(err, "chown ssh agent sock file")
+	}
+
 	// run commands
 	log.Debugf("Run post create commands...")
 	err = PostCreateCommands(setupInfo, log)
@@ -206,6 +212,20 @@ func PatchEtcEnvironment(mergedConfig *config.MergedDevContainerConfig, log log.
 	return nil
 }
 
+func ChownAgentSock(setupInfo *config.Result, log log.Logger) error {
+	user := config.GetRemoteUser(setupInfo)
+
+	agentSockFile := os.Getenv("SSH_AUTH_SOCK")
+	if agentSockFile != "" {
+		err := copy2.ChownR(filepath.Dir(agentSockFile), user)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 func PostCreateCommands(setupInfo *config.Result, log log.Logger) error {
 	remoteUser := config.GetRemoteUser(setupInfo)
 	mergedConfig := setupInfo.MergedConfig