acc: converge ssh/connect-serverless-gpu tunnel test onto a local web…

.github/workflows/push.yml

@@ -262,9 +262,9 @@ 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 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

acceptance/ssh/connect-serverless-gpu/known_hosts

@@ -0,0 +1 @@
+# not actually checked by tests; accept-new appends the ephemeral host key here

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

@@ -0,0 +1 @@
+Connection successful

acceptance/ssh/connect-serverless-gpu/script

@@ -7,16 +7,23 @@ 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.
+  # 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 --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

acceptance/ssh/connect-serverless-gpu/test.toml

@@ -12,6 +12,14 @@ 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
+
 # Serverless GPU is not available in GCP yet
 [CloudEnvs]
   gcp = false

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

acceptance/ssh/connection/out.stdout.txt

@@ -0,0 +1 @@
+Connection successful

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

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"]

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,28 @@ 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.
+//
+// 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 {
+		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)
+		})
+	}
+}