Skip to content

Commit 90d18dd

Browse files
mafredrimafredri
authored
fix(agent): Close stdin and stdout separately to fix pty output loss (coder#6862)
Fixes coder#6656 Closes coder#6840
1 parent 349bfad commit 90d18dd

File tree

5 files changed

+62
-30
lines changed

5 files changed

+62
-30
lines changed

‎agent/agent.go

+25-6
Original file line numberDiff line numberDiff line change
@@ -1115,7 +1115,8 @@ func (a *agent) handleSSHSession(session ssh.Session) (retErr error) {
11151115
go func() {
11161116
for win := range windowSize {
11171117
resizeErr := ptty.Resize(uint16(win.Height), uint16(win.Width))
1118-
if resizeErr != nil {
1118+
// If the pty is closed, then command has exited, no need to log.
1119+
if resizeErr != nil && !errors.Is(resizeErr, pty.ErrClosed) {
11191120
a.logger.Warn(ctx, "failed to resize tty", slog.Error(resizeErr))
11201121
}
11211122
}
@@ -1131,19 +1132,32 @@ func (a *agent) handleSSHSession(session ssh.Session) (retErr error) {
11311132
// output being lost. To avoid this, we wait for the output copy to
11321133
// start before waiting for the command to exit. This ensures that the
11331134
// output copy goroutine will be scheduled before calling close on the
1134-
// pty. There is still a risk of data loss if a command produces a lot
1135-
// of output, see TestAgent_Session_TTY_HugeOutputIsNotLost (skipped).
1135+
// pty. This shouldn't be needed because of `pty.Dup()` below, but it
1136+
// may not be supported on all platforms.
11361137
outputCopyStarted := make(chan struct{})
1137-
ptyOutput := func() io.Reader {
1138+
ptyOutput := func() io.ReadCloser {
11381139
defer close(outputCopyStarted)
1139-
return ptty.Output()
1140+
// Try to dup so we can separate stdin and stdout closure.
1141+
// Once the original pty is closed, the dup will return
1142+
// input/output error once the buffered data has been read.
1143+
stdout, err := ptty.Dup()
1144+
if err == nil {
1145+
return stdout
1146+
}
1147+
// If we can't dup, we shouldn't close
1148+
// the fd since it's tied to stdin.
1149+
return readNopCloser{ptty.Output()}
11401150
}
11411151
wg.Add(1)
11421152
go func() {
11431153
// Ensure data is flushed to session on command exit, if we
11441154
// close the session too soon, we might lose data.
11451155
defer wg.Done()
1146-
_, _ = io.Copy(session, ptyOutput())
1156+
1157+
stdout := ptyOutput()
1158+
defer stdout.Close()
1159+
1160+
_, _ = io.Copy(session, stdout)
11471161
}()
11481162
<-outputCopyStarted
11491163

@@ -1176,6 +1190,11 @@ func (a *agent) handleSSHSession(session ssh.Session) (retErr error) {
11761190
return cmd.Wait()
11771191
}
11781192

1193+
type readNopCloser struct{ io.Reader }
1194+
1195+
// Close implements io.Closer.
1196+
func (readNopCloser) Close() error { return nil }
1197+
11791198
func (a *agent) handleReconnectingPTY(ctx context.Context, logger slog.Logger, msg codersdk.WorkspaceAgentReconnectingPTYInit, conn net.Conn) (retErr error) {
11801199
defer conn.Close()
11811200

‎agent/agent_test.go

+8-22
Original file line numberDiff line numberDiff line change
@@ -350,15 +350,8 @@ func TestAgent_Session_TTY_Hushlogin(t *testing.T) {
350350

351351
func TestAgent_Session_TTY_FastCommandHasOutput(t *testing.T) {
352352
t.Parallel()
353-
if runtime.GOOS == "windows" {
354-
// This might be our implementation, or ConPTY itself.
355-
// It's difficult to find extensive tests for it, so
356-
// it seems like it could be either.
357-
t.Skip("ConPTY appears to be inconsistent on Windows.")
358-
}
359-
360353
// This test is here to prevent regressions where quickly executing
361-
// commands (with TTY) don't flush their output to the SSH session.
354+
// commands (with TTY) don't sync their output to the SSH session.
362355
//
363356
// See: https://github.com/coder/coder/issues/6656
364357
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
@@ -404,20 +397,13 @@ func TestAgent_Session_TTY_FastCommandHasOutput(t *testing.T) {
404397

405398
func TestAgent_Session_TTY_HugeOutputIsNotLost(t *testing.T) {
406399
t.Parallel()
407-
if runtime.GOOS == "windows" {
408-
// This might be our implementation, or ConPTY itself.
409-
// It's difficult to find extensive tests for it, so
410-
// it seems like it could be either.
411-
t.Skip("ConPTY appears to be inconsistent on Windows.")
412-
}
413-
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.")
414-
415-
// This test is here to prevent prove we have a bug where quickly executing
416-
// commands (with TTY) don't flush their output to the SSH session. This is
417-
// due to the pty being closed before all the output has been copied, but
418-
// protecting against this requires a non-trivial rewrite of the output
419-
// processing (or figuring out a way to put the pty in a mode where this
420-
// does not happen).
400+
401+
// This test is here to prevent regressions where a command (with or
402+
// without) a large amount of output would not be fully copied to the
403+
// SSH session. On unix systems, this was fixed by duplicating the file
404+
// descriptor of the PTY master and using it for copying the output.
405+
//
406+
// See: https://github.com/coder/coder/issues/6656
421407
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
422408
defer cancel()
423409
//nolint:dogsled

‎pty/pty.go

+9
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,12 @@ import (
66
"os"
77

88
"github.com/gliderlabs/ssh"
9+
"golang.org/x/xerrors"
910
)
1011

12+
// ErrClosed is returned when a PTY is used after it has been closed.
13+
var ErrClosed = xerrors.New("pty: closed")
14+
1115
// PTY is a minimal interface for interacting with a TTY.
1216
type PTY interface {
1317
io.Closer
@@ -31,6 +35,11 @@ type PTY interface {
3135
// The same stream would be used to provide user input: pty.Input().Write(...)
3236
Input() ReadWriter
3337

38+
// Dup returns a new file descriptor for the PTY.
39+
//
40+
// This is useful for closing stdin and stdout separately.
41+
Dup() (*os.File, error)
42+
3443
// Resize sets the size of the PTY.
3544
Resize(height uint16, width uint16) error
3645
}

‎pty/pty_other.go

+16-2
Original file line numberDiff line numberDiff line change
@@ -7,11 +7,11 @@ import (
77
"os/exec"
88
"runtime"
99
"sync"
10+
"syscall"
1011

1112
"github.com/creack/pty"
1213
"github.com/u-root/u-root/pkg/termios"
1314
"golang.org/x/sys/unix"
14-
"golang.org/x/xerrors"
1515
)
1616

1717
func newPty(opt ...Option) (retPTY *otherPty, err error) {
@@ -113,6 +113,20 @@ func (p *otherPty) Resize(height uint16, width uint16) error {
113113
})
114114
}
115115

116+
func (p *otherPty) Dup() (*os.File, error) {
117+
var newfd int
118+
err := p.control(p.pty, func(fd uintptr) error {
119+
var err error
120+
newfd, err = syscall.Dup(int(fd))
121+
return err
122+
})
123+
if err != nil {
124+
return nil, err
125+
}
126+
127+
return os.NewFile(uintptr(newfd), p.pty.Name()), nil
128+
}
129+
116130
func (p *otherPty) Close() error {
117131
p.mutex.Lock()
118132
defer p.mutex.Unlock()
@@ -131,7 +145,7 @@ func (p *otherPty) Close() error {
131145
if err != nil {
132146
p.err = err
133147
} else {
134-
p.err = xerrors.New("pty: closed")
148+
p.err = ErrClosed
135149
}
136150

137151
return err

‎pty/pty_windows.go

+4
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,10 @@ func (p *ptyWindows) Resize(height uint16, width uint16) error {
123123
return nil
124124
}
125125

126+
func (p *ptyWindows) Dup() (*os.File, error) {
127+
return nil, xerrors.Errorf("not implemented")
128+
}
129+
126130
func (p *ptyWindows) Close() error {
127131
p.closeMutex.Lock()
128132
defer p.closeMutex.Unlock()

0 commit comments

Comments
 (0)