diff --git a/.nextchanges/cli/ssh-connect-pins-host-key.md b/.nextchanges/cli/ssh-connect-pins-host-key.md new file mode 100644 index 00000000000..5f1aceff7dc --- /dev/null +++ b/.nextchanges/cli/ssh-connect-pins-host-key.md @@ -0,0 +1 @@ +* `databricks ssh connect` and `ssh setup` now verify the tunnel's SSH host key against the key the workspace published for the connection, recorded in `~/.databricks/ssh-tunnel-known-hosts/` instead of `~/.ssh/known_hosts`. Reconnecting with a name used before no longer fails with `Host key verification failed` when the compute behind that name changed, and no longer needs a manual `ssh-keygen -R`; host blocks written by an earlier `databricks ssh setup` pick this up once you re-run it. ([#6557](https://github.com/databricks/cli/pull/6557)) diff --git a/acceptance/ssh/connect-serverless-cpu/script b/acceptance/ssh/connect-serverless-cpu/script index 40bbb6debc2..02b0068f3da 100644 --- a/acceptance/ssh/connect-serverless-cpu/script +++ b/acceptance/ssh/connect-serverless-cpu/script @@ -22,6 +22,6 @@ else CLI_RELEASES_DIR="$PWD/releases" title "CLI ssh connect (local sshd)\n" - errcode $CLI ssh connect --releases-dir=$CLI_RELEASES_DIR --user-known-hosts-file=known_hosts -- "echo 'Connection successful'" >out.stdout.txt 2>LOG.stderr + errcode $CLI ssh connect --releases-dir=$CLI_RELEASES_DIR --known-hosts-dir=known-hosts -- "echo 'Connection successful'" >out.stdout.txt 2>LOG.stderr trace print_requests.py //api/2.2/jobs/runs/submit fi diff --git a/acceptance/ssh/connect-serverless-cpu/test.toml b/acceptance/ssh/connect-serverless-cpu/test.toml index 105d62b9ae9..ef09c03dec4 100644 --- a/acceptance/ssh/connect-serverless-cpu/test.toml +++ b/acceptance/ssh/connect-serverless-cpu/test.toml @@ -5,9 +5,11 @@ Cloud = false # Record requests so we can assert the bootstrap job the CLI submits. RecordRequests = true -# Stand-in release archives staged by the local run; not golden output. +# Stand-in release archives staged by the local run, plus the host key the CLI pins under +# known-hosts; not golden output. Ignore = [ "releases", + "known-hosts", ] # The local run drives a real sshd, reliable only on Linux. diff --git a/acceptance/ssh/connect-serverless-gpu/script b/acceptance/ssh/connect-serverless-gpu/script index 0943ec9d2b7..b2e5e3d40fa 100644 --- a/acceptance/ssh/connect-serverless-gpu/script +++ b/acceptance/ssh/connect-serverless-gpu/script @@ -24,6 +24,6 @@ else 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 + errcode $CLI ssh connect --name serverless-gpu-test --accelerator=GPU_1xA10 --releases-dir=$CLI_RELEASES_DIR --known-hosts-dir=known-hosts -- "echo 'Connection successful'" >out.stdout.txt 2>LOG.stderr 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 9ab5c69e8a2..8617cae3cde 100644 --- a/acceptance/ssh/connect-serverless-gpu/test.toml +++ b/acceptance/ssh/connect-serverless-gpu/test.toml @@ -5,9 +5,11 @@ Cloud = false # 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. +# Local runs stage stand-in release archives here, and the CLI pins the server's host key +# under known-hosts; neither is part of the golden output. Ignore = [ "releases", + "known-hosts", ] # Linux-only: the local run drives a real sshd, reliable only there (and where diff --git a/acceptance/ssh/connection/known_hosts b/acceptance/ssh/connection/known_hosts deleted file mode 100644 index 197ec818881..00000000000 --- a/acceptance/ssh/connection/known_hosts +++ /dev/null @@ -1 +0,0 @@ -# not actually checked by tests; accept-new appends the ephemeral host key here diff --git a/acceptance/ssh/connection/script b/acceptance/ssh/connection/script index c69babb6dd0..3858cd1b6ba 100644 --- a/acceptance/ssh/connection/script +++ b/acceptance/ssh/connection/script @@ -10,7 +10,7 @@ fi if [ -n "${CLOUD_ENV:-}" ]; then # 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 + errcode $CLI ssh connect --cluster=$TEST_DEFAULT_CLUSTER_ID --releases-dir=$CLI_RELEASES_DIR --known-hosts-dir=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 @@ -24,6 +24,6 @@ else 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 + errcode $CLI ssh connect --cluster=$TEST_DEFAULT_CLUSTER_ID --releases-dir=$CLI_RELEASES_DIR --known-hosts-dir=known-hosts -- "echo 'Connection successful'" >out.stdout.txt 2>LOG.stderr trace print_requests.py //api/2.2/jobs/runs/submit fi diff --git a/acceptance/ssh/connection/test.toml b/acceptance/ssh/connection/test.toml index 3f67650471d..e09207c3d86 100644 --- a/acceptance/ssh/connection/test.toml +++ b/acceptance/ssh/connection/test.toml @@ -6,9 +6,11 @@ Cloud = false # Assert the dedicated-cluster 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. +# Local runs stage stand-in release archives here, and the CLI pins the server's host key +# under known-hosts; neither is part of the golden output. Ignore = [ "releases", + "known-hosts", ] # Linux-only: the local run drives a real sshd, reliable only there (and where diff --git a/acceptance/ssh/reconnect-reused-name/out.connect-1.txt b/acceptance/ssh/reconnect-reused-name/out.connect-1.txt new file mode 100644 index 00000000000..41cae5e7d16 --- /dev/null +++ b/acceptance/ssh/reconnect-reused-name/out.connect-1.txt @@ -0,0 +1 @@ +Connection successful diff --git a/acceptance/ssh/reconnect-reused-name/out.connect-2.txt b/acceptance/ssh/reconnect-reused-name/out.connect-2.txt new file mode 100644 index 00000000000..41cae5e7d16 --- /dev/null +++ b/acceptance/ssh/reconnect-reused-name/out.connect-2.txt @@ -0,0 +1 @@ +Connection successful diff --git a/acceptance/ssh/reconnect-reused-name/out.test.toml b/acceptance/ssh/reconnect-reused-name/out.test.toml new file mode 100644 index 00000000000..c826e1c76b6 --- /dev/null +++ b/acceptance/ssh/reconnect-reused-name/out.test.toml @@ -0,0 +1,5 @@ +Cloud = false +GOOS.darwin = false +GOOS.linux = true +GOOS.windows = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/acceptance/ssh/reconnect-reused-name/output.txt b/acceptance/ssh/reconnect-reused-name/output.txt new file mode 100644 index 00000000000..43e837cd723 --- /dev/null +++ b/acceptance/ssh/reconnect-reused-name/output.txt @@ -0,0 +1,10 @@ + +=== CLI ssh connect over a stale host key for the connection name + +=== The stale key is replaced by the one the workspace published +stale key replaced + +>>> grep -c ^reused-name ssh- known-hosts/reused-name +1 + +=== Reconnecting to the running server keeps working diff --git a/acceptance/ssh/reconnect-reused-name/script b/acceptance/ssh/reconnect-reused-name/script new file mode 100644 index 00000000000..b07e4304746 --- /dev/null +++ b/acceptance/ssh/reconnect-reused-name/script @@ -0,0 +1,42 @@ +# Reconnecting with a name used before must not fail on the host key +# recorded for that name earlier. The tunnel's host key lives in the workspace, and the +# CLI pins it for the connection name, so a key recorded by an earlier instance (or by +# the same name in another workspace) is replaced rather than tripping strict checking. +# +# The test server backs the /ssh websocket with a real sshd whose host key comes from the +# session's secret scope, so this asserts a full handshake against a pinned key. Skip when +# sshd is absent (only task test-exp-ssh provisions it). +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 + +# No release artifacts locally and they're never executed; stand in empty archives. +mkdir -p releases +: >releases/databricks_cli_linux_amd64.zip +: >releases/databricks_cli_linux_arm64.zip +CLI_RELEASES_DIR="$PWD/releases" + +# The state a user reconnecting with a reused name is in: an entry for this connection +# name, pinned days earlier, whose key no longer belongs to anything. The CLI names the file +# after the connection, so --known-hosts-dir is enough to keep the pin inside the test. +STALE_KEY=AAAAC3NzaC1lZDI1NTE5AAAAICTzJweKcgiNoBUAyuvCY2Qu1Od8mKBON5aJeA03Nl+D +mkdir -p known-hosts +echo "reused-name ssh-ed25519 $STALE_KEY" >known-hosts/reused-name + +# The remote command's output is the assertion (see out.connect-*.txt); the tunnel's own +# progress and teardown messages go to stderr, where the order of the last few lines +# depends on which side closes the websocket first. +title "CLI ssh connect over a stale host key for the connection name\n" +errcode $CLI ssh connect --name reused-name --releases-dir=$CLI_RELEASES_DIR --known-hosts-dir=known-hosts -- "echo 'Connection successful'" >out.connect-1.txt 2>LOG.connect-1 + +title "The stale key is replaced by the one the workspace published\n" +if grep -q "$STALE_KEY" known-hosts/reused-name; then + echo "stale key still pinned" +else + echo "stale key replaced" +fi +trace grep -c "^reused-name ssh-" known-hosts/reused-name + +title "Reconnecting to the running server keeps working\n" +errcode $CLI ssh connect --name reused-name --releases-dir=$CLI_RELEASES_DIR --known-hosts-dir=known-hosts -- "echo 'Connection successful'" >out.connect-2.txt 2>LOG.connect-2 diff --git a/acceptance/ssh/reconnect-reused-name/test.toml b/acceptance/ssh/reconnect-reused-name/test.toml new file mode 100644 index 00000000000..e1577f367a1 --- /dev/null +++ b/acceptance/ssh/reconnect-reused-name/test.toml @@ -0,0 +1,18 @@ +Cloud = false + +# The connection name is explicit, so nothing here depends on the local host name. + +# Local runs stage stand-in release archives here; not part of the golden output. +Ignore = [ + "releases", + "known-hosts", +] + +# 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 +GOOS.darwin = false +GOOS.windows = false + +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/experimental/ssh/README.md b/experimental/ssh/README.md index 769006ef2df..660f1dc6c76 100644 --- a/experimental/ssh/README.md +++ b/experimental/ssh/README.md @@ -17,6 +17,20 @@ B. Spawn an ssh session directly: databricks ssh connect --cluster=id ``` +### Connection names and host keys + +`--name` is a stable handle for serverless compute, not a live session id: connecting with a +name reuses the SSH server still running under it, and starts a new one on the same name once +the previous server has shut down (`--shutdown-delay` after its last client disconnects). + +Both cases verify the server's host key. The server generates the key on first use, keeps it in +the connection's secret scope in the workspace and reuses it for every `sshd` it launches, so +the workspace is the authority on the key. The client reads it from there and pins it in +`~/.databricks/ssh-tunnel-known-hosts/` before each connection, and points ssh at that +file with `StrictHostKeyChecking yes`. Tunnel host keys therefore never land in +`~/.ssh/known_hosts`, where a name - unique only within one workspace - would collide with an +entry left by other compute. + ## Development ```shell ./task build snapshot-release diff --git a/experimental/ssh/cmd/connect.go b/experimental/ssh/cmd/connect.go index adf12e46db6..bb965d8306b 100644 --- a/experimental/ssh/cmd/connect.go +++ b/experimental/ssh/cmd/connect.go @@ -53,7 +53,7 @@ Connect to a dedicated cluster: var handoverTimeout time.Duration var releasesDir string var autoStartCluster bool - var userKnownHostsFile string + var knownHostsDir string var liteswap string var skipSettingsCheck bool var environmentVersion int @@ -82,8 +82,8 @@ Connect to a dedicated cluster: cmd.Flags().StringVar(&releasesDir, "releases-dir", "", "Directory for local SSH tunnel development releases") cmd.Flags().MarkHidden("releases-dir") - cmd.Flags().StringVar(&userKnownHostsFile, "user-known-hosts-file", "", "Path to user known hosts file for SSH client") - cmd.Flags().MarkHidden("user-known-hosts-file") + cmd.Flags().StringVar(&knownHostsDir, "known-hosts-dir", "", "Directory for the known hosts files the CLI maintains for tunnel connections") + cmd.Flags().MarkHidden("known-hosts-dir") cmd.Flags().StringVar(&liteswap, "liteswap", "", "Liteswap header value for traffic routing (dev/test only)") cmd.Flags().MarkHidden("liteswap") @@ -145,7 +145,8 @@ Connect to a dedicated cluster: AutoStartCluster: autoStartCluster, ClientPublicKeyName: clientPublicKeyName, ClientPrivateKeyName: clientPrivateKeyName, - UserKnownHostsFile: userKnownHostsFile, + ServerPublicKeyName: serverPublicKeyName, + KnownHostsDir: knownHostsDir, Liteswap: liteswap, SkipSettingsCheck: skipSettingsCheck, EnvironmentVersion: environmentVersion, diff --git a/experimental/ssh/internal/client/client.go b/experimental/ssh/internal/client/client.go index 3e09bfdc80d..e95351c2e45 100644 --- a/experimental/ssh/internal/client/client.go +++ b/experimental/ssh/internal/client/client.go @@ -97,18 +97,24 @@ type ClientOptions struct { ReleasesDir string // Directory for local SSH keys. Defaults to ~/.databricks/ssh-tunnel-keys SSHKeysDir string + // Directory for the known hosts files the CLI maintains for tunnel connections. + // Defaults to ~/.databricks/ssh-tunnel-known-hosts. The file within it is always named + // after the session, so this relocates a directory the CLI owns and rewrites; it never + // names a file to overwrite. + KnownHostsDir string // Client public key name located in the ssh-tunnel secrets scope. ClientPublicKeyName string // Client private key name located in the ssh-tunnel secrets scope. ClientPrivateKeyName string + // Server public key name located in the ssh-tunnel secrets scope. The server publishes + // its host key there; the client pins it so ssh can verify the server. + ServerPublicKeyName string // If true, the CLI will attempt to start the cluster if it is not running. AutoStartCluster bool // Optional auth profile name. If present, will be added as --profile flag to the ProxyCommand while spawning ssh client. Profile string // Additional arguments to pass to the SSH client in the non proxy mode. AdditionalArgs []string - // Optional path to the user known hosts file. - UserKnownHostsFile string // Liteswap header value for traffic routing (dev/test only). Liteswap string // If true, skip checking and updating IDE settings. @@ -469,6 +475,15 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt log.Infof(ctx, "Cluster ID (from serverless job): %s", clusterID) } + // Pin the running server's host key for this session name. In proxy mode this happens + // before the tunnel carries a single byte, so it is in place well before the ssh client + // that spawned us verifies the key during the key exchange. + knownHostsPath, err := pinServerHostKey(ctx, client, sessionID, secretScopeName, opts) + if err != nil { + outcome.errorCategory = protos.SshTunnelErrorCategoryKeyGenerationFailed + return err + } + if !opts.ProxyMode { cmdio.LogString(ctx, "Connected!") } @@ -480,14 +495,41 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt if opts.ProxyMode { return runSSHProxy(ctx, client, serverPort, clusterID, opts) } else if opts.IDE != "" { - return runIDE(ctx, client, userName, keyPath, serverPort, clusterID, opts) + return runIDE(ctx, client, userName, keyPath, knownHostsPath, serverPort, clusterID, opts) } else { log.Infof(ctx, "Additional SSH arguments: %v", opts.AdditionalArgs) - return spawnSSHClient(ctx, client, userName, keyPath, serverPort, clusterID, opts) + return spawnSSHClient(ctx, client, userName, keyPath, knownHostsPath, serverPort, clusterID, opts) } } -func runIDE(ctx context.Context, client *databricks.WorkspaceClient, userName, keyPath string, serverPort int, clusterID string, opts ClientOptions) error { +// pinServerHostKey records the running server's host key as the entry ssh accepts for this +// session, and returns the path of the known_hosts file it wrote. +// +// The server publishes its host key to the session's secret scope when it starts and reuses +// it for every sshd it launches, which makes the scope the authority on the key. Reading it +// here replaces whatever was recorded for this session name before, so a key left by an +// earlier instance - or by the same name in another workspace, since a name is unique only +// within one - can no longer fail an otherwise valid connection. +func pinServerHostKey(ctx context.Context, client *databricks.WorkspaceClient, sessionID, secretScopeName string, opts ClientOptions) (string, error) { + knownHostsPath, err := sshconfig.GetKnownHostsPath(ctx, sessionID, opts.KnownHostsDir) + if err != nil { + return "", err + } + + publicKey, err := keys.GetSecret(ctx, client, secretScopeName, opts.ServerPublicKeyName) + if err != nil { + return "", fmt.Errorf("failed to read the SSH server host key: %w", err) + } + + if err := sshconfig.PinHostKey(knownHostsPath, sessionID, publicKey); err != nil { + return "", err + } + + log.Infof(ctx, "Pinned the SSH server host key for %s in %s", sessionID, knownHostsPath) + return knownHostsPath, nil +} + +func runIDE(ctx context.Context, client *databricks.WorkspaceClient, userName, keyPath, knownHostsPath string, serverPort int, clusterID string, opts ClientOptions) error { connectionName := opts.SessionIdentifier() if connectionName == "" { return errors.New("connection name is required for IDE integration") @@ -505,7 +547,7 @@ func runIDE(ctx context.Context, client *databricks.WorkspaceClient, userName, k return fmt.Errorf("failed to get SSH config path: %w", err) } - err = ensureSSHConfigEntry(ctx, configPath, connectionName, userName, keyPath, serverPort, clusterID, opts) + err = ensureSSHConfigEntry(ctx, configPath, connectionName, userName, keyPath, knownHostsPath, serverPort, clusterID, opts) if err != nil { return fmt.Errorf("failed to ensure SSH config entry: %w", err) } @@ -513,7 +555,7 @@ func runIDE(ctx context.Context, client *databricks.WorkspaceClient, userName, k return vscode.LaunchIDE(ctx, opts.IDE, connectionName, currentUser.UserName) } -func ensureSSHConfigEntry(ctx context.Context, configPath, hostName, userName, keyPath string, serverPort int, clusterID string, opts ClientOptions) error { +func ensureSSHConfigEntry(ctx context.Context, configPath, hostName, userName, keyPath, knownHostsPath string, serverPort int, clusterID string, opts ClientOptions) error { // Ensure the Include directive exists in the main SSH config err := sshconfig.EnsureIncludeDirective(ctx, configPath) if err != nil { @@ -529,7 +571,10 @@ func ensureSSHConfigEntry(ctx context.Context, configPath, hostName, userName, k return fmt.Errorf("failed to generate ProxyCommand: %w", err) } - hostConfig := sshconfig.GenerateHostConfig(hostName, userName, keyPath, proxyCommand) + // The host key is pinned under the session ID (see pinServerHostKey), so emit it as + // HostKeyAlias to keep the key lookup matching the pinned entry (DECO-27882). Here the + // host alias already is the session ID, but passing it explicitly keeps the two in step. + hostConfig := sshconfig.GenerateHostConfig(hostName, userName, keyPath, knownHostsPath, opts.SessionIdentifier(), proxyCommand) _, err = sshconfig.CreateOrUpdateHostConfig(ctx, hostName, hostConfig, true) if err != nil { @@ -847,19 +892,21 @@ func buildRemoteShellArgs(opts ClientOptions, wsHome string) []string { // allocation (-t) for the interactive case is added before the host: ssh stops // parsing options at the destination, so a -t placed after the host would be // treated as part of the remote command rather than as ssh's force-PTY flag. -func buildSSHArgs(userName, privateKeyPath, proxyCommand, hostName, wsHome string, opts ClientOptions) []string { +// +// Host key checking is strict rather than accept-new because the caller has already +// pinned the server's key in knownHostsPath (see pinServerHostKey), so there is nothing +// left to accept on trust. +func buildSSHArgs(userName, privateKeyPath, knownHostsPath, proxyCommand, hostName, wsHome string, opts ClientOptions) []string { sshArgs := []string{ "-l", userName, "-i", privateKeyPath, "-o", "IdentitiesOnly=yes", - "-o", "StrictHostKeyChecking=accept-new", + "-o", "StrictHostKeyChecking=yes", + "-o", "UserKnownHostsFile=" + knownHostsPath, "-o", "ConnectTimeout=360", "-o", "ServerAliveInterval=" + strconv.Itoa(sshconfig.ServerAliveIntervalSeconds), "-o", "ProxyCommand=" + proxyCommand, } - if opts.UserKnownHostsFile != "" { - sshArgs = append(sshArgs, "-o", "UserKnownHostsFile="+opts.UserKnownHostsFile) - } if len(opts.AdditionalArgs) == 0 { sshArgs = append(sshArgs, "-t") } @@ -868,7 +915,7 @@ func buildSSHArgs(userName, privateKeyPath, proxyCommand, hostName, wsHome strin return sshArgs } -func spawnSSHClient(ctx context.Context, client *databricks.WorkspaceClient, userName, privateKeyPath string, serverPort int, clusterID string, opts ClientOptions) error { +func spawnSSHClient(ctx context.Context, client *databricks.WorkspaceClient, userName, privateKeyPath, knownHostsPath string, serverPort int, clusterID string, opts ClientOptions) error { // Create a copy with metadata for the ProxyCommand optsWithMetadata := opts optsWithMetadata.ServerMetadata = FormatMetadata(userName, serverPort, clusterID) @@ -892,7 +939,7 @@ func spawnSSHClient(ctx context.Context, client *databricks.WorkspaceClient, use } } - sshArgs := buildSSHArgs(userName, privateKeyPath, proxyCommand, hostName, wsHome, opts) + sshArgs := buildSSHArgs(userName, privateKeyPath, knownHostsPath, proxyCommand, hostName, wsHome, opts) log.Debugf(ctx, "Launching SSH client: ssh %s", strings.Join(sshArgs, " ")) sshCmd := exec.CommandContext(ctx, "ssh", sshArgs...) @@ -911,7 +958,7 @@ func spawnSSHClient(ctx context.Context, client *databricks.WorkspaceClient, use // own logs — fetch them from the /logs endpoint and show them instead of leaving the user // with ssh's opaque "Connection closed" message. if exitErr, ok := errors.AsType[*exec.ExitError](err); ok && exitErr.ExitCode() == 255 { - if hint := hostKeyChangedHint(stderrTail.String(), hostName, opts.UserKnownHostsFile); hint != "" { + if hint := hostKeyChangedHint(stderrTail.String(), hostName, knownHostsPath); hint != "" { cmdio.LogString(ctx, cmdio.Yellow(ctx, hint)) } else if logs := fetchServerErrorLogs(ctx, client, clusterID, serverPort, opts.Liteswap); logs != "" { cmdio.LogString(ctx, cmdio.Yellow(ctx, "The SSH connection closed unexpectedly. Recent SSH server errors:")) @@ -1186,10 +1233,11 @@ func (w *tailWriter) String() string { return string(w.buf) } -// hostKeyChangedHint returns advice for clearing a stale known_hosts entry when ssh's stderr -// shows a host-key-verification failure, or "" if the failure was something else. A cluster that -// has been recreated keeps the same connection name but gets a new host key, so the old entry no -// longer matches and ssh aborts the connection. +// hostKeyChangedHint returns advice when ssh's stderr shows a host-key-verification failure, +// or "" if the failure was something else. The CLI pins the key the workspace published for +// this session into knownHostsFile immediately before connecting, so the entry ssh rejected +// is not a leftover from an earlier session: the server that answered is presenting a +// different key than the one the workspace recorded for it. func hostKeyChangedHint(stderr, hostName, knownHostsFile string) string { // "Host key verification failed." is OpenSSH's fixed message for this case; matching it is the // only signal ssh gives (the "don't branch on err.Error()" rule is about Go errors, not the @@ -1197,13 +1245,10 @@ func hostKeyChangedHint(stderr, hostName, knownHostsFile string) string { if !strings.Contains(stderr, "Host key verification failed") { return "" } - cmd := "ssh-keygen -R " + hostName - if knownHostsFile != "" { - // ssh-keygen -R defaults to ~/.ssh/known_hosts, so name the custom file explicitly. - cmd += " -f " + knownHostsFile - } - return "The host key for " + hostName + " has changed. " + - "Remove the stale entry and reconnect:\n " + cmd + return "The SSH server answering for " + hostName + " presented a host key that does not match the one " + + "the workspace published for it.\nThe CLI refreshed that key in " + knownHostsFile + + " just before connecting, so this is not a stale local entry: the tunnel reached a server " + + "the workspace does not know about" } func usagePolicyMatches(storedPolicy, requestedPolicy string) bool { diff --git a/experimental/ssh/internal/client/client_internal_test.go b/experimental/ssh/internal/client/client_internal_test.go index ba15384bb46..4729ef30e82 100644 --- a/experimental/ssh/internal/client/client_internal_test.go +++ b/experimental/ssh/internal/client/client_internal_test.go @@ -319,20 +319,24 @@ func TestHostKeyChangedHint(t *testing.T) { hostName string knownHostsFile string wantContains []string + wantOmits []string wantEmpty bool }{ { - name: "host key failure", - stderr: hostKeyFailureStderr, - hostName: "databricks-cpu-6e7644d0", - wantContains: []string{"databricks-cpu-6e7644d0", "ssh-keygen -R databricks-cpu-6e7644d0"}, + name: "host key failure names the host and the pinned file", + stderr: hostKeyFailureStderr, + hostName: "databricks-cpu-6e7644d0", + knownHostsFile: "/tmp/known_hosts", + wantContains: []string{"databricks-cpu-6e7644d0", "/tmp/known_hosts"}, }, { - name: "host key failure with custom known_hosts file", + // The stale-entry advice this hint used to give no longer applies: the CLI + // rewrites the entry from the workspace before every connection. + name: "host key failure does not blame a stale local entry", stderr: hostKeyFailureStderr, hostName: "databricks-cpu-6e7644d0", knownHostsFile: "/tmp/known_hosts", - wantContains: []string{"ssh-keygen -R databricks-cpu-6e7644d0 -f /tmp/known_hosts"}, + wantOmits: []string{"ssh-keygen -R"}, }, { name: "unrelated failure", @@ -351,6 +355,9 @@ func TestHostKeyChangedHint(t *testing.T) { for _, want := range tt.wantContains { assert.Contains(t, got, want) } + for _, unwanted := range tt.wantOmits { + assert.NotContains(t, got, unwanted) + } }) } } @@ -378,7 +385,7 @@ func TestBuildRemoteShellArgs(t *testing.T) { } func TestBuildSSHArgsSetsServerAliveInterval(t *testing.T) { - args := buildSSHArgs("user", "/key", "proxy command", "myhost", "", ClientOptions{}) + args := buildSSHArgs("user", "/key", "/pins/myhost", "proxy command", "myhost", "", ClientOptions{}) // ssh stops parsing options at the destination, so an option placed after the host would be // treated as part of the remote command rather than as an ssh option. @@ -388,6 +395,21 @@ func TestBuildSSHArgsSetsServerAliveInterval(t *testing.T) { assert.Less(t, optIdx, slices.Index(args, "myhost"), "the option must precede the destination host") } +func TestBuildSSHArgsPinsHostKey(t *testing.T) { + args := buildSSHArgs("user", "/key", "/pins/myhost", "proxy command", "myhost", "", ClientOptions{}) + + // The pinned file is the whole point of strict checking here: without it ssh would + // fall back to ~/.ssh/known_hosts, where an entry for this name may be left over from + // other compute (DECO-27882). + hostIdx := slices.Index(args, "myhost") + for _, want := range []string{"StrictHostKeyChecking=yes", "UserKnownHostsFile=/pins/myhost"} { + optIdx := slices.Index(args, want) + require.NotEqual(t, -1, optIdx, "%s must be passed to ssh", want) + require.Equal(t, "-o", args[optIdx-1]) + assert.Less(t, optIdx, hostIdx, "the option must precede the destination host") + } +} + func TestBuildSSHArgsPTYPlacement(t *testing.T) { indexOf := func(args []string, want string) int { for i, a := range args { @@ -399,7 +421,7 @@ func TestBuildSSHArgsPTYPlacement(t *testing.T) { } t.Run("interactive forces a PTY before the destination", func(t *testing.T) { - args := buildSSHArgs("user", "/key", "proxy command", "myhost", "/Workspace/Users/me@example.com", ClientOptions{}) + args := buildSSHArgs("user", "/key", "/pins/myhost", "proxy command", "myhost", "/Workspace/Users/me@example.com", ClientOptions{}) ptyIdx := indexOf(args, "-t") hostIdx := indexOf(args, "myhost") require.NotEqual(t, -1, ptyIdx, "-t must be present for interactive sessions") @@ -411,7 +433,7 @@ func TestBuildSSHArgsPTYPlacement(t *testing.T) { }) t.Run("non-interactive does not force a PTY", func(t *testing.T) { - args := buildSSHArgs("user", "/key", "proxy command", "myhost", "", ClientOptions{AdditionalArgs: []string{"ls", "-la"}}) + args := buildSSHArgs("user", "/key", "/pins/myhost", "proxy command", "myhost", "", ClientOptions{AdditionalArgs: []string{"ls", "-la"}}) assert.Equal(t, -1, indexOf(args, "-t"), "no PTY for non-interactive passthrough") hostIdx := indexOf(args, "myhost") require.NotEqual(t, -1, hostIdx) diff --git a/experimental/ssh/internal/setup/setup.go b/experimental/ssh/internal/setup/setup.go index aa77e87555d..2f55da1ca46 100644 --- a/experimental/ssh/internal/setup/setup.go +++ b/experimental/ssh/internal/setup/setup.go @@ -46,7 +46,20 @@ func generateHostConfig(ctx context.Context, opts SetupOptions, proxyCommand str return "", fmt.Errorf("failed to get local keys folder: %w", err) } - hostConfig := sshconfig.GenerateHostConfig(opts.HostName, "root", identityFilePath, proxyCommand) + // The ProxyCommand writes this file before the connection reaches host key + // verification, so it does not have to exist yet. It carries no directory override, so + // resolve the default one here as well. + knownHostsPath, err := sshconfig.GetKnownHostsPath(ctx, opts.ClusterID, "") + if err != nil { + return "", err + } + + // The ProxyCommand pins the server's key under the cluster ID (the session ID for a + // dedicated cluster), but the block is written as `Host `. When the + // user-facing name differs from the cluster ID, ssh would look the key up under the name + // and fail strict checking, so pass the cluster ID as HostKeyAlias to match the pinned + // entry (DECO-27882). + hostConfig := sshconfig.GenerateHostConfig(opts.HostName, "root", identityFilePath, knownHostsPath, opts.ClusterID, proxyCommand) return hostConfig, nil } diff --git a/experimental/ssh/internal/setup/setup_test.go b/experimental/ssh/internal/setup/setup_test.go index 7ce1097dcdc..5dc04e1b366 100644 --- a/experimental/ssh/internal/setup/setup_test.go +++ b/experimental/ssh/internal/setup/setup_test.go @@ -9,6 +9,7 @@ import ( "time" "github.com/databricks/cli/experimental/ssh/internal/client" + "github.com/databricks/cli/experimental/ssh/internal/sshconfig" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/experimental/mocks" @@ -103,13 +104,24 @@ func TestGenerateHostConfig_Valid(t *testing.T) { assert.Contains(t, result, "Host test-host") assert.Contains(t, result, "User root") - assert.Contains(t, result, "StrictHostKeyChecking accept-new") assert.Contains(t, result, "--cluster=cluster-123") assert.Contains(t, result, "--shutdown-delay=30s") assert.Contains(t, result, "--profile=test-profile") expectedKeyPath := filepath.Join(tmpDir, "cluster-123") assert.Contains(t, result, fmt.Sprintf(`IdentityFile %q`, expectedKeyPath)) + + // `ssh ` reaches ssh through this block and nothing else, so the host key the + // ProxyCommand pins has to be the one it verifies against (DECO-27882). + assert.Contains(t, result, "StrictHostKeyChecking yes") + expectedKnownHostsPath, err := sshconfig.GetKnownHostsPath(t.Context(), "cluster-123", "") + require.NoError(t, err) + assert.Contains(t, result, fmt.Sprintf(`UserKnownHostsFile %q`, expectedKnownHostsPath)) + + // The host name (test-host) differs from the cluster ID the key is pinned under, so the + // block has to carry HostKeyAlias cluster-123 or strict checking looks the key up under + // test-host and fails (DECO-27882). + assert.Contains(t, result, "\n HostKeyAlias cluster-123\n") } func TestGenerateHostConfig_WithoutProfile(t *testing.T) { @@ -207,6 +219,9 @@ func TestSetup_SuccessfulWithNewConfigFile(t *testing.T) { assert.Contains(t, hostConfigStr, "Host test-host") assert.Contains(t, hostConfigStr, "--cluster=cluster-123") assert.Contains(t, hostConfigStr, "--profile=test-profile") + // The written block pins the key lookup to the cluster ID, which differs from the host + // name test-host (DECO-27882). + assert.Contains(t, hostConfigStr, "HostKeyAlias cluster-123") } func TestSetup_AutoApproveRecreatesExistingHost(t *testing.T) { diff --git a/experimental/ssh/internal/sshconfig/knownhosts.go b/experimental/ssh/internal/sshconfig/knownhosts.go new file mode 100644 index 00000000000..8cfecc2a060 --- /dev/null +++ b/experimental/ssh/internal/sshconfig/knownhosts.go @@ -0,0 +1,89 @@ +package sshconfig + +import ( + "context" + "fmt" + "os" + "path/filepath" + + "github.com/databricks/cli/libs/env" + "golang.org/x/crypto/ssh" +) + +// GetKnownHostsPath returns the known_hosts file the CLI maintains for a session +// (sessionID is the connection name for serverless, the cluster ID otherwise). +// knownHostsDir defaults to ~/.databricks/ssh-tunnel-known-hosts. +// +// The file name is always the session ID, never a path the caller picks, because PinHostKey +// rewrites the file: relocating the directory is all a caller can do, so no shared +// known_hosts can end up being replaced with a single tunnel entry. +// +// Tunnel host keys are deliberately kept out of the user's ~/.ssh/known_hosts. A session +// name identifies compute within one workspace, while ~/.ssh/known_hosts is global and +// keyed by name alone, so the same name used in a second workspace - or against compute +// whose host key was regenerated - collides with the entry left by the first and trips +// strict host key checking on a connection that is perfectly legitimate. +func GetKnownHostsPath(ctx context.Context, sessionID, knownHostsDir string) (string, error) { + if knownHostsDir == "" { + homeDir, err := env.UserHomeDir(ctx) + if err != nil { + return "", fmt.Errorf("failed to get home directory: %w", err) + } + knownHostsDir = filepath.Join(homeDir, ".databricks", "ssh-tunnel-known-hosts") + } + return filepath.Join(knownHostsDir, sessionID), nil +} + +// PinHostKey makes publicKey the entry for hostName in the known_hosts file at path, +// replacing the file's previous contents. +// +// The tunnel server publishes its host key to the workspace, which makes the workspace the +// authority on it: pinning that key before every connection lets ssh verify the server +// with no trust-on-first-use window, and leaves no way for a key recorded under the same +// name earlier to fail a connection. The file belongs to the CLI, which is why it is +// rewritten rather than edited - an entry recorded with HashKnownHosts cannot be matched +// by host name, so selectively dropping the stale one is not possible. +func PinHostKey(path, hostName string, publicKey []byte) error { + parsed, _, _, _, err := ssh.ParseAuthorizedKey(publicKey) + if err != nil { + return fmt.Errorf("failed to parse the SSH server host key: %w", err) + } + // MarshalAuthorizedKey drops any comment and terminates the line with "\n". + line := hostName + " " + string(ssh.MarshalAuthorizedKey(parsed)) + + // Leave an already-correct pin alone, which is the common case: every ssh invocation + // refreshes it through the ProxyCommand, and an IDE opens several at once. Renaming over + // a file another ssh has open fails on Windows. + if existing, err := os.ReadFile(path); err == nil && string(existing) == line { + return nil + } + + dir := filepath.Dir(path) + if err := os.MkdirAll(dir, 0o700); err != nil { + return fmt.Errorf("failed to create known hosts directory: %w", err) + } + + // Write and rename so a connection racing this one (every ssh invocation runs the + // ProxyCommand, which refreshes the pin) never reads a half-written file. + tmp, err := os.CreateTemp(dir, ".known-hosts-*.tmp") + if err != nil { + return fmt.Errorf("failed to create known hosts file: %w", err) + } + defer os.Remove(tmp.Name()) + + _, err = tmp.WriteString(line) + if err == nil { + err = tmp.Chmod(0o600) + } + if closeErr := tmp.Close(); err == nil { + err = closeErr + } + if err != nil { + return fmt.Errorf("failed to write known hosts file: %w", err) + } + + if err := os.Rename(tmp.Name(), path); err != nil { + return fmt.Errorf("failed to replace known hosts file: %w", err) + } + return nil +} diff --git a/experimental/ssh/internal/sshconfig/knownhosts_test.go b/experimental/ssh/internal/sshconfig/knownhosts_test.go new file mode 100644 index 00000000000..1aa3fa19de6 --- /dev/null +++ b/experimental/ssh/internal/sshconfig/knownhosts_test.go @@ -0,0 +1,99 @@ +package sshconfig + +import ( + "os" + "path/filepath" + "runtime" + "testing" + "time" + + "github.com/databricks/cli/libs/env" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const ( + testHostKey = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAICTzJweKcgiNoBUAyuvCY2Qu1Od8mKBON5aJeA03Nl+D" + testOtherHostKey = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIErh/dj+R5gtO5dZ5ToM4KYZTpNHOQTHW/RsjRXeOBt3" +) + +func TestGetKnownHostsPath(t *testing.T) { + tmpDir := t.TempDir() + t.Setenv(env.HomeEnvVar(), tmpDir) + + path, err := GetKnownHostsPath(t.Context(), "databricks-cpu-6e7644d0", "") + require.NoError(t, err) + assert.Equal(t, filepath.Join(tmpDir, ".databricks", "ssh-tunnel-known-hosts", "databricks-cpu-6e7644d0"), path) +} + +func TestGetKnownHostsPathOverrideRelocatesTheDirectoryOnly(t *testing.T) { + // The override names a directory, never the file: PinHostKey replaces the file it is + // given, so a shared known_hosts must stay out of reach of it. + path, err := GetKnownHostsPath(t.Context(), "databricks-cpu-6e7644d0", filepath.Join("/home/me", ".ssh")) + require.NoError(t, err) + assert.Equal(t, filepath.Join("/home/me", ".ssh", "databricks-cpu-6e7644d0"), path) +} + +func TestPinHostKey(t *testing.T) { + path := filepath.Join(t.TempDir(), "known-hosts", "myhost") + + err := PinHostKey(path, "myhost", []byte(testHostKey+" a-comment\n")) + require.NoError(t, err) + + content, err := os.ReadFile(path) + require.NoError(t, err) + // The comment is dropped: known_hosts keeps host and key, and the line is terminated + // so ssh doesn't ignore the last entry. + assert.Equal(t, "myhost "+testHostKey+"\n", string(content)) + + if runtime.GOOS != "windows" { + info, err := os.Stat(path) + require.NoError(t, err) + assert.Equal(t, os.FileMode(0o600), info.Mode().Perm()) + } +} + +func TestPinHostKeyReplacesPreviousEntry(t *testing.T) { + path := filepath.Join(t.TempDir(), "known-hosts", "myhost") + + // A key recorded for this name earlier, in the hashed form OpenSSH writes when + // HashKnownHosts is on - it cannot be matched by host name, which is why the file is + // rewritten rather than edited (DECO-27882). + require.NoError(t, os.MkdirAll(filepath.Dir(path), 0o700)) + require.NoError(t, os.WriteFile(path, []byte("|1|Lvcr5SHFVcEwAhYIbCCuADji6QQ=|rRdETW4qVc96QcC/TtDezxhkfp4= "+testOtherHostKey+"\n"), 0o600)) + + err := PinHostKey(path, "myhost", []byte(testHostKey)) + require.NoError(t, err) + + content, err := os.ReadFile(path) + require.NoError(t, err) + assert.Equal(t, "myhost "+testHostKey+"\n", string(content)) +} + +func TestPinHostKeyLeavesACorrectPinAlone(t *testing.T) { + path := filepath.Join(t.TempDir(), "known-hosts", "myhost") + require.NoError(t, PinHostKey(path, "myhost", []byte(testHostKey))) + + // An IDE opens several connections at once and each refreshes the pin, so an + // already-correct one must not be rewritten. Backdating the file makes the rewrite + // observable without depending on timer resolution. + backdated := time.Now().Add(-time.Hour).Truncate(time.Second) + require.NoError(t, os.Chtimes(path, backdated, backdated)) + + require.NoError(t, PinHostKey(path, "myhost", []byte(testHostKey))) + + info, err := os.Stat(path) + require.NoError(t, err) + assert.Equal(t, backdated, info.ModTime().Truncate(time.Second)) +} + +func TestPinHostKeyRejectsMalformedKey(t *testing.T) { + path := filepath.Join(t.TempDir(), "known-hosts", "myhost") + + err := PinHostKey(path, "myhost", []byte("not-a-key")) + assert.ErrorContains(t, err, "failed to parse the SSH server host key") + + // A malformed key must not leave a file ssh would then fail to verify against. + _, err = os.Stat(path) + assert.ErrorIs(t, err, os.ErrNotExist) +} diff --git a/experimental/ssh/internal/sshconfig/sshconfig.go b/experimental/ssh/internal/sshconfig/sshconfig.go index fbbf3b42fee..884ce9287c8 100644 --- a/experimental/ssh/internal/sshconfig/sshconfig.go +++ b/experimental/ssh/internal/sshconfig/sshconfig.go @@ -213,15 +213,31 @@ func PromptRecreateConfig(ctx context.Context, hostName string) (bool, error) { return response, nil } -func GenerateHostConfig(hostName, userName, identityFile, proxyCommand string) string { +// GenerateHostConfig renders the host block for a tunnel connection. Host key checking is +// strict rather than accept-new: the ProxyCommand pins the server's key (see PinHostKey) +// into knownHostsFile before ssh gets as far as verifying it, so there is no first +// connection that has to be taken on trust. +// +// hostKeyAlias is the name the key is pinned under in knownHostsFile - the session ID, i.e. +// the cluster ID for dedicated compute and the connection name for serverless. When it is +// non-empty it is emitted as HostKeyAlias so ssh looks the key up under that name. This +// matters whenever the user-facing hostName differs from it, as with +// `ssh setup --name --cluster `: without it ssh would look the key up under the +// alias, find no matching entry, and fail strict host key checking (DECO-27882). +func GenerateHostConfig(hostName, userName, identityFile, knownHostsFile, hostKeyAlias, proxyCommand string) string { + hostKeyAliasLine := "" + if hostKeyAlias != "" { + hostKeyAliasLine = fmt.Sprintf(" HostKeyAlias %s\n", hostKeyAlias) + } return fmt.Sprintf(` Host %s User %s ConnectTimeout 360 ServerAliveInterval %d - StrictHostKeyChecking accept-new - IdentitiesOnly yes + StrictHostKeyChecking yes + UserKnownHostsFile %q +%s IdentitiesOnly yes IdentityFile %q ProxyCommand %s -`, hostName, userName, ServerAliveIntervalSeconds, identityFile, proxyCommand) +`, hostName, userName, ServerAliveIntervalSeconds, knownHostsFile, hostKeyAliasLine, identityFile, proxyCommand) } diff --git a/experimental/ssh/internal/sshconfig/sshconfig_test.go b/experimental/ssh/internal/sshconfig/sshconfig_test.go index 23abcbc34d7..7d7f695b12a 100644 --- a/experimental/ssh/internal/sshconfig/sshconfig_test.go +++ b/experimental/ssh/internal/sshconfig/sshconfig_test.go @@ -12,7 +12,7 @@ import ( ) func TestGenerateHostConfigSetsServerAliveInterval(t *testing.T) { - config := GenerateHostConfig("myhost", "root", "/keys/myhost", "databricks ssh connect --proxy") + config := GenerateHostConfig("myhost", "root", "/keys/myhost", "/known-hosts/myhost", "cluster-123", "databricks ssh connect --proxy") // `ssh setup` and `--ide` reach ssh through this block and nothing else, so the option has to // be in it. @@ -26,6 +26,18 @@ func TestGenerateHostConfigSetsServerAliveInterval(t *testing.T) { assert.Greater(t, 3*ServerAliveIntervalSeconds, 30) } +func TestGenerateHostConfigPinsHostKeyAlias(t *testing.T) { + // `ssh setup --name myhost --cluster cluster-123` writes `Host myhost` but pins the server + // key under the cluster ID, so the block has to carry `HostKeyAlias cluster-123` for ssh to + // find the pinned entry under strict checking (DECO-27882). + config := GenerateHostConfig("myhost", "root", "/keys/myhost", "/known-hosts/cluster-123", "cluster-123", "databricks ssh connect --proxy") + assert.Contains(t, config, "\n HostKeyAlias cluster-123\n") + + // An empty alias omits the directive entirely rather than emitting a bare `HostKeyAlias`. + noAlias := GenerateHostConfig("myhost", "root", "/keys/myhost", "/known-hosts/myhost", "", "databricks ssh connect --proxy") + assert.NotContains(t, noAlias, "HostKeyAlias") +} + func TestGetConfigDir(t *testing.T) { dir, err := GetConfigDir(t.Context()) assert.NoError(t, err) diff --git a/libs/testserver/fake_workspace.go b/libs/testserver/fake_workspace.go index 2f077b1b631..bbcefe58ef4 100644 --- a/libs/testserver/fake_workspace.go +++ b/libs/testserver/fake_workspace.go @@ -248,6 +248,13 @@ type FakeWorkspace struct { // clusterVenvs caches Python venvs per existing cluster ID, // matching cloud behavior where libraries are cached on running clusters. clusterVenvs map[string]*clusterEnv + + // sshTunnelHostKeyPEM is the SSH host key every sshd of this workspace's tunnel + // serves, generated on first use. See sshTunnelHostKey. + sshTunnelHostKeyPEM []byte + // sshTunnelHostPublicKey is sshTunnelHostKeyPEM in authorized-key form, published + // to the tunnel's secret scope so a client can pin it. + sshTunnelHostPublicKey []byte } func (s *FakeWorkspace) LockUnlock() func() { diff --git a/libs/testserver/jobs.go b/libs/testserver/jobs.go index 3b16c61b023..c3387be3edb 100644 --- a/libs/testserver/jobs.go +++ b/libs/testserver/jobs.go @@ -593,7 +593,8 @@ const ( ) // writeSSHTunnelMetadata publishes the metadata.json a real tunnel server would -// write next to the bootstrap notebook. Callers must hold the workspace lock. +// write next to the bootstrap notebook, and the host key it would publish to the +// session's secret scope. Callers must hold the workspace lock. func (s *FakeWorkspace) writeSSHTunnelMetadata(request jobs.SubmitRun) { for _, t := range request.Tasks { if t.NotebookTask == nil { @@ -611,9 +612,26 @@ func (s *FakeWorkspace) writeSSHTunnelMetadata(request jobs.SubmitRun) { Info: workspace.ObjectInfo{ObjectType: "FILE", Path: metadataPath}, Data: metadata, } + s.publishSSHTunnelHostKey(t.NotebookTask.BaseParameters["secretScopeName"]) } } +// publishSSHTunnelHostKey stores the tunnel host key's public half in the session's +// secret scope, the way the real server does at startup, so the client can pin it. +// Callers must hold the workspace lock. +func (s *FakeWorkspace) publishSSHTunnelHostKey(scope string) { + if scope == "" { + return + } + if _, err := s.ensureSSHTunnelHostKey(); err != nil { + return + } + if s.Secrets[scope] == nil { + s.Secrets[scope] = make(map[string]string) + } + s.Secrets[scope][sshServerPublicKeySecretKey] = string(s.sshTunnelHostPublicKey) +} + // 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/ssh.go b/libs/testserver/ssh.go index 0cb7eef8846..8efdaa3a851 100644 --- a/libs/testserver/ssh.go +++ b/libs/testserver/ssh.go @@ -32,6 +32,11 @@ func sshTunnelRemoteUser() string { // under (mirrors clientPublicKeyName in experimental/ssh); sshd authorizes it below. const sshClientPublicKeySecretKey = "client-public-key" +// sshServerPublicKeySecretKey is the secret the tunnel server publishes its host key +// under (mirrors serverPublicKeyName in experimental/ssh); the client reads it to pin +// the host key before connecting. +const sshServerPublicKeySecretKey = "server-public-key" + var sshTunnelUpgrader = websocket.Upgrader{} // sshdTerminationTimeout caps how long Cmd.Wait blocks on a stuck sshd once the tunnel closes. @@ -63,6 +68,38 @@ func (s *FakeWorkspace) sshClientAuthorizedKey() []byte { return nil } +// sshTunnelHostKey returns the host key in PEM form that every sshd of this workspace +// serves, generating it on first use. The real tunnel server generates its host key once, +// stores it in the session's secret scope and reuses it for every sshd it launches, so one +// key per workspace is what a client sees across connections. +func (s *FakeWorkspace) sshTunnelHostKey() ([]byte, error) { + defer s.LockUnlock()() + return s.ensureSSHTunnelHostKey() +} + +// ensureSSHTunnelHostKey generates the workspace's tunnel host key if it has none. +// Callers must hold the workspace lock. +func (s *FakeWorkspace) ensureSSHTunnelHostKey() ([]byte, error) { + if s.sshTunnelHostKeyPEM != nil { + return s.sshTunnelHostKeyPEM, nil + } + pub, priv, err := ed25519.GenerateKey(rand.Reader) + if err != nil { + return nil, err + } + block, err := ssh.MarshalPrivateKey(priv, "") + if err != nil { + return nil, err + } + sshPub, err := ssh.NewPublicKey(pub) + if err != nil { + return nil, err + } + s.sshTunnelHostKeyPEM = pem.EncodeToMemory(block) + s.sshTunnelHostPublicKey = ssh.MarshalAuthorizedKey(sshPub) + return s.sshTunnelHostKeyPEM, nil +} + // sshTunnelHandler stands in for a cluster's tunnel server: it upgrades /ssh to a // websocket and drives a real `sshd -i` over it, so `ssh connect` runs a full handshake, // auth, and remote exec locally. Requires sshd; the acceptance test skips when it's absent @@ -83,12 +120,19 @@ func (s *Server) sshTunnelHandler(w http.ResponseWriter, r *http.Request) { // The ws dial carries the CLI's bearer token, so this token-scoped workspace is the // one that stored the uploaded public key. - authorizedKey := s.getWorkspaceForToken(getToken(r)).sshClientAuthorizedKey() + ws := s.getWorkspaceForToken(getToken(r)) + authorizedKey := ws.sshClientAuthorizedKey() if authorizedKey == nil { s.t.Logf("ssh tunnel: client public key secret %q not found", sshClientPublicKeySecretKey) return } + hostKey, err := ws.sshTunnelHostKey() + if err != nil { + s.t.Logf("ssh tunnel: host key: %s", err) + return + } + dir, err := os.MkdirTemp("", "testserver-sshd-") if err != nil { s.t.Logf("ssh tunnel: temp dir: %s", err) @@ -96,7 +140,7 @@ func (s *Server) sshTunnelHandler(w http.ResponseWriter, r *http.Request) { } defer os.RemoveAll(dir) - configPath, err := writeSSHDFiles(dir, authorizedKey) + configPath, err := writeSSHDFiles(dir, authorizedKey, hostKey) if err != nil { s.t.Logf("ssh tunnel: write sshd files: %s", err) return @@ -168,20 +212,12 @@ func (s *Server) sshTunnelHandler(w http.ResponseWriter, r *http.Request) { _ = cmd.Wait() } -// writeSSHDFiles writes an ephemeral host key, authorized_keys, and a minimal sshd_config, +// writeSSHDFiles writes the host key, authorized_keys, and a minimal sshd_config, // returning the config path. StrictModes/UsePAM are off so sshd accepts the temp-dir keys // and doesn't try to switch users when run non-root in inetd mode (-i). -func writeSSHDFiles(dir string, authorizedKey []byte) (string, error) { - _, hostPriv, err := ed25519.GenerateKey(rand.Reader) - if err != nil { - return "", err - } - block, err := ssh.MarshalPrivateKey(hostPriv, "") - if err != nil { - return "", err - } +func writeSSHDFiles(dir string, authorizedKey, hostKeyPEM []byte) (string, error) { hostKeyPath := filepath.Join(dir, "host_key") - if err := os.WriteFile(hostKeyPath, pem.EncodeToMemory(block), 0o600); err != nil { + if err := os.WriteFile(hostKeyPath, hostKeyPEM, 0o600); err != nil { return "", err }