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/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 + } + } +}