diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml index 37c22bb0d3b..8f5f981e6b8 100644 --- a/.github/workflows/push.yml +++ b/.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 diff --git a/acceptance/ssh/connect-serverless-gpu/out.stdout.txt b/acceptance/ssh/connect-serverless-gpu/out.stdout.txt index 41cae5e7d16..e69de29bb2d 100644 --- a/acceptance/ssh/connect-serverless-gpu/out.stdout.txt +++ b/acceptance/ssh/connect-serverless-gpu/out.stdout.txt @@ -1 +0,0 @@ -Connection successful diff --git a/acceptance/ssh/connect-serverless-gpu/out.test.toml b/acceptance/ssh/connect-serverless-gpu/out.test.toml index ab8691ddb9b..c777e3ce206 100644 --- a/acceptance/ssh/connect-serverless-gpu/out.test.toml +++ b/acceptance/ssh/connect-serverless-gpu/out.test.toml @@ -1,4 +1,4 @@ -Local = false +Local = true Cloud = false RequiresUnityCatalog = true CloudEnvs.gcp = false diff --git a/acceptance/ssh/connect-serverless-gpu/output.txt b/acceptance/ssh/connect-serverless-gpu/output.txt index e69de29bb2d..b7013545a8b 100644 --- a/acceptance/ssh/connect-serverless-gpu/output.txt +++ b/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 + } +} diff --git a/acceptance/ssh/connect-serverless-gpu/script b/acceptance/ssh/connect-serverless-gpu/script index 8c8538f6e32..e7e95e443d0 100644 --- a/acceptance/ssh/connect-serverless-gpu/script +++ b/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 diff --git a/acceptance/ssh/connect-serverless-gpu/test.toml b/acceptance/ssh/connect-serverless-gpu/test.toml index fa125d9bb5d..24b0870236f 100644 --- a/acceptance/ssh/connect-serverless-gpu/test.toml +++ b/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 diff --git a/experimental/ssh/internal/proxy/client_server_sshd_test.go b/experimental/ssh/internal/proxy/client_server_sshd_test.go new file mode 100644 index 00000000000..097668d2921 --- /dev/null +++ b/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") + } +} diff --git a/libs/testserver/handlers.go b/libs/testserver/handlers.go index 802baa0118a..338eb7e3e50 100644 --- a/libs/testserver/handlers.go +++ b/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) diff --git a/libs/testserver/jobs.go b/libs/testserver/jobs.go index c3537d00187..2d019c3e496 100644 --- a/libs/testserver/jobs.go +++ b/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. diff --git a/libs/testserver/secrets.go b/libs/testserver/secrets.go index 1f298ea9f0c..1b737be0d61 100644 --- a/libs/testserver/secrets.go +++ b/libs/testserver/secrets.go @@ -4,6 +4,7 @@ import ( "encoding/base64" "encoding/json" "fmt" + "slices" "github.com/databricks/databricks-sdk-go/service/workspace" ) @@ -40,6 +41,38 @@ func (s *FakeWorkspace) SecretsPut(req Request) Response { return Response{} } +// SecretsList models GET /api/2.0/secrets/list. A missing scope must return +// RESOURCE_DOES_NOT_EXIST so callers can branch on apierr.ErrResourceDoesNotExist. +func (s *FakeWorkspace) SecretsList(req Request) Response { + defer s.LockUnlock()() + + scope := req.URL.Query().Get("scope") + + if _, exists := s.SecretScopes[scope]; !exists { + return Response{ + StatusCode: 404, + Body: map[string]string{"error_code": "RESOURCE_DOES_NOT_EXIST", "message": fmt.Sprintf("Scope %s does not exist", scope)}, + } + } + + keys := make([]string, 0, len(s.Secrets[scope])) + for key := range s.Secrets[scope] { + keys = append(keys, key) + } + slices.Sort(keys) + + secrets := make([]workspace.SecretMetadata, 0, len(keys)) + for _, key := range keys { + secrets = append(secrets, workspace.SecretMetadata{Key: key}) + } + + return Response{ + Body: workspace.ListSecretsResponse{ + Secrets: secrets, + }, + } +} + func (s *FakeWorkspace) SecretsGet(req Request) Response { defer s.LockUnlock()()