Repository navigation
ssh: don't report a failed handover when the session ends mid-handover - #6976
Conversation
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 <no-reply@databricks.com>
…comes 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 <no-reply@databricks.com>
…on 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 <no-reply@databricks.com>
Co-authored-by: Isaac <no-reply@databricks.com>
Integration test reportCommit: f7836a7
Top 6 slowest tests (at least 2 minutes):
|
| var closeConnSignal error | ||
| if !websocket.IsCloseError(err, websocket.CloseNormalClosure) { | ||
| switch { | ||
| case ctx.Err() != nil: |
There was a problem hiding this comment.
[P2 / merge blocker] Clean server exit during handover becomes a connection failure
This receiving loop also runs on the server. If sshd's stdout reaches EOF while the server processes the handover's normal close acknowledgment, this new branch replaces that acknowledgment with context.Canceled. acceptHandover then takes its error path and closes the replacement websocket without swapping it in or sending a graceful close. The subsequent closeProxyConnection still targets the old socket, so the client reports ErrWebsocketDropped instead of a clean exit.
Reproduction and results
The test below deterministically pauses the server's close handler after receiving the normal handover acknowledgment, produces EOF on the server's source, waits for teardown, and then releases the handler while the client's input remains open. It uses local WebSockets and pipes, not a real SSH server.
Save it as experimental/ssh/internal/proxy/review_server_eof_test.go and run:
go test -race ./experimental/ssh/internal/proxy -run '^TestReviewServerEOFDuringHandover$' -count=20- PR head
77ea838d70fe87000a9446a1eabeb8eccb77b5dd: 20/20 failures, returningproxy websocket dropped/websocket: close 1006 (abnormal closure): unexpected EOF. - Same regression test with
client.goandproxy.gofrom base5ed93138d9cb6a49776d2f4f0b668ad08a378be9supplied through a Go overlay: 20/20 passes under-race.
Proposed fix
Restrict this cancellation override to the client/initiating side, preserving the server's handling of a normal close acknowledgment. The existing role distinction is whether pc.createWebsocketConnection is non-nil: clients dial; servers are constructed with a nil dialer. reattach already uses this distinction.
Concretely, change this case to:
case pc.createWebsocketConnection != nil && ctx.Err() != nil:Keep the new client.go cancellation handling and the dial-failure precedence unchanged. Update the adjacent comment to explain that the override is client-side; a normal server acknowledgment must still complete the swap so cleanup can gracefully close the replacement connection. Do not simply prioritize normal closures over cancellation on both sides, since the client still needs the teardown handling this PR adds.
I tested this one-line guard through a Go overlay: the regression test plus TestSessionEndDuringHandover, TestHandoverDialCanceledKeepsHandingOver, TestQuickHandover, and TestHandoverDialFailureKeepsSessionAlive all passed 30 runs under -race. This is a tested proposal, not a change pushed to the PR branch.
Complete regression test
package proxy
import (
"context"
"io"
"net/http"
"net/http/httptest"
"sync/atomic"
"testing"
"time"
"github.com/gorilla/websocket"
"github.com/stretchr/testify/require"
)
type reviewServerEOFSource struct {
io.ReadCloser
closed chan struct{}
}
func (source *reviewServerEOFSource) Close() error {
err := source.ReadCloser.Close()
close(source.closed)
return err
}
func TestReviewServerEOFDuringHandover(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 := &reviewServerEOFSource{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(writer http.ResponseWriter, request *http.Request) {
if !accepted.CompareAndSwap(false, true) {
handoverDone <- serverProxy.acceptHandover(ctx, writer, request)
return
}
if err := serverProxy.accept(writer, request); err != nil {
serverDone <- err
return
}
connection := serverProxy.conn.Load()
defaultCloseHandler := connection.CloseHandler()
connection.SetCloseHandler(func(code int, text string) error {
err := defaultCloseHandler(code, text)
close(ackReceived)
select {
case <-releaseAck:
case <-ctx.Done():
}
return err
})
close(serverReady)
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) {
connection, response, err := websocket.DefaultDialer.DialContext(dialCtx, "ws"+server.URL[4:], nil)
if response != nil {
response.Body.Close()
}
return connection, err
})
}()
select {
case <-serverReady:
case <-ctx.Done():
t.Fatal("server did not accept the connection")
}
_, err := serverWriter.Write([]byte("SSH-2.0-review\r\n"))
require.NoError(t, err)
ticks <- time.Now()
select {
case <-ackReceived:
case <-ctx.Done():
t.Fatal("server did not receive the handover close acknowledgment")
}
require.NoError(t, serverWriter.Close())
select {
case <-source.closed:
case <-ctx.Done():
t.Fatal("server did not process EOF")
}
close(releaseAck)
select {
case err := <-handoverDone:
t.Logf("server handover result: %v", err)
case <-ctx.Done():
t.Fatal("handover did not finish")
}
select {
case err := <-serverDone:
require.NoError(t, err)
case <-ctx.Done():
t.Fatal("server did not finish")
}
select {
case err := <-clientDone:
require.NoError(t, err, "a clean server EOF must remain a clean client exit during handover")
case <-ctx.Done():
t.Fatal("client did not finish")
}
}Posted with AI assistance.
There was a problem hiding this comment.
Confirmed and fixed. Your repro failed 20/20 on 77ea838d7 and passed with base client.go/proxy.go; it's now TestServerEOFDuringHandoverIsACleanExit (moved to handover_teardown_test.go so it also runs on Windows).
The final fix differs from the client-only guard you proposed (that was 75dd6d500): in d43fa7f0d the receiving loop checks for a normal close first and completes the swap on both sides, as on main. Only a read that teardown broke passes ctx.Err() on. That covers this case, so the guard no longer changed behavior, and I dropped it.
Posted with AI assistance.
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 <no-reply@databricks.com>
| _, err := serverWriter.Write([]byte("SSH-2.0-test\r\n")) | ||
| require.NoError(t, err) | ||
| ticks <- time.Now() |
There was a problem hiding this comment.
[P2 / merge blocker] Wait for the client to receive the banner before starting handover
The production finding is correctly fixed by the client-only guard. There is still a synchronization gap in this test; it was also present in the reproducer from my earlier comment.
serverWriter.Write only waits for the pipe reader to consume the banner, not for runSendingLoop to send it or the client to receive it. The following interleaving is therefore possible:
- The server's sending loop consumes the banner, then is descheduled before
sendMessage. - This write returns and the test sends the handover tick.
acceptHandoveracquires the server's handover mutex first. - The sending loop resumes and blocks trying to send the banner under that same mutex. It cannot proceed to read the subsequent EOF.
- The test waits for
source.closedbefore closingreleaseAck, but teardown needs that EOF, and the handover holds the mutex while waiting for the paused close acknowledgment.
This is a circular wait until the five-second context deadline, at which point the test fails with the server did not start its teardown even though the production fix is correct.
Reproduce
To force the scheduling window, temporarily add this method to the test's closeSignalingSource. The delay is diagnostic instrumentation, not a proposed fix:
func (source *closeSignalingSource) Read(buffer []byte) (int, error) {
count, err := source.ReadCloser.Read(buffer)
if count > 0 {
time.Sleep(50 * time.Millisecond)
}
return count, err
}Then run:
go test -race ./experimental/ssh/internal/proxy -run '^TestServerEOFDuringHandoverIsACleanExit$' -count=2 -vOn 75dd6d500139a2bef053fdd1a878254560d73331, this failed 2/2, each after five seconds. Ordinary unperturbed runs passed; this deliberately forces a valid interleaving that the test does not synchronize against.
Proposed fix
Capture client output with clientOutput := newTestBuffer(t) and pass it as the destination to RunClientProxy instead of io.Discard. Before sending the handover tick, wait until the client actually observes the banner:
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()This establishes the intended precondition without sleeps. With that synchronization added, the test passed 30/30 under -race with the same 50 ms diagnostic delay still in place. Remove the diagnostic Read override from the final change.
Separately, the original regression, the added regression, and related handover tests passed 30 unperturbed race-enabled runs; the full proxy/client/server package suites passed too. Supplying the pre-fix proxy.go through a Go overlay makes the added regression fail 20/20, so its coverage of the original production bug is sound.
Posted with AI assistance.
There was a problem hiding this comment.
Good catch. Reproduced with your 50 ms read delay (3/3 failures), fixed in c7998d817 as you suggested: the client output goes to a testBuffer and the test waits for the banner before the tick. With the delay still in, it passed 30/30 under -race; the delay isn't in the commit.
Posted with AI assistance.
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 <no-reply@databricks.com>
anton-107
left a comment
anton-107
left a comment
There was a problem hiding this comment.
Please preserve prompt server cleanup when a clean client EOF races with an acknowledged handover. The inline comment includes a deterministic regression test (20/20 failures at this head; 20/20 passes with base production code). The demonstrated impact is bounded delayed cleanup of a resumable session, not data loss.
Posted with AI assistance.
| var closeConnSignal error | ||
| if !websocket.IsCloseError(err, websocket.CloseNormalClosure) { | ||
| switch { | ||
| case pc.createWebsocketConnection != nil && ctx.Err() != nil: |
There was a problem hiding this comment.
[P2] Preserve clean client shutdown after an acknowledged handover
This cancellation branch also overrides a normal handover close acknowledgment when client input reaches EOF. If the server has already swapped to the replacement websocket, signaling context.Canceled here sends initiateHandover through the error path at lines 698–701: it raw-closes newConn without installing it. Deferred client cleanup still targets the old websocket, so the replacement never receives the resumable protocol's "finished" close frame.
The client returns success, but a server whose source remains open treats this as a transient disconnect and enters awaitReattach, keeping the session alive for the default 90-second reattachment grace period instead of shutting down promptly. This is a regression in existing cleanup behavior; the demonstrated impact is delayed cleanup, not data loss or an interrupted active session.
Reproduction: pause the client's close handler after it sends the normal ACK; wait for the server to finish handover; close the client's input and wait for teardown to start; then release the close handler. Assert that successful client exit also ends the server session, rather than entering reattachment grace.
With the test below, run:
go test -race ./experimental/ssh/internal/proxy \
-run '^TestReviewClientEOFDuringHandoverReleasesServer$' -count=20- PR head
c7998d81780d2b61ceebb74b2718fa3d6660067f: 20/20 failures, reportingclean client EOF left the server waiting 1m30s for reattachment. - Identical test with only production
client.goandproxy.goreplaced through a Go build overlay by their versions at base5ed93138d9cb6a49776d2f4f0b668ad08a378be9: 20/20 passes. The current test helpers are retained for this comparison.
Requested change: preserve explicit clean-session completion on the replacement websocket when EOF races with an acknowledged handover. Complete the acknowledged swap before graceful teardown, or otherwise ensure the replacement receives the "finished" notification rather than a raw close that looks recoverable to the server. Keep the cancellation/error-preservation fixes and the server-side fix, and add regression coverage asserting prompt server cleanup. This is a proposed fix direction, not a verified patch.
The test reuses newTestBuffer and closeSignalingSource already present in this PR's client_server_test.go. Its build constraint matches that helper file. Save it as experimental/ssh/internal/proxy/review_client_eof_test.go:
Deterministic regression test
//go:build !windows
package proxy
import (
"context"
"io"
"net/http"
"net/http/httptest"
"sync/atomic"
"testing"
"time"
"github.com/gorilla/websocket"
"github.com/stretchr/testify/require"
)
func TestReviewClientEOFDuringHandoverReleasesServer(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(writer http.ResponseWriter, request *http.Request) {
if !accepted.CompareAndSwap(false, true) {
handoverDone <- serverProxy.acceptHandover(ctx, writer, request)
return
}
if err := serverProxy.accept(writer, request); err != nil {
serverDone <- err
return
}
close(serverReady)
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) {
connection, response, err := websocket.DefaultDialer.DialContext(dialCtx, "ws"+server.URL[4:], nil)
if response != nil {
response.Body.Close()
}
if err == nil && dials.Add(1) == 1 {
defaultCloseHandler := connection.CloseHandler()
connection.SetCloseHandler(func(code int, text string) error {
err := defaultCloseHandler(code, text)
close(ackSent)
select {
case <-releaseAck:
case <-ctx.Done():
}
return err
})
}
return connection, err
})
}()
select {
case <-serverReady:
case <-ctx.Done():
t.Fatal("server did not accept the connection")
}
banner := []byte("SSH-2.0-review\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("client did not acknowledge handover")
}
select {
case err := <-handoverDone:
require.NoError(t, err)
case <-ctx.Done():
t.Fatal("server did not finish the handover")
}
require.NoError(t, clientWriter.Close())
select {
case <-clientSource.closed:
case <-ctx.Done():
t.Fatal("client did not start its teardown")
}
close(releaseAck)
select {
case err := <-clientDone:
require.NoError(t, err)
case <-ctx.Done():
t.Fatal("client did not finish")
}
select {
case err := <-serverDone:
require.NoError(t, err)
case <-serverProxy.resume.parked:
t.Fatalf("clean client EOF left the server waiting %v for reattachment", serverProxy.resume.grace)
case <-ctx.Done():
t.Fatal("server did not finish after clean client EOF")
}
}Posted with AI assistance.
There was a problem hiding this comment.
Confirmed and fixed in d43fa7f0d. Your exact test from this comment passes 20/20 under -race on the current head f7836a7f6 (it failed 20/20 on c7998d817). It's in the PR as TestClientEOFDuringHandoverReleasesServer, in handover_teardown_test.go so it also runs on Windows.
The fix follows your first direction: a normal close now completes the swap first, even when the session is ending, so RunClientProxy's deferred close sends "finished" on the replacement and the server ends at once. Only a read that teardown broke passes ctx.Err() on.
Two related orderings are not fixed here because main has them too: (1) teardown breaks the read before the handover close arrives, so the initiator raw-closes newConn; (2) during a handover a "finished" close is treated as the handover close. Both are under "Known limits" in the description. I'd handle them in a follow-up, e.g. by treating proxySessionFinished as session end in the handover branch. Happy to do that here instead if you'd rather.
Posted with AI assistance.
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 <no-reply@databricks.com>
…dows 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 <no-reply@databricks.com>
Integration test reportCommit: f7e98f7
42 interesting tests: 38 FAIL, 4 flaky
Top 50 slowest tests (at least 2 minutes):
|
4 participants
Why
TestQuickHandoverfailed in the merge queue: https://github.com/databricks/cli/actions/runs/37613006327/job/112764757499The cause is a race at session end. A handover can wait for the old websocket to close when the session ends. Teardown then closes that websocket itself. The receiving loop sent the read error (
use of closed network connection) toinitiateHandover. The client handover goroutine then returnedErrHandoverFailed.This also affects real sessions, not only the test:
sshexit during a handover reportedErrHandoverFailed, because no other error existed.normalizeProxyErrorthen changed it tonil, so the real error was lost.Changes
proxy.go: During a handover, the receiving loop now decides in this order:ctx.Err()to the handover initiator. Before, it sent the read error.client.go: IfinitiateHandoverreturnscontext.Canceled, the handover goroutine returnsnil. The result ofproxy.startthen decides how the session ended. The check comes after the dial-failed branch, because a failed dial never changes the live connection. It comes before the resumable branch, so both modes stop the same way.DeadlineExceeded) or a dropped websocket, still returnsErrHandoverFailed.Tests
TestSessionEndDuringHandover: the handover dial goes to a server that never answers, so the handover stays in progress. Then the session ends. The test has 4 cases: a clean exit and a source error, in non-resumable and resumable mode. Each case checks the session result. It also checks that no new handover starts after the session ends.proxy.gochange, all 4 cases failed in most of 10 runs.client.gochange, the source-error and resumable cases failed in most of 10 runs.TestServerEOFDuringHandoverIsACleanExit: sshd's stdout ends while the server reads the client's close acknowledgment for a handover. The client must exit cleanly.TestClientEOFDuringHandoverReleasesServer: the client's input ends while the client holds the server's handover close. The server must end the session at once, and not wait for a reattach.TestHandoverDialCanceledKeepsHandingOver: a dial that fails withcontext.Canceleddoes not stop later handovers.TestQuickHandover, andTestHandoverDialFailureKeepsSessionAlivepassed in 30 of 30 runs with-race. Theexperimental/ssh/internal/proxy,internal/client, andinternal/serverpackages pass with-race.The test client now uses a
Donechannel. With it, a test can wait forRunClientProxyto return beforeCleanupcloses the input pipe.Known limits
These edge cases exist on
maintoo, and this PR does not change them. Each needs a teardown during the short window of a handover, which runs once every 30 minutes:"finished"close.This pull request and its description were written by Isaac.