Return proper exit code on ssh with TTY (#3192)

* Return proper exit code on ssh with TTY

Signed-off-by: Spike Curtis <spike@coder.com>

* Fix revive lint

Signed-off-by: Spike Curtis <spike@coder.com>

* Fix Windows exit code for missing command

Signed-off-by: Spike Curtis <spike@coder.com>

* Fix close error handling on agent TTY

Signed-off-by: Spike Curtis <spike@coder.com>

agent/agent.go

@@ -43,6 +43,11 @@ const (
 	ProtocolReconnectingPTY = "reconnecting-pty"
 	ProtocolSSH             = "ssh"
 	ProtocolDial            = "dial"
+
+	// MagicSessionErrorCode indicates that something went wrong with the session, rather than the
+	// command just returning a nonzero exit code, and is chosen as an arbitrary, high number
+	// unlikely to shadow other exit codes, which are typically 1, 2, 3, etc.
+	MagicSessionErrorCode = 229
 )
 
 type Options struct {
@@ -273,9 +278,17 @@ func (a *agent) init(ctx context.Context) {
 		},
 		Handler: func(session ssh.Session) {
 			err := a.handleSSHSession(session)
+			var exitError *exec.ExitError
+			if xerrors.As(err, &exitError) {
+				a.logger.Debug(ctx, "ssh session returned", slog.Error(exitError))
+				_ = session.Exit(exitError.ExitCode())
+				return
+			}
 			if err != nil {
 				a.logger.Warn(ctx, "ssh session failed", slog.Error(err))
-				_ = session.Exit(1)
+				// This exit code is designed to be unlikely to be confused for a legit exit code
+				// from the process.
+				_ = session.Exit(MagicSessionErrorCode)
 				return
 			}
 		},
@@ -403,7 +416,7 @@ func (a *agent) createCommand(ctx context.Context, rawCommand string, env []stri
 	return cmd, nil
 }
 
-func (a *agent) handleSSHSession(session ssh.Session) error {
+func (a *agent) handleSSHSession(session ssh.Session) (retErr error) {
 	cmd, err := a.createCommand(session.Context(), session.RawCommand(), session.Environ())
 	if err != nil {
 		return err
@@ -426,14 +439,24 @@ func (a *agent) handleSSHSession(session ssh.Session) error {
 		if err != nil {
 			return xerrors.Errorf("start command: %w", err)
 		}
+		defer func() {
+			closeErr := ptty.Close()
+			if closeErr != nil {
+				a.logger.Warn(context.Background(), "failed to close tty",
+					slog.Error(closeErr))
+				if retErr == nil {
+					retErr = closeErr
+				}
+			}
+		}()
 		err = ptty.Resize(uint16(sshPty.Window.Height), uint16(sshPty.Window.Width))
 		if err != nil {
 			return xerrors.Errorf("resize ptty: %w", err)
 		}
 		go func() {
 			for win := range windowSize {
-				err = ptty.Resize(uint16(win.Height), uint16(win.Width))
-				if err != nil {
+				resizeErr := ptty.Resize(uint16(win.Height), uint16(win.Width))
+				if resizeErr != nil {
 					a.logger.Warn(context.Background(), "failed to resize tty", slog.Error(err))
 				}
 			}
@@ -444,9 +467,15 @@ func (a *agent) handleSSHSession(session ssh.Session) error {
 		go func() {
 			_, _ = io.Copy(session, ptty.Output())
 		}()
-		_, _ = process.Wait()
-		_ = ptty.Close()
-		return nil
+		err = process.Wait()
+		var exitErr *exec.ExitError
+		// ExitErrors just mean the command we run returned a non-zero exit code, which is normal
+		// and not something to be concerned about.  But, if it's something else, we should log it.
+		if err != nil && !xerrors.As(err, &exitErr) {
+			a.logger.Warn(context.Background(), "wait error",
+				slog.Error(err))
+		}
+		return err
 	}
 
 	cmd.Stdout = session
@@ -549,7 +578,7 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, rawID string, conn ne
 		go func() {
 			// If the process dies randomly, we should
 			// close the pty.
-			_, _ = process.Wait()
+			_ = process.Wait()
 			rpty.Close()
 		}()
 		go func() {

agent/agent_test.go

@@ -16,6 +16,8 @@ import (
 	"testing"
 	"time"
 
+	"golang.org/x/xerrors"
+
 	scp "github.com/bramvdbogaerde/go-scp"
 	"github.com/google/uuid"
 	"github.com/pion/udp"
@@ -69,7 +71,7 @@ func TestAgent(t *testing.T) {
 		require.True(t, strings.HasSuffix(strings.TrimSpace(string(output)), "gitssh --"))
 	})
 
-	t.Run("SessionTTY", func(t *testing.T) {
+	t.Run("SessionTTYShell", func(t *testing.T) {
 		t.Parallel()
 		if runtime.GOOS == "windows" {
 			// This might be our implementation, or ConPTY itself.
@@ -103,6 +105,29 @@ func TestAgent(t *testing.T) {
 		require.NoError(t, err)
 	})
 
+	t.Run("SessionTTYExitCode", func(t *testing.T) {
+		t.Parallel()
+		session := setupSSHSession(t, agent.Metadata{})
+		command := "areallynotrealcommand"
+		err := session.RequestPty("xterm", 128, 128, ssh.TerminalModes{})
+		require.NoError(t, err)
+		ptty := ptytest.New(t)
+		require.NoError(t, err)
+		session.Stdout = ptty.Output()
+		session.Stderr = ptty.Output()
+		session.Stdin = ptty.Input()
+		err = session.Start(command)
+		require.NoError(t, err)
+		err = session.Wait()
+		exitErr := &ssh.ExitError{}
+		require.True(t, xerrors.As(err, &exitErr))
+		if runtime.GOOS == "windows" {
+			assert.Equal(t, 1, exitErr.ExitStatus())
+		} else {
+			assert.Equal(t, 127, exitErr.ExitStatus())
+		}
+	})
+
 	t.Run("LocalForwarding", func(t *testing.T) {
 		t.Parallel()
 		random, err := net.Listen("tcp", "127.0.0.1:0")