From 1bd7266afe8e97587cc7bfe9c3a6e5afc422ba02 Mon Sep 17 00:00:00 2001 From: Rada Kamysheva Date: Tue, 14 Jul 2026 08:29:15 +0000 Subject: [PATCH 1/7] acc: converge ssh/connect-serverless-gpu tunnel test onto a local websocket route Follow-up to #5878. That PR covered the SSH tunnel in two styles: an acceptance test asserting the submitted bootstrap job, and a Go test (TestClientServerRealSSHD) driving the tunnel transport. This folds the tunnel coverage into the acceptance test so a single local run exercises both layers. Two enabling changes: - Product fix (experimental/ssh/internal/client/websockets.go): derive the websocket scheme from the workspace host (http -> ws, else wss) instead of hardcoding wss, so the client can dial the plaintext local test server. Adds a unit test for the URL builder. - Test server (libs/testserver): add a raw, connection-hijacking handler registration (Router.HandleRaw) plus 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/connect-serverless-gpu now runs a single `ssh connect --proxy` locally: it submits the same serverless-GPU 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. --- .../ssh/connect-serverless-gpu/out.proxy.txt | 1 + .../ssh/connect-serverless-gpu/out.stdout.txt | 0 acceptance/ssh/connect-serverless-gpu/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/connect-serverless-gpu/out.proxy.txt delete mode 100644 acceptance/ssh/connect-serverless-gpu/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/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 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 906be46ad275d21207d6ca9179d865b9a20b5ebe Mon Sep 17 00:00:00 2001 From: Rada Kamysheva Date: Tue, 14 Jul 2026 08:39:02 +0000 Subject: [PATCH 2/7] acc: tighten SSH tunnel comments --- acceptance/ssh/connect-serverless-gpu/script | 11 ++++------- experimental/ssh/internal/client/websockets.go | 10 +++------- libs/testserver/handlers.go | 4 ++-- libs/testserver/router.go | 9 +++------ libs/testserver/ssh.go | 16 +++++----------- 5 files changed, 17 insertions(+), 33 deletions(-) diff --git a/acceptance/ssh/connect-serverless-gpu/script b/acceptance/ssh/connect-serverless-gpu/script index 1a04f7f1d18..8d030d41dc3 100644 --- a/acceptance/ssh/connect-serverless-gpu/script +++ b/acceptance/ssh/connect-serverless-gpu/script @@ -16,13 +16,10 @@ if [ -n "${CLOUD_ENV:-}" ]; then trace $CLI jobs get-run "$run_id" > LOG.job fi else - # 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. + # No compute locally, but proxy mode still submits the bootstrap job and then + # pipes stdin/stdout over the driver-proxy websocket, which the test server + # echoes. One run asserts both the job and a working ws:// tunnel; proxy mode + # skips the ssh spawn, so it can't hang on a handshake against the echo. 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 diff --git a/experimental/ssh/internal/client/websockets.go b/experimental/ssh/internal/client/websockets.go index 71fec9571f6..fc280674957 100644 --- a/experimental/ssh/internal/client/websockets.go +++ b/experimental/ssh/internal/client/websockets.go @@ -46,14 +46,10 @@ func getProxyURL(ctx context.Context, client *databricks.WorkspaceClient, connID 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. +// buildProxyWebsocketURL builds the driver-proxy websocket URL for an SSH tunnel. // -// 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. +// 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) { u, err := url.Parse(host) if err != nil { diff --git a/libs/testserver/handlers.go b/libs/testserver/handlers.go index b2831aa2259..9826e9e5f00 100644 --- a/libs/testserver/handlers.go +++ b/libs/testserver/handlers.go @@ -753,8 +753,8 @@ 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. + // /ssh upgrades to a websocket, so it hijacks the connection rather than + // returning a normal response. server.HandleRaw("GET", "/driver-proxy-api/o/{workspace_id}/{cluster_id}/{port}/ssh", sshTunnelEchoHandler) // Secrets ACLs: diff --git a/libs/testserver/router.go b/libs/testserver/router.go index 0ddad7f7cef..bcf5493b347 100644 --- a/libs/testserver/router.go +++ b/libs/testserver/router.go @@ -93,12 +93,9 @@ 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). +// HandleRaw registers a handler with the raw ResponseWriter and Request, bypassing +// the JSON serve() pipeline, for endpoints that take over the connection (e.g. a +// websocket upgrade). First registration wins; only wildcard patterns are supported. func (r *Router) HandleRaw(method, path string, handler http.HandlerFunc) { pattern := method + " " + path if r.wildcard[pattern] { diff --git a/libs/testserver/ssh.go b/libs/testserver/ssh.go index e8451268f9d..cf34083ef10 100644 --- a/libs/testserver/ssh.go +++ b/libs/testserver/ssh.go @@ -7,17 +7,12 @@ import ( ) // 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. +// sshTunnelEchoHandler stands in for the tunnel server a real cluster runs behind +// the driver proxy: with no compute locally it just echoes binary frames, enough +// to drive `ssh connect --proxy` end to end. In-process (no `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 { @@ -28,8 +23,7 @@ func sshTunnelEchoHandler(w http.ResponseWriter, r *http.Request) { 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. + // Client closed (its stdin hit EOF), or the connection is gone. return } if messageType != websocket.BinaryMessage { From 51f9252e96818eb6b870bd336470321f841abfbd Mon Sep 17 00:00:00 2001 From: Rada Kamysheva Date: Tue, 14 Jul 2026 09:39:04 +0000 Subject: [PATCH 3/7] acc: back the tunnel /ssh route with a real sshd so the handshake runs locally Deepens the convergence started earlier on this branch. Instead of echoing frames, the test server now upgrades the driver-proxy /ssh websocket to a real `sshd -i` and bridges binary frames to its stdio, so a single local `ssh connect ... -- "echo ..."` completes a genuine SSH handshake, public-key auth, and remote command exec with no Databricks compute. - libs/testserver: replace the echo handler with (*Server).sshTunnelHandler, which reads the CLI's uploaded client-public-key secret (the ws dial carries the bearer token, so the token-scoped workspace is the one that stored it), writes an ephemeral host key + authorized_keys + a minimal sshd_config, spawns sshd, and pumps ws frames <-> sshd stdin/stdout. /metadata now returns the current OS user so `ssh -l ` matches the account sshd runs as. - acceptance/ssh/connect-serverless-gpu: run the full handshake locally, still assert the bootstrap job payload, and self-skip via SKIP_TEST when sshd is absent (the general test job doesn't provision it; task test-exp-ssh does). Gated to linux, matching the removed Go test's rationale. - Delete TestClientServerRealSSHD: the acceptance test now covers the real handshake end to end, so the Go stand-in is redundant. The cat-echo transport tests remain. --- .github/workflows/push.yml | 7 +- .../ssh/connect-serverless-gpu/known_hosts | 1 + .../ssh/connect-serverless-gpu/out.proxy.txt | 1 - .../ssh/connect-serverless-gpu/out.stdout.txt | 1 + .../ssh/connect-serverless-gpu/out.test.toml | 1 + acceptance/ssh/connect-serverless-gpu/script | 18 +- .../ssh/connect-serverless-gpu/test.toml | 6 + .../internal/proxy/client_server_sshd_test.go | 191 ---------------- libs/testserver/handlers.go | 8 +- libs/testserver/jobs.go | 1 - libs/testserver/ssh.go | 212 ++++++++++++++++-- 11 files changed, 228 insertions(+), 219 deletions(-) create mode 100644 acceptance/ssh/connect-serverless-gpu/known_hosts delete mode 100644 acceptance/ssh/connect-serverless-gpu/out.proxy.txt create mode 100644 acceptance/ssh/connect-serverless-gpu/out.stdout.txt delete mode 100644 experimental/ssh/internal/proxy/client_server_sshd_test.go diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml index 8f5f981e6b8..9db62463267 100644 --- a/.github/workflows/push.yml +++ b/.github/workflows/push.yml @@ -262,9 +262,10 @@ jobs: with: cache-key: test-exp-ssh - # The proxy tests exercise a real OpenSSH server, which the runner image - # ships only the client half of. Install it on Linux (the OS the sshd - # test is gated to); other platforms skip that test. + # The ssh acceptance test drives a real OpenSSH server behind the test + # server's tunnel, which the runner image ships only the client half of. + # Install it on Linux (the OS the test is gated to); other platforms and + # the general test job (no install) skip it via SKIP_TEST. - name: Install OpenSSH server if: runner.os == 'Linux' run: sudo apt-get update && sudo apt-get install -y openssh-server diff --git a/acceptance/ssh/connect-serverless-gpu/known_hosts b/acceptance/ssh/connect-serverless-gpu/known_hosts new file mode 100644 index 00000000000..197ec818881 --- /dev/null +++ b/acceptance/ssh/connect-serverless-gpu/known_hosts @@ -0,0 +1 @@ +# not actually checked by tests; accept-new appends the ephemeral host key here diff --git a/acceptance/ssh/connect-serverless-gpu/out.proxy.txt b/acceptance/ssh/connect-serverless-gpu/out.proxy.txt deleted file mode 100644 index b68efaafe5e..00000000000 --- a/acceptance/ssh/connect-serverless-gpu/out.proxy.txt +++ /dev/null @@ -1 +0,0 @@ -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 new file mode 100644 index 00000000000..41cae5e7d16 --- /dev/null +++ b/acceptance/ssh/connect-serverless-gpu/out.stdout.txt @@ -0,0 +1 @@ +Connection successful diff --git a/acceptance/ssh/connect-serverless-gpu/out.test.toml b/acceptance/ssh/connect-serverless-gpu/out.test.toml index c777e3ce206..39c772fe40e 100644 --- a/acceptance/ssh/connect-serverless-gpu/out.test.toml +++ b/acceptance/ssh/connect-serverless-gpu/out.test.toml @@ -1,5 +1,6 @@ Local = true Cloud = false RequiresUnityCatalog = true +GOOS.linux = true CloudEnvs.gcp = false EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/acceptance/ssh/connect-serverless-gpu/script b/acceptance/ssh/connect-serverless-gpu/script index 8d030d41dc3..e458d1bd5a5 100644 --- a/acceptance/ssh/connect-serverless-gpu/script +++ b/acceptance/ssh/connect-serverless-gpu/script @@ -16,10 +16,18 @@ if [ -n "${CLOUD_ENV:-}" ]; then trace $CLI jobs get-run "$run_id" > LOG.job fi else - # No compute locally, but proxy mode still submits the bootstrap job and then - # pipes stdin/stdout over the driver-proxy websocket, which the test server - # echoes. One run asserts both the job and a working ws:// tunnel; proxy mode - # skips the ssh spawn, so it can't hang on a handshake against the echo. - 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 + # No Databricks compute locally, but the test server upgrades the driver-proxy + # /ssh websocket to a real `sshd -i`, so the same full handshake, public-key + # auth, and remote exec run end to end over the ws:// tunnel. This one run + # asserts both the bootstrap job submission and a working tunnel. + # + # Needs a real OpenSSH server; the general test job doesn't provision one (only + # `task test-exp-ssh` does), so skip there. The test server probes sshd the same + # way, so agreement is guaranteed on a given machine. + if [ -z "$(command -v sshd || ls /usr/sbin/sshd /usr/local/sbin/sshd /sbin/sshd 2>/dev/null)" ]; then + echo "SKIP_TEST sshd (openssh-server) not installed" + exit 0 + fi + errcode $CLI ssh connect --name serverless-gpu-test --accelerator=GPU_1xA10 --releases-dir=$CLI_RELEASES_DIR --user-known-hosts-file=known_hosts -- "echo 'Connection successful'" >out.stdout.txt 2>LOG.stderr trace print_requests.py //api/2.2/jobs/runs/submit fi diff --git a/acceptance/ssh/connect-serverless-gpu/test.toml b/acceptance/ssh/connect-serverless-gpu/test.toml index 24b0870236f..1a54fb976dc 100644 --- a/acceptance/ssh/connect-serverless-gpu/test.toml +++ b/acceptance/ssh/connect-serverless-gpu/test.toml @@ -12,6 +12,12 @@ Ignore = [ "releases", ] +# The local run drives a real sshd behind the test server's tunnel; running it as +# a non-root user in inetd mode with a throwaway config is reliable on Linux (which +# is where task test-exp-ssh provisions openssh-server), not on macOS/Windows. +[GOOS] + linux = true + # Serverless GPU is not available in GCP yet [CloudEnvs] gcp = false diff --git a/experimental/ssh/internal/proxy/client_server_sshd_test.go b/experimental/ssh/internal/proxy/client_server_sshd_test.go deleted file mode 100644 index 097668d2921..00000000000 --- a/experimental/ssh/internal/proxy/client_server_sshd_test.go +++ /dev/null @@ -1,191 +0,0 @@ -//go:build linux - -package proxy - -import ( - "context" - "crypto/ed25519" - "crypto/rand" - "encoding/pem" - "fmt" - "io" - "net" - "net/http/httptest" - "os" - "os/exec" - "os/user" - "path/filepath" - "testing" - "time" - - "github.com/databricks/cli/libs/cmdio" - "github.com/gorilla/websocket" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" - "golang.org/x/crypto/ssh" -) - -// findSSHD returns the path to the OpenSSH server binary, or "" if it isn't installed. -func findSSHD() string { - if p, err := exec.LookPath("sshd"); err == nil { - return p - } - // LookPath misses sshd because it lives in sbin, which is usually not on a - // non-root user's PATH; probe the conventional locations directly. - for _, p := range []string{"/usr/sbin/sshd", "/usr/local/sbin/sshd", "/sbin/sshd"} { - if _, err := os.Stat(p); err == nil { - return p - } - } - return "" -} - -// pipeConn adapts a pair of pipes into a net.Conn so an in-process SSH client can speak the SSH -// protocol over the proxy transport (RunClientProxy reads/writes the other ends of these pipes). -type pipeConn struct { - r *io.PipeReader - w *io.PipeWriter -} - -func (c pipeConn) Read(p []byte) (int, error) { return c.r.Read(p) } -func (c pipeConn) Write(p []byte) (int, error) { return c.w.Write(p) } - -func (c pipeConn) Close() error { - _ = c.r.Close() - return c.w.Close() -} - -func (pipeConn) LocalAddr() net.Addr { return pipeAddr{} } -func (pipeConn) RemoteAddr() net.Addr { return pipeAddr{} } -func (pipeConn) SetDeadline(_ time.Time) error { return nil } -func (pipeConn) SetReadDeadline(_ time.Time) error { return nil } -func (pipeConn) SetWriteDeadline(_ time.Time) error { return nil } - -type pipeAddr struct{} - -func (pipeAddr) Network() string { return "pipe" } -func (pipeAddr) String() string { return "pipe" } - -func writeOpenSSHPrivateKey(t *testing.T, path string, key ed25519.PrivateKey) { - block, err := ssh.MarshalPrivateKey(key, "") - require.NoError(t, err) - require.NoError(t, os.WriteFile(path, pem.EncodeToMemory(block), 0o600)) -} - -// writeSSHDConfig writes a minimal sshd_config that lets a real sshd run as a non-root user in -// inetd mode (-i) and authenticate the current user by public key. StrictModes/UsePAM are off so -// sshd doesn't reject the temp-dir key files or try to switch users. -func writeSSHDConfig(t *testing.T, dir, hostKeyPath, authKeysPath string) string { - cfg := fmt.Sprintf( - "HostKey %s\n"+ - "AuthorizedKeysFile %s\n"+ - "PidFile none\n"+ - "StrictModes no\n"+ - "UsePAM no\n"+ - "PubkeyAuthentication yes\n"+ - "PasswordAuthentication no\n"+ - "LogLevel ERROR\n", - hostKeyPath, authKeysPath) - path := filepath.Join(dir, "sshd_config") - require.NoError(t, os.WriteFile(path, []byte(cfg), 0o600)) - return path -} - -// TestClientServerRealSSHD runs a real OpenSSH server (sshd) behind the websocket proxy and a real -// SSH client through RunClientProxy, exercising the actual SSH handshake and a remote command -// end-to-end without any Databricks compute. It is the local stand-in for the ssh-connectivity -// coverage that otherwise requires a cluster: the tunnel must carry the SSH protocol faithfully in -// both directions. Sibling tests in client_server_test.go cover the transport with a `cat` echo -// server; this one adds a genuine sshd so the handshake, auth, and channel exec are validated too. -// -// The test is Linux-only: running sshd -i as a non-root user with a throwaway config is reliable -// there (CI installs openssh-server), but not on macOS. It still skips gracefully when sshd is -// absent so the general `test` job, which doesn't install it, doesn't fail. -func TestClientServerRealSSHD(t *testing.T) { - sshdPath := findSSHD() - if sshdPath == "" { - t.Skip("sshd (openssh-server) not installed; skipping real SSH handshake test") - } - - currentUser, err := user.Current() - require.NoError(t, err) - - dir := t.TempDir() - - // Host key: identifies the server; the client pins it via FixedHostKey. - _, hostPriv, err := ed25519.GenerateKey(rand.Reader) - require.NoError(t, err) - hostKeyPath := filepath.Join(dir, "host_key") - writeOpenSSHPrivateKey(t, hostKeyPath, hostPriv) - hostSigner, err := ssh.NewSignerFromSigner(hostPriv) - require.NoError(t, err) - - // Client key: authorized on the server via authorized_keys, presented by the SSH client to log in. - _, clientPriv, err := ed25519.GenerateKey(rand.Reader) - require.NoError(t, err) - clientSigner, err := ssh.NewSignerFromSigner(clientPriv) - require.NoError(t, err) - authKeysPath := filepath.Join(dir, "authorized_keys") - require.NoError(t, os.WriteFile(authKeysPath, ssh.MarshalAuthorizedKey(clientSigner.PublicKey()), 0o600)) - - sshdConfigPath := writeSSHDConfig(t, dir, hostKeyPath, authKeysPath) - - ctx := cmdio.MockDiscard(t.Context()) - - // Server proxy: each websocket connection spawns `sshd -i`, which speaks SSH over its stdio. - connections := NewConnectionsManager(2, time.Hour) - server := httptest.NewServer(NewProxyServer(ctx, connections, func(ctx context.Context) *exec.Cmd { - return exec.CommandContext(ctx, sshdPath, "-f", sshdConfigPath, "-i", "-e") - })) - 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 - return conn, err - } - - // Wire the SSH client's transport through RunClientProxy: - // ssh client --writes--> clientToProxy --> proxy sends over ws --> sshd stdin - // ssh client <--reads--- proxyToClient <-- proxy recvs from ws <-- sshd stdout - clientToProxyR, clientToProxyW := io.Pipe() - proxyToClientR, proxyToClientW := io.Pipe() - - proxyDone := make(chan error, 1) - go func() { - requestHandoverTick := func() <-chan time.Time { return time.After(time.Hour) } - proxyDone <- RunClientProxy(ctx, clientToProxyR, proxyToClientW, requestHandoverTick, createConn) - }() - - conn := pipeConn{r: proxyToClientR, w: clientToProxyW} - sshConfig := &ssh.ClientConfig{ - User: currentUser.Username, - Auth: []ssh.AuthMethod{ssh.PublicKeys(clientSigner)}, - HostKeyCallback: ssh.FixedHostKey(hostSigner.PublicKey()), - Timeout: 30 * time.Second, - } - - clientConn, chans, reqs, err := ssh.NewClientConn(conn, "pipe", sshConfig) - require.NoError(t, err, "SSH handshake failed") - sshClient := ssh.NewClient(clientConn, chans, reqs) - defer sshClient.Close() - - session, err := sshClient.NewSession() - require.NoError(t, err) - defer session.Close() - - out, err := session.Output("echo 'Connection successful'") - require.NoError(t, err) - assert.Equal(t, "Connection successful\n", string(out)) - - // Tear down the client side and confirm the proxy shuts down cleanly rather than hanging. - require.NoError(t, sshClient.Close()) - _ = conn.Close() - - select { - case err := <-proxyDone: - require.NoError(t, err) - case <-time.After(10 * time.Second): - t.Fatal("client proxy did not shut down after the SSH session ended") - } -} diff --git a/libs/testserver/handlers.go b/libs/testserver/handlers.go index 9826e9e5f00..a76e7c5601b 100644 --- a/libs/testserver/handlers.go +++ b/libs/testserver/handlers.go @@ -744,9 +744,11 @@ func AddDefaultHandlers(server *Server) { }) // SSH tunnel server behind the driver proxy: /metadata returns the remote - // login user, /logs is a best-effort error tail fetched on failure. + // login user, /logs is a best-effort error tail fetched on failure. The user + // must be the OS user the local sshd runs as (the real server reports the same), + // so `ssh -l ` matches the account authorized in authorized_keys. server.Handle("GET", "/driver-proxy-api/o/{workspace_id}/{cluster_id}/{port}/metadata", func(req Request) any { - return Response{Body: sshTunnelRemoteUser} + return Response{Body: sshTunnelRemoteUser()} }) server.Handle("GET", "/driver-proxy-api/o/{workspace_id}/{cluster_id}/{port}/logs", func(req Request) any { @@ -755,7 +757,7 @@ func AddDefaultHandlers(server *Server) { // /ssh upgrades to a websocket, so it hijacks the connection rather than // returning a normal response. - server.HandleRaw("GET", "/driver-proxy-api/o/{workspace_id}/{cluster_id}/{port}/ssh", sshTunnelEchoHandler) + server.HandleRaw("GET", "/driver-proxy-api/o/{workspace_id}/{cluster_id}/{port}/ssh", server.sshTunnelHandler) // Secrets ACLs: server.Handle("GET", "/api/2.0/secrets/acls/get", func(req Request) any { diff --git a/libs/testserver/jobs.go b/libs/testserver/jobs.go index 2d019c3e496..9097d43c086 100644 --- a/libs/testserver/jobs.go +++ b/libs/testserver/jobs.go @@ -534,7 +534,6 @@ const ( sshTunnelBootstrapNotebook = "ssh-server-bootstrap" sshTunnelServerPort = 7772 sshTunnelClusterID = "1234-567890-serverless" - sshTunnelRemoteUser = "spark" ) // writeSSHTunnelMetadata publishes the metadata.json a real tunnel server would diff --git a/libs/testserver/ssh.go b/libs/testserver/ssh.go index cf34083ef10..1f4421bfd78 100644 --- a/libs/testserver/ssh.go +++ b/libs/testserver/ssh.go @@ -1,36 +1,218 @@ package testserver import ( + "context" + "crypto/ed25519" + "crypto/rand" + "encoding/pem" + "fmt" + "io" "net/http" + "os" + "os/exec" + "os/user" + "path/filepath" + "sync" + "time" "github.com/gorilla/websocket" + "golang.org/x/crypto/ssh" ) +// sshTunnelRemoteUser reports the login user the driver-proxy /metadata endpoint +// returns. The local sshd runs as the current OS user, so `ssh -l ` must use +// that same account; the real tunnel server likewise reports its current user. +func sshTunnelRemoteUser() string { + if u, err := user.Current(); err == nil { + return u.Username + } + return "" +} + +// sshClientPublicKeySecretKey is the secret key name the CLI stores the client's +// authorized public key under (mirrors clientPublicKeyName in experimental/ssh). +// The tunnel server on a real cluster reads it to build authorized_keys; the test +// server does the same so the local handshake authenticates the CLI's key. +const sshClientPublicKeySecretKey = "client-public-key" + // sshTunnelUpgrader upgrades the driver-proxy /ssh request to a websocket. var sshTunnelUpgrader = websocket.Upgrader{} -// sshTunnelEchoHandler stands in for the tunnel server a real cluster runs behind -// the driver proxy: with no compute locally it just echoes binary frames, enough -// to drive `ssh connect --proxy` end to end. In-process (no `cat` subprocess), so -// it also works on Windows. -func sshTunnelEchoHandler(w http.ResponseWriter, r *http.Request) { +// sshdTerminationTimeout bounds how long we wait for sshd to exit after the +// tunnel closes before Cmd.Wait gives up, so a stuck child can't wedge the test. +const sshdTerminationTimeout = 10 * time.Second + +// findSSHD returns the path to the OpenSSH server binary, or "" if it isn't installed. +// LookPath usually misses sshd because it lives in sbin (not on a non-root PATH), so +// probe the conventional locations directly. Mirrors the acceptance test's own probe. +func findSSHD() string { + if p, err := exec.LookPath("sshd"); err == nil { + return p + } + for _, p := range []string{"/usr/sbin/sshd", "/usr/local/sbin/sshd", "/sbin/sshd"} { + if _, err := os.Stat(p); err == nil { + return p + } + } + return "" +} + +// sshClientAuthorizedKey returns the client's authorized public key the CLI uploaded +// via secrets/put, scanning every scope since the scope name is derived from the +// session and current user. Returns nil if no such secret exists. +func (s *FakeWorkspace) sshClientAuthorizedKey() []byte { + defer s.LockUnlock()() + for _, scope := range s.Secrets { + if v, ok := scope[sshClientPublicKeySecretKey]; ok { + return []byte(v) + } + } + return nil +} + +// sshTunnelHandler stands in for the tunnel server a real cluster runs behind the +// driver proxy. With no Databricks compute locally, it upgrades the /ssh request to +// a websocket and drives a real `sshd -i` over that transport, so `ssh connect` +// completes an actual SSH handshake, public-key auth, and remote command exec end to +// end. It is the local stand-in for cluster connectivity: the tunnel must carry the +// SSH protocol faithfully in both directions. +// +// Requires a real OpenSSH server; callers gate the test to skip where sshd is absent +// (see acceptance/ssh/connect-serverless-gpu). sshd -i as a non-root user with a +// throwaway config is reliable on Linux, which is where that test runs. +func (s *Server) sshTunnelHandler(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 { - // Client closed (its stdin hit EOF), or the connection is gone. - return - } - if messageType != websocket.BinaryMessage { - continue + sshdPath := findSSHD() + if sshdPath == "" { + return + } + + // The ws dial carries the CLI's bearer token, so the token-scoped workspace here + // is the same one the connect flow uploaded the client's public key to. + authorizedKey := s.getWorkspaceForToken(getToken(r)).sshClientAuthorizedKey() + if authorizedKey == nil { + return + } + + dir, err := os.MkdirTemp("", "testserver-sshd-") + if err != nil { + return + } + defer os.RemoveAll(dir) + + configPath, err := writeSSHDFiles(dir, authorizedKey) + if err != nil { + return + } + + ctx, cancel := context.WithCancel(r.Context()) + defer cancel() + + cmd := exec.CommandContext(ctx, sshdPath, "-f", configPath, "-i", "-e") + cmd.WaitDelay = sshdTerminationTimeout + cmd.Stderr = io.Discard + sshdStdin, err := cmd.StdinPipe() + if err != nil { + return + } + sshdStdout, err := cmd.StdoutPipe() + if err != nil { + return + } + if err := cmd.Start(); err != nil { + return + } + + var closeOnce sync.Once + closeConn := func() { closeOnce.Do(func() { _ = conn.Close() }) } + + var wg sync.WaitGroup + wg.Add(2) + + // sshd stdout -> websocket binary frames. + go func() { + defer wg.Done() + buf := make([]byte, 4096) + for { + n, readErr := sshdStdout.Read(buf) + if n > 0 { + if writeErr := conn.WriteMessage(websocket.BinaryMessage, buf[:n]); writeErr != nil { + break + } + } + if readErr != nil { + break + } } - if err := conn.WriteMessage(websocket.BinaryMessage, data); err != nil { - return + // sshd exited (or the client is gone): unblock the reader below. + closeConn() + }() + + // websocket binary frames -> sshd stdin. + go func() { + defer wg.Done() + for { + messageType, data, readErr := conn.ReadMessage() + if readErr != nil { + break + } + if messageType != websocket.BinaryMessage { + continue + } + if _, writeErr := sshdStdin.Write(data); writeErr != nil { + break + } } + // The client closed its side (ssh session ended): EOF sshd's stdin so it exits. + _ = sshdStdin.Close() + }() + + wg.Wait() + cancel() + _ = cmd.Wait() +} + +// writeSSHDFiles lays out the host key, authorized_keys, and a minimal sshd_config +// in dir and returns the config path. StrictModes/UsePAM are off so sshd doesn't +// reject the temp-dir key files or try to switch users when run as a non-root user +// in inetd mode (-i). +func writeSSHDFiles(dir string, authorizedKey []byte) (string, error) { + _, hostPriv, err := ed25519.GenerateKey(rand.Reader) + if err != nil { + return "", err + } + block, err := ssh.MarshalPrivateKey(hostPriv, "") + if err != nil { + return "", err + } + hostKeyPath := filepath.Join(dir, "host_key") + if err := os.WriteFile(hostKeyPath, pem.EncodeToMemory(block), 0o600); err != nil { + return "", err + } + + authKeysPath := filepath.Join(dir, "authorized_keys") + if err := os.WriteFile(authKeysPath, authorizedKey, 0o600); err != nil { + return "", err + } + + cfg := fmt.Sprintf( + "HostKey %s\n"+ + "AuthorizedKeysFile %s\n"+ + "PidFile none\n"+ + "StrictModes no\n"+ + "UsePAM no\n"+ + "PubkeyAuthentication yes\n"+ + "PasswordAuthentication no\n"+ + "LogLevel ERROR\n", + hostKeyPath, authKeysPath) + configPath := filepath.Join(dir, "sshd_config") + if err := os.WriteFile(configPath, []byte(cfg), 0o600); err != nil { + return "", err } + return configPath, nil } From 904e41d0de23e5f13788dbbbb8ccd82ff6d437be Mon Sep 17 00:00:00 2001 From: Rada Kamysheva Date: Tue, 14 Jul 2026 09:58:46 +0000 Subject: [PATCH 4/7] acc: fix lint and make the ssh tunnel test truly Linux-only CI caught two issues in the previous commit: - lint: use sync.OnceFunc and wg.Go instead of sync.Once/wg.Add+Done in the test server's ssh handler. - GOOS keys absent from the map default to enabled, so `linux = true` alone did not disable macOS/Windows. They ran the handshake: macOS found its system sshd and hung, Windows found sshd.exe and failed fast. Disable darwin and windows explicitly to keep the test Linux-only. --- .../ssh/connect-serverless-gpu/out.test.toml | 2 ++ acceptance/ssh/connect-serverless-gpu/test.toml | 6 +++++- libs/testserver/ssh.go | 14 +++++--------- 3 files changed, 12 insertions(+), 10 deletions(-) diff --git a/acceptance/ssh/connect-serverless-gpu/out.test.toml b/acceptance/ssh/connect-serverless-gpu/out.test.toml index 39c772fe40e..1dc1ba1a690 100644 --- a/acceptance/ssh/connect-serverless-gpu/out.test.toml +++ b/acceptance/ssh/connect-serverless-gpu/out.test.toml @@ -1,6 +1,8 @@ Local = true Cloud = false RequiresUnityCatalog = true +GOOS.darwin = false GOOS.linux = true +GOOS.windows = false CloudEnvs.gcp = false EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/acceptance/ssh/connect-serverless-gpu/test.toml b/acceptance/ssh/connect-serverless-gpu/test.toml index 1a54fb976dc..d93ed96b515 100644 --- a/acceptance/ssh/connect-serverless-gpu/test.toml +++ b/acceptance/ssh/connect-serverless-gpu/test.toml @@ -14,9 +14,13 @@ Ignore = [ # The local run drives a real sshd behind the test server's tunnel; running it as # a non-root user in inetd mode with a throwaway config is reliable on Linux (which -# is where task test-exp-ssh provisions openssh-server), not on macOS/Windows. +# is where task test-exp-ssh provisions openssh-server), not on macOS/Windows. GOOS +# keys absent from this map default to enabled, so the other OSes must be disabled +# explicitly to keep the test Linux-only. [GOOS] linux = true + darwin = false + windows = false # Serverless GPU is not available in GCP yet [CloudEnvs] diff --git a/libs/testserver/ssh.go b/libs/testserver/ssh.go index 1f4421bfd78..c190a92e859 100644 --- a/libs/testserver/ssh.go +++ b/libs/testserver/ssh.go @@ -128,15 +128,12 @@ func (s *Server) sshTunnelHandler(w http.ResponseWriter, r *http.Request) { return } - var closeOnce sync.Once - closeConn := func() { closeOnce.Do(func() { _ = conn.Close() }) } + closeConn := sync.OnceFunc(func() { _ = conn.Close() }) var wg sync.WaitGroup - wg.Add(2) // sshd stdout -> websocket binary frames. - go func() { - defer wg.Done() + wg.Go(func() { buf := make([]byte, 4096) for { n, readErr := sshdStdout.Read(buf) @@ -151,11 +148,10 @@ func (s *Server) sshTunnelHandler(w http.ResponseWriter, r *http.Request) { } // sshd exited (or the client is gone): unblock the reader below. closeConn() - }() + }) // websocket binary frames -> sshd stdin. - go func() { - defer wg.Done() + wg.Go(func() { for { messageType, data, readErr := conn.ReadMessage() if readErr != nil { @@ -170,7 +166,7 @@ func (s *Server) sshTunnelHandler(w http.ResponseWriter, r *http.Request) { } // The client closed its side (ssh session ended): EOF sshd's stdin so it exits. _ = sshdStdin.Close() - }() + }) wg.Wait() cancel() From 04c8e743404f4b505a1130d92d32410fd94b9a40 Mon Sep 17 00:00:00 2001 From: Rada Kamysheva Date: Tue, 14 Jul 2026 10:10:02 +0000 Subject: [PATCH 5/7] acc: tighten SSH tunnel comments --- .github/workflows/push.yml | 7 ++- acceptance/ssh/connect-serverless-gpu/script | 12 ++--- .../ssh/connect-serverless-gpu/test.toml | 8 ++- libs/testserver/handlers.go | 9 ++-- libs/testserver/ssh.go | 50 +++++++------------ 5 files changed, 31 insertions(+), 55 deletions(-) diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml index 9db62463267..eb0ea35e05b 100644 --- a/.github/workflows/push.yml +++ b/.github/workflows/push.yml @@ -262,10 +262,9 @@ jobs: with: cache-key: test-exp-ssh - # The ssh acceptance test drives a real OpenSSH server behind the test - # server's tunnel, which the runner image ships only the client half of. - # Install it on Linux (the OS the test is gated to); other platforms and - # the general test job (no install) skip it via SKIP_TEST. + # The ssh acceptance test drives a real sshd behind the tunnel; the runner + # ships only the client. Install it on Linux (the OS the test is gated to); + # elsewhere and in the general test job it skips via SKIP_TEST. - name: Install OpenSSH server if: runner.os == 'Linux' run: sudo apt-get update && sudo apt-get install -y openssh-server diff --git a/acceptance/ssh/connect-serverless-gpu/script b/acceptance/ssh/connect-serverless-gpu/script index e458d1bd5a5..0943ec9d2b7 100644 --- a/acceptance/ssh/connect-serverless-gpu/script +++ b/acceptance/ssh/connect-serverless-gpu/script @@ -16,14 +16,10 @@ if [ -n "${CLOUD_ENV:-}" ]; then trace $CLI jobs get-run "$run_id" > LOG.job fi else - # No Databricks compute locally, but the test server upgrades the driver-proxy - # /ssh websocket to a real `sshd -i`, so the same full handshake, public-key - # auth, and remote exec run end to end over the ws:// tunnel. This one run - # asserts both the bootstrap job submission and a working tunnel. - # - # Needs a real OpenSSH server; the general test job doesn't provision one (only - # `task test-exp-ssh` does), so skip there. The test server probes sshd the same - # way, so agreement is guaranteed on a given machine. + # No compute locally: the test server backs the driver-proxy /ssh websocket with a + # real sshd, so this one run asserts both the bootstrap job and a full handshake + + # remote exec over the ws:// tunnel. Needs sshd (only task test-exp-ssh provisions + # it); the test server probes the same way, so skip where it's absent. if [ -z "$(command -v sshd || ls /usr/sbin/sshd /usr/local/sbin/sshd /sbin/sshd 2>/dev/null)" ]; then echo "SKIP_TEST sshd (openssh-server) not installed" exit 0 diff --git a/acceptance/ssh/connect-serverless-gpu/test.toml b/acceptance/ssh/connect-serverless-gpu/test.toml index d93ed96b515..a08b994b74e 100644 --- a/acceptance/ssh/connect-serverless-gpu/test.toml +++ b/acceptance/ssh/connect-serverless-gpu/test.toml @@ -12,11 +12,9 @@ Ignore = [ "releases", ] -# The local run drives a real sshd behind the test server's tunnel; running it as -# a non-root user in inetd mode with a throwaway config is reliable on Linux (which -# is where task test-exp-ssh provisions openssh-server), not on macOS/Windows. GOOS -# keys absent from this map default to enabled, so the other OSes must be disabled -# explicitly to keep the test Linux-only. +# Linux-only: the local run drives a real sshd, reliable only there (and where +# task test-exp-ssh provisions it). Absent GOOS keys default to enabled, so the +# other OSes must be disabled explicitly. [GOOS] linux = true darwin = false diff --git a/libs/testserver/handlers.go b/libs/testserver/handlers.go index a76e7c5601b..a51a13b4afe 100644 --- a/libs/testserver/handlers.go +++ b/libs/testserver/handlers.go @@ -743,10 +743,9 @@ func AddDefaultHandlers(server *Server) { return req.Workspace.SecretsList(req) }) - // SSH tunnel server behind the driver proxy: /metadata returns the remote - // login user, /logs is a best-effort error tail fetched on failure. The user - // must be the OS user the local sshd runs as (the real server reports the same), - // so `ssh -l ` matches the account authorized in authorized_keys. + // SSH tunnel server behind the driver proxy: /metadata returns the remote login + // user, /logs is a best-effort error tail, /ssh hijacks the connection for a + // websocket (so it's raw, not a JSON response). server.Handle("GET", "/driver-proxy-api/o/{workspace_id}/{cluster_id}/{port}/metadata", func(req Request) any { return Response{Body: sshTunnelRemoteUser()} }) @@ -755,8 +754,6 @@ func AddDefaultHandlers(server *Server) { return Response{Body: ""} }) - // /ssh upgrades to a websocket, so it hijacks the connection rather than - // returning a normal response. server.HandleRaw("GET", "/driver-proxy-api/o/{workspace_id}/{cluster_id}/{port}/ssh", server.sshTunnelHandler) // Secrets ACLs: diff --git a/libs/testserver/ssh.go b/libs/testserver/ssh.go index c190a92e859..283e4016b55 100644 --- a/libs/testserver/ssh.go +++ b/libs/testserver/ssh.go @@ -19,9 +19,8 @@ import ( "golang.org/x/crypto/ssh" ) -// sshTunnelRemoteUser reports the login user the driver-proxy /metadata endpoint -// returns. The local sshd runs as the current OS user, so `ssh -l ` must use -// that same account; the real tunnel server likewise reports its current user. +// sshTunnelRemoteUser is the login user /metadata reports. It must be the OS user +// the local sshd runs as, so `ssh -l ` matches; the real server reports the same. func sshTunnelRemoteUser() string { if u, err := user.Current(); err == nil { return u.Username @@ -29,22 +28,17 @@ func sshTunnelRemoteUser() string { return "" } -// sshClientPublicKeySecretKey is the secret key name the CLI stores the client's -// authorized public key under (mirrors clientPublicKeyName in experimental/ssh). -// The tunnel server on a real cluster reads it to build authorized_keys; the test -// server does the same so the local handshake authenticates the CLI's key. +// sshClientPublicKeySecretKey is the secret the CLI stores its authorized public key +// under (mirrors clientPublicKeyName in experimental/ssh); sshd authorizes it below. const sshClientPublicKeySecretKey = "client-public-key" -// sshTunnelUpgrader upgrades the driver-proxy /ssh request to a websocket. var sshTunnelUpgrader = websocket.Upgrader{} -// sshdTerminationTimeout bounds how long we wait for sshd to exit after the -// tunnel closes before Cmd.Wait gives up, so a stuck child can't wedge the test. +// sshdTerminationTimeout caps how long Cmd.Wait blocks on a stuck sshd once the tunnel closes. const sshdTerminationTimeout = 10 * time.Second -// findSSHD returns the path to the OpenSSH server binary, or "" if it isn't installed. -// LookPath usually misses sshd because it lives in sbin (not on a non-root PATH), so -// probe the conventional locations directly. Mirrors the acceptance test's own probe. +// findSSHD returns the sshd path, or "" if absent. LookPath misses it in sbin, so also +// probe the usual locations. Mirrors the acceptance test's probe. func findSSHD() string { if p, err := exec.LookPath("sshd"); err == nil { return p @@ -57,9 +51,8 @@ func findSSHD() string { return "" } -// sshClientAuthorizedKey returns the client's authorized public key the CLI uploaded -// via secrets/put, scanning every scope since the scope name is derived from the -// session and current user. Returns nil if no such secret exists. +// sshClientAuthorizedKey returns the public key the CLI uploaded via secrets/put, or nil. +// Scans all scopes since the scope name embeds the session and current user. func (s *FakeWorkspace) sshClientAuthorizedKey() []byte { defer s.LockUnlock()() for _, scope := range s.Secrets { @@ -70,16 +63,10 @@ func (s *FakeWorkspace) sshClientAuthorizedKey() []byte { return nil } -// sshTunnelHandler stands in for the tunnel server a real cluster runs behind the -// driver proxy. With no Databricks compute locally, it upgrades the /ssh request to -// a websocket and drives a real `sshd -i` over that transport, so `ssh connect` -// completes an actual SSH handshake, public-key auth, and remote command exec end to -// end. It is the local stand-in for cluster connectivity: the tunnel must carry the -// SSH protocol faithfully in both directions. -// -// Requires a real OpenSSH server; callers gate the test to skip where sshd is absent -// (see acceptance/ssh/connect-serverless-gpu). sshd -i as a non-root user with a -// throwaway config is reliable on Linux, which is where that test runs. +// sshTunnelHandler stands in for a cluster's tunnel server: it upgrades /ssh to a +// websocket and drives a real `sshd -i` over it, so `ssh connect` runs a full handshake, +// auth, and remote exec locally. Requires sshd; the acceptance test skips when it's absent +// and is pinned to Linux, where sshd -i as a non-root user is reliable. func (s *Server) sshTunnelHandler(w http.ResponseWriter, r *http.Request) { conn, err := sshTunnelUpgrader.Upgrade(w, r, nil) if err != nil { @@ -92,8 +79,8 @@ func (s *Server) sshTunnelHandler(w http.ResponseWriter, r *http.Request) { return } - // The ws dial carries the CLI's bearer token, so the token-scoped workspace here - // is the same one the connect flow uploaded the client's public key to. + // The ws dial carries the CLI's bearer token, so this token-scoped workspace is the + // one that stored the uploaded public key. authorizedKey := s.getWorkspaceForToken(getToken(r)).sshClientAuthorizedKey() if authorizedKey == nil { return @@ -173,10 +160,9 @@ func (s *Server) sshTunnelHandler(w http.ResponseWriter, r *http.Request) { _ = cmd.Wait() } -// writeSSHDFiles lays out the host key, authorized_keys, and a minimal sshd_config -// in dir and returns the config path. StrictModes/UsePAM are off so sshd doesn't -// reject the temp-dir key files or try to switch users when run as a non-root user -// in inetd mode (-i). +// writeSSHDFiles writes an ephemeral host key, authorized_keys, and a minimal sshd_config, +// returning the config path. StrictModes/UsePAM are off so sshd accepts the temp-dir keys +// and doesn't try to switch users when run non-root in inetd mode (-i). func writeSSHDFiles(dir string, authorizedKey []byte) (string, error) { _, hostPriv, err := ed25519.GenerateKey(rand.Reader) if err != nil { From 945a942aa7c6bddd3cd2355a12ca8174ef093262 Mon Sep 17 00:00:00 2001 From: Rada Kamysheva Date: Wed, 15 Jul 2026 06:50:17 +0000 Subject: [PATCH 6/7] acc: log ssh tunnel handler setup failures The testserver /ssh handler returned silently on every setup failure (missing secret, pipe/Start errors, etc.), so a broken sshd setup only surfaced as an empty-stdout golden diff. Log each early-return reason so a future breakage is self-explanatory. --- libs/testserver/ssh.go | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/libs/testserver/ssh.go b/libs/testserver/ssh.go index 283e4016b55..0cb7eef8846 100644 --- a/libs/testserver/ssh.go +++ b/libs/testserver/ssh.go @@ -70,12 +70,14 @@ func (s *FakeWorkspace) sshClientAuthorizedKey() []byte { func (s *Server) sshTunnelHandler(w http.ResponseWriter, r *http.Request) { conn, err := sshTunnelUpgrader.Upgrade(w, r, nil) if err != nil { + s.t.Logf("ssh tunnel: websocket upgrade failed: %s", err) return } defer conn.Close() sshdPath := findSSHD() if sshdPath == "" { + s.t.Logf("ssh tunnel: sshd not found; tunnel cannot start") return } @@ -83,17 +85,20 @@ func (s *Server) sshTunnelHandler(w http.ResponseWriter, r *http.Request) { // one that stored the uploaded public key. authorizedKey := s.getWorkspaceForToken(getToken(r)).sshClientAuthorizedKey() if authorizedKey == nil { + s.t.Logf("ssh tunnel: client public key secret %q not found", sshClientPublicKeySecretKey) return } dir, err := os.MkdirTemp("", "testserver-sshd-") if err != nil { + s.t.Logf("ssh tunnel: temp dir: %s", err) return } defer os.RemoveAll(dir) configPath, err := writeSSHDFiles(dir, authorizedKey) if err != nil { + s.t.Logf("ssh tunnel: write sshd files: %s", err) return } @@ -105,13 +110,16 @@ func (s *Server) sshTunnelHandler(w http.ResponseWriter, r *http.Request) { cmd.Stderr = io.Discard sshdStdin, err := cmd.StdinPipe() if err != nil { + s.t.Logf("ssh tunnel: stdin pipe: %s", err) return } sshdStdout, err := cmd.StdoutPipe() if err != nil { + s.t.Logf("ssh tunnel: stdout pipe: %s", err) return } if err := cmd.Start(); err != nil { + s.t.Logf("ssh tunnel: start sshd: %s", err) return } From 82d7be311266be4ddce7d7b9a129c2bf47138c9d Mon Sep 17 00:00:00 2001 From: Rada Kamysheva Date: Wed, 15 Jul 2026 08:52:41 +0000 Subject: [PATCH 7/7] acc: converge ssh/connection onto the real-sshd tunnel route The testserver /ssh route now drives a real sshd, so ssh/connection's local run completes a genuine handshake and remote exec instead of producing nothing. Mirror connect-serverless-gpu: skip when sshd is absent, gate to Linux (macOS hangs on system sshd, Windows differs), and record "Connection successful" as the golden stdout. --- acceptance/ssh/connection/known_hosts | 3 +-- acceptance/ssh/connection/out.stdout.txt | 1 + acceptance/ssh/connection/out.test.toml | 3 +++ acceptance/ssh/connection/script | 17 ++++++++++++----- acceptance/ssh/connection/test.toml | 8 ++++++++ 5 files changed, 25 insertions(+), 7 deletions(-) diff --git a/acceptance/ssh/connection/known_hosts b/acceptance/ssh/connection/known_hosts index b254044c77a..197ec818881 100644 --- a/acceptance/ssh/connection/known_hosts +++ b/acceptance/ssh/connection/known_hosts @@ -1,2 +1 @@ -# not actually checked by tests -cluster ssh-rsa key +# not actually checked by tests; accept-new appends the ephemeral host key here diff --git a/acceptance/ssh/connection/out.stdout.txt b/acceptance/ssh/connection/out.stdout.txt index e69de29bb2d..41cae5e7d16 100644 --- a/acceptance/ssh/connection/out.stdout.txt +++ b/acceptance/ssh/connection/out.stdout.txt @@ -0,0 +1 @@ +Connection successful diff --git a/acceptance/ssh/connection/out.test.toml b/acceptance/ssh/connection/out.test.toml index 9d444f6ae49..b614b395da0 100644 --- a/acceptance/ssh/connection/out.test.toml +++ b/acceptance/ssh/connection/out.test.toml @@ -1,4 +1,7 @@ Local = true Cloud = false RequiresCluster = true +GOOS.darwin = false +GOOS.linux = true +GOOS.windows = false EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/acceptance/ssh/connection/script b/acceptance/ssh/connection/script index 3e206bfbb01..c69babb6dd0 100644 --- a/acceptance/ssh/connection/script +++ b/acceptance/ssh/connection/script @@ -7,16 +7,23 @@ 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 the 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. + # No compute locally: the test server backs the driver-proxy /ssh websocket with a + # real sshd, so this one run asserts both the bootstrap job and a full handshake + + # remote exec over the ws:// tunnel. Needs sshd (only task test-exp-ssh provisions + # it); the test server probes the same way, so skip where it's absent. + if [ -z "$(command -v sshd || ls /usr/sbin/sshd /usr/local/sbin/sshd /sbin/sshd 2>/dev/null)" ]; then + echo "SKIP_TEST sshd (openssh-server) not installed" + exit 0 + 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 trace print_requests.py //api/2.2/jobs/runs/submit fi diff --git a/acceptance/ssh/connection/test.toml b/acceptance/ssh/connection/test.toml index e2a33c0a5b9..74d8761c89a 100644 --- a/acceptance/ssh/connection/test.toml +++ b/acceptance/ssh/connection/test.toml @@ -13,5 +13,13 @@ Ignore = [ "releases", ] +# Linux-only: the local run drives a real sshd, reliable only there (and where +# task test-exp-ssh provisions it). Absent GOOS keys default to enabled, so the +# other OSes must be disabled explicitly. +[GOOS] + linux = true + darwin = false + windows = false + [EnvMatrix] DATABRICKS_BUNDLE_ENGINE = ["direct"]