acc: converge ssh/connection tunnel test onto a local websocket route

acceptance/ssh/connect-serverless-gpu/out.proxy.txt

@@ -0,0 +1 @@
+hello over the ssh tunnel

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

acceptance/ssh/connection/out.proxy.txt

@@ -0,0 +1 @@
+hello over the ssh tunnel

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

experimental/ssh/internal/client/websockets.go

@@ -4,19 +4,20 @@ import (
 	"context"
 	"fmt"
 	"net/http"
+	"net/url"
 
 	"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) {
-	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/<workspace-id>/... 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
 }

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

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)

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) {

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