From 7cb21208b23f095d6ea91233e7e053b592baa30e Mon Sep 17 00:00:00 2001 From: Anton Nekipelov <226657+anton-107@users.noreply.github.com> Date: Mon, 7 Sep 2026 14:52:35 +0000 Subject: [PATCH 1/4] experimental/ssh: survive a dropped tunnel connection instead of ending the session A single TCP reset anywhere between the client and the workspace ended an `ssh connect` session, and both ends destroyed their state within milliseconds, so nothing survived to reconnect to. It happens several times a session for some users, badly enough that they avoid the terminal entirely. The proxy already knew how to swap a websocket underneath a running sshd - that is what the periodic auth handover does - but only for a connection that still works: the close frame is written after the last data frame, and TCP ordering is the whole safety argument. An abrupt reset provides no such barrier, so this adds the byte accounting that replaces it. Three changes, smallest first. ## A failed handover dial no longer ends the session The handover dials a replacement websocket every `--handover-timeout` (30m by default). Any failure of that dial ended the session, even though nothing had been swapped yet and the connection being replaced was still carrying traffic - so one transient blip dropped a healthy tunnel on a 30-minute clock. The pre-swap failure is now tagged so the client stays on its current connection and waits for the next tick. Deferring the refresh is safe: the driver proxy authenticates a websocket at upgrade time, so a live connection is not re-checked. Retrying the dial in place would not be safe. A dial can fail after the server has already accepted it and begun `acceptHandover`, and a second dial would then race the first for the same connection. ## A dropped session is no longer reported as a successful one `isSuccess` is set before the proxy loop starts and `category()` short-circuited on it, so every mid-session drop reached telemetry as a clean, successful session: the drop rate was unmeasurable and a fix for it unverifiable. The sites that end a session now carry sentinels (`ErrConnectFailed`, `ErrWebsocketDropped`, `ErrHandoverFailed`) mapped onto three new categories. `isSuccess` keeps its meaning - the tunnel was established - so `is_success` separates a failed connection attempt from a session that connected and was later cut short, and `error_category` gives the cause of either. `category()` no longer discards a category just because the tunnel came up, but still refuses to invent one: an established session that ended unattributed stays `TYPE_UNSPECIFIED`, because the ssh client and the user's own remote command both exit non-zero there and neither is a tunnel failure. Only a failed connection attempt falls back to `UNKNOWN`, as before. ## A dropped connection is reattached rather than mourned Both ends now count the payload they have handed to the websocket and the payload they have written to their destination, and keep a bounded tail of the former for replay. The buffer, not the socket, is the source of truth: bytes are appended before they are written, so a write that fails on a dying connection loses nothing. On a drop the client redials with the offset it has delivered, the server answers with its own, and each side replays exactly the difference. SSH verifies a MAC over the byte stream, so this has to be exact - one lost or duplicated byte would disconnect the session rather than repair it. Delivered counts are acknowledged every 64 KB so neither replay buffer grows. Server side, a dropped connection parks the session rather than reaping it: sshd stays alive, the client slot stays held (which is also what keeps the shutdown timer cancelled) and a grace timer tears it down exactly as before if nobody comes back. The three timings are each derived from a constraint rather than picked: the client gives up after 60s because ssh itself does at roughly 90s (`ServerAliveInterval` 30 x OpenSSH's default `ServerAliveCountMax` 3); the server's grace is 90s so it never reaps a client still trying; the 1 MB replay cap is far above the unacknowledged window that acks every 64 KB can produce, so reaching it means the peer stopped acknowledging and the connection is already beyond saving. ### Three ordering hazards, each handled - The write mutex cannot be held across the wait for the peer, because that is the lock the other side needs to complete the reattach - it deadlocks the server. `sendGate` throttles the sending loop instead; correctness rests on the write mutex plus appending before writing. - A reattach must not report a delivered count that can still move. The client often notices a drop first, and the server may still be draining data buffered on the dying connection, so `acceptReattach` retires that connection and waits for the receiving loop to stop delivering before it answers. Otherwise the client replays from a stale offset and the server writes those bytes to sshd twice. - A handover tick landing during a client reattach would wait out its own timeout for a receiving loop that is busy reattaching, and end the session. The client holds the write mutex across its whole reattach so the two serialize. ### Both ends must agree, and version skew is real The generated ssh config pins `--metadata` into the `ProxyCommand`, which skips the version-scoped server lookup, so an upgraded client can reach a server an older CLI started. The client therefore probes a new `/capabilities` endpoint first; an older server has no such route, 404s, and the session behaves exactly as it does today. A resume-capable server also refuses a reattach for a session it no longer holds (410) rather than starting a fresh sshd, since replaying into that would fail the SSH stream with a corrupted MAC - a worse error than the drop. The test server reports `resume: false`, because it drives sshd directly over the websocket rather than running the CLI's own proxy server and has none of the session bookkeeping a reattach needs. ## Tests Byte-exact delivery across single and repeated resets, driven through a TCP relay that sends real RSTs (`httptest`'s `CloseClientConnections` is no use: it does not touch the hijacked connections a websocket upgrade leaves behind). The reattach exchange is asserted frame by frame with a deliberately non-empty replay, so the handshake ordering is pinned rather than incidentally satisfied. Plus the replay arithmetic in isolation, the non-resumable path unchanged, the grace period releasing an abandoned session, and 410/400 on malformed reattach requests. `TestADropIsNeverReportedAsACleanExit` guards the errgroup race fixed in 75fcc0045 from the drop side rather than the keepalive side. That fix keeps a cancellation from masking the error that caused it; reverting it fails this test within a few of its 25 attempts. It was found independently here while chasing an intermittent failure of the new drop test, and landed upstream first, so only the test remains. `category()` needed a real merge with the extension-category split in 3adaa48d0. Its precedence was cyclic across the two changes: an interruption had to outrank the category recorded at the failure site, `isSuccess` had to outrank an interruption, and a session-end category has to outrank `isSuccess`. Broken by making the split explicit instead - once the tunnel is up, the only reportable category is a session-end one, so an established session no longer consults the connection-attempt rules at all. Still to validate against real compute: whether the driver proxy passes the reattach query parameter through a websocket upgrade unchanged. `?id=` already shows query parameters survive, so this is a confirmation rather than a risk, but it is load-bearing for the handshake. Co-authored-by: Isaac --- .../cli/ssh-tunnel-handover-dial-failure.md | 1 + .nextchanges/cli/ssh-tunnel-resume.md | 1 + experimental/ssh/internal/client/client.go | 95 ++++- .../internal/client/client_internal_test.go | 71 ++++ .../ssh/internal/client/websockets.go | 23 +- .../ssh/internal/client/websockets_test.go | 28 +- experimental/ssh/internal/proxy/client.go | 32 +- .../ssh/internal/proxy/client_server_test.go | 80 +++- experimental/ssh/internal/proxy/drop_test.go | 163 ++++++++ .../ssh/internal/proxy/keepalive_test.go | 10 +- experimental/ssh/internal/proxy/proxy.go | 227 ++++++++++- experimental/ssh/internal/proxy/proxy_test.go | 2 +- experimental/ssh/internal/proxy/resume.go | 371 ++++++++++++++++++ .../ssh/internal/proxy/resume_e2e_test.go | 203 ++++++++++ .../ssh/internal/proxy/resume_test.go | 89 +++++ experimental/ssh/internal/proxy/server.go | 70 +++- experimental/ssh/internal/server/server.go | 13 + libs/telemetry/protos/ssh_tunnel.go | 31 +- libs/testserver/handlers.go | 9 + 19 files changed, 1461 insertions(+), 58 deletions(-) create mode 100644 .nextchanges/cli/ssh-tunnel-handover-dial-failure.md create mode 100644 .nextchanges/cli/ssh-tunnel-resume.md create mode 100644 experimental/ssh/internal/proxy/drop_test.go create mode 100644 experimental/ssh/internal/proxy/resume.go create mode 100644 experimental/ssh/internal/proxy/resume_e2e_test.go create mode 100644 experimental/ssh/internal/proxy/resume_test.go diff --git a/.nextchanges/cli/ssh-tunnel-handover-dial-failure.md b/.nextchanges/cli/ssh-tunnel-handover-dial-failure.md new file mode 100644 index 00000000000..58733738ec4 --- /dev/null +++ b/.nextchanges/cli/ssh-tunnel-handover-dial-failure.md @@ -0,0 +1 @@ +* Keep an `ssh connect` session alive when the periodic auth handover cannot open its replacement websocket. A transient dial failure previously ended a healthy session, even though the connection being replaced was still working. diff --git a/.nextchanges/cli/ssh-tunnel-resume.md b/.nextchanges/cli/ssh-tunnel-resume.md new file mode 100644 index 00000000000..95d0e76d86d --- /dev/null +++ b/.nextchanges/cli/ssh-tunnel-resume.md @@ -0,0 +1 @@ +* `ssh connect` sessions now survive an unexpected connection drop. When the tunnel websocket is reset mid-session, the CLI reattaches to the running session and replays the bytes that were lost, so the shell and everything running in it stay intact instead of the session ending with `failed to read from websocket`. Requires an SSH server started by a CLI that supports it; against an older server the connection behaves as before. diff --git a/experimental/ssh/internal/client/client.go b/experimental/ssh/internal/client/client.go index 97cc5334126..5551ca4d2e5 100644 --- a/experimental/ssh/internal/client/client.go +++ b/experimental/ssh/internal/client/client.go @@ -6,6 +6,7 @@ import ( _ "embed" "encoding/base64" "encoding/hex" + "encoding/json" "errors" "fmt" "io" @@ -450,7 +451,12 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt outcome.isSuccess = true if opts.ProxyMode { - return runSSHProxy(ctx, client, serverPort, clusterID, opts) + proxyErr := runSSHProxy(ctx, client, serverPort, clusterID, opts) + // isSuccess stays true - the tunnel was established - so the category is what says + // whether the session ran to completion or was cut short, and why. Without it a + // mid-session drop is indistinguishable from a clean exit in telemetry. + outcome.errorCategory = proxySessionEndCategory(proxyErr) + return proxyErr } else if opts.IDE != "" { return runIDE(ctx, client, userName, keyPath, serverPort, clusterID, opts) } else { @@ -898,13 +904,54 @@ func spawnSSHClient(ctx context.Context, client *databricks.WorkspaceClient, use } func runSSHProxy(ctx context.Context, client *databricks.WorkspaceClient, serverPort int, clusterID string, opts ClientOptions) error { - createConn := func(ctx context.Context, connID string) (*websocket.Conn, error) { - return createWebsocketConnection(ctx, client, connID, clusterID, serverPort, opts.Liteswap) + resumable := serverSupportsResume(ctx, client, clusterID, serverPort, opts.Liteswap) + if !resumable { + log.Infof(ctx, "The SSH server does not support session resume, a dropped connection will end the session") + } + createConn := func(ctx context.Context, req proxy.DialRequest) (*websocket.Conn, error) { + req.ResumeCapable = resumable + return createWebsocketConnection(ctx, client, req, clusterID, serverPort, opts.Liteswap) } requestHandoverTick := func() <-chan time.Time { return time.After(opts.HandoverTimeout) } - return proxy.RunClientProxy(ctx, os.Stdin, os.Stdout, requestHandoverTick, opts.KeepaliveInterval, createConn) + return proxy.RunClientProxy(ctx, os.Stdin, os.Stdout, requestHandoverTick, opts.KeepaliveInterval, resumable, createConn) +} + +// serverSupportsResume reports whether the running SSH server speaks the resume protocol. +// +// This must be established before the session starts, not discovered when a connection drops. +// A server from an older CLI has no /capabilities endpoint and, worse, would answer a reattach +// request by starting a fresh sshd: replaying into that fails the SSH stream with a corrupted MAC +// instead of the clear error the user gets today. A 404 (or any other failure) therefore means +// "assume not supported". Version skew is real here because the generated ssh config pins +// --metadata into the ProxyCommand, which skips the version-scoped server lookup, so an upgraded +// client can reach a server an older CLI started. +func serverSupportsResume(ctx context.Context, client *databricks.WorkspaceClient, clusterID string, serverPort int, liteswap string) bool { + req, err := newDriverProxyRequest(ctx, client, clusterID, serverPort, "capabilities", liteswap) + if err != nil { + log.Debugf(ctx, "Failed to build the server capabilities request: %v", err) + return false + } + httpClient := &http.Client{Transport: client.Config.HTTPTransport} + resp, err := httpClient.Do(req) + if err != nil { + log.Debugf(ctx, "Failed to query the server capabilities: %v", err) + return false + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + log.Debugf(ctx, "The server does not serve /capabilities (status %d)", resp.StatusCode) + return false + } + var capabilities struct { + Resume bool `json:"resume"` + } + if err := json.NewDecoder(resp.Body).Decode(&capabilities); err != nil { + log.Debugf(ctx, "Failed to decode the server capabilities: %v", err) + return false + } + return capabilities.Resume } // accessModeUILabel maps a cluster's access mode to the name shown in the Databricks UI. @@ -1278,11 +1325,22 @@ func sshExtensionErrorCategory(err error) protos.SshTunnelErrorCategory { return protos.SshTunnelErrorCategoryUnknown } -// category returns the error category to report. An interrupted attempt means the user gave -// up, whichever call happened to observe it first, so it wins over the category recorded at -// the failure site. An unattributed failure is reported as UNKNOWN so that it stays countable. +// category returns the error category to report. Once the tunnel is up nothing that follows is a +// connection failure, so an established session reports only how it ended, and only the proxy can +// say that. For a connection attempt, an interruption means the user gave up, whichever call +// happened to observe it first, so it wins over the category recorded at the failure site; an +// unattributed attempt is reported as UNKNOWN so that it stays countable. func (o connectOutcome) category() protos.SshTunnelErrorCategory { - if o.isSuccess || o.err == nil { + if o.err == nil { + return protos.SshTunnelErrorCategoryUnspecified + } + if o.isSuccess { + // proxySessionEndCategory is the only thing that sets a category this late, so an empty + // one means the session simply ended: an interruption, an ordinary exit, or a non-zero + // exit from the ssh client or the user's own remote command. None is a tunnel failure. + if o.errorCategory != "" { + return o.errorCategory + } return protos.SshTunnelErrorCategoryUnspecified } if errors.Is(o.ctxErr, context.Canceled) || errors.Is(o.err, context.Canceled) { @@ -1294,6 +1352,27 @@ func (o connectOutcome) category() protos.SshTunnelErrorCategory { return o.errorCategory } +// proxySessionEndCategory attributes how a proxy-mode session ended. A dropped websocket is +// checked first on purpose: a drop that lands during a handover surfaces from either the +// receiving loop or the handover goroutine, whichever the errgroup records first, and it +// should be counted as a drop in both cases. HANDOVER_FAILED is then only the handover's own +// failures. An unrecognised error is left unattributed - normalizeProxyError already maps a +// clean finish and a user interrupt to nil. +func proxySessionEndCategory(err error) protos.SshTunnelErrorCategory { + switch { + case err == nil: + return "" + case errors.Is(err, proxy.ErrWebsocketDropped): + return protos.SshTunnelErrorCategoryWebsocketDropped + case errors.Is(err, proxy.ErrHandoverFailed): + return protos.SshTunnelErrorCategoryHandoverFailed + case errors.Is(err, proxy.ErrConnectFailed): + return protos.SshTunnelErrorCategoryWebsocketConnectFailed + default: + return "" + } +} + func logSshTunnelEvent(ctx context.Context, opts ClientOptions, outcome connectOutcome) { telemetry.Log(ctx, protos.DatabricksCliLog{ SshTunnelEvent: buildSshTunnelEvent(opts, outcome), diff --git a/experimental/ssh/internal/client/client_internal_test.go b/experimental/ssh/internal/client/client_internal_test.go index ba15384bb46..33b15e18336 100644 --- a/experimental/ssh/internal/client/client_internal_test.go +++ b/experimental/ssh/internal/client/client_internal_test.go @@ -11,6 +11,7 @@ import ( "testing" "time" + "github.com/databricks/cli/experimental/ssh/internal/proxy" "github.com/databricks/cli/experimental/ssh/internal/sshconfig" "github.com/databricks/cli/experimental/ssh/internal/vscode" "github.com/databricks/cli/libs/cmdio" @@ -525,6 +526,17 @@ func TestConnectOutcomeCategory(t *testing.T) { outcome: connectOutcome{isSuccess: true, err: errFailed}, want: protos.SshTunnelErrorCategoryUnspecified, }, + { + // A session end the proxy did attribute must survive isSuccess, or a mid-session + // drop is indistinguishable from a clean exit. + name: "attributed session end after a successful connection keeps its category", + outcome: connectOutcome{ + isSuccess: true, + errorCategory: protos.SshTunnelErrorCategoryWebsocketDropped, + err: errFailed, + }, + want: protos.SshTunnelErrorCategoryWebsocketDropped, + }, { name: "attributed failure keeps its category", outcome: connectOutcome{errorCategory: protos.SshTunnelErrorCategoryIDECommandNotOnPath, err: errFailed}, @@ -582,6 +594,18 @@ func TestConnectOutcomeCategory(t *testing.T) { outcome: connectOutcome{isSuccess: true, ctxErr: context.Canceled, err: errFailed}, want: protos.SshTunnelErrorCategoryUnspecified, }, + { + // ...but a session end the proxy did attribute outranks the interruption, or a drop + // that happened to coincide with the user giving up would be lost. + name: "an attributed session end wins over an interruption", + outcome: connectOutcome{ + isSuccess: true, + ctxErr: context.Canceled, + errorCategory: protos.SshTunnelErrorCategoryWebsocketDropped, + err: errFailed, + }, + want: protos.SshTunnelErrorCategoryWebsocketDropped, + }, } for _, tt := range tests { @@ -635,6 +659,53 @@ func TestSshExtensionErrorCategory(t *testing.T) { } } +func TestProxySessionEndCategory(t *testing.T) { + tests := []struct { + name string + err error + want protos.SshTunnelErrorCategory + }{ + { + name: "clean finish is not attributed", + err: nil, + want: "", + }, + { + name: "dropped websocket", + err: fmt.Errorf("wrapped: %w", proxy.ErrWebsocketDropped), + want: protos.SshTunnelErrorCategoryWebsocketDropped, + }, + { + name: "handover failure", + err: fmt.Errorf("wrapped: %w", proxy.ErrHandoverFailed), + want: protos.SshTunnelErrorCategoryHandoverFailed, + }, + { + name: "connect failure", + err: fmt.Errorf("wrapped: %w", proxy.ErrConnectFailed), + want: protos.SshTunnelErrorCategoryWebsocketConnectFailed, + }, + { + // A drop landing during a handover surfaces from either the receiving loop or the + // handover goroutine, so it must be counted as a drop either way. + name: "a drop during a handover counts as a drop", + err: errors.Join(proxy.ErrHandoverFailed, proxy.ErrWebsocketDropped), + want: protos.SshTunnelErrorCategoryWebsocketDropped, + }, + { + name: "an unrecognised error is left unattributed", + err: errors.New("something else"), + want: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.want, proxySessionEndCategory(tt.err)) + }) + } +} + func TestBuildSshTunnelEventReportsErrorCategory(t *testing.T) { got := buildSshTunnelEvent(ClientOptions{ConnectionName: "my-conn", IDE: "vscode"}, connectOutcome{ errorCategory: protos.SshTunnelErrorCategoryIDECommandNotOnPath, diff --git a/experimental/ssh/internal/client/websockets.go b/experimental/ssh/internal/client/websockets.go index fc280674957..7f0d30aa04a 100644 --- a/experimental/ssh/internal/client/websockets.go +++ b/experimental/ssh/internal/client/websockets.go @@ -5,14 +5,16 @@ import ( "fmt" "net/http" "net/url" + "strconv" + "github.com/databricks/cli/experimental/ssh/internal/proxy" "github.com/databricks/cli/libs/auth" "github.com/databricks/databricks-sdk-go" "github.com/gorilla/websocket" ) -func createWebsocketConnection(ctx context.Context, client *databricks.WorkspaceClient, connID, clusterID string, serverPort int, liteswap string) (*websocket.Conn, error) { - proxyURL, err := getProxyURL(ctx, client, connID, clusterID, serverPort) +func createWebsocketConnection(ctx context.Context, client *databricks.WorkspaceClient, dial proxy.DialRequest, clusterID string, serverPort int, liteswap string) (*websocket.Conn, error) { + proxyURL, err := getProxyURL(ctx, client, dial, clusterID, serverPort) if err != nil { return nil, fmt.Errorf("failed to get proxy URL: %w", err) } @@ -38,19 +40,19 @@ func createWebsocketConnection(ctx context.Context, client *databricks.Workspace return conn, nil } -func getProxyURL(ctx context.Context, client *databricks.WorkspaceClient, connID, clusterID string, serverPort int) (string, error) { +func getProxyURL(ctx context.Context, client *databricks.WorkspaceClient, dial proxy.DialRequest, clusterID string, serverPort int) (string, error) { workspaceID, err := auth.ResolveWorkspaceID(ctx, client) if err != nil { return "", fmt.Errorf("failed to get current workspace ID: %w", err) } - return buildProxyWebsocketURL(client.Config.Host, workspaceID, clusterID, serverPort, connID) + return buildProxyWebsocketURL(client.Config.Host, workspaceID, clusterID, serverPort, dial) } // buildProxyWebsocketURL builds the driver-proxy websocket URL for an SSH tunnel. // // The scheme follows the host (http -> ws, else wss) instead of being hardcoded to // wss, so the tunnel is also diallable against the plaintext local test server. -func buildProxyWebsocketURL(host, workspaceID, clusterID string, serverPort int, connID string) (string, error) { +func buildProxyWebsocketURL(host, workspaceID, clusterID string, serverPort int, dial proxy.DialRequest) (string, error) { u, err := url.Parse(host) if err != nil { return "", fmt.Errorf("failed to parse host %q: %w", host, err) @@ -65,6 +67,15 @@ func buildProxyWebsocketURL(host, workspaceID, clusterID string, serverPort int, // the driver-proxy endpoint and uses an "o" path segment regardless of // whether the workspace ID itself is the legacy or new shape. u.Path = fmt.Sprintf("/driver-proxy-api/o/%s/%s/%d/ssh", workspaceID, clusterID, serverPort) - u.RawQuery = url.Values{"id": {connID}}.Encode() + query := url.Values{"id": {dial.ConnID}} + if dial.ResumeCapable { + // Sending "delivered" at all is what tells the server this client speaks the resume + // protocol, so it buffers its own output for replay from the start of the session. + query.Set("delivered", strconv.FormatInt(dial.Delivered, 10)) + if dial.Reattach { + query.Set("reattach", "1") + } + } + u.RawQuery = query.Encode() return u.String(), nil } diff --git a/experimental/ssh/internal/client/websockets_test.go b/experimental/ssh/internal/client/websockets_test.go index 601ce5e2097..2c8c106282f 100644 --- a/experimental/ssh/internal/client/websockets_test.go +++ b/experimental/ssh/internal/client/websockets_test.go @@ -3,6 +3,7 @@ package client import ( "testing" + "github.com/databricks/cli/experimental/ssh/internal/proxy" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -11,23 +12,48 @@ func TestBuildProxyWebsocketURL(t *testing.T) { tests := []struct { name string host string + dial proxy.DialRequest want string }{ { name: "https host is dialed over wss", host: "https://my-workspace.cloud.databricks.test", + dial: proxy.DialRequest{ConnID: "conn-1"}, want: "wss://my-workspace.cloud.databricks.test/driver-proxy-api/o/900800700600/1234-567890-abc/7772/ssh?id=conn-1", }, { name: "plaintext http host is dialed over ws", host: "http://127.0.0.1:8080", + dial: proxy.DialRequest{ConnID: "conn-1"}, want: "ws://127.0.0.1:8080/driver-proxy-api/o/900800700600/1234-567890-abc/7772/ssh?id=conn-1", }, + { + // A server that does not speak the resume protocol must see the URL it always saw, + // so no resume parameters leak out when the capability probe said no. + name: "a non-resumable dial carries no resume parameters", + host: "http://127.0.0.1:8080", + dial: proxy.DialRequest{ConnID: "conn-1", Delivered: 4096, Reattach: true}, + want: "ws://127.0.0.1:8080/driver-proxy-api/o/900800700600/1234-567890-abc/7772/ssh?id=conn-1", + }, + { + // The offset travels on every dial, not just a reattach: its presence is what tells + // the server to start buffering its own output for replay. + name: "a resumable dial always carries the delivered offset", + host: "http://127.0.0.1:8080", + dial: proxy.DialRequest{ConnID: "conn-1", ResumeCapable: true}, + want: "ws://127.0.0.1:8080/driver-proxy-api/o/900800700600/1234-567890-abc/7772/ssh?delivered=0&id=conn-1", + }, + { + name: "a reattach states its intent and its offset", + host: "http://127.0.0.1:8080", + dial: proxy.DialRequest{ConnID: "conn-1", ResumeCapable: true, Delivered: 4096, Reattach: true}, + want: "ws://127.0.0.1:8080/driver-proxy-api/o/900800700600/1234-567890-abc/7772/ssh?delivered=4096&id=conn-1&reattach=1", + }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - got, err := buildProxyWebsocketURL(tt.host, "900800700600", "1234-567890-abc", 7772, "conn-1") + got, err := buildProxyWebsocketURL(tt.host, "900800700600", "1234-567890-abc", 7772, tt.dial) require.NoError(t, err) assert.Equal(t, tt.want, got) }) diff --git a/experimental/ssh/internal/proxy/client.go b/experimental/ssh/internal/proxy/client.go index c9cb405aedc..53aca3880bd 100644 --- a/experimental/ssh/internal/proxy/client.go +++ b/experimental/ssh/internal/proxy/client.go @@ -42,8 +42,8 @@ func (f *firstByteWriter) Write(p []byte) (int, error) { // one a handover creates — logs the pongs coming back for our keepalive pings. Debug visibility only: // the receiving loop stays the only judge of whether a connection is alive. func logPongs(ctx context.Context, createConn createWebsocketConnectionFunc) createWebsocketConnectionFunc { - return func(connCtx context.Context, connID string) (*websocket.Conn, error) { - conn, err := createConn(connCtx, connID) + return func(connCtx context.Context, req DialRequest) (*websocket.Conn, error) { + conn, err := createConn(connCtx, req) if err != nil { return nil, err } @@ -55,13 +55,23 @@ func logPongs(ctx context.Context, createConn createWebsocketConnectionFunc) cre } } -func RunClientProxy(ctx context.Context, src io.ReadCloser, dst io.Writer, requestHandoverTick func() <-chan time.Time, keepaliveInterval time.Duration, createConn createWebsocketConnectionFunc) error { - proxy := newProxyConnection(logPongs(ctx, createConn)) +// RunClientProxy proxies the SSH byte stream over a websocket to the tunnel server. +// +// resumable turns on the resume protocol, which lets a session survive an unexpected disconnect. +// It must only be set when the server is known to speak it: an older server answers a reattach +// request by starting a fresh sshd, and replaying into that corrupts the SSH stream instead of +// repairing it. +func RunClientProxy(ctx context.Context, src io.ReadCloser, dst io.Writer, requestHandoverTick func() <-chan time.Time, keepaliveInterval time.Duration, resumable bool, createConn createWebsocketConnectionFunc) error { + newConnection := newProxyConnection + if resumable { + newConnection = newResumableProxyConnection + } + proxy := newConnection(logPongs(ctx, createConn)) log.Infof(ctx, "Establishing SSH proxy connection...") ctx, cancel := context.WithCancel(ctx) defer cancel() if err := proxy.connect(ctx); err != nil { - return fmt.Errorf("failed to connect to proxy: %w", err) + return errors.Join(ErrConnectFailed, fmt.Errorf("failed to connect to proxy: %w", err)) } defer proxy.close() log.Infof(ctx, "SSH proxy connection established") @@ -86,7 +96,17 @@ func RunClientProxy(ctx context.Context, src io.ReadCloser, dst io.Writer, reque return nil case <-requestHandoverTick(): if err := proxy.initiateHandover(gCtx); err != nil { - return err + // 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 + // a drop every handover interval. The next tick tries again. Deferring + // the auth refresh is safe: the driver proxy authenticates a websocket + // at upgrade time, so a live connection is not re-checked. + if errors.Is(err, errHandoverDialFailed) { + log.Warnf(gCtx, "Could not open a replacement connection for the auth handover, staying on the current one: %v", err) + continue + } + 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 7aa45d35e81..82c07635297 100644 --- a/experimental/ssh/internal/proxy/client_server_test.go +++ b/experimental/ssh/internal/proxy/client_server_test.go @@ -12,6 +12,7 @@ import ( "net/http/httptest" "os/exec" "sync" + "sync/atomic" "testing" "time" @@ -39,21 +40,27 @@ type testClient struct { } func createTestClient(t *testing.T, serverURL string, requestHandoverTick func() <-chan time.Time, keepaliveInterval time.Duration, errChan chan error) *testClient { - ctx := cmdio.MockDiscard(t.Context()) - clientInput, clientInputWriter := io.Pipe() - clientOutput := newTestBuffer(t) wsURL := "ws" + serverURL[4:] - createConn := func(ctx context.Context, connID string) (*websocket.Conn, error) { - url := fmt.Sprintf("%s?id=%s", wsURL, connID) + createConn := func(ctx context.Context, dial DialRequest) (*websocket.Conn, error) { + url := fmt.Sprintf("%s?id=%s", wsURL, dial.ConnID) conn, _, err := websocket.DefaultDialer.Dial(url, nil) // nolint:bodyclose return conn, err } + return createTestClientWithDialer(t, createConn, requestHandoverTick, keepaliveInterval, false, errChan) +} + +// createTestClientWithDialer is createTestClient with the websocket dialer supplied by the caller, +// so a test can control which dials succeed - the initial connection's or a handover's. +func createTestClientWithDialer(t *testing.T, createConn createWebsocketConnectionFunc, requestHandoverTick func() <-chan time.Time, keepaliveInterval time.Duration, resumable bool, errChan chan error) *testClient { + ctx := cmdio.MockDiscard(t.Context()) + clientInput, clientInputWriter := io.Pipe() + clientOutput := newTestBuffer(t) if requestHandoverTick == nil { requestHandoverTick = neverTick } wg := sync.WaitGroup{} wg.Go(func() { - err := RunClientProxy(ctx, clientInput, clientOutput, requestHandoverTick, keepaliveInterval, createConn) + 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 { errChan <- err @@ -241,6 +248,55 @@ func TestQuickHandover(t *testing.T) { assert.Equal(t, string(expectedOutput), client.Output.String()) } +// A handover that fails while dialing its replacement connection must not end the session: the +// connection it was meant to replace is still live and carrying traffic. Until this was handled, +// a single transient dial failure - a token refresh, a DNS blip, a proxy stumble - dropped an +// otherwise healthy session once every handover interval. +func TestHandoverDialFailureKeepsSessionAlive(t *testing.T) { + server := createTestServer(t, 2, time.Hour) + defer server.Close() + + wsURL := "ws" + server.URL[4:] + var dials atomic.Int32 + createConn := func(ctx context.Context, dial DialRequest) (*websocket.Conn, error) { + // Let the initial connection through and fail every handover dial after it. + if dials.Add(1) > 1 { + return nil, errors.New("simulated transient dial failure") + } + 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) + defer client.Cleanup() + + beforeMsg := []byte("before handover\n") + _, err := client.InputWriter.Write(beforeMsg) + require.NoError(t, err) + require.NoError(t, client.Output.AssertWrite(beforeMsg)) + + handoverChan <- time.Now() + + // The original connection must still be proxying both ways. sendMessage blocks on the + // handover mutex, so this write cannot overtake the failed handover. + afterMsg := []byte("after failed handover\n") + _, err = client.InputWriter.Write(afterMsg) + require.NoError(t, err) + require.NoError(t, client.Output.AssertWrite(afterMsg)) + + select { + case err := <-errChan: + t.Fatalf("session ended after a failed handover dial: %v", err) + default: + } + assert.Equal(t, int32(2), dials.Load(), "expected the initial dial plus exactly one handover dial") +} + // 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 @@ -255,8 +311,8 @@ func TestClientExitsWhenServerCommandFails(t *testing.T) { defer server.Close() wsURL := "ws" + server.URL[4:] - createConn := func(ctx context.Context, connID string) (*websocket.Conn, error) { - conn, _, err := websocket.DefaultDialer.Dial(fmt.Sprintf("%s?id=%s", wsURL, connID), nil) // nolint:bodyclose + createConn := func(ctx context.Context, dial DialRequest) (*websocket.Conn, error) { + conn, _, err := websocket.DefaultDialer.Dial(fmt.Sprintf("%s?id=%s", wsURL, dial.ConnID), nil) // nolint:bodyclose return conn, err } // Source is never closed by the test; only the server-side close must drive the client to exit. @@ -265,7 +321,7 @@ func TestClientExitsWhenServerCommandFails(t *testing.T) { done := make(chan error, 1) go func() { - done <- RunClientProxy(ctx, src, io.Discard, requestHandoverTick, time.Hour, createConn) + done <- RunClientProxy(ctx, src, io.Discard, requestHandoverTick, time.Hour, false, createConn) }() select { @@ -303,8 +359,8 @@ func TestClientTimesOutWhenServerSendsNothing(t *testing.T) { defer server.Close() wsURL := "ws" + server.URL[4:] - createConn := func(ctx context.Context, connID string) (*websocket.Conn, error) { - conn, _, err := websocket.DefaultDialer.Dial(fmt.Sprintf("%s?id=%s", wsURL, connID), nil) // nolint:bodyclose + createConn := func(ctx context.Context, dial DialRequest) (*websocket.Conn, error) { + conn, _, err := websocket.DefaultDialer.Dial(fmt.Sprintf("%s?id=%s", wsURL, dial.ConnID), nil) // nolint:bodyclose return conn, err } src, _ := io.Pipe() @@ -312,7 +368,7 @@ func TestClientTimesOutWhenServerSendsNothing(t *testing.T) { done := make(chan error, 1) go func() { - done <- RunClientProxy(ctx, src, io.Discard, requestHandoverTick, time.Hour, createConn) + done <- RunClientProxy(ctx, src, io.Discard, requestHandoverTick, time.Hour, false, createConn) }() select { diff --git a/experimental/ssh/internal/proxy/drop_test.go b/experimental/ssh/internal/proxy/drop_test.go new file mode 100644 index 00000000000..19bef72d0c6 --- /dev/null +++ b/experimental/ssh/internal/proxy/drop_test.go @@ -0,0 +1,163 @@ +//go:build !windows + +package proxy + +import ( + "context" + "fmt" + "io" + "net" + "sync" + "testing" + "time" + + "github.com/gorilla/websocket" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// tcpRelay stands in for the workspace front door: it forwards TCP between the client and the +// SSH proxy server, and can reset the client leg the way a load balancer recycling a target +// does. httptest's own CloseClientConnections is no use here - it does not touch the hijacked +// connections a websocket upgrade leaves behind. +type tcpRelay struct { + listener net.Listener + upstream string + mu sync.Mutex + clientConns []*net.TCPConn +} + +func newTCPRelay(t *testing.T, upstream string) *tcpRelay { + listener, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + relay := &tcpRelay{listener: listener, upstream: upstream} + t.Cleanup(func() { listener.Close() }) + go relay.serve() + return relay +} + +func (r *tcpRelay) serve() { + for { + downstream, err := r.listener.Accept() + if err != nil { + return + } + upstream, err := net.Dial("tcp", r.upstream) + if err != nil { + downstream.Close() + return + } + r.mu.Lock() + r.clientConns = append(r.clientConns, downstream.(*net.TCPConn)) + r.mu.Unlock() + // Both directions end with the connection being torn down, so a copy error is the + // expected way for these to finish. + go func() { + _, _ = io.Copy(upstream, downstream) + upstream.Close() + }() + go func() { + _, _ = io.Copy(downstream, upstream) + downstream.Close() + }() + } +} + +// resetClients sends a TCP RST on every client leg, so the client's next read fails with +// "connection reset by peer" rather than seeing a clean close. SetLinger is asserted: without +// it the close is graceful and the test would exercise the wrong path. Connections are taken off +// the list as they are reset, so a later call only touches legs opened since. +func (r *tcpRelay) resetClients(t *testing.T) { + r.mu.Lock() + conns := r.clientConns + r.clientConns = nil + r.mu.Unlock() + for _, conn := range conns { + require.NoError(t, conn.SetLinger(0)) + require.NoError(t, conn.Close()) + } +} + +// createResumableTestClient builds a client with the resume protocol enabled, dialing through a +// URL that mirrors the production one: the delivered offset rides along on every dial, and a +// reattach says so explicitly. +func createResumableTestClient(t *testing.T, serverURL string, errChan chan error) *testClient { + wsURL := "ws" + serverURL[4:] + createConn := func(ctx context.Context, dial DialRequest) (*websocket.Conn, error) { + url := fmt.Sprintf("%s?id=%s", wsURL, dial.ConnID) + if dial.ResumeCapable { + url += fmt.Sprintf("&delivered=%d", dial.Delivered) + if dial.Reattach { + url += "&reattach=1" + } + } + conn, _, err := websocket.DefaultDialer.Dial(url, nil) // nolint:bodyclose + return conn, err + } + return createTestClientWithDialer(t, createConn, nil, time.Hour, true, errChan) +} + +// URL returns the relay's address in the http form createTestClient expects. +func (r *tcpRelay) URL() string { + return "http://" + r.listener.Addr().String() +} + +// A mid-session reset on a connection without resume negotiated must end the session and surface +// as ErrWebsocketDropped. Both halves matter: telemetry needs to tell a dropped session apart from +// a clean exit, and this is also the fallback a client talking to an older server relies on, so it +// has to keep working unchanged. +func TestMidSessionResetIsAttributedAsADrop(t *testing.T) { + server := createTestServer(t, 2, time.Hour) + defer server.Close() + + relay := newTCPRelay(t, server.Listener.Addr().String()) + + errChan := make(chan error, 1) + client := createTestClient(t, relay.URL(), nil, time.Hour, errChan) + defer client.Cleanup() + + msg := []byte("before drop\n") + _, err := client.InputWriter.Write(msg) + require.NoError(t, err) + require.NoError(t, client.Output.AssertWrite(msg)) + + relay.resetClients(t) + + select { + case err := <-errChan: + assert.ErrorIs(t, err, ErrWebsocketDropped) + case <-time.After(10 * time.Second): + t.Fatal("the client proxy did not report the dropped connection") + } +} + +// The error that ends a session must survive the cancellation it triggers. proxy.start cancels the +// context before errgroup records its error, so the handover and keepalive goroutines wake up and +// return first; if they reported the cancellation as their own error, errgroup would keep that one +// and normalizeProxyError would turn a dropped session into a clean exit. The window is small, so +// this runs the drop repeatedly rather than once. +func TestADropIsNeverReportedAsACleanExit(t *testing.T) { + for attempt := range 25 { + server := createTestServer(t, 2, time.Hour) + relay := newTCPRelay(t, server.Listener.Addr().String()) + + errChan := make(chan error, 1) + client := createTestClient(t, relay.URL(), nil, time.Hour, errChan) + + msg := []byte("before drop\n") + _, err := client.InputWriter.Write(msg) + require.NoError(t, err) + require.NoError(t, client.Output.AssertWrite(msg)) + + relay.resetClients(t) + + select { + case err := <-errChan: + require.ErrorIs(t, err, ErrWebsocketDropped, "attempt %d", attempt) + case <-time.After(10 * time.Second): + t.Fatalf("attempt %d: the drop was reported as a clean exit", attempt) + } + client.Cleanup() + server.Close() + } +} diff --git a/experimental/ssh/internal/proxy/keepalive_test.go b/experimental/ssh/internal/proxy/keepalive_test.go index c972624f6f1..ac1178a1be3 100644 --- a/experimental/ssh/internal/proxy/keepalive_test.go +++ b/experimental/ssh/internal/proxy/keepalive_test.go @@ -125,8 +125,8 @@ func keepaliveTestDialer(serverURL string, onNetConn func(*pausableConn)) create return wrapped, nil }, } - return func(ctx context.Context, connID string) (*websocket.Conn, error) { - conn, _, err := dialer.DialContext(ctx, fmt.Sprintf("%s?id=%s", wsURL, connID), nil) // nolint:bodyclose + return func(ctx context.Context, dial DialRequest) (*websocket.Conn, error) { + conn, _, err := dialer.DialContext(ctx, fmt.Sprintf("%s?id=%s", wsURL, dial.ConnID), nil) // nolint:bodyclose return conn, err } } @@ -143,7 +143,7 @@ func TestKeepalivePingReachesServer(t *testing.T) { src, _ := io.Pipe() done := make(chan error, 1) go func() { - done <- RunClientProxy(ctx, src, io.Discard, neverTick, 20*time.Millisecond, keepaliveTestDialer(server.URL, nil)) + done <- RunClientProxy(ctx, src, io.Discard, neverTick, 20*time.Millisecond, false, keepaliveTestDialer(server.URL, nil)) }() select { @@ -167,7 +167,7 @@ func TestKeepalivePingFailureDoesNotEndSession(t *testing.T) { src, _ := io.Pipe() done := make(chan error, 1) go func() { - done <- RunClientProxy(ctx, src, io.Discard, neverTick, 20*time.Millisecond, + done <- RunClientProxy(ctx, src, io.Discard, neverTick, 20*time.Millisecond, false, keepaliveTestDialer(server.URL, func(c *pausableConn) { socket.Store(c) })) }() @@ -241,7 +241,7 @@ func TestKeepalivePingFailureDoesNotHangTheSession(t *testing.T) { src, srcWriter := io.Pipe() done := make(chan error, 1) go func() { - done <- RunClientProxy(ctx, src, io.Discard, neverTick, 20*time.Millisecond, + done <- RunClientProxy(ctx, src, io.Discard, neverTick, 20*time.Millisecond, false, keepaliveTestDialer(server.URL, func(c *pausableConn) { socket.Store(c) })) }() diff --git a/experimental/ssh/internal/proxy/proxy.go b/experimental/ssh/internal/proxy/proxy.go index 685b02a9b46..665478cd19b 100644 --- a/experimental/ssh/internal/proxy/proxy.go +++ b/experimental/ssh/internal/proxy/proxy.go @@ -2,6 +2,7 @@ package proxy import ( "context" + "encoding/json" "errors" "fmt" "io" @@ -12,17 +13,42 @@ import ( "sync/atomic" "time" + "github.com/databricks/cli/libs/log" "github.com/google/uuid" "github.com/gorilla/websocket" "golang.org/x/sync/errgroup" ) +// Sentinels for how a proxy session ended, so callers can attribute it for telemetry without +// matching on error text. Joined onto the error at the site that detected it. +var ( + // ErrConnectFailed marks a failure to establish the initial proxy websocket. + ErrConnectFailed = errors.New("proxy websocket could not be established") + // ErrWebsocketDropped marks an established proxy websocket that stopped carrying traffic. + ErrWebsocketDropped = errors.New("proxy websocket dropped") + // ErrHandoverFailed marks a handover that ended the session. A handover that only failed + // to dial its replacement does not end the session and is not reported with this. + ErrHandoverFailed = errors.New("proxy handover failed") +) + var ( errProxyEOF = errors.New("proxy EOF error") errSendingLoopStopped = errors.New("sending loop stopped") errReceivingLoopStopped = errors.New("receiving loop stopped") + // Marks a handover that failed while opening its replacement connection, before any + // connection state changed. The current connection is still the one both proxy loops + // use, so the session can carry on with it instead of ending. + errHandoverDialFailed = errors.New("handover dial failed") + // Marks a write that failed on a resumable connection. The payload is already buffered for + // replay, so the sending loop treats it as a pause rather than the end of the session. + errSendFailedResumable = errors.New("send failed on a resumable connection") ) +// proxyResumeGrace is how long the server holds a dropped session - sshd, its client slot and its +// buffered output - waiting for the client to come back. Longer than proxyResumeBudget so the +// server never reaps a client that is still trying. It is a var so tests can shorten it. +var proxyResumeGrace = 90 * time.Second + const ( // Same as gorilla/websocket default read/write buffer sizes. Bigger payloads will be split into multiple ws frames. proxyBufferSize = 4 * 1024 @@ -34,8 +60,59 @@ const ( // connection parks a write until the kernel gives up retransmitting (~15 minutes with Linux // defaults), and close() and the sending loop need that same lock, so the ping caps its wait. proxyPingWriteTimeout = 5 * time.Second + + // How long the client keeps trying to reattach to its session after the connection drops. + // It has to stay clear of ssh's own ceiling: ServerAliveInterval 30 (see + // sshconfig.ServerAliveIntervalSeconds) times OpenSSH's default ServerAliveCountMax of 3 + // means ssh gives up on an unresponsive tunnel after about 90 seconds, and a resume that + // outlasts that repairs a session ssh has already abandoned. + proxyResumeBudget = 60 * time.Second + // Backoff between resume dials. The first attempt is immediate: a reset often clears at once. + proxyResumeRetryBackoff = 2 * time.Second + // Bounds the wait for the peer's first frame on a reattached connection, which carries the + // offset to replay from. The connection is new, but the peer may be wedged. + proxyResumeHandshakeTimeout = 10 * time.Second + // Cap on payload held for replay, per direction. The unacknowledged window is bounded by + // the TCP send buffer plus whatever the driver proxy holds, and acks land every + // proxyAckThreshold bytes, so this is never approached in practice. Reaching it means the + // peer stopped acknowledging, i.e. the connection is already beyond saving. + proxyResumeBufferLimit = 1 << 20 + // How much payload may be delivered before we tell the peer about it, so it can release + // its replay buffer. Small enough to keep the window far below proxyResumeBufferLimit. + proxyAckThreshold = 64 << 10 ) +// resumeState is the per-connection bookkeeping a resumable transport needs. It is nil unless +// both ends negotiated resume, in which case the proxy behaves exactly as it did before. +type resumeState struct { + // Outgoing payload that may still have to be replayed. + sendBuf *sendBuffer + // Total payload bytes written to the destination. The peer replays from this offset, so it + // only advances after a successful write. + delivered atomic.Int64 + // The delivered count we last told the peer about, so an ack is only sent once the number + // has actually moved. + acked atomic.Int64 + // Carries the replacement connection to a parked receiving loop. Only the server uses it: + // it cannot dial, so it waits here for the client's inbound reattach request. Buffered so a + // client that reattaches before this side has noticed the drop is picked up rather than missed. + resumed chan *websocket.Conn + // Signalled by the receiving loop once it has stopped delivering, so a reattach can report a + // delivered count that cannot move under it. Server side only, and buffered for the same + // reason as resumed. + parked chan struct{} + // Throttles the sending loop while a reattach is in progress. + gate sendGate +} + +func newResumeState() *resumeState { + return &resumeState{ + sendBuf: newSendBuffer(proxyResumeBufferLimit), + resumed: make(chan *websocket.Conn, 1), + parked: make(chan struct{}, 1), + } +} + // handoverCoordination holds the context and channels used to coordinate a single handover operation // between the receiving loop and the handover initiator (initiateHandover or acceptHandover). type handoverCoordination struct { @@ -107,9 +184,31 @@ type proxyConnection struct { // Channel that is closed when the initial connection is established (or failed). // Prevents race conditions where handover is accepted before the initial connection is ready. ready chan struct{} + // Byte accounting for reattaching to this session after an unexpected disconnect, or nil + // when resume was not negotiated. Immutable after construction. + resume *resumeState } -type createWebsocketConnectionFunc func(ctx context.Context, connID string) (*websocket.Conn, error) +// DialRequest describes the connection a client is asking the server for. +type DialRequest struct { + // Identifies the session. A server that already has a connection under this ID treats the + // dial as a handover or a reattach rather than a new session. + ConnID string + // Whether this client speaks the resume protocol. When set, the dial carries the delivered + // offset below, and that parameter's presence is how the server learns it must buffer its + // own output for replay too. + ResumeCapable bool + // How many payload bytes this side has written to its destination. Sent on every dial, not + // just a reattach, so the offset is always current when a drop does happen. + Delivered int64 + // Asks the server to reattach this connection to an existing session whose previous + // connection dropped, and to replay what was lost. Stated explicitly rather than inferred + // from the server's own view of the connection, because the client often notices the drop + // first and would otherwise race the server into treating a reattach as a handover. + Reattach bool +} + +type createWebsocketConnectionFunc func(ctx context.Context, req DialRequest) (*websocket.Conn, error) func newProxyConnection(createConn createWebsocketConnectionFunc) *proxyConnection { return &proxyConnection{ @@ -119,6 +218,21 @@ func newProxyConnection(createConn createWebsocketConnectionFunc) *proxyConnecti } } +// newResumableProxyConnection is newProxyConnection with the byte accounting that lets the +// session survive an unexpected disconnect. Both ends must agree: a server that does not speak +// the protocol tears the session down on the first dropped connection regardless, and a client +// must not attempt a resume against one (it would replay into a freshly spawned sshd). +func newResumableProxyConnection(createConn createWebsocketConnectionFunc) *proxyConnection { + pc := newProxyConnection(createConn) + pc.resume = newResumeState() + return pc +} + +// resumable reports whether this connection can reattach to its session after a drop. +func (pc *proxyConnection) resumable() bool { + return pc.resume != nil +} + func (pc *proxyConnection) start(ctx context.Context, src io.ReadCloser, dst io.Writer) error { g, gCtx := errgroup.WithContext(ctx) g.Go(func() error { @@ -150,7 +264,9 @@ func (pc *proxyConnection) start(ctx context.Context, src io.ReadCloser, dst io. func (pc *proxyConnection) connect(ctx context.Context) error { defer close(pc.ready) - conn, err := pc.createWebsocketConnection(ctx, pc.connID) + // Nothing has been delivered yet, so the initial dial reports offset zero. Sending it at all + // is what tells a resume-capable server that this client speaks the protocol. + conn, err := pc.createWebsocketConnection(ctx, DialRequest{ConnID: pc.connID, ResumeCapable: pc.resumable()}) if err != nil { return err } @@ -185,11 +301,27 @@ func (pc *proxyConnection) runSendingLoop(ctx context.Context, src io.Reader) er b := make([]byte, proxyBufferSize) n, readErr := src.Read(b) if n > 0 { + // Wait out any reattach in progress, so a connection that is down does not fill the + // whole replay window before it comes back. src stays blocked on the OS side + // meanwhile, which is the backpressure we want. + if pc.resumable() { + if err := pc.resume.gate.wait(ctx); err != nil { + return err + } + } // This will block during handover - we stop sending anything except the close message. // Meanwhile the "src" (sshd server stdout or ssh client stdin) will be buffered/blocked on the OS side until we start reading from it again. err := pc.sendMessage(websocket.BinaryMessage, b[:n]) - if err != nil { - return fmt.Errorf("failed to send message: %w", err) + switch { + case errors.Is(err, errSendFailedResumable): + // Buffered for replay, and sendMessage has closed the connection so the + // receiving loop starts the reattach. Carry on reading src: its bytes accumulate + // in the replay buffer, and the gate above holds the next write until the + // connection is back. Falls through to readErr rather than continuing the loop, + // so a read that returned data together with an error still reports it. + log.Debugf(ctx, "Send failed on a resumable connection, waiting for the reattach: %v", err) + case err != nil: + return errors.Join(ErrWebsocketDropped, fmt.Errorf("failed to send message: %w", err)) } } if readErr != nil { @@ -205,8 +337,50 @@ func (pc *proxyConnection) runSendingLoop(ctx context.Context, src io.Reader) er func (pc *proxyConnection) sendMessage(mt int, data []byte) error { pc.handoverMutex.Lock() defer pc.handoverMutex.Unlock() + // Record the payload before writing it, and under the same lock a resume swaps the + // connection with: that way the buffer always holds every byte the peer may still be + // missing, and a resume can never replay a range the sending loop is still appending to. + if pc.resumable() && mt == websocket.BinaryMessage { + if err := pc.resume.sendBuf.append(data); err != nil { + return err + } + } conn := pc.conn.Load() - return conn.WriteMessage(mt, data) + err := conn.WriteMessage(mt, data) + if err != nil && pc.resumable() && mt == websocket.BinaryMessage { + // The payload is buffered, so this failure costs no data. gorilla latches a permanent + // write error after any failed write, so this connection can never send again: close it + // to fail the receiving loop's read now and start the resume, rather than let the + // sending loop fill the whole window first. + conn.Close() + return errors.Join(errSendFailedResumable, err) + } + return err +} + +// sendControlMessage tells the peer how much payload we have written to our destination, so it +// can release that much of its replay buffer. +func (pc *proxyConnection) sendControlMessage(delivered int64) error { + payload, err := json.Marshal(controlMessage{Delivered: delivered}) + if err != nil { + return err + } + return pc.sendMessage(websocket.TextMessage, payload) +} + +// ackDelivered reports our delivered count to the peer once it has moved far enough to be worth +// a frame. A failure is not fatal: the ack is only an optimisation that keeps the peer's replay +// buffer small, and a genuinely broken connection is reported by the loops themselves. +func (pc *proxyConnection) ackDelivered(ctx context.Context) { + delivered := pc.resume.delivered.Load() + if delivered-pc.resume.acked.Load() < proxyAckThreshold { + return + } + if err := pc.sendControlMessage(delivered); err != nil { + log.Debugf(ctx, "Failed to acknowledge %d delivered bytes: %v", delivered, err) + return + } + pc.resume.acked.Store(delivered) } // sendPing writes a keepalive ping on the current connection. Unlike sendMessage it takes neither @@ -230,7 +404,7 @@ func (pc *proxyConnection) runReceivingLoop(ctx context.Context, dst io.Writer) if handover := pc.handoverState.Load(); handover != nil { var closeConnSignal error if !websocket.IsCloseError(err, websocket.CloseNormalClosure) { - closeConnSignal = fmt.Errorf("failed to read from websocket during handover: %w", err) + 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). if err := handover.signalConnectionClosed(closeConnSignal); err != nil { @@ -247,18 +421,39 @@ func (pc *proxyConnection) runReceivingLoop(ctx context.Context, dst io.Writer) } else { if errors.Is(err, io.EOF) || websocket.IsCloseError(err, websocket.CloseNormalClosure) { return errors.Join(errProxyEOF, err) - } else { - return fmt.Errorf("failed to read from websocket: %w", err) } + // An unexpected drop. With resume negotiated the session state on both ends + // outlives the connection, so reattach instead of ending the session. + if pc.resumable() { + if resumeErr := pc.reattach(ctx); resumeErr != nil { + return errors.Join(ErrWebsocketDropped, fmt.Errorf("failed to reattach after the connection dropped: %w", resumeErr)) + } + continue + } + return errors.Join(ErrWebsocketDropped, fmt.Errorf("failed to read from websocket: %w", err)) } } + if mt == websocket.TextMessage && pc.resumable() { + var msg controlMessage + if err := json.Unmarshal(data, &msg); err != nil { + return fmt.Errorf("failed to decode control message: %w", err) + } + pc.resume.sendBuf.ack(msg.Delivered) + continue + } if mt != websocket.BinaryMessage { return errors.New("received non-binary websocket message") } if _, err := dst.Write(data); err != nil { return fmt.Errorf("failed to copy to writer: %w", err) } + if pc.resumable() { + // Only count what actually reached the destination: this is the offset the peer + // replays from, so counting an unwritten byte would silently lose it. + pc.resume.delivered.Add(int64(len(data))) + pc.ackDelivered(ctx) + } } } @@ -315,9 +510,21 @@ func (pc *proxyConnection) initiateHandover(ctx context.Context) error { // Create a new websocket connection by sending an /ssh?id= request to the server. // When server realises it's an ID of an existing connection, it will start AcceptHandover process. - newConn, err := pc.createWebsocketConnection(handoverCtx, pc.connID) + newConn, err := pc.createWebsocketConnection(handoverCtx, DialRequest{ + ConnID: pc.connID, + ResumeCapable: pc.resumable(), + // A handover replaces a connection that still works, so the close-frame barrier keeps + // the byte stream intact and nothing needs replaying. The offset still travels, so the + // server keeps buffering for the drop that may come later. + Delivered: pc.deliveredCount(), + }) if err != nil { - return fmt.Errorf("failed to create new websocket connection: %w", err) + // Nothing has been swapped yet: pc.conn is still live and the receiving loop is still + // reading it. Tag the error so the caller can keep the session on it - see + // errHandoverDialFailed. Retrying the dial here instead would be unsafe: a dial can + // fail after the server already accepted it and began its side of the handover, and a + // second dial would then race the first one's acceptHandover for the same connection. + return errors.Join(errHandoverDialFailed, fmt.Errorf("failed to create new websocket connection: %w", err)) } // Wait for the server to close the old connection diff --git a/experimental/ssh/internal/proxy/proxy_test.go b/experimental/ssh/internal/proxy/proxy_test.go index 1e3af34f40d..5abb7a7186b 100644 --- a/experimental/ssh/internal/proxy/proxy_test.go +++ b/experimental/ssh/internal/proxy/proxy_test.go @@ -155,7 +155,7 @@ func setupTestClientWithDialHook(ctx context.Context, t *testing.T, serverURL st clientInput, clientInputWriter := io.Pipe() clientOutput := newTestBuffer(t) wsURL := "ws" + serverURL[4:] - clientProxy := newProxyConnection(func(ctx context.Context, connID string) (*websocket.Conn, error) { + clientProxy := newProxyConnection(func(ctx context.Context, dial DialRequest) (*websocket.Conn, error) { if onDial != nil { onDial() } diff --git a/experimental/ssh/internal/proxy/resume.go b/experimental/ssh/internal/proxy/resume.go new file mode 100644 index 00000000000..3634cb92c2b --- /dev/null +++ b/experimental/ssh/internal/proxy/resume.go @@ -0,0 +1,371 @@ +package proxy + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "sync" + "time" + + "github.com/databricks/cli/libs/log" + "github.com/gorilla/websocket" +) + +// errSendWindowExhausted means the peer has not acknowledged anything for a full buffer's worth +// of payload. The connection cannot be resumed past this point, because the bytes the peer would +// need replayed are the ones we would have to discard to make room. +var errSendWindowExhausted = errors.New("resume send buffer is full: the peer stopped acknowledging") + +// errReplayUnavailable means the peer asked to resume from an offset we no longer hold. It can +// only happen if the peer acknowledged those bytes and then asked for them again, so it is a +// protocol violation rather than a condition to recover from. +var errReplayUnavailable = errors.New("the bytes needed to resume have already been acknowledged and discarded") + +// sendBuffer holds the tail of the outgoing payload stream so it can be replayed after an +// unexpected disconnect. +// +// The buffer, not the websocket, is the source of truth for what has been sent: bytes are +// appended before they are written, so a write that fails on a dying connection loses nothing. +// SSH runs its own sequence numbers and MACs over the byte stream (RFC 4253 section 6), so a +// resumed connection has to deliver exactly the bytes the peer missed, once, in order - a +// single lost or duplicated byte disconnects the session with a corrupted MAC. +type sendBuffer struct { + mu sync.Mutex + // Total payload bytes appended since the session began. + sent int64 + // Total payload bytes the peer has confirmed writing to its destination. buf holds + // exactly the range [acked, sent). + acked int64 + buf []byte + limit int +} + +func newSendBuffer(limit int) *sendBuffer { + return &sendBuffer{limit: limit} +} + +// append records payload as sent. It fails once the unacknowledged window would outgrow the +// limit, which means the peer has stopped acknowledging and the stream can no longer be repaired. +func (b *sendBuffer) append(payload []byte) error { + b.mu.Lock() + defer b.mu.Unlock() + if len(b.buf)+len(payload) > b.limit { + return fmt.Errorf("%w: %d unacknowledged bytes, limit %d", errSendWindowExhausted, len(b.buf), b.limit) + } + b.buf = append(b.buf, payload...) + b.sent += int64(len(payload)) + return nil +} + +// ack discards everything the peer has confirmed delivering. A stale or duplicated +// acknowledgement is ignored rather than treated as an error: acks are sent periodically and +// may arrive out of order relative to a resume. +func (b *sendBuffer) ack(delivered int64) { + b.mu.Lock() + defer b.mu.Unlock() + if delivered <= b.acked || delivered > b.sent { + return + } + b.buf = b.buf[delivered-b.acked:] + b.acked = delivered +} + +// replayFrom returns the bytes the peer is missing: everything from the offset it last delivered +// up to what we have sent. The returned slice is a copy, so the caller can write it without +// holding the lock. +func (b *sendBuffer) replayFrom(delivered int64) ([]byte, error) { + b.mu.Lock() + defer b.mu.Unlock() + if delivered < b.acked { + return nil, fmt.Errorf("%w: peer asked for offset %d, buffer starts at %d", errReplayUnavailable, delivered, b.acked) + } + if delivered > b.sent { + return nil, fmt.Errorf("peer claims to have delivered %d bytes but only %d were sent", delivered, b.sent) + } + missing := b.buf[delivered-b.acked:] + return append([]byte(nil), missing...), nil +} + +// controlMessage is exchanged as a websocket text frame; payload always stays binary. It carries +// the sender's delivered count, both as the periodic acknowledgement and as the first frame of a +// resumed connection, where it tells the peer where to replay from. +// +// Text frames are only ever sent once resume has been negotiated: a server from an older CLI +// treats any non-binary frame as a protocol error and ends the session. +type controlMessage struct { + Delivered int64 `json:"delivered"` +} + +// deliveredCount is how many payload bytes this side has written to its destination, or zero when +// resume was not negotiated. +func (pc *proxyConnection) deliveredCount() int64 { + if !pc.resumable() { + return 0 + } + return pc.resume.delivered.Load() +} + +// sendGate throttles the sending loop while a reattach is in progress. +// +// It is deliberately not the write mutex. A reattach has to wait for the peer - the client for its +// dial to be answered, the server for the client to come back - and the write mutex is exactly +// what the other side needs to finish the reattach, so holding it across that wait deadlocks the +// server. Correctness rests on the write mutex plus appending before writing; this only stops the +// sending loop from filling the whole replay window while the connection is down. +type sendGate struct { + mu sync.Mutex + waitCh chan struct{} +} + +func (g *sendGate) park() { + g.mu.Lock() + defer g.mu.Unlock() + if g.waitCh == nil { + g.waitCh = make(chan struct{}) + } +} + +func (g *sendGate) release() { + g.mu.Lock() + defer g.mu.Unlock() + if g.waitCh != nil { + close(g.waitCh) + g.waitCh = nil + } +} + +// wait blocks until the gate is released, and returns immediately when it is already open. +func (g *sendGate) wait(ctx context.Context) error { + g.mu.Lock() + waitCh := g.waitCh + g.mu.Unlock() + if waitCh == nil { + return nil + } + select { + case <-waitCh: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +// reattach repairs a dropped connection without the SSH session noticing: both ends replay +// whatever the other is missing, so the byte stream continues exactly where it left off. +// +// The receiving loop owns this, because it is the goroutine that must stop reading the dead +// connection. The sending loop is only throttled: its payload is already in the replay buffer. +func (pc *proxyConnection) reattach(ctx context.Context) error { + pc.resume.gate.park() + defer pc.resume.gate.release() + + // The server cannot dial - its client is behind the driver proxy - so it waits for an inbound + // reattach request instead. + if pc.createWebsocketConnection == nil { + return pc.awaitReattach(ctx) + } + return pc.dialReattach(ctx) +} + +// dialReattach reattaches from the client side: redial the session, learn where the server got +// to, and replay what it is missing. +func (pc *proxyConnection) dialReattach(ctx context.Context) error { + // Held for the whole reattach, dial included, so a handover tick cannot start one while this is + // in flight. A handover needs the receiving loop to answer it, and the receiving loop is the + // goroutine running this - the handover would wait out its own timeout and end a session that + // was about to be repaired. Safe to hold across the dial on this side: only the client + // initiates handovers, and nothing the reattach waits on needs this lock. The server's side + // cannot do the same, which is what sendGate is for. + pc.handoverMutex.Lock() + defer pc.handoverMutex.Unlock() + + budgetCtx, cancel := context.WithTimeout(ctx, proxyResumeBudget) + defer cancel() + + log.Warnf(ctx, "SSH tunnel connection dropped, reattaching to the session...") + conn, err := pc.redial(budgetCtx) + if err != nil { + return err + } + + // The server's first frame says how much of our output it wrote, which is where we replay + // from. It replays what we are missing right after, and those frames wait in the socket + // until the receiving loop picks the new connection up. + serverDelivered, err := readResumeHandshake(conn) + if err != nil { + conn.Close() + return err + } + if err := pc.replayTo(conn, serverDelivered); err != nil { + conn.Close() + return err + } + pc.conn.Store(conn) + log.Warnf(ctx, "SSH tunnel connection reattached, the session continues") + return nil +} + +func (pc *proxyConnection) redial(ctx context.Context) (*websocket.Conn, error) { + var lastErr error + for { + conn, err := pc.createWebsocketConnection(ctx, DialRequest{ + ConnID: pc.connID, + Delivered: pc.resume.delivered.Load(), + Reattach: true, + ResumeCapable: true, + }) + if err == nil { + return conn, nil + } + lastErr = err + log.Debugf(ctx, "Reattach dial failed, retrying: %v", err) + select { + case <-ctx.Done(): + return nil, fmt.Errorf("gave up reattaching after %v: %w", proxyResumeBudget, lastErr) + case <-time.After(proxyResumeRetryBackoff): + } + } +} + +// awaitReattach reattaches from the server side by waiting for the client to come back. The +// session - sshd, the client slot, and the buffered output - is held for the grace period. +// +// It first announces that it has stopped delivering, which is what makes its delivered count +// stable for acceptReattach to report. Both channels are buffered, so a client that reattaches +// before this side has even noticed the drop is picked up rather than missed. +func (pc *proxyConnection) awaitReattach(ctx context.Context) error { + log.Infof(ctx, "Connection dropped, holding the session for up to %v for the client to reattach", proxyResumeGrace) + select { + case pc.resume.parked <- struct{}{}: + default: + // A previous park is still queued, which means acceptReattach has not consumed it yet. + // Nothing to add: it is about to read a delivered count that is already stable. + } + select { + case <-pc.resume.resumed: + // acceptReattach has already replayed and installed the new connection. + log.Info(ctx, "Client reattached to the session") + return nil + case <-time.After(proxyResumeGrace): + return fmt.Errorf("the client did not reattach within %v", proxyResumeGrace) + case <-ctx.Done(): + return ctx.Err() + } +} + +// awaitParked waits until the receiving loop has stopped delivering to sshd, so this side's +// delivered count cannot move while a reattach reports it. +// +// Without this the client can reattach while the server is still draining data buffered on the +// dying connection: the greeting would carry a stale offset, the client would replay from it, and +// the server would write those bytes to sshd twice. SSH would then fail on a corrupted MAC - the +// exact failure the replay accounting exists to prevent. +func (pc *proxyConnection) awaitParked(ctx context.Context) error { + select { + case <-pc.resume.parked: + return nil + case <-time.After(proxyResumeHandshakeTimeout): + return fmt.Errorf("the receiving loop did not stop delivering within %v", proxyResumeHandshakeTimeout) + case <-ctx.Done(): + return ctx.Err() + } +} + +// acceptReattach installs a client's replacement connection on the server side: announce where we +// got to, replay what the client is missing, and hand the connection to the parked receiving loop. +// +// Called from the HTTP handler goroutine, so it takes the write lock the loops use. +func (pc *proxyConnection) acceptReattach(ctx context.Context, w http.ResponseWriter, r *http.Request, clientDelivered int64) error { + select { + case <-pc.ready: + case <-ctx.Done(): + return ctx.Err() + } + + // Retire the dying connection before anything else. The receiving loop may still be draining + // data buffered on it - a reset that only tore down the client's leg leaves this side's read + // succeeding for a while, then hanging - and every byte it delivers moves the offset the + // greeting below is about to report. + previous := pc.conn.Load() + if previous != nil { + previous.Close() + } + if err := pc.awaitParked(ctx); err != nil { + return err + } + + pc.handoverMutex.Lock() + defer pc.handoverMutex.Unlock() + + conn, err := pc.acceptWebsocketConnection(w, r) + if err != nil { + return fmt.Errorf("failed to accept the reattached connection: %w", err) + } + // The offset first, then the payload: the client reads exactly one control frame before it + // hands the connection to its own receiving loop. + if err := pc.sendResumeHandshake(conn); err != nil { + conn.Close() + return err + } + if err := pc.replayTo(conn, clientDelivered); err != nil { + conn.Close() + return err + } + pc.conn.Store(conn) + + select { + case pc.resume.resumed <- conn: + default: + // The loop already gave up, or a previous signal is still queued. Either way the session + // is past saving; the loop's own grace timer reports it. + log.Warnf(ctx, "Reattach signal could not be delivered to the receiving loop") + } + return nil +} + +// replayTo sends the peer everything it is missing, given how much it says it delivered. +func (pc *proxyConnection) replayTo(conn *websocket.Conn, peerDelivered int64) error { + missing, err := pc.resume.sendBuf.replayFrom(peerDelivered) + if err != nil { + return err + } + if len(missing) == 0 { + return nil + } + return conn.WriteMessage(websocket.BinaryMessage, missing) +} + +// sendResumeHandshake announces our delivered count as the first frame of a reattached +// connection, so the peer knows where to replay from. Written directly rather than through +// sendMessage: the connection is not installed yet, and this must not be recorded as payload. +func (pc *proxyConnection) sendResumeHandshake(conn *websocket.Conn) error { + payload, err := json.Marshal(controlMessage{Delivered: pc.resume.delivered.Load()}) + if err != nil { + return err + } + return conn.WriteMessage(websocket.TextMessage, payload) +} + +// readResumeHandshake reads the delivered count the peer sends as the first frame of a reattached +// connection. The deadline bounds the wait: the connection is new, but the peer may be wedged. +func readResumeHandshake(conn *websocket.Conn) (int64, error) { + if err := conn.SetReadDeadline(time.Now().Add(proxyResumeHandshakeTimeout)); err != nil { + return 0, err + } + defer func() { _ = conn.SetReadDeadline(time.Time{}) }() + + mt, data, err := conn.ReadMessage() + if err != nil { + return 0, fmt.Errorf("failed to read the reattach handshake: %w", err) + } + if mt != websocket.TextMessage { + return 0, fmt.Errorf("expected a reattach handshake control frame, got websocket message type %d", mt) + } + var msg controlMessage + if err := json.Unmarshal(data, &msg); err != nil { + return 0, fmt.Errorf("failed to decode the reattach handshake: %w", err) + } + return msg.Delivered, nil +} diff --git a/experimental/ssh/internal/proxy/resume_e2e_test.go b/experimental/ssh/internal/proxy/resume_e2e_test.go new file mode 100644 index 00000000000..3c4f4dbc292 --- /dev/null +++ b/experimental/ssh/internal/proxy/resume_e2e_test.go @@ -0,0 +1,203 @@ +//go:build !windows + +package proxy + +import ( + "context" + "encoding/json" + "fmt" + "net" + "net/http" + "net/http/httptest" + "os/exec" + "testing" + "time" + + "github.com/databricks/cli/libs/cmdio" + "github.com/gorilla/websocket" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// The invariant the whole resume protocol exists to hold: a reset must not cost or duplicate a +// single byte. SSH verifies a MAC over the byte stream (RFC 4253 section 6), so a resume that got +// this wrong would disconnect the session rather than repair it - a worse failure than the drop. +// +// The reset lands immediately after a burst of writes, while the echo of that burst is still in +// flight, so the server has bytes it must genuinely replay. +func TestResumeLosesNoBytesWhenResetMidStream(t *testing.T) { + server := createTestServer(t, 2, time.Hour) + defer server.Close() + + relay := newTCPRelay(t, server.Listener.Addr().String()) + + errChan := make(chan error, 1) + client := createResumableTestClient(t, relay.URL(), errChan) + defer client.Cleanup() + + const lines = 500 + var expected []byte + for i := range lines { + line := fmt.Appendf(nil, "line %d\n", i) + _, err := client.InputWriter.Write(line) + require.NoError(t, err) + expected = append(expected, line...) + } + + relay.resetClients(t) + + require.NoError(t, client.Output.WaitForWrite(fmt.Appendf(nil, "line %d\n", lines-1)), + "the session did not survive the reset") + assert.Equal(t, string(expected), client.Output.String()) + + select { + case err := <-errChan: + t.Fatalf("session ended despite being resumable: %v", err) + default: + } +} + +// Customers report drops several times a session, so surviving one reset is not enough. +func TestResumeSurvivesRepeatedResets(t *testing.T) { + server := createTestServer(t, 2, time.Hour) + defer server.Close() + + relay := newTCPRelay(t, server.Listener.Addr().String()) + + errChan := make(chan error, 1) + client := createResumableTestClient(t, relay.URL(), errChan) + defer client.Cleanup() + + const rounds = 4 + const linesPerRound = 100 + var expected []byte + for round := range rounds { + for i := range linesPerRound { + line := fmt.Appendf(nil, "round %d line %d\n", round, i) + _, err := client.InputWriter.Write(line) + require.NoError(t, err) + expected = append(expected, line...) + } + relay.resetClients(t) + last := fmt.Appendf(nil, "round %d line %d\n", round, linesPerRound-1) + require.NoError(t, client.Output.WaitForWrite(last), "round %d did not survive its reset", round) + } + + assert.Equal(t, string(expected), client.Output.String()) + + select { + case err := <-errChan: + t.Fatalf("session ended despite being resumable: %v", err) + default: + } +} + +// A client that never comes back must not pin sshd and a client slot forever. Releasing the slot +// is also what restarts the shutdown timer, so without this a single dropped session would keep +// the whole server alive until its own timeout. +func TestServerReleasesASessionThatIsNeverReattached(t *testing.T) { + originalGrace := proxyResumeGrace + proxyResumeGrace = 300 * time.Millisecond + defer func() { proxyResumeGrace = originalGrace }() + + ctx := cmdio.MockDiscard(t.Context()) + connections := NewConnectionsManager(2, time.Hour) + proxyServer := NewProxyServer(ctx, connections, func(ctx context.Context) *exec.Cmd { + return exec.CommandContext(ctx, "cat", "-u") + }) + server := httptest.NewServer(proxyServer) + defer server.Close() + + wsURL := "ws" + server.URL[4:] + conn, _, err := websocket.DefaultDialer.Dial(wsURL+"?id=abandoned&delivered=0", nil) // nolint:bodyclose + require.NoError(t, err) + + require.NoError(t, conn.WriteMessage(websocket.BinaryMessage, []byte("hello\n"))) + _, _, err = conn.ReadMessage() + require.NoError(t, err) + require.Equal(t, 1, connections.Count()) + + tcpConn, ok := conn.UnderlyingConn().(*net.TCPConn) + require.True(t, ok) + require.NoError(t, tcpConn.SetLinger(0)) + require.NoError(t, tcpConn.Close()) + + require.Eventually(t, func() bool { + return connections.Count() == 0 + }, 10*time.Second, 20*time.Millisecond, "the server held the session past its resume grace period") +} + +// A reattach for a session the server no longer holds must be refused. Starting a fresh session +// instead would hand the client a new sshd, and replaying into that fails the SSH stream with a +// corrupted MAC instead of a clear error. +func TestReattachToAnUnknownSessionIsRefused(t *testing.T) { + server := createTestServer(t, 2, time.Hour) + defer server.Close() + + wsURL := "ws" + server.URL[4:] + _, resp, err := websocket.DefaultDialer.Dial(wsURL+"?id=never-existed&delivered=0&reattach=1", nil) // nolint:bodyclose + require.Error(t, err) + require.NotNil(t, resp) + defer resp.Body.Close() + assert.Equal(t, http.StatusGone, resp.StatusCode) +} + +// Drives the wire protocol by hand to pin the reattach exchange itself, with a replay that is +// deliberately non-empty: the client under-reports its delivered offset as zero, so the server has +// to replay the whole session. Asserts the ordering the client depends on - the delivered offset +// arrives as a text frame first, and the replayed payload follows it. +func TestReattachReplaysFromTheOffsetTheClientReports(t *testing.T) { + server := createTestServer(t, 2, time.Hour) + defer server.Close() + + wsURL := "ws" + server.URL[4:] + const sessionID = "replay-session" + + conn, _, err := websocket.DefaultDialer.Dial(wsURL+"?id="+sessionID+"&delivered=0", nil) // nolint:bodyclose + require.NoError(t, err) + + const payload = "hello\n" + require.NoError(t, conn.WriteMessage(websocket.BinaryMessage, []byte(payload))) + mt, echo, err := conn.ReadMessage() + require.NoError(t, err) + require.Equal(t, websocket.BinaryMessage, mt) + require.Equal(t, payload, string(echo)) + + // Kill the connection without a close handshake, so the server treats it as a drop and parks + // the session instead of tearing it down. + tcpConn, ok := conn.UnderlyingConn().(*net.TCPConn) + require.True(t, ok) + require.NoError(t, tcpConn.SetLinger(0)) + require.NoError(t, tcpConn.Close()) + + // Reattach claiming to have delivered nothing, so the replay covers the whole echo. + resumed, _, err := websocket.DefaultDialer.Dial(wsURL+"?id="+sessionID+"&delivered=0&reattach=1", nil) // nolint:bodyclose + require.NoError(t, err) + defer resumed.Close() + + mt, greeting, err := resumed.ReadMessage() + require.NoError(t, err) + require.Equal(t, websocket.TextMessage, mt, "the reattach handshake must arrive before any payload") + var msg controlMessage + require.NoError(t, json.Unmarshal(greeting, &msg)) + assert.Equal(t, int64(len(payload)), msg.Delivered, "the server wrote our payload to sshd, so that is its delivered count") + + mt, replayed, err := resumed.ReadMessage() + require.NoError(t, err) + require.Equal(t, websocket.BinaryMessage, mt) + assert.Equal(t, payload, string(replayed), "the server must replay the echo the client claimed not to have") +} + +// A reattach without the offset it replays from is meaningless, and silently treating it as a new +// session is the failure this protocol is designed to avoid. +func TestReattachWithoutADeliveredOffsetIsRejected(t *testing.T) { + server := createTestServer(t, 2, time.Hour) + defer server.Close() + + wsURL := "ws" + server.URL[4:] + _, resp, err := websocket.DefaultDialer.Dial(wsURL+"?id=some-session&reattach=1", nil) // nolint:bodyclose + require.Error(t, err) + require.NotNil(t, resp) + defer resp.Body.Close() + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) +} diff --git a/experimental/ssh/internal/proxy/resume_test.go b/experimental/ssh/internal/proxy/resume_test.go new file mode 100644 index 00000000000..052ab4cc11d --- /dev/null +++ b/experimental/ssh/internal/proxy/resume_test.go @@ -0,0 +1,89 @@ +package proxy + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestSendBufferReplaysWhatThePeerMissed(t *testing.T) { + b := newSendBuffer(64) + require.NoError(t, b.append([]byte("hello "))) + require.NoError(t, b.append([]byte("world"))) + + // The peer wrote only the first 6 bytes before the connection broke. + missing, err := b.replayFrom(6) + require.NoError(t, err) + assert.Equal(t, "world", string(missing)) +} + +func TestSendBufferReplayFromCurrentOffsetIsEmpty(t *testing.T) { + b := newSendBuffer(64) + require.NoError(t, b.append([]byte("hello"))) + + missing, err := b.replayFrom(5) + require.NoError(t, err) + assert.Empty(t, missing) +} + +func TestSendBufferAckDiscardsThePrefix(t *testing.T) { + b := newSendBuffer(64) + require.NoError(t, b.append([]byte("aaaabbbb"))) + b.ack(4) + + // Everything from the acknowledged offset is still replayable. + missing, err := b.replayFrom(4) + require.NoError(t, err) + assert.Equal(t, "bbbb", string(missing)) + + // The discarded prefix is not. + _, err = b.replayFrom(3) + assert.ErrorIs(t, err, errReplayUnavailable) +} + +func TestSendBufferAckFreesTheWindow(t *testing.T) { + b := newSendBuffer(8) + require.NoError(t, b.append([]byte("12345678"))) + require.ErrorIs(t, b.append([]byte("9")), errSendWindowExhausted) + + b.ack(8) + assert.NoError(t, b.append([]byte("9"))) +} + +func TestSendBufferIgnoresStaleAndImpossibleAcks(t *testing.T) { + b := newSendBuffer(64) + require.NoError(t, b.append([]byte("abcdef"))) + b.ack(4) + + // A stale ack must not rewind the buffer, and an ack beyond what we sent must not + // discard bytes the peer cannot have received. + b.ack(2) + b.ack(99) + + missing, err := b.replayFrom(4) + require.NoError(t, err) + assert.Equal(t, "ef", string(missing)) +} + +func TestSendBufferRejectsAnImpossibleReplayOffset(t *testing.T) { + b := newSendBuffer(64) + require.NoError(t, b.append([]byte("abc"))) + + _, err := b.replayFrom(4) + require.Error(t, err) + assert.NotErrorIs(t, err, errReplayUnavailable) +} + +func TestSendBufferReplayIsACopy(t *testing.T) { + b := newSendBuffer(64) + require.NoError(t, b.append([]byte("abcdef"))) + + missing, err := b.replayFrom(0) + require.NoError(t, err) + missing[0] = 'z' + + again, err := b.replayFrom(0) + require.NoError(t, err) + assert.Equal(t, "abcdef", string(again), "mutating a replay must not corrupt the buffer") +} diff --git a/experimental/ssh/internal/proxy/server.go b/experimental/ssh/internal/proxy/server.go index c709e85a72c..550aaede9eb 100644 --- a/experimental/ssh/internal/proxy/server.go +++ b/experimental/ssh/internal/proxy/server.go @@ -2,10 +2,12 @@ package proxy import ( "context" + "errors" "fmt" "net/http" "os" "os/exec" + "strconv" "time" "github.com/databricks/cli/libs/log" @@ -36,12 +38,63 @@ func (server *proxyServer) ServeHTTP(w http.ResponseWriter, r *http.Request) { http.Error(w, "Missing 'id' query parameter", http.StatusBadRequest) return } + req, err := parseDialRequest(r) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } ctx := log.NewContext(server.ctx, log.GetLogger(server.ctx).With("session", id)) - if conn, exists := server.connections.Get(id); exists && conn != nil { + conn, exists := server.connections.Get(id) + switch { + case exists && conn != nil && req.Reattach: + server.handleReattach(ctx, w, r, conn, req) + case exists && conn != nil: server.handleExistingConnection(ctx, w, r, conn) - } else { - server.handleNewConnection(ctx, w, r, id) + case req.Reattach: + // The session is gone: its grace period expired, or the server restarted. Say so instead + // of starting a fresh one. A new sshd would answer the client's replayed bytes with a + // fresh SSH handshake, and ssh would fail on a corrupted MAC rather than a clear error. + log.Info(ctx, "Reattach requested for a session that no longer exists") + http.Error(w, "Session no longer exists", http.StatusGone) + default: + server.handleNewConnection(ctx, w, r, id, req) + } +} + +// parseDialRequest reads the resume protocol's query parameters. Both are absent for a client +// that does not speak it, which leaves the session non-resumable on this side too. +func parseDialRequest(r *http.Request) (DialRequest, error) { + query := r.URL.Query() + req := DialRequest{ + ConnID: query.Get("id"), + Reattach: query.Get("reattach") == "1", + } + if raw := query.Get("delivered"); raw != "" { + delivered, err := strconv.ParseInt(raw, 10, 64) + if err != nil || delivered < 0 { + return DialRequest{}, fmt.Errorf("invalid 'delivered' query parameter: %q", raw) + } + req.Delivered = delivered + req.ResumeCapable = true + } + if req.Reattach && !req.ResumeCapable { + return DialRequest{}, errors.New("'reattach' requires the 'delivered' query parameter") } + return req, nil +} + +func (server *proxyServer) handleReattach(ctx context.Context, w http.ResponseWriter, r *http.Request, conn *proxyConnection, req DialRequest) { + if !conn.resumable() { + log.Info(ctx, "Reattach requested for a session that was not started as resumable") + http.Error(w, "Session is not resumable", http.StatusConflict) + return + } + log.Info(ctx, "Client reattaching to a dropped connection") + if err := conn.acceptReattach(ctx, w, r, req.Delivered); err != nil { + log.Errorf(ctx, "Failed to accept the reattach: %v", err) + return + } + log.Info(ctx, "Reattach accepted") } func (server *proxyServer) handleExistingConnection(ctx context.Context, w http.ResponseWriter, r *http.Request, conn *proxyConnection) { @@ -55,8 +108,15 @@ func (server *proxyServer) handleExistingConnection(ctx context.Context, w http. } } -func (server *proxyServer) handleNewConnection(ctx context.Context, w http.ResponseWriter, r *http.Request, id string) { - conn := newProxyConnection(nil) +func (server *proxyServer) handleNewConnection(ctx context.Context, w http.ResponseWriter, r *http.Request, id string, req DialRequest) { + // The server never dials, so it passes no connection factory: reattaching, for it, means + // waiting for the client to come back. + var conn *proxyConnection + if req.ResumeCapable { + conn = newResumableProxyConnection(nil) + } else { + conn = newProxyConnection(nil) + } if !server.connections.TryAdd(id, conn) { log.Info(ctx, "Maximum clients reached, rejecting connection") http.Error(w, "Maximum clients reached", http.StatusServiceUnavailable) diff --git a/experimental/ssh/internal/server/server.go b/experimental/ssh/internal/server/server.go index 4356700d87d..72c74b258b3 100644 --- a/experimental/ssh/internal/server/server.go +++ b/experimental/ssh/internal/server/server.go @@ -4,6 +4,7 @@ import ( "bytes" "context" _ "embed" + "encoding/json" "errors" "fmt" "io" @@ -107,10 +108,12 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ServerOpt http.Handle("/ssh", proxy.NewProxyServer(ctx, connections, createServerCommand)) http.HandleFunc("/metadata", serveMetadata) http.HandleFunc("/logs", logBuf.serveHTTP) + http.HandleFunc("/capabilities", serveCapabilities) http.Handle("/driver-proxy-http/ssh", proxy.NewProxyServer(ctx, connections, createServerCommand)) http.HandleFunc("/driver-proxy-http/metadata", serveMetadata) http.HandleFunc("/driver-proxy-http/logs", logBuf.serveHTTP) + http.HandleFunc("/driver-proxy-http/capabilities", serveCapabilities) go handleTimeout(ctx, connections.TimedOut, opts.ShutdownDelay) @@ -121,6 +124,16 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ServerOpt })) } +// serveCapabilities tells the client which optional parts of the tunnel protocol this server +// speaks. A server from an older CLI has no such route and returns 404, which the client reads as +// "none of them" - the negotiation this endpoint exists for. +func serveCapabilities(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(w).Encode(map[string]bool{"resume": true}); err != nil { + http.Error(w, "Failed to write capabilities", http.StatusInternalServerError) + } +} + func serveMetadata(w http.ResponseWriter, r *http.Request) { currentUser, err := user.Current() if err != nil { diff --git a/libs/telemetry/protos/ssh_tunnel.go b/libs/telemetry/protos/ssh_tunnel.go index 9fb19a7f702..3240c7e9efe 100644 --- a/libs/telemetry/protos/ssh_tunnel.go +++ b/libs/telemetry/protos/ssh_tunnel.go @@ -22,6 +22,11 @@ const ( // be attributed without logging the error text, which carries cluster names, paths and user // names. // +// It also classifies how an established session ended, for the ends the CLI can attribute +// (the WEBSOCKET_* and HANDOVER_* categories). Those rows carry is_success = true, because +// the tunnel was established: use is_success to separate a failed connection attempt from a +// session that connected and was later cut short, and this field for the cause of either. +// // IDE_SSH_EXTENSION_MISSING was retired in favour of the four IDE_SSH_EXTENSION_* categories // below: it reported all four outcomes as one, and they call for different fixes. Rows written // before the split still carry it, so a query spanning that release has to accept both. @@ -80,6 +85,21 @@ const ( // out to the IDE and see only a killed child process. SshTunnelErrorCategoryUserAborted SshTunnelErrorCategory = "USER_ABORTED" + // The proxy could not establish its websocket to the SSH server at all. Distinguished + // from WEBSOCKET_DROPPED so "never connected" is not counted as a mid-session drop. + SshTunnelErrorCategoryWebsocketConnectFailed SshTunnelErrorCategory = "WEBSOCKET_CONNECT_FAILED" + + // An established proxy websocket stopped carrying traffic mid-session, ending the SSH + // session with it. Typically a TCP reset from somewhere on the path between the client + // and the workspace; the CLI cannot yet resume from it, so every occurrence is a + // user-visible dropped session. + SshTunnelErrorCategoryWebsocketDropped SshTunnelErrorCategory = "WEBSOCKET_DROPPED" + + // The periodic auth handover failed in a way that ended the session. A handover that + // only failed to dial its replacement is not reported here: the session continues on + // its existing connection. + SshTunnelErrorCategoryHandoverFailed SshTunnelErrorCategory = "HANDOVER_FAILED" + // A failure that does not correspond to any of the categories above. The connect path // attributes every per-environment blocker, so a rise here points at a CLI bug (or a new // failure mode that needs its own category) rather than a user's setup. @@ -130,9 +150,12 @@ type SshTunnelEvent struct { // Only the presence is recorded, not the policy ID itself. HasUsagePolicy bool `json:"has_usage_policy"` - // Why the connection attempt failed, or TYPE_UNSPECIFIED on success. Deliberately - // without omitempty: the field is what identifies a failure's cause, so an empty value - // must not be silently dropped into an indistinguishable NULL. Every failure path sets - // a category, falling back to UNKNOWN. + // Why the connection attempt failed, how an established session ended, or + // TYPE_UNSPECIFIED when neither applies. Deliberately without omitempty: the field is + // what identifies a failure's cause, so an empty value must not be silently dropped into + // an indistinguishable NULL. Every failed connection attempt sets a category, falling + // back to UNKNOWN; an established session sets one only for the ends the CLI can + // attribute, since the ssh client and the user's own remote command also exit non-zero + // here and neither is a tunnel failure. ErrorCategory SshTunnelErrorCategory `json:"error_category"` } diff --git a/libs/testserver/handlers.go b/libs/testserver/handlers.go index 8c44d97eae7..b6510e8ed4b 100644 --- a/libs/testserver/handlers.go +++ b/libs/testserver/handlers.go @@ -791,6 +791,15 @@ func AddDefaultHandlers(server *Server) { return Response{Body: ""} }) + // /capabilities reports which optional parts of the tunnel protocol the server speaks. + // This fake drives sshd directly over the websocket rather than running the CLI's own + // proxy server, so it has none of the session bookkeeping a resume needs and says so. + // The resume protocol itself is covered by the proxy package's tests, which run the + // real server implementation. + server.Handle("GET", "/driver-proxy-api/o/{workspace_id}/{cluster_id}/{port}/capabilities", func(req Request) any { + return Response{Body: map[string]bool{"resume": false}} + }) + server.HandleRaw("GET", "/driver-proxy-api/o/{workspace_id}/{cluster_id}/{port}/ssh", server.sshTunnelHandler) // Secrets ACLs: From d6f8a0cc2836d7f1ae0a18aee4d62c06825399f5 Mon Sep 17 00:00:00 2001 From: Anton Nekipelov <226657+anton-107@users.noreply.github.com> Date: Tue, 8 Sep 2026 10:21:06 +0000 Subject: [PATCH 2/4] experimental/ssh: address review on tunnel resume/handover Three blocking review fixes on the resume/handover machinery: - proxy.go: a failed write on a resumable connection now closes the connection and returns errSendFailedResumable for ANY message type, not only binary payloads. A failed ack (a text control frame) previously only logged, so on a one-way-from-server flow the connection stayed write-poisoned: the receiving loop kept running, the peer stopped getting acks, its 1 MB replay buffer filled, and the session ended instead of reattaching. It now drives the same reattach path as a failed binary write. Buffering semantics are unchanged (only binary payloads are appended to the replay buffer). Adds a test covering the failed-ack case. - client_server_test.go: TestHandoverDialFailureKeepsSessionAlive no longer relies on the handover-tick send to imply the dial started. The fake dialer now signals when the handover dial is attempted (it runs while initiateHandover holds handoverMutex), and the test waits for that signal before sending post-handover traffic and asserting the dial count. This removes the race behind the macOS "dials == 1" flake. - .nextchanges: append the required trailing PR link to both new changelog fragments so the changelog preview check passes. Co-authored-by: Isaac --- .../cli/ssh-tunnel-handover-dial-failure.md | 2 +- .nextchanges/cli/ssh-tunnel-resume.md | 2 +- .../ssh/internal/proxy/client_server_test.go | 19 +++++++++ experimental/ssh/internal/proxy/proxy.go | 17 +++++--- experimental/ssh/internal/proxy/proxy_test.go | 42 +++++++++++++++++++ 5 files changed, 75 insertions(+), 7 deletions(-) diff --git a/.nextchanges/cli/ssh-tunnel-handover-dial-failure.md b/.nextchanges/cli/ssh-tunnel-handover-dial-failure.md index 58733738ec4..4a7a58523d5 100644 --- a/.nextchanges/cli/ssh-tunnel-handover-dial-failure.md +++ b/.nextchanges/cli/ssh-tunnel-handover-dial-failure.md @@ -1 +1 @@ -* Keep an `ssh connect` session alive when the periodic auth handover cannot open its replacement websocket. A transient dial failure previously ended a healthy session, even though the connection being replaced was still working. +* Keep an `ssh connect` session alive when the periodic auth handover cannot open its replacement websocket. A transient dial failure previously ended a healthy session, even though the connection being replaced was still working. ([#6558](https://github.com/databricks/cli/pull/6558)) diff --git a/.nextchanges/cli/ssh-tunnel-resume.md b/.nextchanges/cli/ssh-tunnel-resume.md index 95d0e76d86d..fd4ed6edc27 100644 --- a/.nextchanges/cli/ssh-tunnel-resume.md +++ b/.nextchanges/cli/ssh-tunnel-resume.md @@ -1 +1 @@ -* `ssh connect` sessions now survive an unexpected connection drop. When the tunnel websocket is reset mid-session, the CLI reattaches to the running session and replays the bytes that were lost, so the shell and everything running in it stay intact instead of the session ending with `failed to read from websocket`. Requires an SSH server started by a CLI that supports it; against an older server the connection behaves as before. +* `ssh connect` sessions now survive an unexpected connection drop. When the tunnel websocket is reset mid-session, the CLI reattaches to the running session and replays the bytes that were lost, so the shell and everything running in it stay intact instead of the session ending with `failed to read from websocket`. Requires an SSH server started by a CLI that supports it; against an older server the connection behaves as before. ([#6558](https://github.com/databricks/cli/pull/6558)) diff --git a/experimental/ssh/internal/proxy/client_server_test.go b/experimental/ssh/internal/proxy/client_server_test.go index 82c07635297..11a051e74c6 100644 --- a/experimental/ssh/internal/proxy/client_server_test.go +++ b/experimental/ssh/internal/proxy/client_server_test.go @@ -258,9 +258,18 @@ func TestHandoverDialFailureKeepsSessionAlive(t *testing.T) { wsURL := "ws" + server.URL[4:] var dials atomic.Int32 + // Signalled the instant a handover dial is attempted (and made to fail). initiateHandover + // already holds handoverMutex by the time it dials, so a receive here proves the handover + // goroutine has entered the dial and taken the mutex. Buffered and sent non-blockingly so the + // dialer never stalls on it even if more dials than expected occur. + handoverDialAttempted := make(chan struct{}, 1) createConn := func(ctx context.Context, dial DialRequest) (*websocket.Conn, error) { // Let the initial connection through and fail every handover dial after it. if dials.Add(1) > 1 { + select { + case handoverDialAttempted <- struct{}{}: + default: + } return nil, errors.New("simulated transient dial failure") } url := fmt.Sprintf("%s?id=%s", wsURL, dial.ConnID) @@ -282,6 +291,16 @@ func TestHandoverDialFailureKeepsSessionAlive(t *testing.T) { handoverChan <- time.Now() + // Completing the tick send only proves the handover goroutine received the tick; it does not + // prove it acquired handoverMutex and reached the dial. Wait for the dial to actually be + // attempted before sending more traffic - otherwise the payload below can traverse the + // original connection before the handover even starts, which is the macOS "dials == 1" flake. + select { + case <-handoverDialAttempted: + case <-time.After(10 * time.Second): + t.Fatal("the handover never attempted its replacement dial") + } + // The original connection must still be proxying both ways. sendMessage blocks on the // handover mutex, so this write cannot overtake the failed handover. afterMsg := []byte("after failed handover\n") diff --git a/experimental/ssh/internal/proxy/proxy.go b/experimental/ssh/internal/proxy/proxy.go index 665478cd19b..eba1d1b0fba 100644 --- a/experimental/ssh/internal/proxy/proxy.go +++ b/experimental/ssh/internal/proxy/proxy.go @@ -347,11 +347,18 @@ func (pc *proxyConnection) sendMessage(mt int, data []byte) error { } conn := pc.conn.Load() err := conn.WriteMessage(mt, data) - if err != nil && pc.resumable() && mt == websocket.BinaryMessage { - // The payload is buffered, so this failure costs no data. gorilla latches a permanent - // write error after any failed write, so this connection can never send again: close it - // to fail the receiving loop's read now and start the resume, rather than let the - // sending loop fill the whole window first. + if err != nil && pc.resumable() { + // This failure costs no data whatever the message type: a binary payload was buffered + // above, and control/ack messages are regenerated after the resume. gorilla latches a + // permanent write error after any failed write, so this connection can never send again - + // close it to fail the receiving loop's read now and drive the reattach, rather than let + // the sending loop fill the whole window first. This must cover a failed ack (a text + // control frame) too, not only a binary payload: when traffic is one-way from the server + // the receiving side never writes a binary frame, so a poisoned connection would otherwise + // only ever surface as a failed ack. Left as a bare log, the read loop kept running while + // the peer stopped getting acks, its replay buffer filled to the limit, and the session + // ended instead of reattaching. A failed close message reaches here only during teardown, + // where closing the connection is what happens next anyway. conn.Close() return errors.Join(errSendFailedResumable, err) } diff --git a/experimental/ssh/internal/proxy/proxy_test.go b/experimental/ssh/internal/proxy/proxy_test.go index 5abb7a7186b..aebf43b9330 100644 --- a/experimental/ssh/internal/proxy/proxy_test.go +++ b/experimental/ssh/internal/proxy/proxy_test.go @@ -6,6 +6,7 @@ import ( "errors" "fmt" "io" + "net" "net/http" "net/http/httptest" "sync" @@ -268,3 +269,44 @@ func TestConnectionHandover(t *testing.T) { require.NoError(t, err) } } + +// A failed acknowledgement write on a resumable connection must be treated exactly like a failed +// binary write: close the connection and report errSendFailedResumable, so the receiving loop's +// next read fails and drives the reattach. When traffic is one-way from the server the receiving +// side never writes a binary frame, so a poisoned connection surfaces only as a failed ack; left +// as a bare log it would let the peer's replay buffer fill and end the session instead. +func TestFailedAckWriteClosesResumableConnectionToDriveReattach(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if err != nil { + return + } + // A live peer to dial; drain until the client goes away. + for { + if _, _, err := conn.ReadMessage(); err != nil { + return + } + } + })) + defer server.Close() + + wsURL := "ws" + server.URL[4:] + conn, err := createTestWebsocketConnection(wsURL) + require.NoError(t, err) + + pc := newResumableProxyConnection(nil) + pc.conn.Store(conn) + + // Poison the write side the way gorilla latches it after any failed write, without disturbing + // reads - the one-way-from-server case where only the ack ever fails. + require.NoError(t, conn.SetWriteDeadline(time.Now().Add(-time.Hour))) + + // sendControlMessage is the ack path (a text control frame), not a binary payload. + err = pc.sendControlMessage(42) + require.ErrorIs(t, err, errSendFailedResumable, "a failed ack write on a resumable connection must report the resumable-send failure that drives the reattach") + + // It must have closed the connection: a second close returns net.ErrClosed. With the + // connection closed, the receiving loop's next read fails and reattaches, rather than the + // failure being silently swallowed. + require.ErrorIs(t, conn.Close(), net.ErrClosed, "sendMessage must close the poisoned connection so the receiving loop reattaches") +} From 6553a809029e9d3f4dc2e2a4618351009e77e705 Mon Sep 17 00:00:00 2001 From: Anton Nekipelov <226657+anton-107@users.noreply.github.com> Date: Tue, 8 Sep 2026 13:26:51 +0000 Subject: [PATCH 3/4] experimental/ssh: do not warn about a dropped connection on a clean exit Every successful `ssh connect` session ended with a false warning, no --debug needed: $ databricks ssh connect --name bv-prewarm-lf -- true # exits 0 Connected! Received termination signal, cleaning up... Warn: SSH tunnel connection dropped, reattaching to the session... On teardown the context is cancelled, start's context watcher closes the websocket to unblock the receiving loop's read, and that read then fails with something other than a normal closure. The loop's own check for cancellation sits at the top, before the read, so a cancellation that lands while the read is in flight was missed and the failure was taken for an unexpected drop: with resume negotiated it entered reattach(), logged the warning, and only then failed fast on the already-cancelled context. Presentation was the whole impact - exit code and teardown timing were unaffected - but it told users their connection had dropped on every single session, which would erode trust in the genuine warning. The receiving loop now returns the cancellation instead. Placed before the resumable branch, because neither branch fits a teardown: the reattach cannot succeed on a cancelled context anyway, and ErrWebsocketDropped would bill an ordinary exit to a tunnel failure in telemetry - which matters now that an established session no longer discards its error category. The test pins the ordering that makes this exact rather than incidental: it waits for a ping handler to fire from inside ReadMessage, which proves the read cannot return until the connection is closed, then cancels and closes in the same order the watcher does. Reverting the guard fails it, warning and all. Co-authored-by: Isaac --- experimental/ssh/internal/proxy/proxy.go | 9 +++ experimental/ssh/internal/proxy/proxy_test.go | 74 +++++++++++++++++++ 2 files changed, 83 insertions(+) diff --git a/experimental/ssh/internal/proxy/proxy.go b/experimental/ssh/internal/proxy/proxy.go index eba1d1b0fba..e1e295071b1 100644 --- a/experimental/ssh/internal/proxy/proxy.go +++ b/experimental/ssh/internal/proxy/proxy.go @@ -429,6 +429,15 @@ func (pc *proxyConnection) runReceivingLoop(ctx context.Context, dst io.Writer) if errors.Is(err, io.EOF) || websocket.IsCloseError(err, websocket.CloseNormalClosure) { return errors.Join(errProxyEOF, err) } + // A read that fails once our own context is cancelled is the teardown, not a drop: + // start's context watcher closes the connection to unblock this very read, and + // only after the context is done, so cancellation is always visible here first. + // Neither branch below fits - a reattach would warn the user about a drop on every + // clean exit and could not succeed anyway (its redial budget comes from this same + // context), and ErrWebsocketDropped would bill an ordinary exit to a tunnel failure. + if ctx.Err() != nil { + return ctx.Err() + } // An unexpected drop. With resume negotiated the session state on both ends // outlives the connection, so reattach instead of ending the session. if pc.resumable() { diff --git a/experimental/ssh/internal/proxy/proxy_test.go b/experimental/ssh/internal/proxy/proxy_test.go index aebf43b9330..d31cf70eae5 100644 --- a/experimental/ssh/internal/proxy/proxy_test.go +++ b/experimental/ssh/internal/proxy/proxy_test.go @@ -10,6 +10,7 @@ import ( "net/http" "net/http/httptest" "sync" + "sync/atomic" "testing" "time" @@ -310,3 +311,76 @@ func TestFailedAckWriteClosesResumableConnectionToDriveReattach(t *testing.T) { // failure being silently swallowed. require.ErrorIs(t, conn.Close(), net.ErrClosed, "sendMessage must close the poisoned connection so the receiving loop reattaches") } + +// A read that fails during teardown must not be mistaken for a drop. start's context watcher +// closes the websocket to unblock this read, so on every clean exit the read fails with the +// context already cancelled - and reattaching there told the user their connection had dropped on +// every single session, having no chance of succeeding on a cancelled context either. +func TestTeardownIsNotTreatedAsADrop(t *testing.T) { + // Signalled from inside the client's ReadMessage: gorilla dispatches control frames from + // there and keeps reading, so this proves the read is in flight and cannot return until the + // connection is closed. Without it the loop could be between iterations instead, where the + // check at the top of the loop handles the cancellation and the test passes vacuously. + readInFlight := make(chan struct{}) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if err != nil { + return + } + if err := conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(time.Minute)); err != nil { + return + } + // Send no payload: the client's read stays blocked until the test closes the connection. + for { + if _, _, err := conn.ReadMessage(); err != nil { + return + } + } + })) + defer server.Close() + + wsURL := "ws" + server.URL[4:] + conn, err := createTestWebsocketConnection(wsURL) + require.NoError(t, err) + conn.SetPingHandler(func(string) error { + close(readInFlight) + return nil + }) + + var dials atomic.Int32 + pc := newResumableProxyConnection(func(ctx context.Context, dial DialRequest) (*websocket.Conn, error) { + dials.Add(1) + return nil, errors.New("a teardown must never dial a reattach") + }) + pc.conn.Store(conn) + + ctx, cancel := context.WithCancel(t.Context()) + loopErr := make(chan error, 1) + go func() { + loopErr <- pc.runReceivingLoop(ctx, io.Discard) + }() + + select { + case <-readInFlight: + case <-time.After(10 * time.Second): + t.Fatal("the receiving loop never reached its read") + } + + // Cancel first and close second, exactly as start's context watcher does it, so the read + // error always surfaces with the cancellation already visible. + cancel() + require.NoError(t, conn.Close()) + + select { + case err := <-loopErr: + require.ErrorIs(t, err, context.Canceled, "a read that fails after cancellation is the teardown, so the loop must report the cancellation") + require.NotErrorIs(t, err, ErrWebsocketDropped, "a clean teardown must not be attributed to a dropped websocket") + case <-time.After(10 * time.Second): + t.Fatal("the receiving loop never returned") + } + + // No dial means no reattach was started, and the warning that prompted this test lives inside + // the reattach, so it cannot have been printed either. + require.Zero(t, dials.Load(), "the reattach must not be attempted once the context is cancelled") +} From 873e42afc5d0b4bc69ac8141853e06adfebfd36e Mon Sep 17 00:00:00 2001 From: Anton Nekipelov <226657+anton-107@users.noreply.github.com> Date: Tue, 8 Sep 2026 15:13:39 +0000 Subject: [PATCH 4/4] experimental/ssh: address review on logging and the changelog - proxy/client.go: log a failed handover dial at debug, not warn. The session carries on with the connection it already has, so there is nothing for the user to act on, and a warn writes into their interactive terminal on top of whatever the session is showing. - client.go: drop the rationale block on serverSupportsResume. Why an older server cannot resume needs no explaining, and the hazard it described - an old server answering a reattach by starting a fresh sshd, which a replay then corrupts - is already documented on RunClientProxy's resumable parameter, where it governs the decision. - .nextchanges: one fragment for the whole change instead of two, stated in terms of what the user gets rather than the two mechanisms behind it. Co-authored-by: Isaac --- .nextchanges/cli/ssh-tunnel-connection-drops.md | 1 + .nextchanges/cli/ssh-tunnel-handover-dial-failure.md | 1 - .nextchanges/cli/ssh-tunnel-resume.md | 1 - experimental/ssh/internal/client/client.go | 8 -------- experimental/ssh/internal/proxy/client.go | 6 ++++-- 5 files changed, 5 insertions(+), 12 deletions(-) create mode 100644 .nextchanges/cli/ssh-tunnel-connection-drops.md delete mode 100644 .nextchanges/cli/ssh-tunnel-handover-dial-failure.md delete mode 100644 .nextchanges/cli/ssh-tunnel-resume.md diff --git a/.nextchanges/cli/ssh-tunnel-connection-drops.md b/.nextchanges/cli/ssh-tunnel-connection-drops.md new file mode 100644 index 00000000000..fd804834d1e --- /dev/null +++ b/.nextchanges/cli/ssh-tunnel-connection-drops.md @@ -0,0 +1 @@ +* `ssh connect` sessions no longer end when the tunnel's websocket connection is lost. The CLI reattaches to the running session and replays the bytes that were missed, so the shell and everything running in it stay intact, and a transient failure to open a replacement connection for the periodic auth refresh is retried rather than ending the session. Reattaching requires an SSH server started by a CLI that supports it; against an older server the connection behaves as before. ([#6558](https://github.com/databricks/cli/pull/6558)) diff --git a/.nextchanges/cli/ssh-tunnel-handover-dial-failure.md b/.nextchanges/cli/ssh-tunnel-handover-dial-failure.md deleted file mode 100644 index 4a7a58523d5..00000000000 --- a/.nextchanges/cli/ssh-tunnel-handover-dial-failure.md +++ /dev/null @@ -1 +0,0 @@ -* Keep an `ssh connect` session alive when the periodic auth handover cannot open its replacement websocket. A transient dial failure previously ended a healthy session, even though the connection being replaced was still working. ([#6558](https://github.com/databricks/cli/pull/6558)) diff --git a/.nextchanges/cli/ssh-tunnel-resume.md b/.nextchanges/cli/ssh-tunnel-resume.md deleted file mode 100644 index fd4ed6edc27..00000000000 --- a/.nextchanges/cli/ssh-tunnel-resume.md +++ /dev/null @@ -1 +0,0 @@ -* `ssh connect` sessions now survive an unexpected connection drop. When the tunnel websocket is reset mid-session, the CLI reattaches to the running session and replays the bytes that were lost, so the shell and everything running in it stay intact instead of the session ending with `failed to read from websocket`. Requires an SSH server started by a CLI that supports it; against an older server the connection behaves as before. ([#6558](https://github.com/databricks/cli/pull/6558)) diff --git a/experimental/ssh/internal/client/client.go b/experimental/ssh/internal/client/client.go index 5551ca4d2e5..a5fe61f6388 100644 --- a/experimental/ssh/internal/client/client.go +++ b/experimental/ssh/internal/client/client.go @@ -919,14 +919,6 @@ func runSSHProxy(ctx context.Context, client *databricks.WorkspaceClient, server } // serverSupportsResume reports whether the running SSH server speaks the resume protocol. -// -// This must be established before the session starts, not discovered when a connection drops. -// A server from an older CLI has no /capabilities endpoint and, worse, would answer a reattach -// request by starting a fresh sshd: replaying into that fails the SSH stream with a corrupted MAC -// instead of the clear error the user gets today. A 404 (or any other failure) therefore means -// "assume not supported". Version skew is real here because the generated ssh config pins -// --metadata into the ProxyCommand, which skips the version-scoped server lookup, so an upgraded -// client can reach a server an older CLI started. func serverSupportsResume(ctx context.Context, client *databricks.WorkspaceClient, clusterID string, serverPort int, liteswap string) bool { req, err := newDriverProxyRequest(ctx, client, clusterID, serverPort, "capabilities", liteswap) if err != nil { diff --git a/experimental/ssh/internal/proxy/client.go b/experimental/ssh/internal/proxy/client.go index 53aca3880bd..1ad9ebe08c2 100644 --- a/experimental/ssh/internal/proxy/client.go +++ b/experimental/ssh/internal/proxy/client.go @@ -101,9 +101,11 @@ func RunClientProxy(ctx context.Context, src io.ReadCloser, dst io.Writer, reque // would throw away a working tunnel - the failure mode customers see as // a drop every handover interval. The next tick tries again. Deferring // the auth refresh is safe: the driver proxy authenticates a websocket - // at upgrade time, so a live connection is not re-checked. + // at upgrade time, so a live connection is not re-checked. Logged at + // debug because nothing changed for the user, and this would otherwise + // write into their interactive terminal. if errors.Is(err, errHandoverDialFailed) { - log.Warnf(gCtx, "Could not open a replacement connection for the auth handover, staying on the current one: %v", err) + log.Debugf(gCtx, "Could not open a replacement connection for the auth handover, staying on the current one: %v", err) continue } return errors.Join(ErrHandoverFailed, err)