From 08ca4ebe7741f07eeadc0745b732ae798ced27fd Mon Sep 17 00:00:00 2001 From: Anton Nekipelov <226657+anton-107@users.noreply.github.com> Date: Fri, 11 Sep 2026 14:32:44 +0000 Subject: [PATCH 1/4] Fix SSH session resume with bounded backpressure DECO-28433: restore dropped-session reattachment without the full-window failure introduced by #6558 or the protocol degradation attempted in #6609. Wait for acknowledgments outside the write lock, send acknowledgments on a separate loop, and resume reads before replaying full windows. Negotiate a new protocol version to avoid enabling resume against the faulty server. Preserve final payload before EOF, distinguish session completion from a dropped socket, retry dropped handshakes, and stop retrying expired sessions. Regression tests reproduce the original 1 MiB truncation, an acknowledgment deadlock during handover, and a simultaneous replay deadlock. They verify 8 MiB transfers in both directions through handover and repeated resets. Validation: SSH unit and local acceptance suites; race tests for changed SSH and telemetry packages; targeted lint; Windows test compilation. Live driver-proxy validation remains required before release. --- .nextchanges/cli/ssh-resume-backpressure.md | 1 + experimental/ssh/internal/client/client.go | 87 +++- .../internal/client/client_internal_test.go | 71 +++ .../ssh/internal/client/websockets.go | 33 +- .../ssh/internal/client/websockets_test.go | 85 +++- .../ssh/internal/proxy/backpressure_test.go | 186 ++++++++ experimental/ssh/internal/proxy/client.go | 34 +- .../ssh/internal/proxy/client_server_test.go | 99 +++- experimental/ssh/internal/proxy/drop_test.go | 163 +++++++ .../ssh/internal/proxy/keepalive_test.go | 10 +- experimental/ssh/internal/proxy/proxy.go | 357 +++++++++++++- experimental/ssh/internal/proxy/proxy_test.go | 118 ++++- experimental/ssh/internal/proxy/resume.go | 436 ++++++++++++++++++ .../ssh/internal/proxy/resume_e2e_test.go | 200 ++++++++ .../ssh/internal/proxy/resume_test.go | 343 ++++++++++++++ experimental/ssh/internal/proxy/server.go | 82 +++- experimental/ssh/internal/server/server.go | 13 + libs/telemetry/protos/ssh_tunnel.go | 31 +- libs/testserver/handlers.go | 9 + 19 files changed, 2293 insertions(+), 65 deletions(-) create mode 100644 .nextchanges/cli/ssh-resume-backpressure.md create mode 100644 experimental/ssh/internal/proxy/backpressure_test.go 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-resume-backpressure.md b/.nextchanges/cli/ssh-resume-backpressure.md new file mode 100644 index 00000000000..3ba1785757c --- /dev/null +++ b/.nextchanges/cli/ssh-resume-backpressure.md @@ -0,0 +1 @@ +* Preserve SSH sessions across temporary tunnel disconnects, with bounded replay and backpressure for large transfers. diff --git a/experimental/ssh/internal/client/client.go b/experimental/ssh/internal/client/client.go index e95351c2e45..5cf0aedaaf9 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" @@ -493,7 +494,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, knownHostsPath, serverPort, clusterID, opts) } else { @@ -973,13 +979,46 @@ 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. +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 { + ResumeVersion int `json:"resume_version"` + } + 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.ResumeVersion == proxy.ResumeProtocolVersion } // accessModeUILabel maps a cluster's access mode to the name shown in the Databricks UI. @@ -1351,11 +1390,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) { @@ -1367,6 +1417,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 4729ef30e82..ae70cb3e85a 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" @@ -547,6 +548,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}, @@ -604,6 +616,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 { @@ -657,6 +681,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..b3b65c2a526 100644 --- a/experimental/ssh/internal/client/websockets.go +++ b/experimental/ssh/internal/client/websockets.go @@ -2,17 +2,20 @@ package client import ( "context" + "errors" "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) } @@ -30,27 +33,33 @@ func createWebsocketConnection(ctx context.Context, client *databricks.Workspace } // websocket connection manages lifecycle of the response object, no need to close the body - conn, _, err := websocket.DefaultDialer.Dial(req.URL.String(), req.Header) // nolint:bodyclose + conn, resp, err := websocket.DefaultDialer.DialContext(ctx, req.URL.String(), req.Header) + if resp != nil { + resp.Body.Close() + } if err != nil { + if dial.Reattach && resp != nil && resp.StatusCode >= 400 && resp.StatusCode < 500 && resp.StatusCode != http.StatusTooManyRequests { + return nil, errors.Join(proxy.ErrReattachRejected, fmt.Errorf("reattach failed (HTTP %d): %w", resp.StatusCode, err)) + } return nil, fmt.Errorf("failed to establish websocket connection: %w", err) } 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 +74,16 @@ 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 { + query.Set(proxy.ResumeVersionParameter, strconv.Itoa(proxy.ResumeProtocolVersion)) + // 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..b3689a40ec6 100644 --- a/experimental/ssh/internal/client/websockets_test.go +++ b/experimental/ssh/internal/client/websockets_test.go @@ -1,33 +1,116 @@ package client import ( + "context" + "net/http" + "net/http/httptest" "testing" + "time" + "github.com/databricks/cli/experimental/ssh/internal/proxy" + "github.com/databricks/databricks-sdk-go" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) +func TestServerSupportsResume(t *testing.T) { + for _, tc := range []struct { + name string + status int + body string + want bool + }{ + {"old server", http.StatusNotFound, "not found", false}, + {"faulty v1 server", http.StatusOK, `{"resume":true}`, false}, + {"corrected protocol", http.StatusOK, `{"resume_version":2}`, true}, + {"unknown protocol", http.StatusOK, `{"resume_version":3}`, false}, + {"invalid response", http.StatusOK, `{`, false}, + } { + t.Run(tc.name, func(t *testing.T) { + mux := http.NewServeMux() + mux.HandleFunc("GET /driver-proxy-api/o/123/cluster/7772/capabilities", func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(tc.status) + _, _ = w.Write([]byte(tc.body)) + }) + server := httptest.NewServer(mux) + defer server.Close() + client, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "test-token", WorkspaceID: "123", AuthType: "pat"}) + require.NoError(t, err) + assert.Equal(t, tc.want, serverSupportsResume(t.Context(), client, "cluster", 7772, "")) + }) + } +} + +func TestCreateWebsocketConnectionReattachRejected(t *testing.T) { + for _, status := range []int{http.StatusGone, http.StatusConflict, http.StatusUnauthorized, http.StatusTooManyRequests, http.StatusServiceUnavailable} { + t.Run(http.StatusText(status), func(t *testing.T) { + mux := http.NewServeMux() + mux.HandleFunc("GET /driver-proxy-api/o/123/cluster/7772/ssh", func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(status) + }) + server := httptest.NewServer(mux) + defer server.Close() + client, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "test-token", WorkspaceID: "123", AuthType: "pat"}) + require.NoError(t, err) + ctx, cancel := context.WithTimeout(t.Context(), time.Second) + defer cancel() + _, err = createWebsocketConnection(ctx, client, proxy.DialRequest{ConnID: "test", ResumeCapable: true, Reattach: true}, "cluster", 7772, "") + require.Error(t, err) + if status < 500 && status != http.StatusTooManyRequests { + assert.ErrorIs(t, err, proxy.ErrReattachRejected) + } else { + assert.NotErrorIs(t, err, proxy.ErrReattachRejected) + } + }) + } +} + 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&resume_version=2", + }, + { + 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&resume_version=2", + }, } 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/backpressure_test.go b/experimental/ssh/internal/proxy/backpressure_test.go new file mode 100644 index 00000000000..81c0fad07ea --- /dev/null +++ b/experimental/ssh/internal/proxy/backpressure_test.go @@ -0,0 +1,186 @@ +package proxy_test + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "sync" + "testing" + "time" + + "github.com/databricks/cli/experimental/ssh/internal/proxy" + "github.com/gorilla/websocket" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestResumeBackpressure covers a live peer whose acknowledgments arrive after the +// replay window fills, as happens during IDE startup through the driver proxy. +func TestResumeBackpressure(t *testing.T) { + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + const window = 1 << 20 + payload := bytes.Repeat([]byte("0123456789abcdef"), window/2) + received := make(chan []byte, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upgrader := websocket.Upgrader{} + conn, err := upgrader.Upgrade(w, r, nil) + if !assert.NoError(t, err) { + return + } + defer conn.Close() + stop := context.AfterFunc(ctx, func() { conn.Close() }) + defer stop() + if !assert.NoError(t, conn.WriteMessage(websocket.BinaryMessage, []byte("SSH-2.0-test\r\n"))) { + return + } + var data []byte + defer func() { received <- data }() + acked := 0 + for len(data) < len(payload) { + mt, frame, err := conn.ReadMessage() + if err != nil { + return + } + if mt == websocket.TextMessage { + var ack struct { + Delivered int `json:"delivered"` + } + if !assert.NoError(t, json.Unmarshal(frame, &ack)) { + return + } + assert.Equal(t, len("SSH-2.0-test\r\n"), ack.Delivered) + continue + } + if !assert.Equal(t, websocket.BinaryMessage, mt) { + return + } + data = append(data, frame...) + if len(data)-acked >= window { + // Give the sender time to try to exceed its window before acknowledging. + select { + case <-time.After(50 * time.Millisecond): + case <-ctx.Done(): + return + } + if err := conn.WriteJSON(map[string]int{"delivered": len(data)}); err != nil { + return + } + acked = len(data) + } + } + _ = conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "finished")) + })) + defer server.Close() + input, writer := io.Pipe() + defer input.Close() + defer writer.Close() + go func() { _, _ = writer.Write(payload) }() + err := proxy.RunClientProxy(ctx, input, io.Discard, func() <-chan time.Time { return nil }, time.Hour, true, + func(ctx context.Context, req proxy.DialRequest) (*websocket.Conn, error) { + conn, resp, err := websocket.DefaultDialer.DialContext(ctx, "ws"+server.URL[4:], nil) + if resp != nil { + resp.Body.Close() + } + return conn, err + }) + assert.NoError(t, err) + cancel() + select { + case data := <-received: + require.Len(t, data, len(payload), "the live peer received a truncated stream") + assert.Equal(t, sha256.Sum256(payload), sha256.Sum256(data)) + case <-time.After(time.Second): + t.Fatal("the peer did not finish") + } +} + +type lockedBuffer struct { + mu sync.Mutex + b bytes.Buffer +} + +func (b *lockedBuffer) Write(p []byte) (int, error) { + b.mu.Lock() + defer b.mu.Unlock() + return b.b.Write(p) +} + +func (b *lockedBuffer) bytes() []byte { + b.mu.Lock() + defer b.mu.Unlock() + return bytes.Clone(b.b.Bytes()) +} + +// TestResumeHandoverWithPendingAck exercises data already in flight when the +// handover owns the write lock. Reading the close frame must not wait on an ack. +func TestResumeHandoverWithPendingAck(t *testing.T) { + ctx, cancel := context.WithTimeout(t.Context(), 3*time.Second) + defer cancel() + initial := make(chan *websocket.Conn, 1) + ready := make(chan struct{}) + payload := bytes.Repeat([]byte("x"), 64<<10) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if !assert.NoError(t, err) { + return + } + defer conn.Close() + stop := context.AfterFunc(ctx, func() { conn.Close() }) + defer stop() + select { + case old := <-initial: + if !assert.NoError(t, old.WriteMessage(websocket.BinaryMessage, payload)) { + return + } + if !assert.NoError(t, old.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "handover"))) { + return + } + if !assert.NoError(t, conn.WriteMessage(websocket.BinaryMessage, []byte("after"))) { + return + } + default: + if !assert.NoError(t, conn.WriteMessage(websocket.BinaryMessage, []byte("before"))) { + return + } + initial <- conn + close(ready) + } + for { + if _, _, err := conn.ReadMessage(); err != nil { + return + } + } + })) + defer server.Close() + input, writer := io.Pipe() + defer input.Close() + defer writer.Close() + var output lockedBuffer + ticks := make(chan time.Time, 1) + done := make(chan error, 1) + go func() { + done <- proxy.RunClientProxy(ctx, input, &output, func() <-chan time.Time { return ticks }, time.Hour, true, + func(ctx context.Context, req proxy.DialRequest) (*websocket.Conn, error) { + conn, resp, err := websocket.DefaultDialer.DialContext(ctx, "ws"+server.URL[4:], nil) + if resp != nil { + resp.Body.Close() + } + return conn, err + }) + }() + <-ready + ticks <- time.Now() + want := append(append([]byte("before"), payload...), []byte("after")...) + assert.Eventually(t, func() bool { return bytes.Equal(want, output.bytes()) }, 2*time.Second, time.Millisecond) + cancel() + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("client did not stop after cancellation") + } +} diff --git a/experimental/ssh/internal/proxy/client.go b/experimental/ssh/internal/proxy/client.go index c9cb405aedc..1ad9ebe08c2 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,19 @@ 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. Logged at + // debug because nothing changed for the user, and this would otherwise + // write into their interactive terminal. + if errors.Is(err, errHandoverDialFailed) { + 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) } } } diff --git a/experimental/ssh/internal/proxy/client_server_test.go b/experimental/ssh/internal/proxy/client_server_test.go index 7aa45d35e81..11a051e74c6 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,74 @@ 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 + // 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) + 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() + + // 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") + _, 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 +330,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 +340,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 +378,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 +387,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..a5162879182 --- /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("&resume_version=2&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 11f8fcca720..5ce9c34d375 100644 --- a/experimental/ssh/internal/proxy/keepalive_test.go +++ b/experimental/ssh/internal/proxy/keepalive_test.go @@ -113,7 +113,7 @@ func startKeepaliveTestServer(t *testing.T) (*httptest.Server, <-chan struct{}) // websocket handshake has completed, so a test can control how its writes behave. func keepaliveTestDialer(serverURL string, onNetConn func(*pausableConn)) createWebsocketConnectionFunc { wsURL := "ws" + serverURL[4:] - return func(ctx context.Context, connID string) (*websocket.Conn, error) { + return func(ctx context.Context, dial DialRequest) (*websocket.Conn, error) { var wrapped *pausableConn dialer := websocket.Dialer{ NetDialContext: func(ctx context.Context, network, addr string) (net.Conn, error) { @@ -125,7 +125,7 @@ func keepaliveTestDialer(serverURL string, onNetConn func(*pausableConn)) create return wrapped, nil }, } - conn, _, err := dialer.DialContext(ctx, fmt.Sprintf("%s?id=%s", wsURL, connID), nil) // nolint:bodyclose + conn, _, err := dialer.DialContext(ctx, fmt.Sprintf("%s?id=%s", wsURL, dial.ConnID), nil) // nolint:bodyclose if err != nil { return nil, err } @@ -152,7 +152,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 { @@ -176,7 +176,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) })) }() @@ -250,7 +250,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..48864984993 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,49 @@ 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") + // ErrReattachRejected marks a permanent rejection, such as an expired session. + ErrReattachRejected = errors.New("the server rejected session reattachment") +) + +// ResumeProtocolVersion requires flow control and explicit session completion. +// Version 1 could terminate a healthy session when its replay window filled. +const ResumeProtocolVersion = 2 + +const ResumeVersionParameter = "resume_version" + 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 outlasts the client's retry budget so the server does not reap +// a session while its client is still trying to reconnect. +const 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 +67,74 @@ 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. A full window pauses the source + // until the peer acknowledges delivery; bursts through the driver proxy can fill it. + 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 + // Acknowledge the tail of an idle stream, including the final bytes before EOF. + proxyAckInterval = 100 * time.Millisecond + proxyCloseWriteTimeout = time.Second + proxySessionFinished = "finished" ) +// 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 + // Coalesces acknowledgments so the receiving loop never waits for a write lock. + ackNeeded chan struct{} + // Protected by handoverMutex. Written before any new payload on a replacement + // socket, after the receiving loop is free to drain the peer's replay. + replay []byte + // 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 + // Serializes inbound reattach attempts, including retiring the old socket. + reattachMu sync.Mutex + // Closed when the session stops, so an HTTP handler cannot revive it. + done chan struct{} + grace time.Duration +} + +func newResumeState() *resumeState { + return &resumeState{ + sendBuf: newSendBuffer(proxyResumeBufferLimit), + ackNeeded: make(chan struct{}, 1), + resumed: make(chan *websocket.Conn, 1), + parked: make(chan struct{}, 1), + done: make(chan struct{}), + grace: proxyResumeGrace, + } +} + // 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 +206,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 +} + +// 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, connID string) (*websocket.Conn, error) +type createWebsocketConnectionFunc func(ctx context.Context, req DialRequest) (*websocket.Conn, error) func newProxyConnection(createConn createWebsocketConnectionFunc) *proxyConnection { return &proxyConnection{ @@ -119,15 +240,43 @@ 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) + var finished atomic.Bool + if pc.resumable() { + g.Go(func() error { + pc.runAckLoop(gCtx) + return nil + }) + } g.Go(func() error { err := pc.runSendingLoop(gCtx, src) + if errors.Is(err, errProxyEOF) { + finished.Store(true) + } // Always return a non nil error to cancel the errgroup context return errors.Join(err, errSendingLoopStopped) }) g.Go(func() error { err := pc.runReceivingLoop(gCtx, dst) + if errors.Is(err, errProxyEOF) { + finished.Store(true) + } // Always return a non nil error to cancel the errgroup context return errors.Join(err, errReceivingLoopStopped) }) @@ -139,7 +288,14 @@ func (pc *proxyConnection) start(ctx context.Context, src io.ReadCloser, dst io. // Both loops can still be stuck on conn.ReadMessage or src.Read and won't notice context cancellation, // so we close the connection and the source (sshd stdout pipe or ssh client stdio) to unblock them. <-gCtx.Done() - return errors.Join(pc.close(), pc.closeConnection(), pc.closeSource(src)) + if pc.resumable() { + close(pc.resume.done) + } + var closeErr error + if finished.Load() || ctx.Err() != nil { + closeErr = pc.close() + } + return errors.Join(closeErr, pc.closeConnection(), pc.closeSource(src)) }) err := g.Wait() if err == nil || isNormalClosure(err) { @@ -150,7 +306,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,15 +343,41 @@ 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 + } + if err := pc.resume.sendBuf.waitForSpace(ctx, n); 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 { if errors.Is(readErr, io.EOF) { + if pc.resumable() { + // EOF may accompany the last payload of a failed write. Keep the + // session alive until replay delivers and acknowledges those bytes. + if err := pc.resume.sendBuf.waitForSpace(ctx, proxyResumeBufferLimit); err != nil { + return err + } + } return errors.Join(errProxyEOF, readErr) } else { return fmt.Errorf("failed to read from source: %w", readErr) @@ -205,8 +389,93 @@ 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) + var err error + if pc.resumable() && len(pc.resume.replay) > 0 { + err = conn.WriteMessage(websocket.BinaryMessage, pc.resume.replay) + if err == nil { + pc.resume.replay = nil + } + } + if err == nil { + err = conn.WriteMessage(mt, data) + } + 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) + } + 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 queues an acknowledgment without blocking reads. A synchronous +// write here deadlocks with a handover waiting for this loop to read a close frame. +func (pc *proxyConnection) ackDelivered() { + delivered := pc.resume.delivered.Load() + if delivered-pc.resume.acked.Load() < proxyAckThreshold { + return + } + pc.requestAck() +} + +// requestAck coalesces wakeups; the writer reads the latest delivered count. +func (pc *proxyConnection) requestAck() { + select { + case pc.resume.ackNeeded <- struct{}{}: + default: + } +} + +// runAckLoop keeps acknowledgment writes independent of payload reads and source +// backpressure. A failed write closes the socket and lets the receiver reattach. +func (pc *proxyConnection) runAckLoop(ctx context.Context) { + ticker := time.NewTicker(proxyAckInterval) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + if pc.resume.delivered.Load() == pc.resume.acked.Load() { + continue + } + case <-pc.resume.ackNeeded: + } + delivered := pc.resume.delivered.Load() + if err := pc.sendControlMessage(delivered); err != nil { + log.Debugf(ctx, "Failed to acknowledge %d delivered bytes: %v", delivered, err) + continue + } + pc.resume.acked.Store(delivered) + } } // sendPing writes a keepalive ping on the current connection. Unlike sendMessage it takes neither @@ -215,7 +484,11 @@ func (pc *proxyConnection) sendMessage(mt int, data []byte) error { // close() and the sending loop also need. func (pc *proxyConnection) sendPing() error { conn := pc.conn.Load() - return conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(proxyPingWriteTimeout)) + err := conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(proxyPingWriteTimeout)) + if err != nil && pc.resumable() { + conn.Close() + } + return err } func (pc *proxyConnection) runReceivingLoop(ctx context.Context, dst io.Writer) error { @@ -230,7 +503,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 { @@ -245,26 +518,68 @@ func (pc *proxyConnection) runReceivingLoop(ctx context.Context, dst io.Writer) // Continue with the receiving loop, pc.conn is now the new connection. continue } else { - if errors.Is(err, io.EOF) || websocket.IsCloseError(err, websocket.CloseNormalClosure) { + closeErr, closed := errors.AsType[*websocket.CloseError](err) + finished := closed && closeErr.Code == websocket.CloseNormalClosure && closeErr.Text == proxySessionFinished + if finished || (!pc.resumable() && (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) } + // 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() { + conn.Close() + 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 { + n, err := dst.Write(data) + if err != nil { return fmt.Errorf("failed to copy to writer: %w", err) } + if n != len(data) { + return fmt.Errorf("failed to copy to writer: %w", io.ErrShortWrite) + } + 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() + } } } func (pc *proxyConnection) close() error { // Keep in mind that pc.sendMessage blocks during handover - err := pc.sendMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")) + reason := "" + if pc.resumable() { + reason = proxySessionFinished + } + // WriteControl can run alongside a blocked data write or handover. + err := pc.conn.Load().WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, reason), time.Now().Add(proxyCloseWriteTimeout)) if err != nil { if isNormalClosure(err) || errors.Is(err, websocket.ErrCloseSent) { return nil @@ -315,9 +630,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..d31cf70eae5 100644 --- a/experimental/ssh/internal/proxy/proxy_test.go +++ b/experimental/ssh/internal/proxy/proxy_test.go @@ -6,9 +6,11 @@ import ( "errors" "fmt" "io" + "net" "net/http" "net/http/httptest" "sync" + "sync/atomic" "testing" "time" @@ -155,7 +157,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() } @@ -268,3 +270,117 @@ 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") +} + +// 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") +} diff --git a/experimental/ssh/internal/proxy/resume.go b/experimental/ssh/internal/proxy/resume.go new file mode 100644 index 00000000000..504b6b0343b --- /dev/null +++ b/experimental/ssh/internal/proxy/resume.go @@ -0,0 +1,436 @@ +package proxy + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "sync" + "time" + + "github.com/databricks/cli/libs/log" + "github.com/gorilla/websocket" +) + +// errSendWindowExhausted prevents overwriting unacknowledged bytes. The sending +// loop waits for room before appending, without holding the websocket write lock. +var errSendWindowExhausted = errors.New("resume send buffer is full") + +// 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 + space chan struct{} +} + +func newSendBuffer(limit int) *sendBuffer { + return &sendBuffer{limit: limit, space: make(chan struct{}, 1)} +} + +// waitForSpace applies backpressure to the sole payload producer. Acknowledgments +// can still arrive, and a reattach can still acquire the websocket write lock. +func (b *sendBuffer) waitForSpace(ctx context.Context, size int) error { + for { + if err := ctx.Err(); err != nil { + return err + } + b.mu.Lock() + available := b.limit - len(b.buf) + b.mu.Unlock() + if size <= available { + return nil + } + select { + case <-b.space: + case <-ctx.Done(): + return ctx.Err() + } + } +} + +// append records payload before its websocket write, while the caller holds the +// same lock used to install a replacement connection and its replay. +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 + select { + case b.space <- struct{}{}: + default: + } +} + +// 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, serverDelivered, 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. + if err := pc.prepareReplay(serverDelivered); err != nil { + conn.Close() + return err + } + pc.conn.Store(conn) + select { + case <-pc.resume.done: + conn.Close() + return ErrReattachRejected + default: + } + pc.requestAck() + log.Warnf(ctx, "SSH tunnel connection reattached, the session continues") + return nil +} + +func (pc *proxyConnection) redial(ctx context.Context) (*websocket.Conn, int64, 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 { + var delivered int64 + delivered, err = readResumeHandshake(ctx, conn) + if err == nil { + return conn, delivered, nil + } + conn.Close() + } + if errors.Is(err, ErrReattachRejected) { + return nil, 0, err + } + lastErr = err + log.Debugf(ctx, "Reattach dial failed, retrying: %v", err) + select { + case <-ctx.Done(): + return nil, 0, fmt.Errorf("gave up reattaching: %w", errors.Join(lastErr, ctx.Err())) + 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", pc.resume.grace) + 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(pc.resume.grace): + return fmt.Errorf("the client did not reattach within %v", pc.resume.grace) + 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 <-pc.resume.done: + return ErrReattachRejected + 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 { + pc.resume.reattachMu.Lock() + defer pc.resume.reattachMu.Unlock() + select { + case <-pc.ready: + case <-pc.resume.done: + return ErrReattachRejected + 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 + } + installed := false + defer func() { + if !installed { + // A failed upgrade or handshake leaves the receiver parked. Let the + // next attempt use that same stable delivered offset. + pc.resume.parked <- struct{}{} + } + }() + + pc.handoverMutex.Lock() + defer pc.handoverMutex.Unlock() + select { + case <-pc.resume.done: + return ErrReattachRejected + default: + } + + 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.prepareReplay(clientDelivered); err != nil { + conn.Close() + return err + } + pc.conn.Store(conn) + // Teardown can run while the handshake holds the write lock. If it closed + // the old socket before this store, this handler must close the replacement. + select { + case <-pc.resume.done: + conn.Close() + return ErrReattachRejected + default: + } + pc.requestAck() + + select { + case pc.resume.resumed <- conn: + installed = true + case <-pc.resume.done: + conn.Close() + return ErrReattachRejected + } + return nil +} + +// prepareReplay queues missing bytes under the write lock. Replaying here would +// deadlock when both peers fill their socket buffers before returning to reads. +func (pc *proxyConnection) prepareReplay(peerDelivered int64) error { + missing, err := pc.resume.sendBuf.replayFrom(peerDelivered) + if err != nil { + return err + } + pc.resume.sendBuf.ack(peerDelivered) + pc.resume.replay = missing + return nil +} + +// 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(ctx context.Context, conn *websocket.Conn) (int64, error) { + stop := context.AfterFunc(ctx, func() { conn.Close() }) + defer stop() + 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..966d0cf53c9 --- /dev/null +++ b/experimental/ssh/internal/proxy/resume_e2e_test.go @@ -0,0 +1,200 @@ +//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) { + 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") + }) + proxyServer.resumeGrace = 300 * time.Millisecond + server := httptest.NewServer(proxyServer) + defer server.Close() + + wsURL := "ws" + server.URL[4:] + conn, _, err := websocket.DefaultDialer.Dial(wsURL+"?id=abandoned&resume_version=2&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&resume_version=2&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+"&resume_version=2&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+"&resume_version=2&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..f78c1c5010e --- /dev/null +++ b/experimental/ssh/internal/proxy/resume_test.go @@ -0,0 +1,343 @@ +package proxy + +import ( + "bytes" + "context" + "errors" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + "github.com/gorilla/websocket" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestResumeRetriesADroppedHandshake(t *testing.T) { + ctx, cancel := context.WithTimeout(t.Context(), 3*time.Second) + defer cancel() + var attempts atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if !assert.NoError(t, err) { + return + } + defer conn.Close() + if attempts.Add(1) == 1 { + return + } + assert.NoError(t, conn.WriteJSON(controlMessage{})) + <-ctx.Done() + })) + defer server.Close() + defer cancel() + pc := newResumableProxyConnection(func(ctx context.Context, req DialRequest) (*websocket.Conn, error) { + conn, resp, err := websocket.DefaultDialer.DialContext(ctx, "ws"+server.URL[4:], nil) + if resp != nil { + resp.Body.Close() + } + return conn, err + }) + require.NoError(t, pc.dialReattach(ctx)) + defer func() { _ = pc.closeConnection() }() + assert.Equal(t, int32(2), attempts.Load()) +} + +func TestResumeDoesNotRetryARejectedSession(t *testing.T) { + var attempts int + pc := newResumableProxyConnection(func(ctx context.Context, req DialRequest) (*websocket.Conn, error) { + attempts++ + return nil, ErrReattachRejected + }) + ctx, cancel := context.WithTimeout(t.Context(), time.Second) + defer cancel() + err := pc.dialReattach(ctx) + assert.ErrorIs(t, err, ErrReattachRejected) + assert.NoError(t, ctx.Err()) + assert.Equal(t, 1, attempts) +} + +func TestSendBufferWaitForSpace(t *testing.T) { + for _, cancelWait := range []bool{false, true} { + t.Run(fmt.Sprintf("cancel=%t", cancelWait), func(t *testing.T) { + b := newSendBuffer(8) + require.NoError(t, b.append([]byte("12345678"))) + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + done := make(chan error, 1) + go func() { done <- b.waitForSpace(ctx, 1) }() + select { + case err := <-done: + t.Fatalf("full window did not block: %v", err) + case <-time.After(20 * time.Millisecond): + } + if cancelWait { + cancel() + } else { + b.ack(1) + } + select { + case err := <-done: + assert.Equal(t, cancelWait, errors.Is(err, context.Canceled)) + case <-time.After(time.Second): + t.Fatal("waiting producer did not wake up") + } + }) + } +} + +func TestParseDialRequestRequiresCorrectedResumeProtocol(t *testing.T) { + for _, tc := range []struct { + query string + wantResume bool + wantError bool + }{ + {"id=test", false, false}, + {"id=test&delivered=0", false, true}, + {"id=test&resume_version=1&delivered=0", false, true}, + {"id=test&resume_version=2&delivered=0", true, false}, + {"id=test&resume_version=2&delivered=-1", false, true}, + {"id=test&resume_version=2", false, true}, + {"id=test&reattach=1", false, true}, + } { + t.Run(tc.query, func(t *testing.T) { + req, err := parseDialRequest(httptest.NewRequest(http.MethodGet, "/?"+tc.query, nil)) + if tc.wantError { + require.Error(t, err) + return + } + require.NoError(t, err) + assert.Equal(t, tc.wantResume, req.ResumeCapable) + }) + } +} + +func TestResumeCancellationInterruptsHandshake(t *testing.T) { + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + ready := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if !assert.NoError(t, err) { + return + } + defer conn.Close() + close(ready) + _, _, _ = conn.ReadMessage() + })) + defer server.Close() + pc := newResumableProxyConnection(func(ctx context.Context, req DialRequest) (*websocket.Conn, error) { + conn, resp, err := websocket.DefaultDialer.DialContext(ctx, "ws"+server.URL[4:], nil) + if resp != nil { + resp.Body.Close() + } + return conn, err + }) + done := make(chan error, 1) + go func() { done <- pc.dialReattach(ctx) }() + <-ready + cancel() + select { + case err := <-done: + assert.ErrorIs(t, err, context.Canceled) + case <-time.After(time.Second): + t.Fatal("cancellation did not interrupt the handshake") + } +} + +func TestResumeReplaysFullWindowsInBothDirections(t *testing.T) { + ctx, cancel := context.WithTimeout(t.Context(), 4*time.Second) + defer cancel() + serverProxy := newResumableProxyConnection(nil) + var accepted atomic.Bool + server := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if accepted.CompareAndSwap(false, true) { + assert.NoError(t, serverProxy.accept(w, r)) + return + } + req, err := parseDialRequest(r) + if assert.NoError(t, err) { + if req.Reattach { + _ = serverProxy.acceptReattach(ctx, w, r, req.Delivered) + } else { + _ = serverProxy.acceptHandover(ctx, w, r) + } + } + })) + server.Config.ConnState = func(conn net.Conn, state http.ConnState) { + if state == http.StateHijacked { + assert.NoError(t, conn.(*net.TCPConn).SetWriteBuffer(4096)) + } + } + server.Start() + defer server.Close() + defer cancel() + sessionCtx := ctx + clientProxy := newResumableProxyConnection(func(ctx context.Context, req DialRequest) (*websocket.Conn, error) { + url := fmt.Sprintf("ws%s?id=test&resume_version=2&delivered=%d", server.URL[4:], req.Delivered) + if req.Reattach { + url += "&reattach=1" + } + dialer := websocket.Dialer{NetDialContext: func(ctx context.Context, network, addr string) (net.Conn, error) { + conn, err := (&net.Dialer{}).DialContext(ctx, network, addr) + if err == nil { + err = conn.(*net.TCPConn).SetWriteBuffer(4096) + } + return conn, err + }} + conn, resp, err := dialer.DialContext(ctx, url, nil) + if resp != nil { + resp.Body.Close() + } + if err == nil { + context.AfterFunc(sessionCtx, func() { conn.Close() }) + } + return conn, err + }) + require.NoError(t, clientProxy.connect(ctx)) + <-serverProxy.ready + defer func() { _ = serverProxy.closeConnection() }() + defer func() { _ = clientProxy.closeConnection() }() + + // Both writes failed before reaching the peer. Neither direction has room for + // new payload, and neither can finish replay unless the peer resumes reading. + toClient := bytes.Repeat([]byte("s"), proxyResumeBufferLimit) + toServer := bytes.Repeat([]byte("c"), proxyResumeBufferLimit) + require.NoError(t, serverProxy.resume.sendBuf.append(toClient)) + require.NoError(t, clientProxy.resume.sendBuf.append(toServer)) + require.NoError(t, clientProxy.closeConnection()) + + clientInput, clientWriter := io.Pipe() + serverInput, serverWriter := io.Pipe() + defer clientWriter.Close() + defer serverWriter.Close() + clientOutput := newTestBuffer(t) + serverOutput := newTestBuffer(t) + clientOutput.OnWrite = make(chan []byte, 4096) + serverOutput.OnWrite = make(chan []byte, 4096) + done := make(chan error, 2) + go func() { done <- serverProxy.start(ctx, serverInput, serverOutput) }() + go func() { done <- clientProxy.start(ctx, clientInput, clientOutput) }() + require.Eventually(t, func() bool { + return clientOutput.Contains(toClient) && serverOutput.Contains(toServer) + }, 2*time.Second, time.Millisecond, "both peers must read while replaying") + + // Resume must remain available after a transfer larger than the replay window, + // including an auth handover and another reset while data is in flight. + moreClient := bytes.Repeat([]byte("download"), proxyResumeBufferLimit) + moreServer := bytes.Repeat([]byte("upload!!"), proxyResumeBufferLimit) + writes := make(chan error, 2) + go func() { _, err := serverWriter.Write(moreClient); writes <- err }() + go func() { _, err := clientWriter.Write(moreServer); writes <- err }() + require.NoError(t, clientProxy.initiateHandover(ctx)) + require.NoError(t, clientProxy.conn.Load().UnderlyingConn().Close()) + for range 2 { + select { + case err := <-writes: + require.NoError(t, err) + case <-ctx.Done(): + t.Fatal("transfer stalled after reattach") + } + } + toClient = append(toClient, moreClient...) + toServer = append(toServer, moreServer...) + require.Eventually(t, func() bool { + return clientOutput.Contains(toClient) && serverOutput.Contains(toServer) + }, 2*time.Second, time.Millisecond, "large transfers must remain byte-exact after reconnecting") + cancel() + for range 2 { + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("proxy did not stop") + } + } + assert.Len(t, clientOutput.String(), len(toClient)) + assert.Len(t, serverOutput.String(), len(toServer)) +} + +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..1d0d9232f75 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" @@ -20,6 +22,7 @@ type proxyServer struct { ctx context.Context connections *ConnectionsManager createServerCommand createServerCommandFunc + resumeGrace time.Duration } func NewProxyServer(ctx context.Context, connections *ConnectionsManager, createServerCommand createServerCommandFunc) *proxyServer { @@ -27,6 +30,7 @@ func NewProxyServer(ctx context.Context, connections *ConnectionsManager, create ctx: ctx, connections: connections, createServerCommand: createServerCommand, + resumeGrace: proxyResumeGrace, } } @@ -36,12 +40,69 @@ 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 != "" { + if query.Get(ResumeVersionParameter) != strconv.Itoa(ResumeProtocolVersion) { + return DialRequest{}, fmt.Errorf("session resume requires %s=%d", ResumeVersionParameter, ResumeProtocolVersion) + } + 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") } + if query.Has(ResumeVersionParameter) && !req.ResumeCapable { + return DialRequest{}, errors.New("'resume_version' 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 +116,16 @@ 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) + conn.resume.grace = server.resumeGrace + } 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) @@ -106,7 +175,8 @@ func runServerProxy(ctx context.Context, proxy *proxyConnection, createServerCom g, gCtx := errgroup.WithContext(ctx) g.Go(func() error { - defer closeProxyConnection(ctx, proxy) + // Let the sending loop drain stdout and receive its final acknowledgment + // before closing the websocket, even if the process has already exited. // Waiting on the underlying Process, not the command itself. // Command.Wait needs to be called to release all resources, // but it's only safe to do so after we've finished reading from stdout, diff --git a/experimental/ssh/internal/server/server.go b/experimental/ssh/internal/server/server.go index 4356700d87d..3c28254bb3f 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]int{proxy.ResumeVersionParameter: proxy.ResumeProtocolVersion}); 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 ce2e6411102..03607853be6 100644 --- a/libs/testserver/handlers.go +++ b/libs/testserver/handlers.go @@ -848,6 +848,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 445ec66b18d2ea4a4c2278d597068f826439cd59 Mon Sep 17 00:00:00 2001 From: Anton Nekipelov <226657+anton-107@users.noreply.github.com> Date: Fri, 11 Sep 2026 16:15:09 +0000 Subject: [PATCH 2/4] Link SSH resume changelog to PR #6650 --- .nextchanges/cli/ssh-resume-backpressure.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.nextchanges/cli/ssh-resume-backpressure.md b/.nextchanges/cli/ssh-resume-backpressure.md index 3ba1785757c..75fabe3f6fa 100644 --- a/.nextchanges/cli/ssh-resume-backpressure.md +++ b/.nextchanges/cli/ssh-resume-backpressure.md @@ -1 +1 @@ -* Preserve SSH sessions across temporary tunnel disconnects, with bounded replay and backpressure for large transfers. +* Preserve SSH sessions across temporary tunnel disconnects, with bounded replay and backpressure for large transfers. ([#6650](https://github.com/databricks/cli/pull/6650)) From c2b8ea275d0803d6d2944775c9a9bf140ea418a0 Mon Sep 17 00:00:00 2001 From: Anton Nekipelov <226657+anton-107@users.noreply.github.com> Date: Sat, 12 Sep 2026 09:54:53 +0000 Subject: [PATCH 3/4] Bound the resume waits that could outlive a session, and cover them Validation of #6650 found three unbounded or unreported paths, each with a test that fails without the fix: - EOF on the source waited indefinitely for the peer to acknowledge the tail. A peer that reads without acknowledging (its own receiving loop wedged, or a half-open connection) kept the CLI alive after the ssh client that spawned it was gone; on the server side it kept sshd, a client slot and the compute alive with nobody attached. Bounded by proxyEOFDrainTimeout. - The capabilities probe had no timeout of its own, so a driver proxy that accepts the request and never answers blocked the connect path before the tunnel was even dialled. - Every session that ended before the SSH server's first byte was reported as a clean exit: proxy.start's deferred cancel always fires before g.Wait delivers the outcome, so the exit path returned nil. Deterministic, not a race - 40 of 40 runs. The exit path now waits for the outcome, bounded. Also report a tunnel that kept dropping as a drop rather than as a missing openssh-server, which is what the handshake timeout's message claims. Added tests: byte-exactness across a reset storm, no goroutine growth across repeated reattaches, client-slot retention for abandoned sessions, and both readings of the close frame's "finished" reason. Co-authored-by: Isaac --- .../ssh/internal/client/capabilities_test.go | 46 +++ experimental/ssh/internal/client/client.go | 8 +- experimental/ssh/internal/proxy/client.go | 24 +- experimental/ssh/internal/proxy/proxy.go | 38 ++- .../ssh/internal/proxy/resume_limits_test.go | 270 ++++++++++++++++++ .../ssh/internal/proxy/resume_stress_test.go | 243 ++++++++++++++++ 6 files changed, 623 insertions(+), 6 deletions(-) create mode 100644 experimental/ssh/internal/client/capabilities_test.go create mode 100644 experimental/ssh/internal/proxy/resume_limits_test.go create mode 100644 experimental/ssh/internal/proxy/resume_stress_test.go diff --git a/experimental/ssh/internal/client/capabilities_test.go b/experimental/ssh/internal/client/capabilities_test.go new file mode 100644 index 00000000000..16e253ac5f4 --- /dev/null +++ b/experimental/ssh/internal/client/capabilities_test.go @@ -0,0 +1,46 @@ +package client + +import ( + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/databricks/databricks-sdk-go" + "github.com/stretchr/testify/require" +) + +// The capabilities probe runs before the tunnel is dialled, on every proxy-mode connect. A driver +// proxy that accepts the request and never answers must not hold the connect path open: resume is +// an optimisation, and failing to learn about it has to cost bounded time, not the whole session. +func TestServerSupportsResumeGivesUpOnAStalledServer(t *testing.T) { + restore := capabilitiesProbeTimeout + capabilitiesProbeTimeout = time.Second + defer func() { capabilitiesProbeTimeout = restore }() + + stalled := make(chan struct{}) + + mux := http.NewServeMux() + mux.HandleFunc("GET /driver-proxy-api/o/123/cluster/7772/capabilities", func(w http.ResponseWriter, r *http.Request) { + <-stalled + }) + server := httptest.NewServer(mux) + // LIFO: release the stalled handler first, so Close does not wait on it. + defer server.Close() + defer close(stalled) + + client, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "test-token", WorkspaceID: "123", AuthType: "pat"}) + require.NoError(t, err) + + done := make(chan bool, 1) + go func() { + done <- serverSupportsResume(t.Context(), client, "cluster", 7772, "") + }() + + select { + case resumable := <-done: + require.False(t, resumable, "a server that never answered cannot be assumed to support resume") + case <-time.After(30 * time.Second): + t.Fatal("the capabilities probe never returned: it has no timeout of its own, so a stalled driver proxy blocks the connect path indefinitely") + } +} diff --git a/experimental/ssh/internal/client/client.go b/experimental/ssh/internal/client/client.go index 5cf0aedaaf9..048e7f06cd3 100644 --- a/experimental/ssh/internal/client/client.go +++ b/experimental/ssh/internal/client/client.go @@ -993,6 +993,9 @@ func runSSHProxy(ctx context.Context, client *databricks.WorkspaceClient, server return proxy.RunClientProxy(ctx, os.Stdin, os.Stdout, requestHandoverTick, opts.KeepaliveInterval, resumable, createConn) } +// capabilitiesProbeTimeout caps the pre-connect capabilities probe. A var so tests can shorten it. +var capabilitiesProbeTimeout = 10 * time.Second + // serverSupportsResume reports whether the running SSH server speaks the resume protocol. func serverSupportsResume(ctx context.Context, client *databricks.WorkspaceClient, clusterID string, serverPort int, liteswap string) bool { req, err := newDriverProxyRequest(ctx, client, clusterID, serverPort, "capabilities", liteswap) @@ -1000,7 +1003,10 @@ func serverSupportsResume(ctx context.Context, client *databricks.WorkspaceClien log.Debugf(ctx, "Failed to build the server capabilities request: %v", err) return false } - httpClient := &http.Client{Transport: client.Config.HTTPTransport} + // Bounded on its own: this probe runs before the tunnel is dialled, and resume is an + // optimisation, so a driver proxy that accepts the request and never answers must cost a + // bounded wait rather than the whole connect path. + httpClient := &http.Client{Transport: client.Config.HTTPTransport, Timeout: capabilitiesProbeTimeout} resp, err := httpClient.Do(req) if err != nil { log.Debugf(ctx, "Failed to query the server capabilities: %v", err) diff --git a/experimental/ssh/internal/proxy/client.go b/experimental/ssh/internal/proxy/client.go index 1ad9ebe08c2..8bab4d9501a 100644 --- a/experimental/ssh/internal/proxy/client.go +++ b/experimental/ssh/internal/proxy/client.go @@ -163,12 +163,34 @@ func RunClientProxy(ctx context.Context, src io.ReadCloser, dst io.Writer, reque case <-time.After(clientHandshakeTimeout): // cancel() (deferred) unblocks what it can; the process exits and reclaims any goroutine // still stuck on os.Stdin. ssh then fails fast instead of hanging on its ConnectTimeout. + if proxy.droppedConnections() > 0 { + // The tunnel connected and then kept dropping, so sshd is not the suspect: pointing + // the user at a missing openssh-server here sends them to the wrong place, and bills + // a network failure to the wrong telemetry category. + return errors.Join(ErrWebsocketDropped, + fmt.Errorf("the tunnel connection dropped %d time(s) before the SSH server responded", proxy.droppedConnections())) + } return errHandshakeTimeout case <-ctx.Done(): - return nil + // proxy.start cancels this context from its own deferred cancel, so it always fires just + // before g.Wait delivers the session's outcome. Without waiting for that outcome, every + // session that ends before the SSH server's first byte - a drop during the handshake, a + // reattach the server refused - is reported as a clean exit: the user sees no reason and + // telemetry records no category. Bounded so a wedged loop still cannot hang the exit. + select { + case err := <-done: + return normalizeProxyError(err) + case <-time.After(proxyExitGrace): + return nil + } } } +// proxyExitGrace bounds how long the exit path waits for the session's outcome once the context is +// cancelled. The outcome is already in flight by then, so this only matters if a proxy loop is +// wedged, and then reporting nothing beats hanging. +const proxyExitGrace = 5 * time.Second + // normalizeProxyError treats a clean finish or a context cancellation (our own exit signal, or the // user interrupting) as success; anything else is a real proxy error. func normalizeProxyError(err error) error { diff --git a/experimental/ssh/internal/proxy/proxy.go b/experimental/ssh/internal/proxy/proxy.go index 48864984993..e8d6a2f4ee3 100644 --- a/experimental/ssh/internal/proxy/proxy.go +++ b/experimental/ssh/internal/proxy/proxy.go @@ -91,6 +91,15 @@ const ( proxySessionFinished = "finished" ) +// proxyEOFDrainTimeout bounds how long the source's EOF waits for the peer to acknowledge the tail +// of the stream. A peer whose own receiving loop is wedged reads without ever acknowledging, and +// waiting for it indefinitely outlives the ssh client that spawned us - on the server side it also +// pins sshd and a client slot, which is what keeps the compute alive. The tail is at most one +// replay window and acknowledgements are sent every proxyAckInterval, so this only expires when the +// peer has genuinely stopped delivering. Long enough to cover a reattach that lands right at EOF. +// A var so tests can shorten it. +var proxyEOFDrainTimeout = proxyResumeGrace + // 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 { @@ -119,6 +128,9 @@ type resumeState struct { gate sendGate // Serializes inbound reattach attempts, including retiring the old socket. reattachMu sync.Mutex + // How many times the connection has dropped, so the exit path can tell a tunnel that kept + // dropping apart from an SSH server that never answered. The two need different messages. + drops atomic.Int64 // Closed when the session stops, so an HTTP handler cannot revive it. done chan struct{} grace time.Duration @@ -255,6 +267,14 @@ func (pc *proxyConnection) resumable() bool { return pc.resume != nil } +// droppedConnections reports how many times this session's connection has dropped. +func (pc *proxyConnection) droppedConnections() int64 { + if !pc.resumable() { + return 0 + } + return pc.resume.drops.Load() +} + func (pc *proxyConnection) start(ctx context.Context, src io.ReadCloser, dst io.Writer) error { g, gCtx := errgroup.WithContext(ctx) var finished atomic.Bool @@ -372,10 +392,19 @@ func (pc *proxyConnection) runSendingLoop(ctx context.Context, src io.Reader) er if readErr != nil { if errors.Is(readErr, io.EOF) { if pc.resumable() { - // EOF may accompany the last payload of a failed write. Keep the - // session alive until replay delivers and acknowledges those bytes. - if err := pc.resume.sendBuf.waitForSpace(ctx, proxyResumeBufferLimit); err != nil { - return err + // EOF may accompany the last payload of a failed write. Keep the session alive + // until replay delivers and acknowledges those bytes - but only for + // proxyEOFDrainTimeout: a peer that reads without ever acknowledging would + // otherwise hold the session open forever, long after the ssh client that + // spawned us is gone. Our own cancellation still ends it immediately. + drainCtx, cancel := context.WithTimeout(ctx, proxyEOFDrainTimeout) + err := pc.resume.sendBuf.waitForSpace(drainCtx, proxyResumeBufferLimit) + cancel() + if err != nil && ctx.Err() != nil { + return ctx.Err() + } + if err != nil { + log.Warnf(ctx, "The peer never acknowledged the last bytes of the session: %v", err) } } return errors.Join(errProxyEOF, readErr) @@ -536,6 +565,7 @@ func (pc *proxyConnection) runReceivingLoop(ctx context.Context, dst io.Writer) // outlives the connection, so reattach instead of ending the session. if pc.resumable() { conn.Close() + pc.resume.drops.Add(1) if resumeErr := pc.reattach(ctx); resumeErr != nil { return errors.Join(ErrWebsocketDropped, fmt.Errorf("failed to reattach after the connection dropped: %w", resumeErr)) } diff --git a/experimental/ssh/internal/proxy/resume_limits_test.go b/experimental/ssh/internal/proxy/resume_limits_test.go new file mode 100644 index 00000000000..0b3077b0e6d --- /dev/null +++ b/experimental/ssh/internal/proxy/resume_limits_test.go @@ -0,0 +1,270 @@ +//go:build !windows + +package proxy + +import ( + "context" + "net" + "net/http" + "net/http/httptest" + "os/exec" + "sync/atomic" + "testing" + "time" + + "github.com/databricks/cli/libs/cmdio" + "github.com/gorilla/websocket" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// blackHoleServer accepts a resume-capable websocket, drains every frame the client sends and never +// acknowledges any of it. That is what a peer whose own receiving loop is wedged looks like from +// here: the socket is alive and writes succeed, but its delivered offset never moves. +func blackHoleServer(t *testing.T) *httptest.Server { + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if err != nil { + return + } + defer conn.Close() + for { + if _, _, err := conn.ReadMessage(); err != nil { + return + } + } + })) +} + +// dropAtOnceServer accepts the upgrade and immediately resets the connection, so the session fails +// before the SSH server's first byte could ever arrive. +func dropAtOnceServer(t *testing.T) *httptest.Server { + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if err != nil { + return + } + if tcpConn, ok := conn.UnderlyingConn().(*net.TCPConn); ok { + _ = tcpConn.SetLinger(0) + } + conn.Close() + })) +} + +// serverWithConnections is createTestServer with the connections manager exposed, so a test can +// assert the server released the session, and with the command it runs supplied by the caller. +func serverWithConnections(t *testing.T, resumeGrace time.Duration, command createServerCommandFunc) (*httptest.Server, *ConnectionsManager) { + ctx := cmdio.MockDiscard(t.Context()) + connections := NewConnectionsManager(1, time.Hour) + proxyServer := NewProxyServer(ctx, connections, command) + proxyServer.resumeGrace = resumeGrace + return httptest.NewServer(proxyServer), connections +} + +// dropWebsocket ends a websocket the way a reset does, with no close handshake, so the peer treats +// it as a drop rather than an orderly shutdown. +func dropWebsocket(t *testing.T, conn *websocket.Conn) { + tcpConn, ok := conn.UnderlyingConn().(*net.TCPConn) + require.True(t, ok) + require.NoError(t, tcpConn.SetLinger(0)) + require.NoError(t, tcpConn.Close()) +} + +// EOF on the source is how every clean exit begins: ssh closes the proxy command's stdin. Waiting +// for the peer to acknowledge the tail is right, but it has to be bounded - a peer that reads +// without ever acknowledging would otherwise keep the CLI alive after the ssh client that spawned +// it is gone. +func TestEOFEndsTheSessionWhenThePeerStopsAcknowledging(t *testing.T) { + restore := proxyEOFDrainTimeout + proxyEOFDrainTimeout = 500 * time.Millisecond + t.Cleanup(func() { proxyEOFDrainTimeout = restore }) + + server := blackHoleServer(t) + defer server.Close() + + srcWriter, done := runResumableClientTo(t, server.URL, &sink{}) + _, err := srcWriter.Write([]byte("hello\n")) + require.NoError(t, err) + // Let the write reach the peer, so the bytes are unacknowledged rather than unsent. + time.Sleep(200 * time.Millisecond) + require.NoError(t, srcWriter.Close()) + + select { + case <-done: + case <-time.After(30 * time.Second): + t.Fatal("the client never exited after EOF: it is waiting for an acknowledgment that will never arrive") + } +} + +// The same wait on the server side, where it costs more: the session holds a client slot and keeps +// the shutdown timer cancelled, so a session that is never released keeps the compute alive with +// nobody attached to it. +func TestServerReleasesASessionWhenTheClientStopsAcknowledging(t *testing.T) { + restore := proxyEOFDrainTimeout + proxyEOFDrainTimeout = 500 * time.Millisecond + t.Cleanup(func() { proxyEOFDrainTimeout = restore }) + + // sshd's stand-in prints and exits, so the server's sending loop reaches EOF straight away. + server, connections := serverWithConnections(t, 500*time.Millisecond, func(ctx context.Context) *exec.Cmd { + return exec.CommandContext(ctx, "sh", "-c", "echo hello; exit 0") + }) + defer server.Close() + + wsURL := "ws" + server.URL[4:] + conn, _, err := websocket.DefaultDialer.Dial(wsURL+"?id=no-acks&resume_version=2&delivered=0", nil) // nolint:bodyclose + require.NoError(t, err) + defer conn.Close() + + // Read everything so the server's writes never block, but acknowledge nothing. + go func() { + for { + if _, _, err := conn.ReadMessage(); err != nil { + return + } + } + }() + + require.Eventually(t, func() bool { + return connections.Count() == 0 + }, 30*time.Second, 50*time.Millisecond, + "the server never released the session: its sending loop is still waiting for an acknowledgment of the bytes sshd printed before it exited") +} + +// A session that dies before the SSH server's first byte must still report why. That error is what +// the user sees instead of a bare ssh failure, and what telemetry attributes the drop from, so +// losing it turns a failed session into an apparently clean exit. +func TestASessionThatFailsBeforeTheFirstByteStillReportsTheFailure(t *testing.T) { + server := dropAtOnceServer(t) + defer server.Close() + + const attempts = 20 + silent := 0 + for range attempts { + srcWriter, done := runPlainClientTo(t, server.URL, &sink{}) + select { + case err := <-done: + if err == nil { + silent++ + } + case <-time.After(30 * time.Second): + t.Fatal("the session neither ended nor reported a failure") + } + srcWriter.Close() + } + assert.Equal(t, 0, silent, + "%d of %d sessions that died before the first byte returned no error at all", silent, attempts) +} + +// What a user is told when the tunnel connects and then keeps dropping before sshd answers. The +// handshake timeout's own message blames a missing openssh-server, which for a network failure +// sends them to the wrong place - and bills it to the wrong telemetry category. +func TestEarlyDropIsReportedAsADropNotAMissingSshd(t *testing.T) { + server := dropAtOnceServer(t) + defer server.Close() + + srcWriter, done := runResumableClientTo(t, server.URL, &sink{}) + defer srcWriter.Close() + + select { + case err := <-done: + require.Error(t, err) + assert.ErrorIs(t, err, ErrWebsocketDropped) + assert.NotErrorIs(t, err, errHandshakeTimeout) + case <-time.After(3 * time.Minute): + t.Fatal("the client never gave up on a connection that drops on every attempt") + } +} + +// A dropped session holds its client slot for the whole grace period, and a user whose network +// flapped long enough for ssh to give up comes back as a new session. This pins the cost: with the +// slots full, that new session is refused until the abandoned ones expire. +func TestAbandonedSessionsHoldTheirClientSlots(t *testing.T) { + server, _ := serverWithConnections(t, 5*time.Second, func(ctx context.Context) *exec.Cmd { + return exec.CommandContext(ctx, "cat", "-u") + }) + defer server.Close() + + wsURL := "ws" + server.URL[4:] + conn, _, err := websocket.DefaultDialer.Dial(wsURL+"?id=first&resume_version=2&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) + dropWebsocket(t, conn) + + _, resp, err := websocket.DefaultDialer.Dial(wsURL+"?id=second&resume_version=2&delivered=0", nil) // nolint:bodyclose + if resp != nil { + defer resp.Body.Close() + } + require.Error(t, err, "the abandoned session must still hold the only slot") + assert.Equal(t, http.StatusServiceUnavailable, resp.StatusCode) +} + +// closingServer answers the first dial, sends one payload byte so the client is past its handshake +// window, then ends the session with a normal-closure close frame carrying the given reason. A +// later reattach is refused the way a server that has torn the session down refuses one. +func closingServer(t *testing.T, reason string) (*httptest.Server, *atomic.Int64) { + var reattachAttempts atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Query().Get("reattach") == "1" { + reattachAttempts.Add(1) + http.Error(w, "Session no longer exists", http.StatusGone) + return + } + conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if err != nil { + return + } + defer conn.Close() + // A handler cannot fail the test directly; a write that fails here shows up as the client + // never seeing the close frame, which its own assertions report. + if err := conn.WriteMessage(websocket.BinaryMessage, []byte("x")); err != nil { + return + } + if err := conn.WriteMessage(websocket.CloseMessage, + websocket.FormatCloseMessage(websocket.CloseNormalClosure, reason)); err != nil { + return + } + time.Sleep(200 * time.Millisecond) + })) + return server, &reattachAttempts +} + +// On a resumable connection the "finished" reason in a close frame is the only thing separating a +// session that ended from a socket that died: a genuine drop arrives as a read error, never as a +// close frame, so the code cannot tell them apart by anything else. That makes every clean exit +// depend on the reason text surviving the path between the client and the workspace. Both readings +// are pinned here, so a change to either side of that contract fails loudly rather than turning +// every clean exit into a reported drop. +func TestNormalClosureReasonDecidesCleanExitVersusDrop(t *testing.T) { + t.Run("the finished reason ends the session cleanly", func(t *testing.T) { + server, reattachAttempts := closingServer(t, proxySessionFinished) + defer server.Close() + + srcWriter, done := runResumableClientTo(t, server.URL, &sink{}) + defer srcWriter.Close() + select { + case err := <-done: + assert.NoError(t, err) + case <-time.After(30 * time.Second): + t.Fatal("the client did not notice the session had finished") + } + assert.Zero(t, reattachAttempts.Load(), "a finished session must not be reattached to") + }) + + t.Run("without the reason the same close frame is a drop", func(t *testing.T) { + server, reattachAttempts := closingServer(t, "") + defer server.Close() + + srcWriter, done := runResumableClientTo(t, server.URL, &sink{}) + defer srcWriter.Close() + select { + case err := <-done: + assert.ErrorIs(t, err, ErrWebsocketDropped) + case <-time.After(30 * time.Second): + t.Fatal("the client hung on a normal closure") + } + assert.Equal(t, int64(1), reattachAttempts.Load(), + "the client tries to reattach, so a stripped reason would warn the user and count a dropped tunnel on every clean exit") + }) +} diff --git a/experimental/ssh/internal/proxy/resume_stress_test.go b/experimental/ssh/internal/proxy/resume_stress_test.go new file mode 100644 index 00000000000..ff67d559cb3 --- /dev/null +++ b/experimental/ssh/internal/proxy/resume_stress_test.go @@ -0,0 +1,243 @@ +//go:build !windows + +package proxy + +import ( + "bytes" + "context" + "errors" + "fmt" + "io" + "runtime" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/databricks/cli/libs/cmdio" + "github.com/gorilla/websocket" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// sink collects everything the proxy writes to its destination. Unlike testBuffer it has no +// bounded channel, which a multi-megabyte transfer would fill and then block the receiving loop +// on - the stall would be the test's, not the proxy's. +type sink struct { + mu sync.Mutex + buf bytes.Buffer +} + +func (s *sink) Write(p []byte) (int, error) { + s.mu.Lock() + defer s.mu.Unlock() + return s.buf.Write(p) +} + +func (s *sink) Len() int { + s.mu.Lock() + defer s.mu.Unlock() + return s.buf.Len() +} + +func (s *sink) Bytes() []byte { + s.mu.Lock() + defer s.mu.Unlock() + return append([]byte(nil), s.buf.Bytes()...) +} + +// runResumableClientTo runs RunClientProxy against url with resume negotiated, writing what it +// receives to dst. It returns the writer standing in for ssh's stdin and the session's outcome. +func runResumableClientTo(t *testing.T, url string, dst io.Writer) (io.WriteCloser, <-chan error) { + ctx := cmdio.MockDiscard(t.Context()) + wsURL := "ws" + url[4:] + createConn := func(ctx context.Context, dial DialRequest) (*websocket.Conn, error) { + u := fmt.Sprintf("%s?id=%s&resume_version=%d&delivered=%d", wsURL, dial.ConnID, ResumeProtocolVersion, dial.Delivered) + if dial.Reattach { + u += "&reattach=1" + } + conn, resp, err := websocket.DefaultDialer.DialContext(ctx, u, nil) // nolint:bodyclose + if resp != nil { + resp.Body.Close() + } + // Mirrors what createWebsocketConnection does in internal/client: a 4xx answer to a + // reattach is a permanent refusal, so the redial loop stops instead of spending its whole + // budget. Without it a test that refuses a reattach waits out proxyResumeBudget. + if err != nil && dial.Reattach && resp != nil && resp.StatusCode >= 400 && resp.StatusCode < 500 { + return nil, errors.Join(ErrReattachRejected, err) + } + return conn, err + } + src, srcWriter := io.Pipe() + done := make(chan error, 1) + go func() { + done <- RunClientProxy(ctx, src, dst, neverTick, time.Hour, true, createConn) + }() + return srcWriter, done +} + +// runPlainClientTo is runResumableClientTo without resume, the path an older server gets. +func runPlainClientTo(t *testing.T, url string, dst io.Writer) (io.WriteCloser, <-chan error) { + ctx := cmdio.MockDiscard(t.Context()) + wsURL := "ws" + url[4:] + createConn := func(ctx context.Context, dial DialRequest) (*websocket.Conn, error) { + conn, _, err := websocket.DefaultDialer.DialContext(ctx, wsURL+"?id="+dial.ConnID, nil) // nolint:bodyclose + return conn, err + } + src, srcWriter := io.Pipe() + done := make(chan error, 1) + go func() { + done <- RunClientProxy(ctx, src, dst, neverTick, time.Hour, false, createConn) + }() + return srcWriter, done +} + +// firstDifference reports where two byte streams diverge, so a failure names an offset instead of +// dumping megabytes. Returns -1 when want is a prefix of got. +func firstDifference(got, want []byte) int { + for i := range min(len(got), len(want)) { + if got[i] != want[i] { + return i + } + } + if len(got) < len(want) { + return len(got) + } + return -1 +} + +// context reports the bytes around an offset, for a failure message. +func context16(b []byte, at int) string { + return string(b[max(at-16, 0):min(at+16, len(b))]) +} + +// resetClientLegs is resetClients without its assertions: a handover closes connections behind the +// relay's back, so a leg that is already gone cannot be reset and is not a failure. +func (r *tcpRelay) resetClientLegs() { + r.mu.Lock() + conns := r.clientConns + r.clientConns = nil + r.mu.Unlock() + for _, conn := range conns { + if conn.SetLinger(0) == nil { + conn.Close() + } + } +} + +// One reset at a quiet moment is the easy case. Customers report drops several times a session, on +// networks that flap, so resets have to be survivable wherever they land - including while a replay +// is still being written. A single lost or duplicated byte fails SSH's MAC (RFC 4253 section 6) and +// disconnects the session the resume exists to save, so the whole stream is compared, not sampled. +func TestResumeStaysByteExactUnderAResetStorm(t *testing.T) { + server := createTestServer(t, 4, time.Hour) + defer server.Close() + relay := newTCPRelay(t, server.Listener.Addr().String()) + + out := &sink{} + srcWriter, done := runResumableClientTo(t, relay.URL(), out) + + const chunks = 512 + const chunkSize = 8 * 1024 + var sent []byte + for i := range chunks { + sent = append(sent, bytes.Repeat([]byte{byte('A' + i%26)}, chunkSize)...) + } + + storm := make(chan struct{}) + stormDone := make(chan struct{}) + var resets atomic.Int64 + go func() { + defer close(stormDone) + for { + select { + case <-storm: + return + case <-time.After(7 * time.Millisecond): + relay.resetClientLegs() + resets.Add(1) + } + } + }() + stop := func() { + close(storm) + <-stormDone + } + + writeDone := make(chan error, 1) + go func() { + // Paced, so the resets land inside the transfer rather than after it. A terminal or an scp + // stream is not delivered as one instantaneous burst either. + for i := range chunks { + if _, err := srcWriter.Write(sent[i*chunkSize : (i+1)*chunkSize]); err != nil { + writeDone <- err + return + } + time.Sleep(5 * time.Millisecond) + } + writeDone <- nil + }() + + deadline := time.After(60 * time.Second) + for out.Len() < len(sent) { + select { + case err := <-done: + stop() + t.Fatalf("session ended after %d resets with %d/%d bytes echoed: %v", resets.Load(), out.Len(), len(sent), err) + case <-deadline: + stop() + t.Fatalf("transfer stalled after %d resets with %d/%d bytes echoed", resets.Load(), out.Len(), len(sent)) + case <-time.After(20 * time.Millisecond): + } + } + stop() + require.NoError(t, <-writeDone) + + got := out.Bytes() + if diff := firstDifference(got, sent); diff >= 0 { + t.Fatalf("the echoed stream diverges from what was sent at offset %d (after %d resets): got %q, want %q", + diff, resets.Load(), context16(got, diff), context16(sent, diff)) + } + assert.Len(t, got, len(sent), "the echoed stream must not carry extra bytes") +} + +// Every reset spawns a reattach, and a session on a flapping network gets hundreds of them. +// Goroutines that outlive their reattach would accumulate for the life of the session. +func TestRepeatedResetsDoNotLeakGoroutines(t *testing.T) { + server := createTestServer(t, 4, time.Hour) + defer server.Close() + relay := newTCPRelay(t, server.Listener.Addr().String()) + + out := &sink{} + srcWriter, done := runResumableClientTo(t, relay.URL(), out) + + // Settle the session first, so the baseline covers only what the resets add. + _, err := srcWriter.Write([]byte("warmup\n")) + require.NoError(t, err) + require.Eventually(t, func() bool { return out.Len() >= len("warmup\n") }, 10*time.Second, 10*time.Millisecond) + time.Sleep(500 * time.Millisecond) + baseline := runtime.NumGoroutine() + + const rounds = 30 + for i := range rounds { + relay.resetClientLegs() + line := []byte("after reset\n") + want := out.Len() + len(line) + _, err := srcWriter.Write(line) + require.NoError(t, err) + require.Eventually(t, func() bool { return out.Len() >= want }, 20*time.Second, 10*time.Millisecond, + "round %d did not survive its reset", i) + } + + select { + case err := <-done: + t.Fatalf("session ended during the resets: %v", err) + default: + } + + // Reattach goroutines retire asynchronously, so give them a moment. + time.Sleep(time.Second) + assert.Less(t, runtime.NumGoroutine()-baseline, rounds, + "goroutine count grew from %d to %d across %d resets, which is close to one per reset", + baseline, runtime.NumGoroutine(), rounds) +} From 29da0875dab8748735f21164f116e70c81b802e0 Mon Sep 17 00:00:00 2001 From: Anton Nekipelov <226657+anton-107@users.noreply.github.com> Date: Tue, 15 Sep 2026 07:29:55 +0000 Subject: [PATCH 4/4] Fix SSH resume recovery during handover and HTTP timeouts --- .../ssh/internal/client/websockets.go | 4 +- .../ssh/internal/client/websockets_test.go | 65 ++++++++- experimental/ssh/internal/proxy/client.go | 6 + experimental/ssh/internal/proxy/proxy.go | 66 +++++---- experimental/ssh/internal/proxy/proxy_test.go | 129 ++++++++++++++++++ 5 files changed, 235 insertions(+), 35 deletions(-) diff --git a/experimental/ssh/internal/client/websockets.go b/experimental/ssh/internal/client/websockets.go index b3b65c2a526..51cd28c50fe 100644 --- a/experimental/ssh/internal/client/websockets.go +++ b/experimental/ssh/internal/client/websockets.go @@ -38,7 +38,9 @@ func createWebsocketConnection(ctx context.Context, client *databricks.Workspace resp.Body.Close() } if err != nil { - if dial.Reattach && resp != nil && resp.StatusCode >= 400 && resp.StatusCode < 500 && resp.StatusCode != http.StatusTooManyRequests { + // Only the server's explicit session refusals bypass the retry budget. + // Other responses, including 408, can come from an intermediary. + if dial.Reattach && resp != nil && (resp.StatusCode == http.StatusConflict || resp.StatusCode == http.StatusGone) { return nil, errors.Join(proxy.ErrReattachRejected, fmt.Errorf("reattach failed (HTTP %d): %w", resp.StatusCode, err)) } return nil, fmt.Errorf("failed to establish websocket connection: %w", err) diff --git a/experimental/ssh/internal/client/websockets_test.go b/experimental/ssh/internal/client/websockets_test.go index b3689a40ec6..05d15636688 100644 --- a/experimental/ssh/internal/client/websockets_test.go +++ b/experimental/ssh/internal/client/websockets_test.go @@ -1,14 +1,18 @@ package client import ( + "bytes" "context" + "io" "net/http" "net/http/httptest" + "sync/atomic" "testing" "time" "github.com/databricks/cli/experimental/ssh/internal/proxy" "github.com/databricks/databricks-sdk-go" + "github.com/gorilla/websocket" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -42,11 +46,22 @@ func TestServerSupportsResume(t *testing.T) { } func TestCreateWebsocketConnectionReattachRejected(t *testing.T) { - for _, status := range []int{http.StatusGone, http.StatusConflict, http.StatusUnauthorized, http.StatusTooManyRequests, http.StatusServiceUnavailable} { - t.Run(http.StatusText(status), func(t *testing.T) { + for _, tc := range []struct { + status int + rejected bool + }{ + {http.StatusGone, true}, + {http.StatusConflict, true}, + {http.StatusUnauthorized, false}, + {http.StatusForbidden, false}, + {http.StatusRequestTimeout, false}, + {http.StatusTooManyRequests, false}, + {http.StatusServiceUnavailable, false}, + } { + t.Run(http.StatusText(tc.status), func(t *testing.T) { mux := http.NewServeMux() mux.HandleFunc("GET /driver-proxy-api/o/123/cluster/7772/ssh", func(w http.ResponseWriter, r *http.Request) { - w.WriteHeader(status) + w.WriteHeader(tc.status) }) server := httptest.NewServer(mux) defer server.Close() @@ -56,7 +71,7 @@ func TestCreateWebsocketConnectionReattachRejected(t *testing.T) { defer cancel() _, err = createWebsocketConnection(ctx, client, proxy.DialRequest{ConnID: "test", ResumeCapable: true, Reattach: true}, "cluster", 7772, "") require.Error(t, err) - if status < 500 && status != http.StatusTooManyRequests { + if tc.rejected { assert.ErrorIs(t, err, proxy.ErrReattachRejected) } else { assert.NotErrorIs(t, err, proxy.ErrReattachRejected) @@ -65,6 +80,48 @@ func TestCreateWebsocketConnectionReattachRejected(t *testing.T) { } } +func TestCreateWebsocketConnectionRetriesRequestTimeout(t *testing.T) { + var reattachments atomic.Int32 + mux := http.NewServeMux() + mux.HandleFunc("GET /driver-proxy-api/o/123/cluster/7772/ssh", func(w http.ResponseWriter, r *http.Request) { + reattach := r.URL.Query().Get("reattach") == "1" + if reattach && reattachments.Add(1) == 1 { + w.WriteHeader(http.StatusRequestTimeout) + return + } + conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if !assert.NoError(t, err) { + return + } + defer conn.Close() + if !reattach { + assert.NoError(t, conn.WriteMessage(websocket.BinaryMessage, []byte("before"))) + return + } + assert.Equal(t, "6", r.URL.Query().Get("delivered")) + assert.NoError(t, conn.WriteJSON(map[string]int{"delivered": 0})) + assert.NoError(t, conn.WriteMessage(websocket.BinaryMessage, []byte("after"))) + assert.NoError(t, conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "finished"))) + }) + server := httptest.NewServer(mux) + defer server.Close() + client, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "test-token", WorkspaceID: "123", AuthType: "pat"}) + require.NoError(t, err) + ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second) + defer cancel() + input, writer := io.Pipe() + defer writer.Close() + var output bytes.Buffer + err = proxy.RunClientProxy(ctx, input, &output, func() <-chan time.Time { return nil }, time.Hour, true, + func(ctx context.Context, req proxy.DialRequest) (*websocket.Conn, error) { + return createWebsocketConnection(ctx, client, req, "cluster", 7772, "") + }) + require.NoError(t, err) + assert.NoError(t, ctx.Err()) + assert.Equal(t, int32(2), reattachments.Load()) + assert.Equal(t, "beforeafter", output.String()) +} + func TestBuildProxyWebsocketURL(t *testing.T) { tests := []struct { name string diff --git a/experimental/ssh/internal/proxy/client.go b/experimental/ssh/internal/proxy/client.go index 8bab4d9501a..8d6a12fb438 100644 --- a/experimental/ssh/internal/proxy/client.go +++ b/experimental/ssh/internal/proxy/client.go @@ -108,6 +108,12 @@ func RunClientProxy(ctx context.Context, src io.ReadCloser, dst io.Writer, reque log.Debugf(gCtx, "Could not open a replacement connection for the auth handover, staying on the current one: %v", err) continue } + if resumable { + // The failed handover closes its sockets. The receiving loop + // reattaches after the handover releases the write lock. + log.Debugf(gCtx, "Auth handover failed, recovering through session reattachment: %v", err) + continue + } return errors.Join(ErrHandoverFailed, err) } } diff --git a/experimental/ssh/internal/proxy/proxy.go b/experimental/ssh/internal/proxy/proxy.go index e8d6a2f4ee3..0cd5e73308c 100644 --- a/experimental/ssh/internal/proxy/proxy.go +++ b/experimental/ssh/internal/proxy/proxy.go @@ -528,51 +528,54 @@ func (pc *proxyConnection) runReceivingLoop(ctx context.Context, dst io.Writer) conn := pc.conn.Load() mt, data, err := conn.ReadMessage() if err != nil { - // During handover a normal closure is expected, but any other error must stop the read loop (and eventually terminate the ssh session). + // A normal closure completes handover. An interrupted handover needs + // reattachment to recover bytes the close-frame exchange did not drain. if handover := pc.handoverState.Load(); handover != nil { var closeConnSignal error if !websocket.IsCloseError(err, websocket.CloseNormalClosure) { 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 { - return err - } + handoverErr := handover.signalConnectionClosed(closeConnSignal) // Wait for the handover initiator to swap the connection. // While we wait for the handover to complete, the new connection might be getting incoming messages. // They will be buffered by the TCP stack and will be read by us after the handover is complete. - if err := handover.waitForConnectionToSwap(); err != nil { - return err + if handoverErr == nil { + handoverErr = handover.waitForConnectionToSwap() } // Continue with the receiving loop, pc.conn is now the new connection. - continue - } else { - closeErr, closed := errors.AsType[*websocket.CloseError](err) - finished := closed && closeErr.Code == websocket.CloseNormalClosure && closeErr.Text == proxySessionFinished - if finished || (!pc.resumable() && (errors.Is(err, io.EOF) || websocket.IsCloseError(err, websocket.CloseNormalClosure))) { - return errors.Join(errProxyEOF, err) + if handoverErr == nil { + continue } - // 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() + if !pc.resumable() { + return handoverErr } - // 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() { - conn.Close() - pc.resume.drops.Add(1) - if resumeErr := pc.reattach(ctx); resumeErr != nil { - return errors.Join(ErrWebsocketDropped, fmt.Errorf("failed to reattach after the connection dropped: %w", resumeErr)) - } - continue + } + closeErr, closed := errors.AsType[*websocket.CloseError](err) + finished := closed && closeErr.Code == websocket.CloseNormalClosure && closeErr.Text == proxySessionFinished + if finished || (!pc.resumable() && (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() { + conn.Close() + pc.resume.drops.Add(1) + if resumeErr := pc.reattach(ctx); resumeErr != nil { + return errors.Join(ErrWebsocketDropped, fmt.Errorf("failed to reattach after the connection dropped: %w", resumeErr)) } - return errors.Join(ErrWebsocketDropped, fmt.Errorf("failed to read from websocket: %w", err)) + continue } + return errors.Join(ErrWebsocketDropped, fmt.Errorf("failed to read from websocket: %w", err)) } if mt == websocket.TextMessage && pc.resumable() { @@ -681,6 +684,7 @@ func (pc *proxyConnection) initiateHandover(ctx context.Context) error { // (it does so when it receives an /ssh request with known connection ID and starts AcceptHandover process). // Receiving loop will signal about closed connection to the coord.connClosed channel. if err := handoverState.waitForConnectionToClose(); err != nil { + pc.conn.Load().Close() newConn.Close() return err } @@ -733,6 +737,7 @@ func (pc *proxyConnection) acceptHandover(ctx context.Context, w http.ResponseWr } err = currentConn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "handover")) if err != nil { + currentConn.Close() newConn.Close() return fmt.Errorf("failed to send close message to the current connection: %w", err) } @@ -741,6 +746,7 @@ func (pc *proxyConnection) acceptHandover(ctx context.Context, w http.ResponseWr // On the client its done automatically by the websocket library with the default close handler. // On the server we then receive a close error in the RunReceivingLoop and signal about it to the coord.connClosed channel. if err := handoverState.waitForConnectionToClose(); err != nil { + currentConn.Close() newConn.Close() return err } diff --git a/experimental/ssh/internal/proxy/proxy_test.go b/experimental/ssh/internal/proxy/proxy_test.go index d31cf70eae5..352161b9d7c 100644 --- a/experimental/ssh/internal/proxy/proxy_test.go +++ b/experimental/ssh/internal/proxy/proxy_test.go @@ -271,6 +271,135 @@ func TestConnectionHandover(t *testing.T) { } } +// TestResumeAfterHandoverDrop keeps the same byte streams alive when the old +// connection drops before the handover's close-frame exchange finishes. +func TestResumeAfterHandoverDrop(t *testing.T) { + for _, tc := range []struct { + name string + dropDuringUpgrade bool + failDial bool + loseCloseReply bool + }{ + {name: "during upgrade", dropDuringUpgrade: true}, + {name: "failed dial", failDial: true}, + {name: "lost close reply", loseCloseReply: true}, + } { + t.Run(tc.name, func(t *testing.T) { + ctx, cancel := context.WithTimeout(t.Context(), 8*time.Second) + defer cancel() + serverProxy := newResumableProxyConnection(nil) + serverInput, serverWriter := io.Pipe() + defer serverWriter.Close() + serverOutput := newTestBuffer(t) + done := make(chan error, 2) + var accepted atomic.Bool + var reattachments atomic.Int32 + server := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if accepted.CompareAndSwap(false, true) { + if !assert.NoError(t, serverProxy.accept(w, r)) { + return + } + done <- serverProxy.start(ctx, serverInput, serverOutput) + return + } + req, err := parseDialRequest(r) + if !assert.NoError(t, err) { + return + } + if req.Reattach { + reattachments.Add(1) + assert.NoError(t, serverProxy.acceptReattach(ctx, w, r, req.Delivered)) + return + } + _ = serverProxy.acceptHandover(ctx, w, r) + })) + var upgrades atomic.Int32 + server.Config.ConnState = func(_ net.Conn, state http.ConnState) { + if state == http.StateHijacked && upgrades.Add(1) == 2 && tc.dropDuringUpgrade { + // Both handover coordinators exist before the replacement is upgraded. + assert.NoError(t, serverProxy.conn.Load().Close()) + } + } + server.Start() + defer server.Close() + defer cancel() + + clientInput, clientWriter := io.Pipe() + defer clientWriter.Close() + clientOutput := newTestBuffer(t) + ticks := make(chan time.Time, 1) + var initial *websocket.Conn + createConn := func(ctx context.Context, req DialRequest) (*websocket.Conn, error) { + if initial != nil && !req.Reattach && tc.failDial { + assert.NoError(t, initial.Close()) + return nil, errors.New("handover dial interrupted") + } + url := fmt.Sprintf("ws%s?id=%s&resume_version=%d&delivered=%d", server.URL[4:], req.ConnID, ResumeProtocolVersion, req.Delivered) + if req.Reattach { + url += "&reattach=1" + } + conn, resp, err := websocket.DefaultDialer.DialContext(ctx, url, nil) + if resp != nil { + resp.Body.Close() + } + if err == nil && initial == nil { + initial = conn + if tc.loseCloseReply { + conn.SetCloseHandler(func(int, string) error { + _ = conn.Close() + return net.ErrClosed + }) + } + } + return conn, err + } + go func() { + done <- RunClientProxy(ctx, clientInput, clientOutput, func() <-chan time.Time { return ticks }, time.Hour, true, createConn) + }() + + beforeClient := []byte("before handover to client") + beforeServer := []byte("before handover to server") + _, err := serverWriter.Write(beforeClient) + require.NoError(t, err) + require.NoError(t, clientOutput.WaitForWrite(beforeClient)) + _, err = clientWriter.Write(beforeServer) + require.NoError(t, err) + require.NoError(t, serverOutput.WaitForWrite(beforeServer)) + ticks <- time.Now() + require.Eventually(t, func() bool { return reattachments.Load() > 0 }, 3*time.Second, time.Millisecond, + "a drop during handover must reattach the session") + + afterClient := bytes.Repeat([]byte("download"), 8192) + afterServer := bytes.Repeat([]byte("upload!!"), 8192) + _, err = serverWriter.Write(afterClient) + require.NoError(t, err) + _, err = clientWriter.Write(afterServer) + require.NoError(t, err) + require.NoError(t, clientOutput.WaitForWrite(afterClient)) + require.NoError(t, serverOutput.WaitForWrite(afterServer)) + select { + case err := <-done: + t.Fatalf("session ended before cancellation: %v", err) + default: + } + cancel() + for range 2 { + select { + case err := <-done: + // Cancellation closes the source to unblock the sending loop. + if err != nil && !errors.Is(err, io.ErrClosedPipe) { + assert.ErrorIs(t, err, context.Canceled) + } + case <-time.After(time.Second): + t.Fatal("proxy did not stop after cancellation") + } + } + assert.Equal(t, string(append(beforeClient, afterClient...)), clientOutput.String()) + assert.Equal(t, string(append(beforeServer, afterServer...)), serverOutput.String()) + }) + } +} + // 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