diff --git a/cli/ssh.go b/cli/ssh.go index d18ac8909f575..bf1299ea944ad 100644 --- a/cli/ssh.go +++ b/cli/ssh.go @@ -18,6 +18,7 @@ import ( "slices" "strings" "sync" + "sync/atomic" "time" "github.com/gen2brain/beeep" @@ -247,9 +248,28 @@ func (r *RootCmd) ssh() *serpent.Command { // In stdio mode, we can't allow any writes to stdin or stdout // because they are used by the SSH protocol. stdioReader, stdioWriter := inv.Stdin, inv.Stdout + // connected is set once the SSH session is consuming stdin, at + // which point its copy logic handles stdin termination. + var connected atomic.Bool if stdio { inv.Stdin = stdioErrLogReader{inv.Logger} inv.Stdout = inv.Stderr + + // ProxyCommand clients close stdin when they abandon a + // connection attempt. Relay stdin through a pipe so that its + // termination is observed even before the connection exists, + // canceling this process instead of leaking it while it waits + // for the agent or dials. See #27954. + piper, pipew := io.Pipe() + go func(src io.Reader) { + _, err := io.Copy(pipew, src) + _ = pipew.CloseWithError(err) + logger.Debug(ctx, "stdin relay finished", slog.Error(err)) + if !connected.Load() { + cancel() + } + }(stdioReader) + stdioReader = piper } // This WaitGroup solves for a race condition where we were logging @@ -460,6 +480,7 @@ func (r *RootCmd) ssh() *serpent.Command { }) defer closeUsage() } + connected.Store(true) return runCoderConnectStdio(ctx, fmt.Sprintf("%s:22", coderConnectHost), stdioReader, stdioWriter, stack, logger) } } @@ -549,6 +570,7 @@ func (r *RootCmd) ssh() *serpent.Command { return nil }, logger, client, workspace, errCh) }() + connected.Store(true) copier.copy(&wg) return nil } diff --git a/cli/ssh_test.go b/cli/ssh_test.go index 2221a23e7bf9c..c90cc6a0a81a4 100644 --- a/cli/ssh_test.go +++ b/cli/ssh_test.go @@ -2170,8 +2170,17 @@ func TestSSH_CoderConnect(t *testing.T) { ctx = context.WithValue(ctx, "fs", fs) client, workspace, agentToken := setupWorkspaceForAgent(t) + // Provide a live stdin for the command; with no stdin attached, the + // command treats it as an abandoned connection and exits early. + clientOutput, clientInput := io.Pipe() + defer func() { + for _, c := range []io.Closer{clientOutput, clientInput} { + _ = c.Close() + } + }() inv, root := clitest.New(t, "ssh", workspace.Name, "--network-info-dir", "/net", "--stdio") clitest.SetupConfig(t, client, root) + inv.Stdin = clientOutput ctx = cli.WithTestOnlyCoderConnectDialer(ctx, &fakeCoderConnectDialer{}) ctx = withCoderConnectRunning(ctx) @@ -2457,6 +2466,48 @@ func (*fakeCoderConnectDialer) DialContext(ctx context.Context, network, addr st return nil, xerrors.Errorf("dial coder connect host %q over %s", addr, network) } +// TestSSHStdio_ExitWhenStdinClosedBeforeAgentConnects ensures the command exits +// promptly when the client closes stdin before a connection is established, +// instead of leaking the process while waiting for the agent or dialing. +// See https://github.com/coder/coder/issues/27954 +func TestSSHStdio_ExitWhenStdinClosedBeforeAgentConnects(t *testing.T) { + t.Parallel() + + client, workspace, _ := setupWorkspaceForAgent(t) + + clientOutput, clientInput := io.Pipe() + serverOutput, serverInput := io.Pipe() + defer func() { + for _, c := range []io.Closer{clientOutput, clientInput, serverOutput, serverInput} { + _ = c.Close() + } + }() + + ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitMedium) + defer cancel() + + inv, root := clitest.New(t, "ssh", "--stdio", workspace.Name) + clitest.SetupConfig(t, client, root) + inv.Stdin = clientOutput + inv.Stdout = serverInput + inv.Stderr = io.Discard + + cmdDone := tGo(t, func() { + err := inv.WithContext(ctx).Run() + assert.Error(t, err) + }) + + // The client abandons the connection attempt by closing stdin. No agent + // is started, so the command would wait for it forever unless it notices. + require.NoError(t, clientOutput.Close()) + + select { + case <-cmdDone: + case <-time.After(testutil.WaitShort): + require.Fail(t, "coder ssh --stdio did not exit after stdin was closed") + } +} + // tGoContext runs fn in a goroutine passing a context that will be // canceled on test completion and wait until fn has finished executing. // Done and cancel are returned for optionally waiting until completion