From 92abea7637059f5eca28321dca190b264895a5a5 Mon Sep 17 00:00:00 2001 From: Rada Kamysheva Date: Tue, 14 Jul 2026 07:33:45 +0000 Subject: [PATCH 1/2] acc: converge ssh/connection tunnel test onto a local websocket route Follow-up to #5907. The dedicated-cluster acceptance test only asserted the bootstrap job the CLI submits; the actual driver-proxy tunnel was covered separately by the Go proxy echo tests because the client hardcoded a wss:// scheme and the test server had no websocket route. Two changes let the local acceptance run exercise the tunnel end to end: - Product fix: derive the websocket scheme from the workspace host (http -> ws, else wss) instead of forcing wss, so the client can dial the plaintext local test server. Add a unit test for the URL builder. - Test server: add a raw (connection-hijacking) handler registration and a driver-proxy /ssh websocket route that echoes binary frames, standing in for the SSH tunnel server a real cluster runs. In-process, so it also works on Windows. acceptance/ssh/connection now runs a single `ssh connect --proxy` locally: it submits the same dedicated-cluster bootstrap job and pipes stdin/stdout over the tunnel, asserting both the job payload and the echoed bytes. Proxy mode skips the ssh subprocess, so it doesn't hang looping a handshake against the echo. --- acceptance/ssh/connection/out.proxy.txt | 1 + acceptance/ssh/connection/out.stdout.txt | 0 acceptance/ssh/connection/script | 16 ++++--- .../ssh/internal/client/websockets.go | 34 ++++++++++++--- .../ssh/internal/client/websockets_test.go | 35 ++++++++++++++++ libs/testserver/handlers.go | 4 ++ libs/testserver/router.go | 15 +++++++ libs/testserver/ssh.go | 42 +++++++++++++++++++ 8 files changed, 136 insertions(+), 11 deletions(-) create mode 100644 acceptance/ssh/connection/out.proxy.txt delete mode 100644 acceptance/ssh/connection/out.stdout.txt create mode 100644 experimental/ssh/internal/client/websockets_test.go create mode 100644 libs/testserver/ssh.go diff --git a/acceptance/ssh/connection/out.proxy.txt b/acceptance/ssh/connection/out.proxy.txt new file mode 100644 index 00000000000..b68efaafe5e --- /dev/null +++ b/acceptance/ssh/connection/out.proxy.txt @@ -0,0 +1 @@ +hello over the ssh tunnel diff --git a/acceptance/ssh/connection/out.stdout.txt b/acceptance/ssh/connection/out.stdout.txt deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/acceptance/ssh/connection/script b/acceptance/ssh/connection/script index 3e206bfbb01..59e30b65631 100644 --- a/acceptance/ssh/connection/script +++ b/acceptance/ssh/connection/script @@ -7,16 +7,22 @@ if [ -z "${CLOUD_ENV:-}" ]; then CLI_RELEASES_DIR="$PWD/releases" fi -errcode $CLI ssh connect --cluster=$TEST_DEFAULT_CLUSTER_ID --releases-dir=$CLI_RELEASES_DIR --user-known-hosts-file=known_hosts -- "echo 'Connection successful'" >out.stdout.txt 2>LOG.stderr - if [ -n "${CLOUD_ENV:-}" ]; then - # On cloud the connection runs the remote command; dump the run on failure. + # On cloud a real dedicated cluster runs the remote command over a full SSH + # handshake; dump the run on failure. + errcode $CLI ssh connect --cluster=$TEST_DEFAULT_CLUSTER_ID --releases-dir=$CLI_RELEASES_DIR --user-known-hosts-file=known_hosts -- "echo 'Connection successful'" >out.stdout.txt 2>LOG.stderr if ! grep -q "Connection successful" out.stdout.txt; then run_id=$(cat LOG.stderr | grep -o "Job submitted successfully with run ID: [0-9]*" | grep -o "[0-9]*$") trace $CLI jobs get-run "$run_id" > LOG.job fi else - # Locally the handshake can't complete (no compute); just assert the CLI - # submitted the right dedicated-cluster bootstrap job. + # Locally there's no compute for a full SSH handshake, but proxy mode still + # submits the same dedicated-cluster bootstrap job and then pipes stdin/stdout + # over the driver-proxy websocket, which the test server echoes back. A single + # run therefore asserts both the submitted job and a working tunnel (the CLI + # dialing the driver-proxy over ws:// and proxying bytes across it), so we need + # neither the full ssh spawn (it would hang looping its handshake against the + # echo) nor a separate Go proxy test. + printf 'hello over the ssh tunnel\n' | errcode $CLI ssh connect --proxy --cluster=$TEST_DEFAULT_CLUSTER_ID --releases-dir=$CLI_RELEASES_DIR >out.proxy.txt 2>LOG.stderr trace print_requests.py //api/2.2/jobs/runs/submit fi diff --git a/experimental/ssh/internal/client/websockets.go b/experimental/ssh/internal/client/websockets.go index 3241a9be2af..71fec9571f6 100644 --- a/experimental/ssh/internal/client/websockets.go +++ b/experimental/ssh/internal/client/websockets.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "net/http" + "net/url" "github.com/databricks/cli/libs/auth" "github.com/databricks/databricks-sdk-go" @@ -11,12 +12,12 @@ import ( ) func createWebsocketConnection(ctx context.Context, client *databricks.WorkspaceClient, connID, clusterID string, serverPort int, liteswap string) (*websocket.Conn, error) { - url, err := getProxyURL(ctx, client, connID, clusterID, serverPort) + proxyURL, err := getProxyURL(ctx, client, connID, clusterID, serverPort) if err != nil { return nil, fmt.Errorf("failed to get proxy URL: %w", err) } - req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + req, err := http.NewRequestWithContext(ctx, http.MethodGet, proxyURL, nil) if err != nil { return nil, fmt.Errorf("failed to create request: %w", err) } @@ -28,7 +29,6 @@ func createWebsocketConnection(ctx context.Context, client *databricks.Workspace return nil, fmt.Errorf("failed to authenticate: %w", err) } - req.URL.Scheme = "wss" // 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 if err != nil { @@ -43,10 +43,32 @@ func getProxyURL(ctx context.Context, client *databricks.WorkspaceClient, connID if err != nil { return "", fmt.Errorf("failed to get current workspace ID: %w", err) } - host := client.Config.Host + return buildProxyWebsocketURL(client.Config.Host, workspaceID, clusterID, serverPort, connID) +} + +// buildProxyWebsocketURL builds the driver-proxy websocket URL for an SSH tunnel +// connection against the given workspace host. +// +// The websocket scheme is derived from the host scheme rather than hardcoded: a +// plaintext http host (the local acceptance test server) is dialed over ws, and +// everything else over wss. Forcing wss would make the tunnel undiallable against +// a plaintext test server, which is why the local flow can only be exercised +// end-to-end once the scheme follows the host. +func buildProxyWebsocketURL(host, workspaceID, clusterID string, serverPort int, connID string) (string, error) { + u, err := url.Parse(host) + if err != nil { + return "", fmt.Errorf("failed to parse host %q: %w", host, err) + } + switch u.Scheme { + case "http": + u.Scheme = "ws" + default: + u.Scheme = "wss" + } // The /driver-proxy-api/o//... path is a legacy URL form on // the driver-proxy endpoint and uses an "o" path segment regardless of // whether the workspace ID itself is the legacy or new shape. - url := fmt.Sprintf("%s/driver-proxy-api/o/%s/%s/%d/ssh?id=%s", host, workspaceID, clusterID, serverPort, connID) - return url, nil + u.Path = fmt.Sprintf("/driver-proxy-api/o/%s/%s/%d/ssh", workspaceID, clusterID, serverPort) + u.RawQuery = url.Values{"id": {connID}}.Encode() + return u.String(), nil } diff --git a/experimental/ssh/internal/client/websockets_test.go b/experimental/ssh/internal/client/websockets_test.go new file mode 100644 index 00000000000..601ce5e2097 --- /dev/null +++ b/experimental/ssh/internal/client/websockets_test.go @@ -0,0 +1,35 @@ +package client + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestBuildProxyWebsocketURL(t *testing.T) { + tests := []struct { + name string + host string + want string + }{ + { + name: "https host is dialed over wss", + host: "https://my-workspace.cloud.databricks.test", + 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", + want: "ws://127.0.0.1:8080/driver-proxy-api/o/900800700600/1234-567890-abc/7772/ssh?id=conn-1", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := buildProxyWebsocketURL(tt.host, "900800700600", "1234-567890-abc", 7772, "conn-1") + require.NoError(t, err) + assert.Equal(t, tt.want, got) + }) + } +} diff --git a/libs/testserver/handlers.go b/libs/testserver/handlers.go index 338eb7e3e50..b2831aa2259 100644 --- a/libs/testserver/handlers.go +++ b/libs/testserver/handlers.go @@ -753,6 +753,10 @@ func AddDefaultHandlers(server *Server) { return Response{Body: ""} }) + // The tunnel's /ssh endpoint upgrades to a websocket, so it takes over the + // connection itself instead of returning a normal JSON/text response. + server.HandleRaw("GET", "/driver-proxy-api/o/{workspace_id}/{cluster_id}/{port}/ssh", sshTunnelEchoHandler) + // Secrets ACLs: server.Handle("GET", "/api/2.0/secrets/acls/get", func(req Request) any { return req.Workspace.SecretsAclsGet(req) diff --git a/libs/testserver/router.go b/libs/testserver/router.go index 00381eae6a9..0ddad7f7cef 100644 --- a/libs/testserver/router.go +++ b/libs/testserver/router.go @@ -93,6 +93,21 @@ func (r *Router) Handle(method, path string, handler HandlerFunc) { }) } +// HandleRaw registers a handler that receives the raw http.ResponseWriter and +// *http.Request, bypassing the JSON serve() pipeline. It is needed for endpoints +// that take over the connection themselves (e.g. a websocket upgrade), which +// serve()'s buffered write path cannot express. First registration wins, matching +// Handle. Only wildcard patterns are supported, which is all the current callers +// need (the driver-proxy websocket path). +func (r *Router) HandleRaw(method, path string, handler http.HandlerFunc) { + pattern := method + " " + path + if r.wildcard[pattern] { + return + } + r.wildcard[pattern] = true + r.mux.HandleFunc(pattern, handler) +} + // ServeHTTP routes a request to the registered handler, falling back to // NotFound if no route matches. func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) { diff --git a/libs/testserver/ssh.go b/libs/testserver/ssh.go new file mode 100644 index 00000000000..e8451268f9d --- /dev/null +++ b/libs/testserver/ssh.go @@ -0,0 +1,42 @@ +package testserver + +import ( + "net/http" + + "github.com/gorilla/websocket" +) + +// sshTunnelUpgrader upgrades the driver-proxy /ssh request to a websocket. +// The zero value accepts the CLI's same-origin-less dialer without extra checks. +var sshTunnelUpgrader = websocket.Upgrader{} + +// sshTunnelEchoHandler stands in for the SSH tunnel server that a real cluster +// runs behind the driver proxy. There is no compute locally, so it just echoes +// every binary frame back to the client. That is enough to exercise the client's +// raw byte proxy (`ssh connect --proxy`) end-to-end: the CLI dials the tunnel, +// pipes stdin over the websocket, and the frames it reads back are what a live +// `cat`-style loopback would return. It mirrors the echo server used by the +// proxy package's Go tests, but in-process and without a `cat` subprocess, so it +// also works on Windows. +func sshTunnelEchoHandler(w http.ResponseWriter, r *http.Request) { + conn, err := sshTunnelUpgrader.Upgrade(w, r, nil) + if err != nil { + return + } + defer conn.Close() + + for { + messageType, data, err := conn.ReadMessage() + if err != nil { + // A normal closure from the client (its stdin reached EOF) ends the + // loop; any other read error means the connection is gone anyway. + return + } + if messageType != websocket.BinaryMessage { + continue + } + if err := conn.WriteMessage(websocket.BinaryMessage, data); err != nil { + return + } + } +} From 26fab01570e4c378742bdad488e1ed9f4c7c6ea2 Mon Sep 17 00:00:00 2001 From: Rada Kamysheva Date: Tue, 14 Jul 2026 08:10:35 +0000 Subject: [PATCH 2/2] acc: converge ssh/connect-serverless-gpu tunnel test too Apply the same convergence to the serverless-GPU acceptance test (follow-up to #5878, where the review asked to fold the two test styles into one acceptance test). Locally it now runs a single `ssh connect --proxy` that submits the same serverless-GPU bootstrap job and pipes stdin/stdout over the driver-proxy websocket echo route, asserting both the job payload and the echoed bytes. --- .../ssh/connect-serverless-gpu/out.proxy.txt | 1 + .../ssh/connect-serverless-gpu/out.stdout.txt | 0 acceptance/ssh/connect-serverless-gpu/script | 16 +++++++++++----- 3 files changed, 12 insertions(+), 5 deletions(-) create mode 100644 acceptance/ssh/connect-serverless-gpu/out.proxy.txt delete mode 100644 acceptance/ssh/connect-serverless-gpu/out.stdout.txt diff --git a/acceptance/ssh/connect-serverless-gpu/out.proxy.txt b/acceptance/ssh/connect-serverless-gpu/out.proxy.txt new file mode 100644 index 00000000000..b68efaafe5e --- /dev/null +++ b/acceptance/ssh/connect-serverless-gpu/out.proxy.txt @@ -0,0 +1 @@ +hello over the ssh tunnel diff --git a/acceptance/ssh/connect-serverless-gpu/out.stdout.txt b/acceptance/ssh/connect-serverless-gpu/out.stdout.txt deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/acceptance/ssh/connect-serverless-gpu/script b/acceptance/ssh/connect-serverless-gpu/script index e7e95e443d0..1a04f7f1d18 100644 --- a/acceptance/ssh/connect-serverless-gpu/script +++ b/acceptance/ssh/connect-serverless-gpu/script @@ -7,16 +7,22 @@ if [ -z "${CLOUD_ENV:-}" ]; then CLI_RELEASES_DIR="$PWD/releases" fi -errcode $CLI ssh connect --name serverless-gpu-test --accelerator=GPU_1xA10 --releases-dir=$CLI_RELEASES_DIR -- "echo 'Connection successful'" >out.stdout.txt 2>LOG.stderr - if [ -n "${CLOUD_ENV:-}" ]; then - # On cloud the connection runs the remote command; dump the run on failure. + # On cloud serverless GPU compute runs the remote command over a full SSH + # handshake; dump the run on failure. + errcode $CLI ssh connect --name serverless-gpu-test --accelerator=GPU_1xA10 --releases-dir=$CLI_RELEASES_DIR -- "echo 'Connection successful'" >out.stdout.txt 2>LOG.stderr if ! grep -q "Connection successful" out.stdout.txt; then run_id=$(cat LOG.stderr | grep -o "Job submitted successfully with run ID: [0-9]*" | grep -o "[0-9]*$") trace $CLI jobs get-run "$run_id" > LOG.job fi else - # Locally the handshake can't complete (no compute); just assert the CLI - # submitted the right serverless GPU bootstrap job. + # Locally there's no compute for a full SSH handshake, but proxy mode still + # submits the same serverless GPU bootstrap job and then pipes stdin/stdout + # over the driver-proxy websocket, which the test server echoes back. A single + # run therefore asserts both the submitted job and a working tunnel (the CLI + # dialing the driver-proxy over ws:// and proxying bytes across it), so we need + # neither the full ssh spawn (it would hang looping its handshake against the + # echo) nor a separate Go proxy test. + printf 'hello over the ssh tunnel\n' | errcode $CLI ssh connect --proxy --name serverless-gpu-test --accelerator=GPU_1xA10 --releases-dir=$CLI_RELEASES_DIR >out.proxy.txt 2>LOG.stderr trace print_requests.py //api/2.2/jobs/runs/submit fi