From 5c777d259bd4c3e2e8e2d00e5115d7e00cc27262 Mon Sep 17 00:00:00 2001 From: Grigory Panov Date: Wed, 7 Oct 2026 14:26:44 +0200 Subject: [PATCH 1/8] ssh: don't report a failed handover when the session ends mid-handover Teardown closes the websocket while a handover may still be waiting for the old connection to close. The receiving loop passed the resulting "use of closed network connection" read error to initiateHandover, and the client's handover goroutine returned ErrHandoverFailed. That error could win the errgroup race against proxy.start's real outcome, so a clean exit reported a handover failure. This is the TestQuickHandover flake seen in the merge queue. The receiving loop now passes the context cancellation on, and the handover goroutine leaves the outcome to proxy.start when it sees it. Co-authored-by: Isaac --- experimental/ssh/internal/proxy/client.go | 6 ++ .../ssh/internal/proxy/client_server_test.go | 64 +++++++++++++++++++ experimental/ssh/internal/proxy/proxy.go | 8 ++- 3 files changed, 77 insertions(+), 1 deletion(-) diff --git a/experimental/ssh/internal/proxy/client.go b/experimental/ssh/internal/proxy/client.go index 8d6a12fb438..634f070156a 100644 --- a/experimental/ssh/internal/proxy/client.go +++ b/experimental/ssh/internal/proxy/client.go @@ -96,6 +96,12 @@ func RunClientProxy(ctx context.Context, src io.ReadCloser, dst io.Writer, reque return nil case <-requestHandoverTick(): if err := proxy.initiateHandover(gCtx); err != nil { + // The session ended while the handover was in progress. proxy.start + // decides the session's outcome; an error here could be recorded first + // and mask it. + if errors.Is(err, context.Canceled) { + return nil + } // A handover that never got past its dial leaves the current connection // untouched and still carrying traffic, so ending the session over it // would throw away a working tunnel - the failure mode customers see as diff --git a/experimental/ssh/internal/proxy/client_server_test.go b/experimental/ssh/internal/proxy/client_server_test.go index 11a051e74c6..4dea121516e 100644 --- a/experimental/ssh/internal/proxy/client_server_test.go +++ b/experimental/ssh/internal/proxy/client_server_test.go @@ -316,6 +316,70 @@ func TestHandoverDialFailureKeepsSessionAlive(t *testing.T) { assert.Equal(t, int32(2), dials.Load(), "expected the initial dial plus exactly one handover dial") } +// A session that ends while a handover is still waiting for the old connection to close must end +// cleanly. Teardown closes that connection itself, and the handover used to report the resulting +// read error as ErrHandoverFailed, which could win the race against the session's real outcome - +// the TestQuickHandover flake. +func TestSessionEndDuringHandoverIsNotAHandoverFailure(t *testing.T) { + server := createTestServer(t, 2, time.Hour) + defer server.Close() + + // The handover dial lands here instead of on the proxy server, so nobody ever closes the old + // connection and the handover stays in progress until the session ends. + silentServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if err != nil { + return + } + defer conn.Close() + for { + if _, _, err := conn.ReadMessage(); err != nil { + return + } + } + })) + defer silentServer.Close() + + wsURL := "ws" + server.URL[4:] + silentURL := "ws" + silentServer.URL[4:] + var dials atomic.Int32 + handoverDialed := make(chan struct{}, 1) + createConn := func(ctx context.Context, dial DialRequest) (*websocket.Conn, error) { + url := fmt.Sprintf("%s?id=%s", wsURL, dial.ConnID) + if dials.Add(1) > 1 { + url = silentURL + defer func() { handoverDialed <- struct{}{} }() + } + conn, _, err := websocket.DefaultDialer.Dial(url, nil) // nolint:bodyclose + return conn, err + } + + handoverChan := make(chan time.Time) + errChan := make(chan error, 1) + client := createTestClientWithDialer(t, createConn, func() <-chan time.Time { + return handoverChan + }, time.Hour, false, errChan) + + msg := []byte("before handover\n") + _, err := client.InputWriter.Write(msg) + require.NoError(t, err) + require.NoError(t, client.Output.AssertWrite(msg)) + + handoverChan <- time.Now() + select { + case <-handoverDialed: + case <-time.After(10 * time.Second): + t.Fatal("the handover never dialed its replacement connection") + } + + client.Cleanup() + select { + case err := <-errChan: + t.Fatalf("session ended with an error: %v", err) + default: + } +} + // TestClientExitsWhenServerCommandFails reproduces the missing-sshd case: the server accepts the // websocket but can't launch its command, so it closes the connection immediately. The client // proxy must exit promptly instead of hanging on the handover goroutine (which would leave the diff --git a/experimental/ssh/internal/proxy/proxy.go b/experimental/ssh/internal/proxy/proxy.go index 493644ad5ae..a6eed7dee34 100644 --- a/experimental/ssh/internal/proxy/proxy.go +++ b/experimental/ssh/internal/proxy/proxy.go @@ -535,7 +535,13 @@ func (pc *proxyConnection) runReceivingLoop(ctx context.Context, dst io.Writer) // reattachment to recover bytes the close-frame exchange did not drain. if handover := pc.handoverState.Load(); handover != nil { var closeConnSignal error - if !websocket.IsCloseError(err, websocket.CloseNormalClosure) { + switch { + case ctx.Err() != nil: + // The session is ending and teardown closed the connection under the handover + // (see the ctx.Err() check below). Pass the cancellation on, so the handover + // initiator does not report a failed handover for an ordinary exit. + closeConnSignal = ctx.Err() + case !websocket.IsCloseError(err, websocket.CloseNormalClosure): closeConnSignal = errors.Join(ErrWebsocketDropped, fmt.Errorf("failed to read from websocket during handover: %w", err)) } // Signal the current connection is closed to the handover initiator (initiateHandover or acceptHandover). From f2ea5e7f45513656a429017abf4fa3218aade1d4 Mon Sep 17 00:00:00 2001 From: Grigory Panov Date: Wed, 7 Oct 2026 14:42:22 +0200 Subject: [PATCH 2/8] ssh: keep handing over after a canceled dial; cover both teardown outcomes Check context.Canceled only after the dial-failed branch, so a dial that fails with Canceled while the session is live is retried on the next tick instead of stopping all later handovers. Make the regression test table-driven: a clean exit must report no error, and a session that ends with a real source error during a handover must report that error, not ErrHandoverFailed. The second case covers the client.go change. The test client now exposes Done, so the test can wait for the session to end before Cleanup closes the pipe. Co-authored-by: Isaac --- experimental/ssh/internal/proxy/client.go | 12 +- .../ssh/internal/proxy/client_server_test.go | 172 +++++++++++++----- 2 files changed, 133 insertions(+), 51 deletions(-) diff --git a/experimental/ssh/internal/proxy/client.go b/experimental/ssh/internal/proxy/client.go index 634f070156a..47cdbc08787 100644 --- a/experimental/ssh/internal/proxy/client.go +++ b/experimental/ssh/internal/proxy/client.go @@ -96,12 +96,6 @@ func RunClientProxy(ctx context.Context, src io.ReadCloser, dst io.Writer, reque return nil case <-requestHandoverTick(): if err := proxy.initiateHandover(gCtx); err != nil { - // The session ended while the handover was in progress. proxy.start - // decides the session's outcome; an error here could be recorded first - // and mask it. - if errors.Is(err, context.Canceled) { - return nil - } // A handover that never got past its dial leaves the current connection // untouched and still carrying traffic, so ending the session over it // would throw away a working tunnel - the failure mode customers see as @@ -120,6 +114,12 @@ func RunClientProxy(ctx context.Context, src io.ReadCloser, dst io.Writer, reque log.Debugf(gCtx, "Auth handover failed, recovering through session reattachment: %v", err) continue } + // The session ended while the handover was in progress. proxy.start + // decides the session's outcome; an error here could be recorded first + // and mask it. + if errors.Is(err, context.Canceled) { + return nil + } return errors.Join(ErrHandoverFailed, err) } } diff --git a/experimental/ssh/internal/proxy/client_server_test.go b/experimental/ssh/internal/proxy/client_server_test.go index 4dea121516e..758d236e0b8 100644 --- a/experimental/ssh/internal/proxy/client_server_test.go +++ b/experimental/ssh/internal/proxy/client_server_test.go @@ -34,9 +34,11 @@ func createTestServer(t *testing.T, maxClients int, shutdownDelay time.Duration) } type testClient struct { - InputWriter io.WriteCloser + InputWriter *io.PipeWriter Output *testBuffer - Cleanup func() + // Done is closed when RunClientProxy returns. + Done <-chan struct{} + Cleanup func() } func createTestClient(t *testing.T, serverURL string, requestHandoverTick func() <-chan time.Time, keepaliveInterval time.Duration, errChan chan error) *testClient { @@ -58,8 +60,10 @@ func createTestClientWithDialer(t *testing.T, createConn createWebsocketConnecti if requestHandoverTick == nil { requestHandoverTick = neverTick } + done := make(chan struct{}) wg := sync.WaitGroup{} wg.Go(func() { + defer close(done) err := RunClientProxy(ctx, clientInput, clientOutput, requestHandoverTick, keepaliveInterval, resumable, createConn) if err != nil && !errors.Is(err, context.Canceled) && !errors.Is(err, io.ErrClosedPipe) { if errChan != nil { @@ -72,6 +76,7 @@ func createTestClientWithDialer(t *testing.T, createConn createWebsocketConnecti return &testClient{ InputWriter: clientInputWriter, Output: clientOutput, + Done: done, Cleanup: func() { clientInput.Close() clientInputWriter.Close() @@ -316,67 +321,144 @@ func TestHandoverDialFailureKeepsSessionAlive(t *testing.T) { assert.Equal(t, int32(2), dials.Load(), "expected the initial dial plus exactly one handover dial") } -// A session that ends while a handover is still waiting for the old connection to close must end -// cleanly. Teardown closes that connection itself, and the handover used to report the resulting -// read error as ErrHandoverFailed, which could win the race against the session's real outcome - -// the TestQuickHandover flake. -func TestSessionEndDuringHandoverIsNotAHandoverFailure(t *testing.T) { +// A handover dial can fail with context.Canceled from the dialer's own internals while the session +// is still live. That is a failed dial, not the session ending, so later ticks must still hand over. +func TestHandoverDialCanceledKeepsHandingOver(t *testing.T) { server := createTestServer(t, 2, time.Hour) defer server.Close() - // The handover dial lands here instead of on the proxy server, so nobody ever closes the old - // connection and the handover stays in progress until the session ends. - silentServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil) - if err != nil { - return - } - defer conn.Close() - for { - if _, _, err := conn.ReadMessage(); err != nil { - return - } - } - })) - defer silentServer.Close() - wsURL := "ws" + server.URL[4:] - silentURL := "ws" + silentServer.URL[4:] var dials atomic.Int32 - handoverDialed := make(chan struct{}, 1) createConn := func(ctx context.Context, dial DialRequest) (*websocket.Conn, error) { - url := fmt.Sprintf("%s?id=%s", wsURL, dial.ConnID) - if dials.Add(1) > 1 { - url = silentURL - defer func() { handoverDialed <- struct{}{} }() + if dials.Add(1) == 2 { + return nil, context.Canceled } + url := fmt.Sprintf("%s?id=%s", wsURL, dial.ConnID) conn, _, err := websocket.DefaultDialer.Dial(url, nil) // nolint:bodyclose return conn, err } handoverChan := make(chan time.Time) - errChan := make(chan error, 1) client := createTestClientWithDialer(t, createConn, func() <-chan time.Time { return handoverChan - }, time.Hour, false, errChan) + }, time.Hour, false, nil) + defer client.Cleanup() - msg := []byte("before handover\n") - _, err := client.InputWriter.Write(msg) - require.NoError(t, err) - require.NoError(t, client.Output.AssertWrite(msg)) + for i := range 2 { + select { + case handoverChan <- time.Now(): + case <-time.After(10 * time.Second): + t.Fatalf("handover %d was never started: the handover loop has stopped", i+1) + } + // sendMessage waits for the handover mutex, so this round trip also waits out the handover. + msg := fmt.Appendf(nil, "after handover %d\n", i+1) + _, err := client.InputWriter.Write(msg) + require.NoError(t, err) + require.NoError(t, client.Output.AssertWrite(msg)) + } + assert.Equal(t, int32(3), dials.Load(), "expected the initial dial plus two handover dials") +} - handoverChan <- time.Now() - select { - case <-handoverDialed: - case <-time.After(10 * time.Second): - t.Fatal("the handover never dialed its replacement connection") +var errSourceFailed = errors.New("source failed") + +// A session that ends while a handover is still waiting for the old connection to close must end +// with the session's own outcome. Teardown closes that connection itself, and the handover used to +// report the resulting read error as ErrHandoverFailed: on a clean exit that was the only error, and +// otherwise it could win the race against the real one - the TestQuickHandover flake. +func TestSessionEndDuringHandover(t *testing.T) { + tests := []struct { + name string + endSession func(w *io.PipeWriter) + wantErr error + }{ + { + name: "clean exit", + endSession: func(w *io.PipeWriter) { w.Close() }, + }, + { + name: "source error", + endSession: func(w *io.PipeWriter) { w.CloseWithError(errSourceFailed) }, + wantErr: errSourceFailed, + }, } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + server := createTestServer(t, 2, time.Hour) + defer server.Close() + + // The handover dial lands here instead of on the proxy server, so nobody ever closes the + // old connection and the handover stays in progress until the session ends. + silentServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if err != nil { + return + } + defer conn.Close() + for { + if _, _, err := conn.ReadMessage(); err != nil { + return + } + } + })) + defer silentServer.Close() + + wsURL := "ws" + server.URL[4:] + silentURL := "ws" + silentServer.URL[4:] + var dials atomic.Int32 + handoverDialed := make(chan struct{}, 1) + createConn := func(ctx context.Context, dial DialRequest) (*websocket.Conn, error) { + isHandover := dials.Add(1) > 1 + url := fmt.Sprintf("%s?id=%s", wsURL, dial.ConnID) + if isHandover { + url = silentURL + } + conn, _, err := websocket.DefaultDialer.Dial(url, nil) // nolint:bodyclose + // Only a successful dial leaves the handover in progress. + if isHandover && err == nil { + handoverDialed <- struct{}{} + } + return conn, err + } - client.Cleanup() - select { - case err := <-errChan: - t.Fatalf("session ended with an error: %v", err) - default: + handoverChan := make(chan time.Time) + errChan := make(chan error, 1) + client := createTestClientWithDialer(t, createConn, func() <-chan time.Time { + return handoverChan + }, time.Hour, false, errChan) + defer client.Cleanup() + + msg := []byte("before handover\n") + _, err := client.InputWriter.Write(msg) + require.NoError(t, err) + require.NoError(t, client.Output.AssertWrite(msg)) + + handoverChan <- time.Now() + select { + case <-handoverDialed: + case <-time.After(10 * time.Second): + t.Fatal("the handover never dialed its replacement connection") + } + + // Wait for the session to end before Cleanup closes the reader too: a read that + // wakes up after that returns io.ErrClosedPipe instead of what endSession set. + tt.endSession(client.InputWriter) + select { + case <-client.Done: + case <-time.After(10 * time.Second): + t.Fatal("the session did not end") + } + var sessionErr error + select { + case sessionErr = <-errChan: + default: + } + if tt.wantErr == nil { + assert.NoError(t, sessionErr) + return + } + assert.ErrorIs(t, sessionErr, tt.wantErr) + assert.NotErrorIs(t, sessionErr, ErrHandoverFailed) + }) } } From 6687c6b3e17fc601f4362d6c84eb20b833b7cda1 Mon Sep 17 00:00:00 2001 From: Grigory Panov Date: Wed, 7 Oct 2026 14:56:06 +0200 Subject: [PATCH 3/8] ssh: stop the handover loop for resumable sessions too when the session ends Check context.Canceled before the resumable branch. A resumable session that ended mid-handover logged a false "recovering through session reattachment" line and went back into the loop, where a ready tick could start another handover for a session being torn down. Tests: run the session-end cases for resumable sessions too, and queue a second tick to assert no handover starts after the session ended. Use the existing blackHoleServer. Wait for each handover dial in TestHandoverDialCanceledKeepsHandingOver so a write cannot overtake the handover. The test client signals completion with one channel instead of a channel plus a WaitGroup. Co-authored-by: Isaac --- experimental/ssh/internal/proxy/client.go | 10 +-- .../ssh/internal/proxy/client_server_test.go | 79 +++++++++++-------- 2 files changed, 50 insertions(+), 39 deletions(-) diff --git a/experimental/ssh/internal/proxy/client.go b/experimental/ssh/internal/proxy/client.go index 47cdbc08787..b31c928865c 100644 --- a/experimental/ssh/internal/proxy/client.go +++ b/experimental/ssh/internal/proxy/client.go @@ -108,18 +108,16 @@ func RunClientProxy(ctx context.Context, src io.ReadCloser, dst io.Writer, reque log.Debugf(gCtx, "Could not open a replacement connection for the auth handover, staying on the current one: %v", err) continue } + // The session ended during the handover: see the gCtx.Done case above. + if errors.Is(err, context.Canceled) { + return nil + } if resumable { // The failed handover closes its sockets. The receiving loop // reattaches after the handover releases the write lock. log.Debugf(gCtx, "Auth handover failed, recovering through session reattachment: %v", err) continue } - // The session ended while the handover was in progress. proxy.start - // decides the session's outcome; an error here could be recorded first - // and mask it. - if errors.Is(err, context.Canceled) { - return nil - } return errors.Join(ErrHandoverFailed, err) } } diff --git a/experimental/ssh/internal/proxy/client_server_test.go b/experimental/ssh/internal/proxy/client_server_test.go index 758d236e0b8..0736c3ae5f5 100644 --- a/experimental/ssh/internal/proxy/client_server_test.go +++ b/experimental/ssh/internal/proxy/client_server_test.go @@ -61,8 +61,7 @@ func createTestClientWithDialer(t *testing.T, createConn createWebsocketConnecti requestHandoverTick = neverTick } done := make(chan struct{}) - wg := sync.WaitGroup{} - wg.Go(func() { + go func() { defer close(done) err := RunClientProxy(ctx, clientInput, clientOutput, requestHandoverTick, keepaliveInterval, resumable, createConn) if err != nil && !errors.Is(err, context.Canceled) && !errors.Is(err, io.ErrClosedPipe) { @@ -72,7 +71,7 @@ func createTestClientWithDialer(t *testing.T, createConn createWebsocketConnecti t.Errorf("client error: %v", err) } } - }) + }() return &testClient{ InputWriter: clientInputWriter, Output: clientOutput, @@ -80,7 +79,7 @@ func createTestClientWithDialer(t *testing.T, createConn createWebsocketConnecti Cleanup: func() { clientInput.Close() clientInputWriter.Close() - wg.Wait() + <-done }, } } @@ -329,8 +328,13 @@ func TestHandoverDialCanceledKeepsHandingOver(t *testing.T) { wsURL := "ws" + server.URL[4:] var dials atomic.Int32 + handoverDialed := make(chan struct{}, 2) createConn := func(ctx context.Context, dial DialRequest) (*websocket.Conn, error) { - if dials.Add(1) == 2 { + n := dials.Add(1) + if n > 1 { + handoverDialed <- struct{}{} + } + if n == 2 { return nil, context.Canceled } url := fmt.Sprintf("%s?id=%s", wsURL, dial.ConnID) @@ -350,6 +354,13 @@ func TestHandoverDialCanceledKeepsHandingOver(t *testing.T) { case <-time.After(10 * time.Second): t.Fatalf("handover %d was never started: the handover loop has stopped", i+1) } + // The tick only proves the loop received it. Wait for the dial, which runs under the + // handover mutex, so the write below cannot overtake the handover. + select { + case <-handoverDialed: + case <-time.After(10 * time.Second): + t.Fatalf("handover %d never dialed", i+1) + } // sendMessage waits for the handover mutex, so this round trip also waits out the handover. msg := fmt.Appendf(nil, "after handover %d\n", i+1) _, err := client.InputWriter.Write(msg) @@ -359,15 +370,13 @@ func TestHandoverDialCanceledKeepsHandingOver(t *testing.T) { assert.Equal(t, int32(3), dials.Load(), "expected the initial dial plus two handover dials") } -var errSourceFailed = errors.New("source failed") - // A session that ends while a handover is still waiting for the old connection to close must end -// with the session's own outcome. Teardown closes that connection itself, and the handover used to -// report the resulting read error as ErrHandoverFailed: on a clean exit that was the only error, and -// otherwise it could win the race against the real one - the TestQuickHandover flake. +// with the session's own outcome, not ErrHandoverFailed, and must not start another handover. func TestSessionEndDuringHandover(t *testing.T) { + errSourceFailed := errors.New("source failed") tests := []struct { name string + resumable bool endSession func(w *io.PipeWriter) wantErr error }{ @@ -380,26 +389,25 @@ func TestSessionEndDuringHandover(t *testing.T) { endSession: func(w *io.PipeWriter) { w.CloseWithError(errSourceFailed) }, wantErr: errSourceFailed, }, + { + name: "resumable clean exit", + resumable: true, + endSession: func(w *io.PipeWriter) { w.Close() }, + }, + { + name: "resumable source error", + resumable: true, + endSession: func(w *io.PipeWriter) { w.CloseWithError(errSourceFailed) }, + wantErr: errSourceFailed, + }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { server := createTestServer(t, 2, time.Hour) defer server.Close() - // The handover dial lands here instead of on the proxy server, so nobody ever closes the // old connection and the handover stays in progress until the session ends. - silentServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil) - if err != nil { - return - } - defer conn.Close() - for { - if _, _, err := conn.ReadMessage(); err != nil { - return - } - } - })) + silentServer := blackHoleServer(t) defer silentServer.Close() wsURL := "ws" + server.URL[4:] @@ -409,30 +417,34 @@ func TestSessionEndDuringHandover(t *testing.T) { createConn := func(ctx context.Context, dial DialRequest) (*websocket.Conn, error) { isHandover := dials.Add(1) > 1 url := fmt.Sprintf("%s?id=%s", wsURL, dial.ConnID) + if dial.ResumeCapable { + url += fmt.Sprintf("&resume_version=2&delivered=%d", dial.Delivered) + } if isHandover { url = silentURL } conn, _, err := websocket.DefaultDialer.Dial(url, nil) // nolint:bodyclose // Only a successful dial leaves the handover in progress. if isHandover && err == nil { - handoverDialed <- struct{}{} + select { + case handoverDialed <- struct{}{}: + default: + } } return conn, err } - handoverChan := make(chan time.Time) + // The second tick waits in the channel while the first handover is in progress, so a + // handover loop that keeps going after the session ended dials again. + ticks := make(chan time.Time, 2) + ticks <- time.Now() + ticks <- time.Now() errChan := make(chan error, 1) client := createTestClientWithDialer(t, createConn, func() <-chan time.Time { - return handoverChan - }, time.Hour, false, errChan) + return ticks + }, time.Hour, tt.resumable, errChan) defer client.Cleanup() - msg := []byte("before handover\n") - _, err := client.InputWriter.Write(msg) - require.NoError(t, err) - require.NoError(t, client.Output.AssertWrite(msg)) - - handoverChan <- time.Now() select { case <-handoverDialed: case <-time.After(10 * time.Second): @@ -452,6 +464,7 @@ func TestSessionEndDuringHandover(t *testing.T) { case sessionErr = <-errChan: default: } + assert.Equal(t, int32(2), dials.Load(), "expected no handover after the session ended") if tt.wantErr == nil { assert.NoError(t, sessionErr) return From 16d628c0776a5564acff0f9d6e266478f915f4db Mon Sep 17 00:00:00 2001 From: Grigory Panov Date: Wed, 7 Oct 2026 15:05:38 +0200 Subject: [PATCH 4/8] ssh: state the real rule in the canceled-dial test comment Co-authored-by: Isaac --- experimental/ssh/internal/proxy/client_server_test.go | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/experimental/ssh/internal/proxy/client_server_test.go b/experimental/ssh/internal/proxy/client_server_test.go index 0736c3ae5f5..176b4ebf0d7 100644 --- a/experimental/ssh/internal/proxy/client_server_test.go +++ b/experimental/ssh/internal/proxy/client_server_test.go @@ -320,8 +320,9 @@ func TestHandoverDialFailureKeepsSessionAlive(t *testing.T) { assert.Equal(t, int32(2), dials.Load(), "expected the initial dial plus exactly one handover dial") } -// A handover dial can fail with context.Canceled from the dialer's own internals while the session -// is still live. That is a failed dial, not the session ending, so later ticks must still hand over. +// A failed dial never touches the live connection, whatever error it returns, so the handover loop +// must check for it before it treats context.Canceled as the session ending. Later ticks must still +// hand over. func TestHandoverDialCanceledKeepsHandingOver(t *testing.T) { server := createTestServer(t, 2, time.Hour) defer server.Close() From 75dd6d500139a2bef053fdd1a878254560d73331 Mon Sep 17 00:00:00 2001 From: Grigory Panov Date: Wed, 7 Oct 2026 16:40:46 +0200 Subject: [PATCH 5/8] ssh: keep the server's handover swap on a clean exit The receiving loop also runs on the server. When sshd's stdout ended while the server read the client's normal close acknowledgment for a handover, the new ctx.Err() case replaced that acknowledgment with context.Canceled. acceptHandover then closed the replacement websocket without swapping it in, and the client saw ErrWebsocketDropped instead of a clean exit. Apply the cancellation override on the client only, the side that dials, as reattach already does. Add a regression test that holds the server's close handler until its source has ended. Co-authored-by: Isaac --- .../ssh/internal/proxy/client_server_test.go | 111 ++++++++++++++++++ experimental/ssh/internal/proxy/proxy.go | 11 +- 2 files changed, 118 insertions(+), 4 deletions(-) diff --git a/experimental/ssh/internal/proxy/client_server_test.go b/experimental/ssh/internal/proxy/client_server_test.go index 176b4ebf0d7..324a984131a 100644 --- a/experimental/ssh/internal/proxy/client_server_test.go +++ b/experimental/ssh/internal/proxy/client_server_test.go @@ -476,6 +476,117 @@ func TestSessionEndDuringHandover(t *testing.T) { } } +// closeSignalingSource reports when the proxy closes it, which happens during the proxy's teardown. +type closeSignalingSource struct { + io.ReadCloser + closed chan struct{} +} + +func (s *closeSignalingSource) Close() error { + err := s.ReadCloser.Close() + close(s.closed) + return err +} + +// The receiving loop also runs on the server. If sshd's stdout ends while the server reads the +// client's normal close acknowledgment for a handover, the server must still finish the swap, so its +// teardown closes the replacement connection gracefully and the client sees a clean exit. +func TestServerEOFDuringHandoverIsACleanExit(t *testing.T) { + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + serverProxy := newProxyConnection(nil) + serverInput, serverWriter := io.Pipe() + defer serverWriter.Close() + source := &closeSignalingSource{ReadCloser: serverInput, closed: make(chan struct{})} + serverReady := make(chan struct{}) + ackReceived := make(chan struct{}) + releaseAck := make(chan struct{}) + serverDone := make(chan error, 1) + handoverDone := make(chan error, 1) + var accepted atomic.Bool + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !accepted.CompareAndSwap(false, true) { + handoverDone <- serverProxy.acceptHandover(ctx, w, r) + return + } + if err := serverProxy.accept(w, r); err != nil { + serverDone <- err + return + } + // Hold the receiving loop right after it reads the client's close acknowledgment, until + // the server's source has ended and its teardown has started. + conn := serverProxy.conn.Load() + defaultCloseHandler := conn.CloseHandler() + conn.SetCloseHandler(func(code int, text string) error { + err := defaultCloseHandler(code, text) + close(ackReceived) + select { + case <-releaseAck: + case <-ctx.Done(): + } + return err + }) + close(serverReady) + // Mirrors runServerProxy. + err := serverProxy.start(ctx, source, io.Discard) + closeProxyConnection(ctx, serverProxy) + serverDone <- err + })) + defer server.Close() + + clientInput, clientWriter := io.Pipe() + defer clientWriter.Close() + ticks := make(chan time.Time, 1) + clientDone := make(chan error, 1) + go func() { + clientDone <- RunClientProxy(ctx, clientInput, io.Discard, func() <-chan time.Time { return ticks }, time.Hour, false, + func(dialCtx context.Context, _ DialRequest) (*websocket.Conn, error) { + conn, _, err := websocket.DefaultDialer.DialContext(dialCtx, "ws"+server.URL[4:], nil) // nolint:bodyclose + return conn, err + }) + }() + + select { + case <-serverReady: + case <-ctx.Done(): + t.Fatal("the server did not accept the connection") + } + // The client waits for the server's first byte before it treats the session as started. + _, err := serverWriter.Write([]byte("SSH-2.0-test\r\n")) + require.NoError(t, err) + ticks <- time.Now() + select { + case <-ackReceived: + case <-ctx.Done(): + t.Fatal("the server did not receive the handover close acknowledgment") + } + require.NoError(t, serverWriter.Close()) + select { + case <-source.closed: + case <-ctx.Done(): + t.Fatal("the server did not start its teardown") + } + close(releaseAck) + + select { + case <-handoverDone: + case <-ctx.Done(): + t.Fatal("the server handover did not finish") + } + select { + case err := <-serverDone: + require.NoError(t, err) + case <-ctx.Done(): + t.Fatal("the server did not finish") + } + select { + case err := <-clientDone: + require.NoError(t, err, "a clean server exit during a handover must be a clean client exit") + case <-ctx.Done(): + t.Fatal("the client did not finish") + } +} + // TestClientExitsWhenServerCommandFails reproduces the missing-sshd case: the server accepts the // websocket but can't launch its command, so it closes the connection immediately. The client // proxy must exit promptly instead of hanging on the handover goroutine (which would leave the diff --git a/experimental/ssh/internal/proxy/proxy.go b/experimental/ssh/internal/proxy/proxy.go index a6eed7dee34..100d17e39df 100644 --- a/experimental/ssh/internal/proxy/proxy.go +++ b/experimental/ssh/internal/proxy/proxy.go @@ -536,10 +536,13 @@ func (pc *proxyConnection) runReceivingLoop(ctx context.Context, dst io.Writer) if handover := pc.handoverState.Load(); handover != nil { var closeConnSignal error switch { - case ctx.Err() != nil: - // The session is ending and teardown closed the connection under the handover - // (see the ctx.Err() check below). Pass the cancellation on, so the handover - // initiator does not report a failed handover for an ordinary exit. + case pc.createWebsocketConnection != nil && ctx.Err() != nil: + // Client only (only the client dials): the session is ending and teardown + // closed the connection under the handover (see the ctx.Err() check below). + // Pass the cancellation on, so the handover initiator does not report a failed + // handover for an ordinary exit. The server must still complete the swap on a + // normal close acknowledgment, so its teardown closes the replacement + // connection gracefully. closeConnSignal = ctx.Err() case !websocket.IsCloseError(err, websocket.CloseNormalClosure): closeConnSignal = errors.Join(ErrWebsocketDropped, fmt.Errorf("failed to read from websocket during handover: %w", err)) From c7998d81780d2b61ceebb74b2718fa3d6660067f Mon Sep 17 00:00:00 2001 From: Grigory Panov Date: Wed, 7 Oct 2026 18:26:32 +0200 Subject: [PATCH 6/8] ssh: wait for the client to receive the banner before the handover tick serverWriter.Write only proves that the server's sending loop read the banner, not that it sent it. If acceptHandover took the handover mutex first, the sending loop blocked on that mutex and never read the EOF, while the handover waited for the held close acknowledgment. The test then deadlocked until its 5s deadline. Capture the client output and wait for the banner before the tick. Co-authored-by: Isaac --- experimental/ssh/internal/proxy/client_server_test.go | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/experimental/ssh/internal/proxy/client_server_test.go b/experimental/ssh/internal/proxy/client_server_test.go index 324a984131a..c1dd18379e0 100644 --- a/experimental/ssh/internal/proxy/client_server_test.go +++ b/experimental/ssh/internal/proxy/client_server_test.go @@ -536,10 +536,11 @@ func TestServerEOFDuringHandoverIsACleanExit(t *testing.T) { clientInput, clientWriter := io.Pipe() defer clientWriter.Close() + clientOutput := newTestBuffer(t) ticks := make(chan time.Time, 1) clientDone := make(chan error, 1) go func() { - clientDone <- RunClientProxy(ctx, clientInput, io.Discard, func() <-chan time.Time { return ticks }, time.Hour, false, + clientDone <- RunClientProxy(ctx, clientInput, clientOutput, func() <-chan time.Time { return ticks }, time.Hour, false, func(dialCtx context.Context, _ DialRequest) (*websocket.Conn, error) { conn, _, err := websocket.DefaultDialer.DialContext(dialCtx, "ws"+server.URL[4:], nil) // nolint:bodyclose return conn, err @@ -551,9 +552,13 @@ func TestServerEOFDuringHandoverIsACleanExit(t *testing.T) { case <-ctx.Done(): t.Fatal("the server did not accept the connection") } - // The client waits for the server's first byte before it treats the session as started. - _, err := serverWriter.Write([]byte("SSH-2.0-test\r\n")) + // The client waits for the server's first byte before it treats the session as started. Wait + // until the client has it: the write only proves the sending loop read it, and a handover that + // takes the mutex before the sending loop sends it would block the loop from reading the EOF. + banner := []byte("SSH-2.0-test\r\n") + _, err := serverWriter.Write(banner) require.NoError(t, err) + require.NoError(t, clientOutput.WaitForWrite(banner)) ticks <- time.Now() select { case <-ackReceived: From d43fa7f0d82548d7fc07a6e5937c21936cc9d0dc Mon Sep 17 00:00:00 2001 From: Grigory Panov Date: Thu, 8 Oct 2026 09:07:26 +0200 Subject: [PATCH 7/8] ssh: complete an acknowledged handover swap even when the session ends If the client's input ended while its receiving loop held the server's normal handover close, the ctx.Err() case replaced that close with context.Canceled. initiateHandover then raw-closed the replacement websocket without installing it, so the client's exit never sent the resumable "finished" close there. The server took the raw close as a drop and waited 90s for a reattach. A normal close now always completes the swap, on both sides, as before this PR. Only a read that teardown broke passes the cancellation on. This also covers the server case the client-only guard handled, so the guard is dropped. Co-authored-by: Isaac --- .../ssh/internal/proxy/client_server_test.go | 106 ++++++++++++++++++ experimental/ssh/internal/proxy/proxy.go | 17 +-- 2 files changed, 115 insertions(+), 8 deletions(-) diff --git a/experimental/ssh/internal/proxy/client_server_test.go b/experimental/ssh/internal/proxy/client_server_test.go index c1dd18379e0..cc8f41adaea 100644 --- a/experimental/ssh/internal/proxy/client_server_test.go +++ b/experimental/ssh/internal/proxy/client_server_test.go @@ -592,6 +592,112 @@ func TestServerEOFDuringHandoverIsACleanExit(t *testing.T) { } } +// If the client's input ends while its receiving loop holds the server's normal handover close, the +// client must still complete the swap: its exit then sends the resumable "finished" close on the +// replacement connection, and the server ends the session at once instead of waiting for a reattach. +func TestClientEOFDuringHandoverReleasesServer(t *testing.T) { + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + serverProxy := newResumableProxyConnection(nil) + serverInput, serverWriter := io.Pipe() + defer serverWriter.Close() + serverReady := make(chan struct{}) + serverDone := make(chan error, 1) + handoverDone := make(chan error, 1) + var accepted atomic.Bool + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !accepted.CompareAndSwap(false, true) { + handoverDone <- serverProxy.acceptHandover(ctx, w, r) + return + } + if err := serverProxy.accept(w, r); err != nil { + serverDone <- err + return + } + close(serverReady) + // Mirrors runServerProxy. + err := serverProxy.start(ctx, serverInput, io.Discard) + closeProxyConnection(ctx, serverProxy) + serverDone <- err + })) + defer server.Close() + + clientInput, clientWriter := io.Pipe() + defer clientWriter.Close() + clientSource := &closeSignalingSource{ReadCloser: clientInput, closed: make(chan struct{})} + clientOutput := newTestBuffer(t) + ticks := make(chan time.Time, 1) + clientDone := make(chan error, 1) + ackSent := make(chan struct{}) + releaseAck := make(chan struct{}) + var dials atomic.Int32 + go func() { + clientDone <- RunClientProxy(ctx, clientSource, clientOutput, func() <-chan time.Time { return ticks }, time.Hour, true, + func(dialCtx context.Context, _ DialRequest) (*websocket.Conn, error) { + conn, _, err := websocket.DefaultDialer.DialContext(dialCtx, "ws"+server.URL[4:], nil) // nolint:bodyclose + // Hold the client's receiving loop right after it acknowledges the server's + // handover close on the first connection, until the client's teardown has started. + if err == nil && dials.Add(1) == 1 { + defaultCloseHandler := conn.CloseHandler() + conn.SetCloseHandler(func(code int, text string) error { + err := defaultCloseHandler(code, text) + close(ackSent) + select { + case <-releaseAck: + case <-ctx.Done(): + } + return err + }) + } + return conn, err + }) + }() + + select { + case <-serverReady: + case <-ctx.Done(): + t.Fatal("the server did not accept the connection") + } + banner := []byte("SSH-2.0-test\r\n") + _, err := serverWriter.Write(banner) + require.NoError(t, err) + require.NoError(t, clientOutput.WaitForWrite(banner)) + ticks <- time.Now() + select { + case <-ackSent: + case <-ctx.Done(): + t.Fatal("the client did not acknowledge the handover close") + } + select { + case err := <-handoverDone: + require.NoError(t, err) + case <-ctx.Done(): + t.Fatal("the server did not finish the handover") + } + require.NoError(t, clientWriter.Close()) + select { + case <-clientSource.closed: + case <-ctx.Done(): + t.Fatal("the client did not start its teardown") + } + close(releaseAck) + + select { + case err := <-clientDone: + require.NoError(t, err) + case <-ctx.Done(): + t.Fatal("the client did not finish") + } + select { + case err := <-serverDone: + require.NoError(t, err) + case <-serverProxy.resume.parked: + t.Fatalf("a clean client exit left the server waiting %v for a reattach", serverProxy.resume.grace) + case <-ctx.Done(): + t.Fatal("the server did not finish after a clean client exit") + } +} + // TestClientExitsWhenServerCommandFails reproduces the missing-sshd case: the server accepts the // websocket but can't launch its command, so it closes the connection immediately. The client // proxy must exit promptly instead of hanging on the handover goroutine (which would leave the diff --git a/experimental/ssh/internal/proxy/proxy.go b/experimental/ssh/internal/proxy/proxy.go index 100d17e39df..6cdcb3b7f8b 100644 --- a/experimental/ssh/internal/proxy/proxy.go +++ b/experimental/ssh/internal/proxy/proxy.go @@ -536,15 +536,16 @@ func (pc *proxyConnection) runReceivingLoop(ctx context.Context, dst io.Writer) if handover := pc.handoverState.Load(); handover != nil { var closeConnSignal error switch { - case pc.createWebsocketConnection != nil && ctx.Err() != nil: - // Client only (only the client dials): the session is ending and teardown - // closed the connection under the handover (see the ctx.Err() check below). - // Pass the cancellation on, so the handover initiator does not report a failed - // handover for an ordinary exit. The server must still complete the swap on a - // normal close acknowledgment, so its teardown closes the replacement - // connection gracefully. + case websocket.IsCloseError(err, websocket.CloseNormalClosure): + // The peer completed its side of the handover, so complete the swap even when + // the session is ending: teardown then closes the replacement connection + // gracefully, and a resumable peer does not wait for a reattach. + case ctx.Err() != nil: + // The session is ending and teardown closed the connection under the handover + // (see the ctx.Err() check below). Pass the cancellation on, so the handover + // initiator does not report a failed handover for an ordinary exit. closeConnSignal = ctx.Err() - case !websocket.IsCloseError(err, websocket.CloseNormalClosure): + default: closeConnSignal = errors.Join(ErrWebsocketDropped, fmt.Errorf("failed to read from websocket during handover: %w", err)) } // Signal the current connection is closed to the handover initiator (initiateHandover or acceptHandover). From f7836a7f6d94366a17005a495969ad7b17f3f4f8 Mon Sep 17 00:00:00 2001 From: Grigory Panov Date: Thu, 8 Oct 2026 09:24:31 +0200 Subject: [PATCH 8/8] ssh: make the session-end test deterministic and run EOF tests on Windows TestSessionEndDuringHandover dialed the cat-backed test server for the initial connection. If that server's echo of the client's teardown close arrived before teardown closed the socket, the read was a normal close, the swap completed, and a queued tick could start a third dial. Both connections now land on a server that drains raw bytes and never answers, so the teardown read always fails locally. Move TestServerEOFDuringHandoverIsACleanExit and TestClientEOFDuringHandoverReleasesServer unchanged to a file without the !windows build tag: they do not need the cat echo server. State in the receiving loop comment that a normal close is usually, not always, the peer's side of the handover. Co-authored-by: Isaac --- .../ssh/internal/proxy/client_server_test.go | 255 ++---------------- .../internal/proxy/handover_teardown_test.go | 236 ++++++++++++++++ experimental/ssh/internal/proxy/proxy.go | 7 +- 3 files changed, 259 insertions(+), 239 deletions(-) create mode 100644 experimental/ssh/internal/proxy/handover_teardown_test.go diff --git a/experimental/ssh/internal/proxy/client_server_test.go b/experimental/ssh/internal/proxy/client_server_test.go index cc8f41adaea..d215612c151 100644 --- a/experimental/ssh/internal/proxy/client_server_test.go +++ b/experimental/ssh/internal/proxy/client_server_test.go @@ -371,6 +371,19 @@ func TestHandoverDialCanceledKeepsHandingOver(t *testing.T) { assert.Equal(t, int32(3), dials.Load(), "expected the initial dial plus two handover dials") } +// rawDrainServer accepts a websocket and reads its raw bytes without processing any frame, so it never +// answers anything, not even a close frame. +func rawDrainServer(t *testing.T) *httptest.Server { + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if err != nil { + return + } + defer conn.Close() + _, _ = io.Copy(io.Discard, conn.NetConn()) + })) +} + // A session that ends while a handover is still waiting for the old connection to close must end // with the session's own outcome, not ErrHandoverFailed, and must not start another handover. func TestSessionEndDuringHandover(t *testing.T) { @@ -404,27 +417,19 @@ func TestSessionEndDuringHandover(t *testing.T) { } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - server := createTestServer(t, 2, time.Hour) + // Both connections land on a server that never answers, so nobody closes the old + // connection and the handover stays in progress until the session ends. It does not + // even echo the close frame of the client's teardown: an echo that arrived before + // teardown closed the socket would complete the swap instead. + server := rawDrainServer(t) defer server.Close() - // The handover dial lands here instead of on the proxy server, so nobody ever closes the - // old connection and the handover stays in progress until the session ends. - silentServer := blackHoleServer(t) - defer silentServer.Close() wsURL := "ws" + server.URL[4:] - silentURL := "ws" + silentServer.URL[4:] var dials atomic.Int32 handoverDialed := make(chan struct{}, 1) createConn := func(ctx context.Context, dial DialRequest) (*websocket.Conn, error) { isHandover := dials.Add(1) > 1 - url := fmt.Sprintf("%s?id=%s", wsURL, dial.ConnID) - if dial.ResumeCapable { - url += fmt.Sprintf("&resume_version=2&delivered=%d", dial.Delivered) - } - if isHandover { - url = silentURL - } - conn, _, err := websocket.DefaultDialer.Dial(url, nil) // nolint:bodyclose + conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) // nolint:bodyclose // Only a successful dial leaves the handover in progress. if isHandover && err == nil { select { @@ -476,228 +481,6 @@ func TestSessionEndDuringHandover(t *testing.T) { } } -// closeSignalingSource reports when the proxy closes it, which happens during the proxy's teardown. -type closeSignalingSource struct { - io.ReadCloser - closed chan struct{} -} - -func (s *closeSignalingSource) Close() error { - err := s.ReadCloser.Close() - close(s.closed) - return err -} - -// The receiving loop also runs on the server. If sshd's stdout ends while the server reads the -// client's normal close acknowledgment for a handover, the server must still finish the swap, so its -// teardown closes the replacement connection gracefully and the client sees a clean exit. -func TestServerEOFDuringHandoverIsACleanExit(t *testing.T) { - ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) - defer cancel() - serverProxy := newProxyConnection(nil) - serverInput, serverWriter := io.Pipe() - defer serverWriter.Close() - source := &closeSignalingSource{ReadCloser: serverInput, closed: make(chan struct{})} - serverReady := make(chan struct{}) - ackReceived := make(chan struct{}) - releaseAck := make(chan struct{}) - serverDone := make(chan error, 1) - handoverDone := make(chan error, 1) - var accepted atomic.Bool - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if !accepted.CompareAndSwap(false, true) { - handoverDone <- serverProxy.acceptHandover(ctx, w, r) - return - } - if err := serverProxy.accept(w, r); err != nil { - serverDone <- err - return - } - // Hold the receiving loop right after it reads the client's close acknowledgment, until - // the server's source has ended and its teardown has started. - conn := serverProxy.conn.Load() - defaultCloseHandler := conn.CloseHandler() - conn.SetCloseHandler(func(code int, text string) error { - err := defaultCloseHandler(code, text) - close(ackReceived) - select { - case <-releaseAck: - case <-ctx.Done(): - } - return err - }) - close(serverReady) - // Mirrors runServerProxy. - err := serverProxy.start(ctx, source, io.Discard) - closeProxyConnection(ctx, serverProxy) - serverDone <- err - })) - defer server.Close() - - clientInput, clientWriter := io.Pipe() - defer clientWriter.Close() - clientOutput := newTestBuffer(t) - ticks := make(chan time.Time, 1) - clientDone := make(chan error, 1) - go func() { - clientDone <- RunClientProxy(ctx, clientInput, clientOutput, func() <-chan time.Time { return ticks }, time.Hour, false, - func(dialCtx context.Context, _ DialRequest) (*websocket.Conn, error) { - conn, _, err := websocket.DefaultDialer.DialContext(dialCtx, "ws"+server.URL[4:], nil) // nolint:bodyclose - return conn, err - }) - }() - - select { - case <-serverReady: - case <-ctx.Done(): - t.Fatal("the server did not accept the connection") - } - // The client waits for the server's first byte before it treats the session as started. Wait - // until the client has it: the write only proves the sending loop read it, and a handover that - // takes the mutex before the sending loop sends it would block the loop from reading the EOF. - banner := []byte("SSH-2.0-test\r\n") - _, err := serverWriter.Write(banner) - require.NoError(t, err) - require.NoError(t, clientOutput.WaitForWrite(banner)) - ticks <- time.Now() - select { - case <-ackReceived: - case <-ctx.Done(): - t.Fatal("the server did not receive the handover close acknowledgment") - } - require.NoError(t, serverWriter.Close()) - select { - case <-source.closed: - case <-ctx.Done(): - t.Fatal("the server did not start its teardown") - } - close(releaseAck) - - select { - case <-handoverDone: - case <-ctx.Done(): - t.Fatal("the server handover did not finish") - } - select { - case err := <-serverDone: - require.NoError(t, err) - case <-ctx.Done(): - t.Fatal("the server did not finish") - } - select { - case err := <-clientDone: - require.NoError(t, err, "a clean server exit during a handover must be a clean client exit") - case <-ctx.Done(): - t.Fatal("the client did not finish") - } -} - -// If the client's input ends while its receiving loop holds the server's normal handover close, the -// client must still complete the swap: its exit then sends the resumable "finished" close on the -// replacement connection, and the server ends the session at once instead of waiting for a reattach. -func TestClientEOFDuringHandoverReleasesServer(t *testing.T) { - ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) - defer cancel() - serverProxy := newResumableProxyConnection(nil) - serverInput, serverWriter := io.Pipe() - defer serverWriter.Close() - serverReady := make(chan struct{}) - serverDone := make(chan error, 1) - handoverDone := make(chan error, 1) - var accepted atomic.Bool - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if !accepted.CompareAndSwap(false, true) { - handoverDone <- serverProxy.acceptHandover(ctx, w, r) - return - } - if err := serverProxy.accept(w, r); err != nil { - serverDone <- err - return - } - close(serverReady) - // Mirrors runServerProxy. - err := serverProxy.start(ctx, serverInput, io.Discard) - closeProxyConnection(ctx, serverProxy) - serverDone <- err - })) - defer server.Close() - - clientInput, clientWriter := io.Pipe() - defer clientWriter.Close() - clientSource := &closeSignalingSource{ReadCloser: clientInput, closed: make(chan struct{})} - clientOutput := newTestBuffer(t) - ticks := make(chan time.Time, 1) - clientDone := make(chan error, 1) - ackSent := make(chan struct{}) - releaseAck := make(chan struct{}) - var dials atomic.Int32 - go func() { - clientDone <- RunClientProxy(ctx, clientSource, clientOutput, func() <-chan time.Time { return ticks }, time.Hour, true, - func(dialCtx context.Context, _ DialRequest) (*websocket.Conn, error) { - conn, _, err := websocket.DefaultDialer.DialContext(dialCtx, "ws"+server.URL[4:], nil) // nolint:bodyclose - // Hold the client's receiving loop right after it acknowledges the server's - // handover close on the first connection, until the client's teardown has started. - if err == nil && dials.Add(1) == 1 { - defaultCloseHandler := conn.CloseHandler() - conn.SetCloseHandler(func(code int, text string) error { - err := defaultCloseHandler(code, text) - close(ackSent) - select { - case <-releaseAck: - case <-ctx.Done(): - } - return err - }) - } - return conn, err - }) - }() - - select { - case <-serverReady: - case <-ctx.Done(): - t.Fatal("the server did not accept the connection") - } - banner := []byte("SSH-2.0-test\r\n") - _, err := serverWriter.Write(banner) - require.NoError(t, err) - require.NoError(t, clientOutput.WaitForWrite(banner)) - ticks <- time.Now() - select { - case <-ackSent: - case <-ctx.Done(): - t.Fatal("the client did not acknowledge the handover close") - } - select { - case err := <-handoverDone: - require.NoError(t, err) - case <-ctx.Done(): - t.Fatal("the server did not finish the handover") - } - require.NoError(t, clientWriter.Close()) - select { - case <-clientSource.closed: - case <-ctx.Done(): - t.Fatal("the client did not start its teardown") - } - close(releaseAck) - - select { - case err := <-clientDone: - require.NoError(t, err) - case <-ctx.Done(): - t.Fatal("the client did not finish") - } - select { - case err := <-serverDone: - require.NoError(t, err) - case <-serverProxy.resume.parked: - t.Fatalf("a clean client exit left the server waiting %v for a reattach", serverProxy.resume.grace) - case <-ctx.Done(): - t.Fatal("the server did not finish after a clean client exit") - } -} - // TestClientExitsWhenServerCommandFails reproduces the missing-sshd case: the server accepts the // websocket but can't launch its command, so it closes the connection immediately. The client // proxy must exit promptly instead of hanging on the handover goroutine (which would leave the diff --git a/experimental/ssh/internal/proxy/handover_teardown_test.go b/experimental/ssh/internal/proxy/handover_teardown_test.go new file mode 100644 index 00000000000..46580d86b73 --- /dev/null +++ b/experimental/ssh/internal/proxy/handover_teardown_test.go @@ -0,0 +1,236 @@ +package proxy + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + "github.com/gorilla/websocket" + "github.com/stretchr/testify/require" +) + +// closeSignalingSource reports when the proxy closes it, which happens during the proxy's teardown. +type closeSignalingSource struct { + io.ReadCloser + closed chan struct{} +} + +func (s *closeSignalingSource) Close() error { + err := s.ReadCloser.Close() + close(s.closed) + return err +} + +// The receiving loop also runs on the server. If sshd's stdout ends while the server reads the +// client's normal close acknowledgment for a handover, the server must still finish the swap, so its +// teardown closes the replacement connection gracefully and the client sees a clean exit. +func TestServerEOFDuringHandoverIsACleanExit(t *testing.T) { + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + serverProxy := newProxyConnection(nil) + serverInput, serverWriter := io.Pipe() + defer serverWriter.Close() + source := &closeSignalingSource{ReadCloser: serverInput, closed: make(chan struct{})} + serverReady := make(chan struct{}) + ackReceived := make(chan struct{}) + releaseAck := make(chan struct{}) + serverDone := make(chan error, 1) + handoverDone := make(chan error, 1) + var accepted atomic.Bool + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !accepted.CompareAndSwap(false, true) { + handoverDone <- serverProxy.acceptHandover(ctx, w, r) + return + } + if err := serverProxy.accept(w, r); err != nil { + serverDone <- err + return + } + // Hold the receiving loop right after it reads the client's close acknowledgment, until + // the server's source has ended and its teardown has started. + conn := serverProxy.conn.Load() + defaultCloseHandler := conn.CloseHandler() + conn.SetCloseHandler(func(code int, text string) error { + err := defaultCloseHandler(code, text) + close(ackReceived) + select { + case <-releaseAck: + case <-ctx.Done(): + } + return err + }) + close(serverReady) + // Mirrors runServerProxy. + err := serverProxy.start(ctx, source, io.Discard) + closeProxyConnection(ctx, serverProxy) + serverDone <- err + })) + defer server.Close() + + clientInput, clientWriter := io.Pipe() + defer clientWriter.Close() + clientOutput := newTestBuffer(t) + ticks := make(chan time.Time, 1) + clientDone := make(chan error, 1) + go func() { + clientDone <- RunClientProxy(ctx, clientInput, clientOutput, func() <-chan time.Time { return ticks }, time.Hour, false, + func(dialCtx context.Context, _ DialRequest) (*websocket.Conn, error) { + conn, _, err := websocket.DefaultDialer.DialContext(dialCtx, "ws"+server.URL[4:], nil) // nolint:bodyclose + return conn, err + }) + }() + + select { + case <-serverReady: + case <-ctx.Done(): + t.Fatal("the server did not accept the connection") + } + // The client waits for the server's first byte before it treats the session as started. Wait + // until the client has it: the write only proves the sending loop read it, and a handover that + // takes the mutex before the sending loop sends it would block the loop from reading the EOF. + banner := []byte("SSH-2.0-test\r\n") + _, err := serverWriter.Write(banner) + require.NoError(t, err) + require.NoError(t, clientOutput.WaitForWrite(banner)) + ticks <- time.Now() + select { + case <-ackReceived: + case <-ctx.Done(): + t.Fatal("the server did not receive the handover close acknowledgment") + } + require.NoError(t, serverWriter.Close()) + select { + case <-source.closed: + case <-ctx.Done(): + t.Fatal("the server did not start its teardown") + } + close(releaseAck) + + select { + case <-handoverDone: + case <-ctx.Done(): + t.Fatal("the server handover did not finish") + } + select { + case err := <-serverDone: + require.NoError(t, err) + case <-ctx.Done(): + t.Fatal("the server did not finish") + } + select { + case err := <-clientDone: + require.NoError(t, err, "a clean server exit during a handover must be a clean client exit") + case <-ctx.Done(): + t.Fatal("the client did not finish") + } +} + +// If the client's input ends while its receiving loop holds the server's normal handover close, the +// client must still complete the swap: its exit then sends the resumable "finished" close on the +// replacement connection, and the server ends the session at once instead of waiting for a reattach. +func TestClientEOFDuringHandoverReleasesServer(t *testing.T) { + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + serverProxy := newResumableProxyConnection(nil) + serverInput, serverWriter := io.Pipe() + defer serverWriter.Close() + serverReady := make(chan struct{}) + serverDone := make(chan error, 1) + handoverDone := make(chan error, 1) + var accepted atomic.Bool + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !accepted.CompareAndSwap(false, true) { + handoverDone <- serverProxy.acceptHandover(ctx, w, r) + return + } + if err := serverProxy.accept(w, r); err != nil { + serverDone <- err + return + } + close(serverReady) + // Mirrors runServerProxy. + err := serverProxy.start(ctx, serverInput, io.Discard) + closeProxyConnection(ctx, serverProxy) + serverDone <- err + })) + defer server.Close() + + clientInput, clientWriter := io.Pipe() + defer clientWriter.Close() + clientSource := &closeSignalingSource{ReadCloser: clientInput, closed: make(chan struct{})} + clientOutput := newTestBuffer(t) + ticks := make(chan time.Time, 1) + clientDone := make(chan error, 1) + ackSent := make(chan struct{}) + releaseAck := make(chan struct{}) + var dials atomic.Int32 + go func() { + clientDone <- RunClientProxy(ctx, clientSource, clientOutput, func() <-chan time.Time { return ticks }, time.Hour, true, + func(dialCtx context.Context, _ DialRequest) (*websocket.Conn, error) { + conn, _, err := websocket.DefaultDialer.DialContext(dialCtx, "ws"+server.URL[4:], nil) // nolint:bodyclose + // Hold the client's receiving loop right after it acknowledges the server's + // handover close on the first connection, until the client's teardown has started. + if err == nil && dials.Add(1) == 1 { + defaultCloseHandler := conn.CloseHandler() + conn.SetCloseHandler(func(code int, text string) error { + err := defaultCloseHandler(code, text) + close(ackSent) + select { + case <-releaseAck: + case <-ctx.Done(): + } + return err + }) + } + return conn, err + }) + }() + + select { + case <-serverReady: + case <-ctx.Done(): + t.Fatal("the server did not accept the connection") + } + banner := []byte("SSH-2.0-test\r\n") + _, err := serverWriter.Write(banner) + require.NoError(t, err) + require.NoError(t, clientOutput.WaitForWrite(banner)) + ticks <- time.Now() + select { + case <-ackSent: + case <-ctx.Done(): + t.Fatal("the client did not acknowledge the handover close") + } + select { + case err := <-handoverDone: + require.NoError(t, err) + case <-ctx.Done(): + t.Fatal("the server did not finish the handover") + } + require.NoError(t, clientWriter.Close()) + select { + case <-clientSource.closed: + case <-ctx.Done(): + t.Fatal("the client did not start its teardown") + } + close(releaseAck) + + select { + case err := <-clientDone: + require.NoError(t, err) + case <-ctx.Done(): + t.Fatal("the client did not finish") + } + select { + case err := <-serverDone: + require.NoError(t, err) + case <-serverProxy.resume.parked: + t.Fatalf("a clean client exit left the server waiting %v for a reattach", serverProxy.resume.grace) + case <-ctx.Done(): + t.Fatal("the server did not finish after a clean client exit") + } +} diff --git a/experimental/ssh/internal/proxy/proxy.go b/experimental/ssh/internal/proxy/proxy.go index 6cdcb3b7f8b..88cee6eff3a 100644 --- a/experimental/ssh/internal/proxy/proxy.go +++ b/experimental/ssh/internal/proxy/proxy.go @@ -537,9 +537,10 @@ func (pc *proxyConnection) runReceivingLoop(ctx context.Context, dst io.Writer) var closeConnSignal error switch { case websocket.IsCloseError(err, websocket.CloseNormalClosure): - // The peer completed its side of the handover, so complete the swap even when - // the session is ending: teardown then closes the replacement connection - // gracefully, and a resumable peer does not wait for a reattach. + // A normal close ends the old connection cleanly, usually as the peer's side of + // the handover. Complete the swap even when the session is ending: teardown then + // closes the replacement connection gracefully, and a resumable peer does not + // wait for a reattach. case ctx.Err() != nil: // The session is ending and teardown closed the connection under the handover // (see the ctx.Err() check below). Pass the cancellation on, so the handover