fix(pro): wait for outer tunnels to exit and clean up properly

cmd/ssh.go

@@ -44,7 +44,6 @@ type SSHCmd struct {
 	SetEnvVars          []string
 
 	Stdio                     bool
-	JumpContainer             bool
 	AgentForwarding           bool
 	GPGAgentForwarding        bool
 	GitSSHSignatureForwarding bool
@@ -436,7 +435,7 @@ func (cmd *SSHCmd) startTunnel(ctx context.Context, devPodConfig *config.Config,
 		if cmd.Proxy {
 			go func() {
 				if err := cmd.startRunnerServices(ctx, devPodConfig, containerClient, log); err != nil {
-					log.Error(err)
+					log.Debug(err)
 				}
 			}()
 		}

pkg/inject/inject.go

@@ -9,9 +9,11 @@ import (
 	"io"
 	"os"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/loft-sh/devpod/pkg/command"
+	"github.com/loft-sh/devpod/pkg/util"
 	"github.com/loft-sh/log"
 	perrors "github.com/pkg/errors"
 )
@@ -114,7 +116,8 @@ func InjectAndExecute(
 	case err = <-execErrChan:
 		result = <-injectChan
 	case result = <-injectChan:
-		// we don't wait for the command termination here and will just retry on error
+		// give exec some time to properly terminate and clean up
+		util.WaitForChan(execErrChan, 2*time.Second)
 	}
 
 	// prefer result error
@@ -126,7 +129,7 @@ func InjectAndExecute(
 		return result.wasExecuted, nil
 	}
 
-	log.Debugf("Rerun command as binary was injected")
+	log.Debug("Rerun command as binary was injected")
 	delayedStderr.Start()
 	return true, exec(ctx, scriptParams.Command, stdin, stdout, delayedStderr)
 }
@@ -281,14 +284,19 @@ func readLine(reader io.Reader) (string, error) {
 }
 
 func pipe(toStdin io.Writer, fromStdin io.Reader, toStdout io.Writer, fromStdout io.Reader) error {
-	errChan := make(chan error, 2)
+	var err error
+	wg := sync.WaitGroup{}
+	wg.Add(1)
 	go func() {
-		_, err := io.Copy(toStdout, fromStdout)
-		errChan <- err
+		defer wg.Done()
+		_, err = io.Copy(toStdout, fromStdout)
 	}()
+	wg.Add(1)
 	go func() {
-		_, err := io.Copy(toStdin, fromStdin)
-		errChan <- err
+		defer wg.Done()
+		_, err = io.Copy(toStdin, fromStdin)
 	}()
-	return <-errChan
+
+	wg.Wait()
+	return err
 }

pkg/ssh/server/ssh.go

@@ -140,7 +140,7 @@ func (s *Server) handler(sess ssh.Session) {
 	var err error
 	if isPty {
 		s.log.Debugf("Execute SSH server PTY command: %s", strings.Join(cmd.Args, " "))
-		err = HandlePTY(sess, ptyReq, winCh, cmd, nil)
+		err = s.HandlePTY(sess, ptyReq, winCh, cmd, nil)
 	} else {
 		s.log.Debugf("Execute SSH server command: %s", strings.Join(cmd.Args, " "))
 		err = s.HandleNonPTY(sess, cmd)
@@ -201,16 +201,17 @@ func (s *Server) HandleNonPTY(sess ssh.Session, cmd *exec.Cmd) (err error) {
 		}
 	}()
 
-	waitGroup.Wait()
+	// order is important here!
 	err = cmd.Wait()
+	waitGroup.Wait()
 	if err != nil {
 		return err
 	}
 
 	return nil
 }
 
-func HandlePTY(
+func (s *Server) HandlePTY(
 	sess ssh.Session,
 	ptyReq ssh.Pty,
 	winCh <-chan ssh.Window,

pkg/tunnel/container.go

@@ -13,6 +13,7 @@ import (
 	"github.com/loft-sh/devpod/pkg/config"
 	"github.com/loft-sh/devpod/pkg/provider"
 	devssh "github.com/loft-sh/devpod/pkg/ssh"
+	"github.com/loft-sh/devpod/pkg/util"
 	"github.com/loft-sh/log"
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
@@ -128,8 +129,10 @@ func (c *ContainerHandler) Run(ctx context.Context, handler Handler, cfg *config
 	// wait for result
 	select {
 	case err := <-containerChan:
+		util.WaitForChan(tunnelChan, 2*time.Second)
 		return errors.Wrap(err, "tunnel to container")
 	case err := <-tunnelChan:
+		util.WaitForChan(containerChan, 2*time.Second)
 		return errors.Wrap(err, "connect to server")
 	}
 }
@@ -204,9 +207,6 @@ func (c *ContainerHandler) runRunInContainer(ctx context.Context, sshClient *ssh
 		defer stdoutWriter.Close()
 		defer cancel()
 
-		c.log.Debugf("Run container tunnel")
-		defer c.log.Debugf("Container tunnel exited")
-
 		command := fmt.Sprintf("'%s' agent container-tunnel --workspace-info '%s'", c.client.AgentPath(), workspaceInfo)
 		if c.log.GetLevel() == logrus.DebugLevel {
 			command += " --debug"

pkg/tunnel/direct.go

@@ -4,8 +4,10 @@ import (
 	"context"
 	"io"
 	"os"
+	"time"
 
 	devssh "github.com/loft-sh/devpod/pkg/ssh"
+	"github.com/loft-sh/devpod/pkg/util"
 	"github.com/pkg/errors"
 )
 
@@ -52,8 +54,10 @@ func NewTunnel(ctx context.Context, tunnel Tunnel, handler Handler) error {
 	// wait for result
 	select {
 	case err := <-innerTunnelChan:
+		util.WaitForChan(outerTunnelChan, 2*time.Second)
 		return errors.Wrap(err, "inner tunnel")
 	case err := <-outerTunnelChan:
+		util.WaitForChan(innerTunnelChan, 2*time.Second)
 		return errors.Wrap(err, "outer tunnel")
 	}
 }

pkg/util/channel.go

@@ -0,0 +1,13 @@
+package util
+
+import "time"
+
+// WaitForChan races the channel closing against a timeout
+func WaitForChan(channel <-chan error, timeout time.Duration) {
+	select {
+	case <-time.After(timeout):
+		return
+	case <-channel:
+		return
+	}
+}