refactor: PTY & SSH (#7100)

* Add ssh tests for longoutput, orphan Signed-off-by: Spike Curtis <spike@coder.com> * PTY/SSH tests & improvements Signed-off-by: Spike Curtis <spike@coder.com> * Fix some tests Signed-off-by: Spike Curtis <spike@coder.com> * Fix linting Signed-off-by: Spike Curtis <spike@coder.com> * fmt Signed-off-by: Spike Curtis <spike@coder.com> * Fix windows test Signed-off-by: Spike Curtis <spike@coder.com> * Windows copy test Signed-off-by: Spike Curtis <spike@coder.com> * WIP Windows pty handling Signed-off-by: Spike Curtis <spike@coder.com> * Fix truncation tests Signed-off-by: Spike Curtis <spike@coder.com> * Appease linter/fmt Signed-off-by: Spike Curtis <spike@coder.com> * Fix typo Signed-off-by: Spike Curtis <spike@coder.com> * Rework truncation test to not assume OS buffers Signed-off-by: Spike Curtis <spike@coder.com> * Disable orphan test on Windows --- uses sh Signed-off-by: Spike Curtis <spike@coder.com> * agent_test running SSH in pty use ptytest.Start Signed-off-by: Spike Curtis <spike@coder.com> * More detail about closing pseudoconsole on windows Signed-off-by: Spike Curtis <spike@coder.com> * Code review fixes Signed-off-by: Spike Curtis <spike@coder.com> * Rearrange ptytest method order Signed-off-by: Spike Curtis <spike@coder.com> * Protect pty.Resize on windows from races Signed-off-by: Spike Curtis <spike@coder.com> * Fix windows bugs Signed-off-by: Spike Curtis <spike@coder.com> * PTY doesn't extend PTYCmd Signed-off-by: Spike Curtis <spike@coder.com> * Fix windows types Signed-off-by: Spike Curtis <spike@coder.com> --------- Signed-off-by: Spike Curtis <spike@coder.com>
2026-09-24 15:04:27 +08:00 · 2023-04-24 14:53:57 +04:00
parent c000f2ec28
commit daee91c6dc
16 changed files with 803 additions and 288 deletions
@@ -1045,7 +1045,7 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, logger slog.Logger, m
 		if err = a.trackConnGoroutine(func() {
 			buffer := make([]byte, 1024)
 			for {
-				read, err := rpty.ptty.Output().Read(buffer)
+				read, err := rpty.ptty.OutputReader().Read(buffer)
 				if err != nil {
 					// When the PTY is closed, this is triggered.
 					break
@@ -1138,7 +1138,7 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, logger slog.Logger, m
 			logger.Warn(ctx, "read conn", slog.Error(err))
 			return nil
 		}
-		_, err = rpty.ptty.Input().Write([]byte(req.Data))
+		_, err = rpty.ptty.InputWriter().Write([]byte(req.Data))
 		if err != nil {
 			logger.Warn(ctx, "write to pty", slog.Error(err))
 			return nil
@@ -1358,7 +1358,7 @@ type reconnectingPTY struct {
 	circularBuffer      *circbuf.Buffer
 	circularBufferMutex sync.RWMutex
 	timeout             *time.Timer
-	ptty                pty.PTY
+	ptty                pty.PTYCmd
 }

 // Close ends all connections to the reconnecting
@@ -45,6 +45,7 @@ import (
 	"github.com/coder/coder/coderd/httpapi"
 	"github.com/coder/coder/codersdk"
 	"github.com/coder/coder/codersdk/agentsdk"
+	"github.com/coder/coder/pty"
 	"github.com/coder/coder/pty/ptytest"
 	"github.com/coder/coder/tailnet"
 	"github.com/coder/coder/tailnet/tailnettest"
@@ -481,17 +482,10 @@ func TestAgent_TCPLocalForwarding(t *testing.T) {
 		}
 	}()

-	pty := ptytest.New(t)
-
-	cmd := setupSSHCommand(t, []string{"-L", fmt.Sprintf("%d:127.0.0.1:%d", randomPort, remotePort)}, []string{"sleep", "5"})
-	cmd.Stdin = pty.Input()
-	cmd.Stdout = pty.Output()
-	cmd.Stderr = pty.Output()
-	err = cmd.Start()
-	require.NoError(t, err)
+	_, proc := setupSSHCommand(t, []string{"-L", fmt.Sprintf("%d:127.0.0.1:%d", randomPort, remotePort)}, []string{"sleep", "5"})

 	go func() {
-		err := cmd.Wait()
+		err := proc.Wait()
 		select {
 		case <-done:
 		default:
@@ -523,7 +517,7 @@ func TestAgent_TCPLocalForwarding(t *testing.T) {

 	<-done

-	_ = cmd.Process.Kill()
+	_ = proc.Kill()
 }

 //nolint:paralleltest // This test reserves a port.
@@ -562,17 +556,10 @@ func TestAgent_TCPRemoteForwarding(t *testing.T) {
 		}
 	}()

-	pty := ptytest.New(t)
-
-	cmd := setupSSHCommand(t, []string{"-R", fmt.Sprintf("127.0.0.1:%d:127.0.0.1:%d", randomPort, localPort)}, []string{"sleep", "5"})
-	cmd.Stdin = pty.Input()
-	cmd.Stdout = pty.Output()
-	cmd.Stderr = pty.Output()
-	err = cmd.Start()
-	require.NoError(t, err)
+	_, proc := setupSSHCommand(t, []string{"-R", fmt.Sprintf("127.0.0.1:%d:127.0.0.1:%d", randomPort, localPort)}, []string{"sleep", "5"})

 	go func() {
-		err := cmd.Wait()
+		err := proc.Wait()
 		select {
 		case <-done:
 		default:
@@ -604,7 +591,7 @@ func TestAgent_TCPRemoteForwarding(t *testing.T) {

 	<-done

-	_ = cmd.Process.Kill()
+	_ = proc.Kill()
 }

 func TestAgent_UnixLocalForwarding(t *testing.T) {
@@ -641,17 +628,10 @@ func TestAgent_UnixLocalForwarding(t *testing.T) {
 		}
 	}()

-	pty := ptytest.New(t)
-
-	cmd := setupSSHCommand(t, []string{"-L", fmt.Sprintf("%s:%s", localSocketPath, remoteSocketPath)}, []string{"sleep", "5"})
-	cmd.Stdin = pty.Input()
-	cmd.Stdout = pty.Output()
-	cmd.Stderr = pty.Output()
-	err = cmd.Start()
-	require.NoError(t, err)
+	_, proc := setupSSHCommand(t, []string{"-L", fmt.Sprintf("%s:%s", localSocketPath, remoteSocketPath)}, []string{"sleep", "5"})

 	go func() {
-		err := cmd.Wait()
+		err := proc.Wait()
 		select {
 		case <-done:
 		default:
@@ -676,7 +656,7 @@ func TestAgent_UnixLocalForwarding(t *testing.T) {
 	_ = conn.Close()
 	<-done

-	_ = cmd.Process.Kill()
+	_ = proc.Kill()
 }

 func TestAgent_UnixRemoteForwarding(t *testing.T) {
@@ -713,17 +693,10 @@ func TestAgent_UnixRemoteForwarding(t *testing.T) {
 		}
 	}()

-	pty := ptytest.New(t)
-
-	cmd := setupSSHCommand(t, []string{"-R", fmt.Sprintf("%s:%s", remoteSocketPath, localSocketPath)}, []string{"sleep", "5"})
-	cmd.Stdin = pty.Input()
-	cmd.Stdout = pty.Output()
-	cmd.Stderr = pty.Output()
-	err = cmd.Start()
-	require.NoError(t, err)
+	_, proc := setupSSHCommand(t, []string{"-R", fmt.Sprintf("%s:%s", remoteSocketPath, localSocketPath)}, []string{"sleep", "5"})

 	go func() {
-		err := cmd.Wait()
+		err := proc.Wait()
 		select {
 		case <-done:
 		default:
@@ -753,7 +726,7 @@ func TestAgent_UnixRemoteForwarding(t *testing.T) {

 	<-done

-	_ = cmd.Process.Kill()
+	_ = proc.Kill()
 }

 func TestAgent_SFTP(t *testing.T) {
@@ -1648,7 +1621,7 @@ func TestAgent_WriteVSCodeConfigs(t *testing.T) {
 	}, testutil.WaitShort, testutil.IntervalFast)
 }

-func setupSSHCommand(t *testing.T, beforeArgs []string, afterArgs []string) *exec.Cmd {
+func setupSSHCommand(t *testing.T, beforeArgs []string, afterArgs []string) (*ptytest.PTYCmd, pty.Process) {
 	//nolint:dogsled
 	agentConn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
 	listener, err := net.Listen("tcp", "127.0.0.1:0")
@@ -1690,7 +1663,8 @@ func setupSSHCommand(t *testing.T, beforeArgs []string, afterArgs []string) *exe
 		"host",
 	)
 	args = append(args, afterArgs...)
-	return exec.Command("ssh", args...)
+	cmd := exec.Command("ssh", args...)
+	return ptytest.Start(t, cmd)
 }

 func setupSSHSession(t *testing.T, options agentsdk.Manifest) *ssh.Session {
@@ -253,102 +253,12 @@ func (s *Server) sessionStart(session ssh.Session, extraEnv []string) (retErr er

 	sshPty, windowSize, isPty := session.Pty()
 	if isPty {
-		// Disable minimal PTY emulation set by gliderlabs/ssh (NL-to-CRNL).
-		// See https://github.com/coder/coder/issues/3371.
-		session.DisablePTYEmulation()
-
-		if !isQuietLogin(session.RawCommand()) {
-			manifest := s.Manifest.Load()
-			if manifest != nil {
-				err = showMOTD(session, manifest.MOTDFile)
-				if err != nil {
-					s.logger.Error(ctx, "show MOTD", slog.Error(err))
-				}
-			} else {
-				s.logger.Warn(ctx, "metadata lookup failed, unable to show MOTD")
-			}
-		}
-
-		cmd.Env = append(cmd.Env, fmt.Sprintf("TERM=%s", sshPty.Term))
-
-		// The pty package sets `SSH_TTY` on supported platforms.
-		ptty, process, err := pty.Start(cmd, pty.WithPTYOption(
-			pty.WithSSHRequest(sshPty),
-			pty.WithLogger(slog.Stdlib(ctx, s.logger, slog.LevelInfo)),
-		))
-		if err != nil {
-			return xerrors.Errorf("start command: %w", err)
-		}
-		var wg sync.WaitGroup
-		defer func() {
-			defer wg.Wait()
-			closeErr := ptty.Close()
-			if closeErr != nil {
-				s.logger.Warn(ctx, "failed to close tty", slog.Error(closeErr))
-				if retErr == nil {
-					retErr = closeErr
-				}
-			}
-		}()
-		go func() {
-			for win := range windowSize {
-				resizeErr := ptty.Resize(uint16(win.Height), uint16(win.Width))
-				// If the pty is closed, then command has exited, no need to log.
-				if resizeErr != nil && !errors.Is(resizeErr, pty.ErrClosed) {
-					s.logger.Warn(ctx, "failed to resize tty", slog.Error(resizeErr))
-				}
-			}
-		}()
-		// We don't add input copy to wait group because
-		// it won't return until the session is closed.
-		go func() {
-			_, _ = io.Copy(ptty.Input(), session)
-		}()
-
-		// In low parallelism scenarios, the command may exit and we may close
-		// the pty before the output copy has started. This can result in the
-		// output being lost. To avoid this, we wait for the output copy to
-		// start before waiting for the command to exit. This ensures that the
-		// output copy goroutine will be scheduled before calling close on the
-		// pty. This shouldn't be needed because of `pty.Dup()` below, but it
-		// may not be supported on all platforms.
-		outputCopyStarted := make(chan struct{})
-		ptyOutput := func() io.ReadCloser {
-			defer close(outputCopyStarted)
-			// Try to dup so we can separate stdin and stdout closure.
-			// Once the original pty is closed, the dup will return
-			// input/output error once the buffered data has been read.
-			stdout, err := ptty.Dup()
-			if err == nil {
-				return stdout
-			}
-			// If we can't dup, we shouldn't close
-			// the fd since it's tied to stdin.
-			return readNopCloser{ptty.Output()}
-		}
-		wg.Add(1)
-		go func() {
-			// Ensure data is flushed to session on command exit, if we
-			// close the session too soon, we might lose data.
-			defer wg.Done()
-
-			stdout := ptyOutput()
-			defer stdout.Close()
-
-			_, _ = io.Copy(session, stdout)
-		}()
-		<-outputCopyStarted
-
-		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) {
-			s.logger.Warn(ctx, "wait error", slog.Error(err))
-		}
-		return err
+		return s.startPTYSession(session, cmd, sshPty, windowSize)
 	}
+	return startNonPTYSession(session, cmd)
+}

+func startNonPTYSession(session ssh.Session, cmd *exec.Cmd) error {
 	cmd.Stdout = session
 	cmd.Stderr = session.Stderr()
 	// This blocks forever until stdin is received if we don't
@@ -368,10 +278,94 @@ func (s *Server) sessionStart(session ssh.Session, extraEnv []string) (retErr er
 	return cmd.Wait()
 }

-type readNopCloser struct{ io.Reader }
+// ptySession is the interface to the ssh.Session that startPTYSession uses
+// we use an interface here so that we can fake it in tests.
+type ptySession interface {
+	io.ReadWriter
+	Context() ssh.Context
+	DisablePTYEmulation()
+	RawCommand() string
+}

-// Close implements io.Closer.
-func (readNopCloser) Close() error { return nil }
+func (s *Server) startPTYSession(session ptySession, cmd *exec.Cmd, sshPty ssh.Pty, windowSize <-chan ssh.Window) (retErr error) {
+	ctx := session.Context()
+	// Disable minimal PTY emulation set by gliderlabs/ssh (NL-to-CRNL).
+	// See https://github.com/coder/coder/issues/3371.
+	session.DisablePTYEmulation()
+
+	if !isQuietLogin(session.RawCommand()) {
+		manifest := s.Manifest.Load()
+		if manifest != nil {
+			err := showMOTD(session, manifest.MOTDFile)
+			if err != nil {
+				s.logger.Error(ctx, "show MOTD", slog.Error(err))
+			}
+		} else {
+			s.logger.Warn(ctx, "metadata lookup failed, unable to show MOTD")
+		}
+	}
+
+	cmd.Env = append(cmd.Env, fmt.Sprintf("TERM=%s", sshPty.Term))
+
+	// The pty package sets `SSH_TTY` on supported platforms.
+	ptty, process, err := pty.Start(cmd, pty.WithPTYOption(
+		pty.WithSSHRequest(sshPty),
+		pty.WithLogger(slog.Stdlib(ctx, s.logger, slog.LevelInfo)),
+	))
+	if err != nil {
+		return xerrors.Errorf("start command: %w", err)
+	}
+	defer func() {
+		closeErr := ptty.Close()
+		if closeErr != nil {
+			s.logger.Warn(ctx, "failed to close tty", slog.Error(closeErr))
+			if retErr == nil {
+				retErr = closeErr
+			}
+		}
+	}()
+	go func() {
+		for win := range windowSize {
+			resizeErr := ptty.Resize(uint16(win.Height), uint16(win.Width))
+			// If the pty is closed, then command has exited, no need to log.
+			if resizeErr != nil && !errors.Is(resizeErr, pty.ErrClosed) {
+				s.logger.Warn(ctx, "failed to resize tty", slog.Error(resizeErr))
+			}
+		}
+	}()
+
+	go func() {
+		_, _ = io.Copy(ptty.InputWriter(), session)
+	}()
+
+	// We need to wait for the command output to finish copying.  It's safe to
+	// just do this copy on the main handler goroutine because one of two things
+	// will happen:
+	//
+	// 1. The command completes & closes the TTY, which then triggers an error
+	//    after we've Read() all the buffered data from the PTY.
+	// 2. The client hangs up, which cancels the command's Context, and go will
+	//    kill the command's process.  This then has the same effect as (1).
+	n, err := io.Copy(session, ptty.OutputReader())
+	s.logger.Debug(ctx, "copy output done", slog.F("bytes", n), slog.Error(err))
+	if err != nil {
+		return xerrors.Errorf("copy error: %w", err)
+	}
+	// We've gotten all the output, but we need to wait for the process to
+	// complete so that we can get the exit code.  This returns
+	// immediately if the TTY was closed as part of the command exiting.
+	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) {
+		s.logger.Warn(ctx, "wait error", slog.Error(err))
+	}
+	if err != nil {
+		return xerrors.Errorf("process wait: %w", err)
+	}
+	return nil
+}

 func (s *Server) sftpHandler(session ssh.Session) {
 	ctx := session.Context()
@@ -0,0 +1,190 @@
+//go:build !windows
+
+package agentssh
+
+import (
+	"bufio"
+	"context"
+	"io"
+	"net"
+	"os/exec"
+	"testing"
+
+	gliderssh "github.com/gliderlabs/ssh"
+	"github.com/spf13/afero"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/testutil"
+
+	"cdr.dev/slog/sloggers/slogtest"
+)
+
+const longScript = `
+echo "started"
+sleep 30
+echo "done"
+`
+
+// Test_sessionStart_orphan tests running a command that takes a long time to
+// exit normally, and terminate the SSH session context early to verify that we
+// return quickly and don't leave the command running as an "orphan" with no
+// active SSH session.
+func Test_sessionStart_orphan(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitMedium)
+	defer cancel()
+	logger := slogtest.Make(t, nil)
+	s, err := NewServer(ctx, logger, afero.NewMemMapFs(), 0, "")
+	require.NoError(t, err)
+
+	// Here we're going to call the handler directly with a faked SSH session
+	// that just uses io.Pipes instead of a network socket.  There is a large
+	// variation in the time between closing the socket from the client side and
+	// the SSH server canceling the session Context, which would lead to a flaky
+	// test if we did it that way.  So instead, we directly cancel the context
+	// in this test.
+	sessionCtx, sessionCancel := context.WithCancel(ctx)
+	toClient, fromClient, sess := newTestSession(sessionCtx)
+	ptyInfo := gliderssh.Pty{}
+	windowSize := make(chan gliderssh.Window)
+	close(windowSize)
+	// the command gets the session context so that Go will terminate it when
+	// the session expires.
+	cmd := exec.CommandContext(sessionCtx, "sh", "-c", longScript)
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		// we don't really care what the error is here.  In the larger scenario,
+		// the client has disconnected, so we can't return any error information
+		// to them.
+		_ = s.startPTYSession(sess, cmd, ptyInfo, windowSize)
+	}()
+
+	readDone := make(chan struct{})
+	go func() {
+		defer close(readDone)
+		s := bufio.NewScanner(toClient)
+		assert.True(t, s.Scan())
+		txt := s.Text()
+		assert.Equal(t, "started", txt, "output corrupted")
+	}()
+
+	waitForChan(ctx, t, readDone, "read timeout")
+	// process is started, and should be sleeping for ~30 seconds
+
+	sessionCancel()
+
+	// now, we wait for the handler to complete.  If it does so before the
+	// main test timeout, we consider this a pass.  If not, it indicates
+	// that the server isn't properly shutting down sessions when they are
+	// disconnected client side, which could lead to processes hanging around
+	// indefinitely.
+	waitForChan(ctx, t, done, "handler timeout")
+
+	err = fromClient.Close()
+	require.NoError(t, err)
+}
+
+func waitForChan(ctx context.Context, t *testing.T, c <-chan struct{}, msg string) {
+	t.Helper()
+	select {
+	case <-c:
+		// OK!
+	case <-ctx.Done():
+		t.Fatal(msg)
+	}
+}
+
+type testSession struct {
+	ctx testSSHContext
+
+	// c2p is the client -> pty buffer
+	toPty *io.PipeReader
+	// p2c is the pty -> client buffer
+	fromPty *io.PipeWriter
+}
+
+type testSSHContext struct {
+	context.Context
+}
+
+func newTestSession(ctx context.Context) (toClient *io.PipeReader, fromClient *io.PipeWriter, s ptySession) {
+	toClient, fromPty := io.Pipe()
+	toPty, fromClient := io.Pipe()
+
+	return toClient, fromClient, &testSession{
+		ctx:     testSSHContext{ctx},
+		toPty:   toPty,
+		fromPty: fromPty,
+	}
+}
+
+func (s *testSession) Context() gliderssh.Context {
+	return s.ctx
+}
+
+func (*testSession) DisablePTYEmulation() {}
+
+// RawCommand returns "quiet logon" so that the PTY handler doesn't attempt to
+// write the message of the day, which will interfere with our tests.  It writes
+// the message of the day if it's a shell login (zero length RawCommand()).
+func (*testSession) RawCommand() string { return "quiet logon" }
+
+func (s *testSession) Read(p []byte) (n int, err error) {
+	return s.toPty.Read(p)
+}
+
+func (s *testSession) Write(p []byte) (n int, err error) {
+	return s.fromPty.Write(p)
+}
+
+func (testSSHContext) Lock() {
+	panic("not implemented")
+}
+
+func (testSSHContext) Unlock() {
+	panic("not implemented")
+}
+
+// User returns the username used when establishing the SSH connection.
+func (testSSHContext) User() string {
+	panic("not implemented")
+}
+
+// SessionID returns the session hash.
+func (testSSHContext) SessionID() string {
+	panic("not implemented")
+}
+
+// ClientVersion returns the version reported by the client.
+func (testSSHContext) ClientVersion() string {
+	panic("not implemented")
+}
+
+// ServerVersion returns the version reported by the server.
+func (testSSHContext) ServerVersion() string {
+	panic("not implemented")
+}
+
+// RemoteAddr returns the remote address for this connection.
+func (testSSHContext) RemoteAddr() net.Addr {
+	panic("not implemented")
+}
+
+// LocalAddr returns the local address for this connection.
+func (testSSHContext) LocalAddr() net.Addr {
+	panic("not implemented")
+}
+
+// Permissions returns the Permissions object used for this connection.
+func (testSSHContext) Permissions() *gliderssh.Permissions {
+	panic("not implemented")
+}
+
+// SetValue allows you to easily write new values into the underlying context.
+func (testSSHContext) SetValue(_, _ interface{}) {
+	panic("not implemented")
+}