acc: run ssh/connect-serverless-gpu locally

.github/workflows/push.yml

@@ -262,6 +262,13 @@ 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.
+      - name: Install OpenSSH server
+        if: runner.os == 'Linux'
+        run: sudo apt-get update && sudo apt-get install -y openssh-server
+
       - name: Run tests
         run: |
           go tool -modfile=tools/task/go.mod task test-exp-ssh

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

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

acceptance/ssh/connect-serverless-gpu/output.txt

@@ -0,0 +1,40 @@
+
+>>> print_requests.py //api/2.2/jobs/runs/submit
+{
+  "method": "POST",
+  "path": "/api/2.2/jobs/runs/submit",
+  "body": {
+    "environments": [
+      {
+        "environment_key": "ssh_tunnel_serverless",
+        "spec": {
+          "environment_version": "4"
+        }
+      }
+    ],
+    "run_name": "ssh-server-bootstrap-serverless-gpu-test",
+    "tasks": [
+      {
+        "compute": {
+          "hardware_accelerator": "GPU_1xA10"
+        },
+        "environment_key": "ssh_tunnel_serverless",
+        "notebook_task": {
+          "base_parameters": {
+            "authorizedKeySecretName": "client-public-key",
+            "maxClients": "10",
+            "secretScopeName": "[USERNAME]-serverless-gpu-test-ssh-tunnel-keys",
+            "serverless": "true",
+            "sessionId": "serverless-gpu-test",
+            "shutdownDelay": "10m0s",
+            "version": "[CLI_VERSION]"
+          },
+          "notebook_path": "/Workspace/Users/[USERNAME]/.databricks/ssh-tunnel/[CLI_VERSION]/serverless-gpu-test/ssh-server-bootstrap"
+        },
+        "task_key": "start_ssh_server",
+        "timeout_seconds": 86400
+      }
+    ],
+    "timeout_seconds": 86400
+  }
+}

acceptance/ssh/connect-serverless-gpu/script

@@ -1,6 +1,22 @@
+if [ -z "${CLOUD_ENV:-}" ]; then
+  # Local runs have no release artifacts and never execute the uploaded
+  # binaries, so stand in empty archives (dev name omits the version).
+  mkdir -p releases
+  : >releases/databricks_cli_linux_amd64.zip
+  : >releases/databricks_cli_linux_arm64.zip
+  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 ! 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
+if [ -n "${CLOUD_ENV:-}" ]; then
+  # On cloud the connection runs the remote command; dump the run on failure.
+  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.
+  trace print_requests.py //api/2.2/jobs/runs/submit
 fi

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

@@ -1,9 +1,17 @@
-Local = false
+Local = true
 Cloud = false
 
 # Serverless GPU is only available in newer environments
 RequiresUnityCatalog = true
 
+# Assert the serverless GPU bootstrap job the CLI submits on a local run.
+RecordRequests = true
+
+# Local runs stage stand-in release archives here; not part of the golden output.
+Ignore = [
+  "releases",
+]
+
 # Serverless GPU is not available in GCP yet
 [CloudEnvs]
   gcp = false

experimental/ssh/internal/proxy/client_server_sshd_test.go

@@ -0,0 +1,191 @@
+//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")
+	}
+}

libs/testserver/handlers.go

@@ -105,6 +105,15 @@ func AddDefaultHandlers(server *Server) {
 		path := req.URL.Query().Get("path")
 		data := req.Workspace.WorkspaceExport(path)
 
+		// A missing object returns 404, matching the real API; returning the nil
+		// body otherwise trips the response normalizer.
+		if data == nil {
+			return Response{
+				StatusCode: 404,
+				Body:       map[string]string{"message": fmt.Sprintf("Path (%s) doesn't exist.", path)},
+			}
+		}
+
 		// The filer reads the raw object body via ?direct_download=true, while
 		// the SDK's Workspace.Export (used by `databricks workspace export`)
 		// requests JSON and expects the base64-encoded content field.
@@ -730,6 +739,20 @@ func AddDefaultHandlers(server *Server) {
 		return req.Workspace.SecretsGet(req)
 	})
 
+	server.Handle("GET", "/api/2.0/secrets/list", func(req Request) any {
+		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.
+	server.Handle("GET", "/driver-proxy-api/o/{workspace_id}/{cluster_id}/{port}/metadata", func(req Request) any {
+		return Response{Body: sshTunnelRemoteUser}
+	})
+
+	server.Handle("GET", "/driver-proxy-api/o/{workspace_id}/{cluster_id}/{port}/logs", func(req Request) any {
+		return Response{Body: ""}
+	})
+
 	// Secrets ACLs:
 	server.Handle("GET", "/api/2.0/secrets/acls/get", func(req Request) any {
 		return req.Workspace.SecretsAclsGet(req)

libs/testserver/jobs.go

@@ -15,6 +15,7 @@ import (
 
 	"github.com/databricks/databricks-sdk-go/service/compute"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
+	"github.com/databricks/databricks-sdk-go/service/workspace"
 )
 
 const missingJobGitProviderMessage = "git_source.git_provider must be one of: github,gitlab,bitbucketcloud,gitlabenterpriseedition,bitbucketserver,azuredevopsservices,githubenterprise,awscodecommit"
@@ -519,9 +520,45 @@ func (s *FakeWorkspace) JobsSubmit(req Request) Response {
 		Tasks:      tasks,
 	}
 
+	// No tunnel server runs locally, so synthesize the metadata.json it would
+	// publish; `ssh connect` polls for it before connecting.
+	if strings.HasPrefix(runName, sshTunnelBootstrapRunPrefix) {
+		s.writeSSHTunnelMetadata(request)
+	}
+
 	return Response{Body: jobs.SubmitRunResponse{RunId: runId}}
 }
 
+const (
+	sshTunnelBootstrapRunPrefix = "ssh-server-bootstrap-"
+	sshTunnelBootstrapNotebook  = "ssh-server-bootstrap"
+	sshTunnelServerPort         = 7772
+	sshTunnelClusterID          = "1234-567890-serverless"
+	sshTunnelRemoteUser         = "spark"
+)
+
+// writeSSHTunnelMetadata publishes the metadata.json a real tunnel server would
+// write next to the bootstrap notebook. Callers must hold the workspace lock.
+func (s *FakeWorkspace) writeSSHTunnelMetadata(request jobs.SubmitRun) {
+	for _, t := range request.Tasks {
+		if t.NotebookTask == nil {
+			continue
+		}
+		metadataPath := strings.TrimSuffix(t.NotebookTask.NotebookPath, sshTunnelBootstrapNotebook) + "metadata.json"
+		metadata, err := json.Marshal(map[string]any{
+			"port":       sshTunnelServerPort,
+			"cluster_id": sshTunnelClusterID,
+		})
+		if err != nil {
+			continue
+		}
+		s.files[metadataPath] = FileEntry{
+			Info: workspace.ObjectInfo{ObjectType: "FILE", Path: metadataPath},
+			Data: metadata,
+		}
+	}
+}
+
 // executePythonWheelTask runs a python wheel task locally using uv.
 // For tasks using existing_cluster_id, the venv is cached per cluster to match
 // cloud behavior where libraries are cached on running clusters.