fix(agent): Close stdin and stdout separately to fix pty output loss

mafredri · mafredri · commit 197bedcab292 · 2023-03-29T10:33:33.000Z
Fixes #6656 Closes #6840
diff --git a/agent/agent.go b/agent/agent.go
@@ -1131,19 +1131,32 @@ func (a *agent) handleSSHSession(session ssh.Session) (retErr error) {
 		// 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. There is still a risk of data loss if a command produces a lot
-		// of output, see TestAgent_Session_TTY_HugeOutputIsNotLost (skipped).
+		// 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.Reader {
+		ptyOutput := func() io.ReadCloser {
 			defer close(outputCopyStarted)
-			return ptty.Output()
+			// 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()
-			_, _ = io.Copy(session, ptyOutput())
+
+			stdout := ptyOutput()
+			defer stdout.Close()
+
+			_, _ = io.Copy(session, stdout)
 		}()
 		<-outputCopyStarted
 
@@ -1372,6 +1385,11 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, logger slog.Logger, m
 	}
 }
 
+type readNopCloser struct{ io.Reader }
+
+// Close implements io.Closer.
+func (readNopCloser) Close() error { return nil }
+
 // startReportingConnectionStats runs the connection stats reporting goroutine.
 func (a *agent) startReportingConnectionStats(ctx context.Context) {
 	reportStats := func(networkStats map[netlogtype.Connection]netlogtype.Counts) {
diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -350,15 +350,8 @@ func TestAgent_Session_TTY_Hushlogin(t *testing.T) {
 
 func TestAgent_Session_TTY_FastCommandHasOutput(t *testing.T) {
 	t.Parallel()
-	if runtime.GOOS == "windows" {
-		// This might be our implementation, or ConPTY itself.
-		// It's difficult to find extensive tests for it, so
-		// it seems like it could be either.
-		t.Skip("ConPTY appears to be inconsistent on Windows.")
-	}
-
 	// This test is here to prevent regressions where quickly executing
-	// commands (with TTY) don't flush their output to the SSH session.
+	// commands (with TTY) don't sync their output to the SSH session.
 	//
 	// See: https://github.com/coder/coder/issues/6656
 	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
@@ -404,20 +397,13 @@ func TestAgent_Session_TTY_FastCommandHasOutput(t *testing.T) {
 
 func TestAgent_Session_TTY_HugeOutputIsNotLost(t *testing.T) {
 	t.Parallel()
-	if runtime.GOOS == "windows" {
-		// This might be our implementation, or ConPTY itself.
-		// It's difficult to find extensive tests for it, so
-		// it seems like it could be either.
-		t.Skip("ConPTY appears to be inconsistent on Windows.")
-	}
-	t.Skip("This test proves we have a bug where parts of large output on a PTY can be lost after the command exits, skipped to avoid test failures.")
-
-	// This test is here to prevent prove we have a bug where quickly executing
-	// commands (with TTY) don't flush their output to the SSH session. This is
-	// due to the pty being closed before all the output has been copied, but
-	// protecting against this requires a non-trivial rewrite of the output
-	// processing (or figuring out a way to put the pty in a mode where this
-	// does not happen).
+
+	// This test is here to prevent regressions where a command (with or
+	// without) a large amount of output would not be fully copied to the
+	// SSH session. On unix systems, this was fixed by duplicating the file
+	// descriptor of the PTY master and using it for copying the output.
+	//
+	// See: https://github.com/coder/coder/issues/6656
 	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
 	defer cancel()
 	//nolint:dogsled
diff --git a/pty/pty.go b/pty/pty.go
@@ -31,6 +31,11 @@ type PTY interface {
 	// The same stream would be used to provide user input: pty.Input().Write(...)
 	Input() ReadWriter
 
+	// Dup returns a new file descriptor for the PTY.
+	//
+	// This is useful for closing stdin and stdout separately.
+	Dup() (*os.File, error)
+
 	// Resize sets the size of the PTY.
 	Resize(height uint16, width uint16) error
 }
diff --git a/pty/pty_other.go b/pty/pty_other.go
@@ -7,6 +7,7 @@ import (
 	"os/exec"
 	"runtime"
 	"sync"
+	"syscall"
 
 	"github.com/creack/pty"
 	"github.com/u-root/u-root/pkg/termios"
@@ -113,6 +114,20 @@ func (p *otherPty) Resize(height uint16, width uint16) error {
 	})
 }
 
+func (p *otherPty) Dup() (*os.File, error) {
+	var newfd int
+	err := p.control(p.pty, func(fd uintptr) error {
+		var err error
+		newfd, err = syscall.Dup(int(fd))
+		return err
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	return os.NewFile(uintptr(newfd), p.pty.Name()), nil
+}
+
 func (p *otherPty) Close() error {
 	p.mutex.Lock()
 	defer p.mutex.Unlock()
diff --git a/pty/pty_windows.go b/pty/pty_windows.go
@@ -123,6 +123,10 @@ func (p *ptyWindows) Resize(height uint16, width uint16) error {
 	return nil
 }
 
+func (p *ptyWindows) Dup() (*os.File, error) {
+	return nil, xerrors.Errorf("not implemented")
+}
+
 func (p *ptyWindows) Close() error {
 	p.closeMutex.Lock()
 	defer p.closeMutex.Unlock()
