From a8c6f7ad331feb6890215b76932668869b050fa5 Mon Sep 17 00:00:00 2001 From: Dmitry Meyer Date: Mon, 31 Aug 2026 14:47:10 +0000 Subject: [PATCH 1/2] Fix adding authorized keys in Kubernetes and Slurm The Kubernetes jump pod added the user's key with `grep -qvF KEY file`, which asks whether some line does not contain the key, not whether no line contains it. The jump pod's `authorized_keys` holds the project key, so the condition held whether or not the user's key was already there, and every run appended another copy of it. Users keep the same key across runs, so the file of a jump pod, shared by the whole namespace, grew by one duplicate per run. Had the file been empty, the condition would have been false and the key never added at all. Slurm got the condition right, but matched the whole entry, comment included, so the same key submitted with a different comment was added a second time. All three implementations -- SSH fleets, Kubernetes, Slurm -- are replaced with `get_add_authorized_keys_script()`, extracted from the SSH fleet one, the only one of the three that was correct. It returns the script instead of running it, as the call sites execute it in three different ways: paramiko, `ssh` over the jump pod, and the Slurm client. * Keys reach the script as heredoc data instead of being interpolated into its commands, so the shell never parses anything that came from a key. * Entries are rebuilt from `parse_public_key()`, a new helper parsing a public key in OpenSSH disk format, so that nothing unvalidated is written to the file. A key that fails to parse is skipped with a warning: one bad key does not keep the rest out of the file. * Entries are deduplicated on the key blob alone, since neither the comment nor the `command="/bin/false"` options Kubernetes and Slurm prepend are a part of the identity of a key. * Kubernetes and Slurm entries now carry the `# added by dstack` marker the SSH fleet entries have carried since #4179. Co-authored-by: Claude Opus 5 (1M context) --- .../core/backends/base/authorized_keys.py | 80 +++++++++ .../core/backends/kubernetes/compute.py | 11 +- .../_internal/core/backends/slurm/compute.py | 23 +-- .../services/ssh_fleets/provisioning.py | 38 +--- src/dstack/_internal/utils/ssh.py | 74 ++++++++ .../backends/base/test_authorized_keys.py | 165 ++++++++++++++++++ src/tests/_internal/utils/test_ssh.py | 77 +++++++- 7 files changed, 412 insertions(+), 56 deletions(-) create mode 100644 src/dstack/_internal/core/backends/base/authorized_keys.py create mode 100644 src/tests/_internal/core/backends/base/test_authorized_keys.py diff --git a/src/dstack/_internal/core/backends/base/authorized_keys.py b/src/dstack/_internal/core/backends/base/authorized_keys.py new file mode 100644 index 000000000..c804c1975 --- /dev/null +++ b/src/dstack/_internal/core/backends/base/authorized_keys.py @@ -0,0 +1,80 @@ +from textwrap import dedent +from typing import Optional +from uuid import uuid4 + +from dstack._internal.utils.logging import get_logger +from dstack._internal.utils.ssh import parse_public_key + +logger = get_logger(__name__) + + +# Appended to the comment field of every authorized_keys entry added by the server. sshd +# ignores everything after the key blob, so the marker has no effect on authentication; it +# only records that the entry is ours. +# +# The shim adds its own entries, marked with `# added by dstack-shim`, and rewrites only the +# entries carrying that marker, see runner/internal/shim/authorized_keys.go. It matches the +# marker as an exact suffix, and this one is always appended last, therefore the entries added +# by this script are never touched by the shim. Do not change this value to anything that ends +# with the shim's marker. +DSTACK_PUBLIC_KEY_MARKER = "# added by dstack" + + +def get_add_authorized_keys_script( + authorized_keys: list[str], + *, + options: Optional[str] = None, +) -> str: + """ + Builds a POSIX shell script adding the given public keys to `~/.ssh/authorized_keys`. + + The `~/.ssh` directory and the file are created if missing. An entry is added only if its + key blob is not in the file yet, so the script can be run repeatedly; entries already in + the file, whoever added them, are never modified or removed. + + Every entry is rebuilt from the parsed key and marked with DSTACK_PUBLIC_KEY_MARKER, so + that nothing unvalidated reaches the file. Keys that cannot be parsed are skipped with a + warning -- one bad key does not keep the rest out of the file. + + The keys are passed to the script as heredoc data and never interpolated into commands, + therefore the shell does not parse anything that came from a key. + + Args: + authorized_keys: The public keys in OpenSSH disk format. + options: The authorized_keys options to prepend to every entry, e.g. + `command="/bin/false"`. Must be a single line with no tabs. + + Returns: + The script. + """ + entries: list[str] = [] + for authorized_key in authorized_keys: + try: + key = parse_public_key(authorized_key) + except ValueError as e: + logger.warning("Failed to parse authorized key: %r: %s", authorized_key, e) + continue + key.comment = " ".join(filter(None, [key.comment, DSTACK_PUBLIC_KEY_MARKER])) + entry = str(key) + if options is not None: + entry = f"{options} {entry}" + # The blob is the identity of the key -- the comment and the options are not, an entry + # differing only in them is the same key and must not be added twice + entries.append(f"{key.blob_base64}\t{entry}") + eof = f"EOF_{uuid4().hex}" + header = dedent(f"""\ + set -eu + if [ ! -e ~/.ssh/authorized_keys ]; then + mkdir -p ~/.ssh + chmod 700 ~/.ssh + touch ~/.ssh/authorized_keys + chmod 600 ~/.ssh/authorized_keys + elif [ -n "$(tail -c1 ~/.ssh/authorized_keys)" ]; then + echo >> ~/.ssh/authorized_keys + fi + while IFS='\t' read -r blob entry <&3; do + if ! grep -qF "$blob" ~/.ssh/authorized_keys; then + echo "$entry" >> ~/.ssh/authorized_keys + fi + done 3<<'{eof}'""") + return "\n".join([header, *entries, eof]) diff --git a/src/dstack/_internal/core/backends/kubernetes/compute.py b/src/dstack/_internal/core/backends/kubernetes/compute.py index f60c285bd..92ecb7501 100644 --- a/src/dstack/_internal/core/backends/kubernetes/compute.py +++ b/src/dstack/_internal/core/backends/kubernetes/compute.py @@ -14,6 +14,9 @@ from kubernetes import client from typing_extensions import Self +from dstack._internal.core.backends.base.authorized_keys import ( + get_add_authorized_keys_script, +) from dstack._internal.core.backends.base.compute import ( Compute, ComputeWithAllOffersCached, @@ -1035,11 +1038,9 @@ def _check_and_configure_jump_pod_service( ssh_private_key=project_ssh_private_key, # command= in authorized_keys is equivalent to ForceCommand in sshd_config # By forcing the /bin/false command we only allow proxy jumping, no shell access - command=f""" - if grep -qvF '{user_ssh_public_key}' ~/.ssh/authorized_keys; then - echo 'command="/bin/false" {user_ssh_public_key}' >> ~/.ssh/authorized_keys - fi - """, + command=get_add_authorized_keys_script( + [user_ssh_public_key], options='command="/bin/false"' + ), ) if ssh_exit_status != 0: logger.debug( diff --git a/src/dstack/_internal/core/backends/slurm/compute.py b/src/dstack/_internal/core/backends/slurm/compute.py index 4f5baf44b..2083278bc 100644 --- a/src/dstack/_internal/core/backends/slurm/compute.py +++ b/src/dstack/_internal/core/backends/slurm/compute.py @@ -5,6 +5,9 @@ from functools import partial from typing import Optional +from dstack._internal.core.backends.base.authorized_keys import ( + get_add_authorized_keys_script, +) from dstack._internal.core.backends.base.compute import ( Compute, ComputeWithAllOffersCached, @@ -319,20 +322,12 @@ def _run_slurm_job( for job_node in job_nodes ] - res = client.exec(f""" - set -eu - if [ ! -e ~/.ssh/authorized_keys ]; then - mkdir -p ~/.ssh - chmod 700 ~/.ssh - touch ~/.ssh/authorized_keys - chmod 600 ~/.ssh/authorized_keys - fi - for key in {shlex.join(authorized_keys)}; do - if ! grep -qF "$key" ~/.ssh/authorized_keys; then - echo 'command="/bin/false"' "$key" >> ~/.ssh/authorized_keys - fi - done - """) + # command= in authorized_keys is equivalent to ForceCommand in sshd_config. + # By forcing the /bin/false command we only allow proxy jumping through the login + # node, no shell access + res = client.exec( + get_add_authorized_keys_script(authorized_keys, options='command="/bin/false"') + ) if not res.ok: raise ComputeError(f"Failed to add authorized keys: {res}") diff --git a/src/dstack/_internal/server/services/ssh_fleets/provisioning.py b/src/dstack/_internal/server/services/ssh_fleets/provisioning.py index 4ceac4a26..6feae9ca3 100644 --- a/src/dstack/_internal/server/services/ssh_fleets/provisioning.py +++ b/src/dstack/_internal/server/services/ssh_fleets/provisioning.py @@ -1,7 +1,6 @@ import io import json import time -import uuid from contextlib import contextmanager, nullcontext from textwrap import dedent from typing import Any, Dict, Generator, List, Optional @@ -9,6 +8,9 @@ import paramiko from gpuhunt import AcceleratorVendor, correct_gpu_memory_gib +from dstack._internal.core.backends.base.authorized_keys import ( + get_add_authorized_keys_script, +) from dstack._internal.core.backends.base.compute import ( DSTACK_SHIM_RESTART_INTERVAL_SECONDS, GoArchType, @@ -91,39 +93,7 @@ def upload_envs(client: paramiko.SSHClient, working_dir: str, envs: Dict[str, st def add_authorized_keys(client: paramiko.SSHClient, authorized_keys: list[str]) -> None: - heredoc_lines: list[str] = [] - for key in authorized_keys: - # Not using paramiko.pkey.PublicBlob.from_string() to avoid unnecessary parsing/validation - try: - key_type, key_blob, *key_comment_parts = key.split() - except ValueError as e: - logger.warning("Failed to parse authorized key: %r: %s", key, e) - continue - key_comment_parts.append("# added by dstack") - key_comment = " ".join(key_comment_parts) - heredoc_lines.append(f"{key_type}\t{key_blob}\t{key_comment}") - eof = f"EOF_{uuid.uuid4().hex}" - script_parts = [ - f""" - set -eu - if [ ! -e ~/.ssh/authorized_keys ]; then - mkdir -p ~/.ssh - chmod 700 ~/.ssh - touch ~/.ssh/authorized_keys - chmod 600 ~/.ssh/authorized_keys - elif [ -n "$(tail -c1 ~/.ssh/authorized_keys)" ]; then - echo >> ~/.ssh/authorized_keys - fi - while IFS='\t' read -r type blob comment <&3; do - if ! grep -qF "$blob" ~/.ssh/authorized_keys; then - echo "$type $blob $comment" >> ~/.ssh/authorized_keys - fi - done 3<<'{eof}' - """.strip(), - *heredoc_lines, - eof, - ] - script = "\n".join(script_parts) + script = get_add_authorized_keys_script(authorized_keys) try: _, stdout, stderr = client.exec_command(script, timeout=5) out = stdout.read().strip().decode() diff --git a/src/dstack/_internal/utils/ssh.py b/src/dstack/_internal/utils/ssh.py index d8a701233..b55670006 100644 --- a/src/dstack/_internal/utils/ssh.py +++ b/src/dstack/_internal/utils/ssh.py @@ -1,11 +1,14 @@ +import base64 import io import os import re import shutil +import struct import subprocess import sys import tempfile from contextlib import suppress +from dataclasses import dataclass from pathlib import Path from typing import Dict, Optional, Union @@ -33,6 +36,77 @@ def get_public_key_fingerprint(text: str) -> str: return pk.fingerprint +@dataclass +class PublicKey: + """ + A public key in OpenSSH disk format, parsed into fields. + + Converting to str renders the key back into the one-line `type blob [comment]` form. + + Attributes: + type: The key type, e.g. `ssh-ed25519`. + blob_base64: The base64-encoded key blob, as it appears in the key file. + comment: The comment or None if the key has no comment. + """ + + type: str + blob_base64: str + comment: Optional[str] = None + + def __str__(self) -> str: + if not self.comment: + return f"{self.type} {self.blob_base64}" + return f"{self.type} {self.blob_base64} {self.comment}" + + +def parse_public_key(key: str) -> PublicKey: + """ + Parses a public key in OpenSSH disk format into its fields. + + Performs basic validation -- ensures that the key consists of exactly one line, that the blob + is valid base64, and that the type field matches the type encoded in the blob. The key type + itself is not restricted, that is, keys of any type, including types unsupported by dstack, + are accepted. + + Options (an optional field preceding the key type in the authorized_keys format) are not + supported -- the first field is always interpreted as a key type, so a line with options is + rejected as invalid. + + The comment, if present, is normalized -- surrounding whitespaces are removed, adjacent + whitespaces are collapsed into a single space. + + Args: + key: The public key in OpenSSH disk format, a `type blob [comment]` string. + + Returns: + The parsed key. The blob is stored as is, without decoding. + + Raises: + ValueError: Invalid public key. + """ + lines = key.strip().splitlines() + if len(lines) != 1: + raise ValueError("Expected a single line") + try: + type_declared, blob_base64, *comment_parts = lines[0].split() + except ValueError: + raise ValueError("Not enough fields") + # paramiko.pkey.PublicBlob.from_string() performs the same key type check + try: + blob = base64.b64decode(blob_base64, validate=True) + [type_length] = struct.unpack(">I", blob[:4]) + type_parsed = blob[4 : 4 + type_length].decode() + except (ValueError, struct.error) as e: + raise ValueError(f"Failed to parse key: {e}") from e + if type_declared != type_parsed: + raise ValueError(f"Key type mismatch: {type_declared} != {type_parsed}") + if comment_parts: + comment = " ".join(comment_parts) + else: + comment = None + return PublicKey(type=type_declared, blob_base64=blob_base64, comment=comment) + + def get_host_config(hostname: str, ssh_config_path: PathLike = default_ssh_config_path) -> dict: ssh_config_path = os.path.expanduser(ssh_config_path) if os.path.exists(ssh_config_path): diff --git a/src/tests/_internal/core/backends/base/test_authorized_keys.py b/src/tests/_internal/core/backends/base/test_authorized_keys.py new file mode 100644 index 000000000..94e97d673 --- /dev/null +++ b/src/tests/_internal/core/backends/base/test_authorized_keys.py @@ -0,0 +1,165 @@ +import os +import stat +import subprocess +from pathlib import Path + +import pytest + +from dstack._internal.core.backends.base.authorized_keys import ( + DSTACK_PUBLIC_KEY_MARKER, + get_add_authorized_keys_script, +) + +ED25519_BLOB = "AAAAC3NzaC1lZDI1NTE5AAAAINOmx0T+hBRaJ6jCi21ZYe2NW3EZS8e0Mdwl+yZJt+kD" +ED25519_KEY = f"ssh-ed25519 {ED25519_BLOB}" +ECDSA_BLOB = ( + "AAAAE2VjZHNhLXNoYTItbmlzdHAyNTYAAAAIbmlzdHAyNTYAAABBBE8xPz2OD5CXgZyHY6D70lMVk5IZyiWQRw5h" + "9HHFrsqMp0v7SKEp89dBKAEPhq2h/4LfTFuH3aycKs/oZpbFbWM=" +) +ECDSA_KEY = f"ecdsa-sha2-nistp256 {ECDSA_BLOB}" + +SCRIPT_TIMEOUT = 10 + + +def entries(script: str) -> list[str]: + """Returns the authorized_keys entries the script would add, in order.""" + lines = script.split("\n") + heredoc_start = next(i for i, line in enumerate(lines) if line.startswith("done 3<<")) + 1 + # the last line is the heredoc terminator + return [line.split("\t", 1)[1] for line in lines[heredoc_start:-1]] + + +class TestGetAddAuthorizedKeysScript: + def test_appends_marker_to_comment(self): + script = get_add_authorized_keys_script([f"{ED25519_KEY} dev@host"]) + + assert entries(script) == [f"{ED25519_KEY} dev@host {DSTACK_PUBLIC_KEY_MARKER}"] + + def test_marker_is_the_whole_comment_if_the_key_has_none(self): + script = get_add_authorized_keys_script([ED25519_KEY]) + + assert entries(script) == [f"{ED25519_KEY} {DSTACK_PUBLIC_KEY_MARKER}"] + + def test_prepends_options(self): + script = get_add_authorized_keys_script([ED25519_KEY], options='command="/bin/false"') + + assert entries(script) == [ + f'command="/bin/false" {ED25519_KEY} {DSTACK_PUBLIC_KEY_MARKER}' + ] + + def test_normalizes_comment_whitespace(self): + # a tab in a comment would otherwise break the tab-delimited heredoc + script = get_add_authorized_keys_script([f"{ED25519_KEY} my\tlaptop key"]) + + assert entries(script) == [f"{ED25519_KEY} my laptop key {DSTACK_PUBLIC_KEY_MARKER}"] + + @pytest.mark.parametrize( + "invalid_key", + [ + "", + "not a key", + "ssh-ed25519", + f"ssh-rsa {ED25519_BLOB}", + # options are not a part of the on-disk public key format + f'command="/bin/false" {ED25519_KEY}', + # a second key smuggled in via a newline + f"{ED25519_KEY} dev@host\nssh-rsa AAAAB3NzaC1yc2E evil", + ], + ) + def test_skips_invalid_keys(self, invalid_key: str): + script = get_add_authorized_keys_script([invalid_key, ECDSA_KEY]) + + assert entries(script) == [f"{ECDSA_KEY} {DSTACK_PUBLIC_KEY_MARKER}"] + + def test_no_keys(self): + # the script still creates the file, it just adds nothing + assert entries(get_add_authorized_keys_script([])) == [] + + +def run_script(script: str, home: Path) -> None: + result = subprocess.run( + ["sh", "-c", script], + env={"HOME": str(home), "PATH": os.environ["PATH"]}, + capture_output=True, + text=True, + # a broken script must fail the test, not block it on the terminal + stdin=subprocess.DEVNULL, + timeout=SCRIPT_TIMEOUT, + ) + assert (result.returncode, result.stdout, result.stderr) == (0, "", "") + + +def read_authorized_keys(home: Path) -> str: + return (home / ".ssh" / "authorized_keys").read_text() + + +def write_authorized_keys(home: Path, content: str) -> None: + (home / ".ssh").mkdir(exist_ok=True) + (home / ".ssh" / "authorized_keys").write_text(content) + + +class TestAddAuthorizedKeysScriptExecution: + def test_creates_ssh_dir_and_file(self, tmp_path: Path): + run_script(get_add_authorized_keys_script([f"{ED25519_KEY} dev@host"]), tmp_path) + + assert read_authorized_keys(tmp_path) == ( + f"{ED25519_KEY} dev@host {DSTACK_PUBLIC_KEY_MARKER}\n" + ) + assert stat.S_IMODE((tmp_path / ".ssh").stat().st_mode) == 0o700 + assert stat.S_IMODE((tmp_path / ".ssh" / "authorized_keys").stat().st_mode) == 0o600 + + def test_creates_file_without_keys(self, tmp_path: Path): + run_script(get_add_authorized_keys_script([]), tmp_path) + + assert read_authorized_keys(tmp_path) == "" + + def test_keeps_existing_entries(self, tmp_path: Path): + write_authorized_keys(tmp_path, f"{ECDSA_KEY} added-by-hand\n") + + run_script(get_add_authorized_keys_script([f"{ED25519_KEY} dev@host"]), tmp_path) + + assert read_authorized_keys(tmp_path) == ( + f"{ECDSA_KEY} added-by-hand\n{ED25519_KEY} dev@host {DSTACK_PUBLIC_KEY_MARKER}\n" + ) + + def test_adds_newline_to_file_not_ending_with_one(self, tmp_path: Path): + write_authorized_keys(tmp_path, f"{ECDSA_KEY} added-by-hand") + + run_script(get_add_authorized_keys_script([f"{ED25519_KEY} dev@host"]), tmp_path) + + assert read_authorized_keys(tmp_path) == ( + f"{ECDSA_KEY} added-by-hand\n{ED25519_KEY} dev@host {DSTACK_PUBLIC_KEY_MARKER}\n" + ) + + def test_is_idempotent(self, tmp_path: Path): + script = get_add_authorized_keys_script([f"{ED25519_KEY} dev@host", ECDSA_KEY]) + + run_script(script, tmp_path) + run_script(script, tmp_path) + + assert read_authorized_keys(tmp_path) == ( + f"{ED25519_KEY} dev@host {DSTACK_PUBLIC_KEY_MARKER}\n" + f"{ECDSA_KEY} {DSTACK_PUBLIC_KEY_MARKER}\n" + ) + + def test_does_not_add_a_key_already_in_the_file(self, tmp_path: Path): + # the same key with another comment and no marker is still the same key + write_authorized_keys(tmp_path, f"{ED25519_KEY} added-by-hand\n") + + run_script(get_add_authorized_keys_script([f"{ED25519_KEY} dev@host"]), tmp_path) + + assert read_authorized_keys(tmp_path) == f"{ED25519_KEY} added-by-hand\n" + + def test_writes_comment_with_shell_metacharacters_verbatim(self, tmp_path: Path): + # the comment is passed to the script as data, so the shell does not parse it + home = tmp_path / "home" + home.mkdir() + touched = tmp_path / "touched" + comment = f"x';touch {touched};'" + + run_script(get_add_authorized_keys_script([f"{ED25519_KEY} {comment}"]), home) + + assert not touched.exists() + assert read_authorized_keys(home) == ( + f"{ED25519_KEY} {comment} {DSTACK_PUBLIC_KEY_MARKER}\n" + ) diff --git a/src/tests/_internal/utils/test_ssh.py b/src/tests/_internal/utils/test_ssh.py index 0512d9bf0..92ed9bef3 100644 --- a/src/tests/_internal/utils/test_ssh.py +++ b/src/tests/_internal/utils/test_ssh.py @@ -1,6 +1,7 @@ import subprocess import unittest from pathlib import Path +from typing import Optional from unittest.mock import MagicMock, patch import pytest @@ -9,10 +10,12 @@ from dstack._internal.utils import crypto from dstack._internal.utils.path import FilePath from dstack._internal.utils.ssh import ( + PublicKey, check_required_ssh_version, find_ssh_util, include_ssh_config, normalize_path, + parse_public_key, pkey_from_str, resolve_ssh_key, update_ssh_config, @@ -29,9 +32,8 @@ W3EZS8e0Mdwl+yZJt+kDAAAAC3Rlc3RAZHN0YWNrAQI= -----END OPENSSH PRIVATE KEY----- """ -PUBLIC_KEY_NO_COMMENT = ( - "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAINOmx0T+hBRaJ6jCi21ZYe2NW3EZS8e0Mdwl+yZJt+kD" -) +PUBLIC_KEY_BLOB = "AAAAC3NzaC1lZDI1NTE5AAAAINOmx0T+hBRaJ6jCi21ZYe2NW3EZS8e0Mdwl+yZJt+kD" +PUBLIC_KEY_NO_COMMENT = f"ssh-ed25519 {PUBLIC_KEY_BLOB}" PUBLIC_KEY = f"{PUBLIC_KEY_NO_COMMENT} test@dstack\n" # A valid public key of a type paramiko cannot construct a PKey for UNSUPPORTED_TYPE_PUBLIC_KEY = ( @@ -238,3 +240,72 @@ def test_raises_on_invalid_key(self, tmp_path: Path, contents: str): def test_raises_if_key_does_not_exist(self, tmp_path: Path): with pytest.raises(OSError): resolve_ssh_key(tmp_path / "id_ed25519") + + +class TestParsePublicKey: + def test_key_with_comment(self): + assert parse_public_key(PUBLIC_KEY) == PublicKey( + type="ssh-ed25519", blob_base64=PUBLIC_KEY_BLOB, comment="test@dstack" + ) + + def test_key_without_comment(self): + assert parse_public_key(PUBLIC_KEY_NO_COMMENT) == PublicKey( + type="ssh-ed25519", blob_base64=PUBLIC_KEY_BLOB, comment=None + ) + + def test_accepts_unsupported_key_type(self): + # any key type is accepted, even one paramiko cannot construct a PKey for + key = parse_public_key(UNSUPPORTED_TYPE_PUBLIC_KEY) + + assert key.type == "sk-ssh-ed25519@openssh.com" + # the blob is returned as is, with base64 padding preserved + assert key.blob_base64 == UNSUPPORTED_TYPE_PUBLIC_KEY.split()[1] + assert key.comment == "test@dstack" + + def test_normalizes_comment(self): + key = parse_public_key(f"{PUBLIC_KEY_NO_COMMENT} two \t words \n") + + assert key.comment == "two words" + + def test_accepts_crlf_line_ending(self): + key = parse_public_key(f"{PUBLIC_KEY_NO_COMMENT} test@dstack\r\n") + + assert key.comment == "test@dstack" + + @pytest.mark.parametrize("key", [PUBLIC_KEY, PUBLIC_KEY_NO_COMMENT]) + def test_str_round_trips(self, key: str): + parsed = parse_public_key(key) + + assert str(parsed) == key.strip() + assert parse_public_key(str(parsed)) == parsed + + @pytest.mark.parametrize("comment", [None, ""]) + def test_str_omits_absent_comment(self, comment: Optional[str]): + key = PublicKey(type="ssh-ed25519", blob_base64=PUBLIC_KEY_BLOB, comment=comment) + + assert str(key) == PUBLIC_KEY_NO_COMMENT + + @pytest.mark.parametrize( + ("key", "error"), + [ + ("", "Expected a single line"), + (" ", "Expected a single line"), + ("\n\n", "Expected a single line"), + (PUBLIC_KEY + PUBLIC_KEY, "Expected a single line"), + ("ssh-ed25519", "Not enough fields"), + ("ssh-ed25519 not-base64!", "Failed to parse key"), + # a stray non-base64 character that b64decode would otherwise discard silently, + # yielding a valid blob but leaving the character in the returned key + (f"ssh-ed25519 {PUBLIC_KEY_BLOB[:4]}!{PUBLIC_KEY_BLOB[4:]}", "Failed to parse key"), + # the blob is shorter than the 4-byte type length prefix + ("ssh-ed25519 AAA=", "Failed to parse key"), + # the type field in the blob is not valid UTF-8 + ("ssh-ed25519 AAAABP////8=", "Failed to parse key"), + (f"ssh-rsa {PUBLIC_KEY_BLOB}", "Key type mismatch"), + # options are not supported, the first field is always read as a key type + (f'command="/bin/false" {PUBLIC_KEY_NO_COMMENT}', "Failed to parse key"), + ], + ) + def test_raises_on_invalid_key(self, key: str, error: str): + with pytest.raises(ValueError, match=error): + parse_public_key(key) From 4db4bd774efbd3c177e7139c2355d227542acd9d Mon Sep 17 00:00:00 2001 From: Dmitry Meyer Date: Tue, 1 Sep 2026 08:12:29 +0000 Subject: [PATCH 2/2] Make dstack marker optional --- .../_internal/core/backends/base/authorized_keys.py | 13 +++++++++---- .../_internal/core/backends/kubernetes/compute.py | 11 ++++++++--- src/dstack/_internal/core/backends/slurm/compute.py | 11 +++++++---- .../core/backends/base/test_authorized_keys.py | 7 +++++++ 4 files changed, 31 insertions(+), 11 deletions(-) diff --git a/src/dstack/_internal/core/backends/base/authorized_keys.py b/src/dstack/_internal/core/backends/base/authorized_keys.py index c804c1975..15fcad132 100644 --- a/src/dstack/_internal/core/backends/base/authorized_keys.py +++ b/src/dstack/_internal/core/backends/base/authorized_keys.py @@ -23,6 +23,7 @@ def get_add_authorized_keys_script( authorized_keys: list[str], *, + add_dstack_marker: bool = True, options: Optional[str] = None, ) -> str: """ @@ -32,15 +33,18 @@ def get_add_authorized_keys_script( key blob is not in the file yet, so the script can be run repeatedly; entries already in the file, whoever added them, are never modified or removed. - Every entry is rebuilt from the parsed key and marked with DSTACK_PUBLIC_KEY_MARKER, so - that nothing unvalidated reaches the file. Keys that cannot be parsed are skipped with a - warning -- one bad key does not keep the rest out of the file. + Every entry is rebuilt from the parsed key, so that nothing unvalidated reaches the file. + Keys that cannot be parsed are skipped with a warning -- one bad key does not keep the + rest out of the file. The keys are passed to the script as heredoc data and never interpolated into commands, therefore the shell does not parse anything that came from a key. Args: authorized_keys: The public keys in OpenSSH disk format. + add_dstack_marker: Whether to append DSTACK_PUBLIC_KEY_MARKER to every entry. Set to + False only where the whole file is managed by dstack and there are no foreign + entries to tell ours from. options: The authorized_keys options to prepend to every entry, e.g. `command="/bin/false"`. Must be a single line with no tabs. @@ -54,8 +58,9 @@ def get_add_authorized_keys_script( except ValueError as e: logger.warning("Failed to parse authorized key: %r: %s", authorized_key, e) continue - key.comment = " ".join(filter(None, [key.comment, DSTACK_PUBLIC_KEY_MARKER])) entry = str(key) + if add_dstack_marker: + entry = f"{entry} {DSTACK_PUBLIC_KEY_MARKER}" if options is not None: entry = f"{options} {entry}" # The blob is the identity of the key -- the comment and the options are not, an entry diff --git a/src/dstack/_internal/core/backends/kubernetes/compute.py b/src/dstack/_internal/core/backends/kubernetes/compute.py index 92ecb7501..d21784f3e 100644 --- a/src/dstack/_internal/core/backends/kubernetes/compute.py +++ b/src/dstack/_internal/core/backends/kubernetes/compute.py @@ -1036,10 +1036,15 @@ def _check_and_configure_jump_pod_service( port=jump_pod_port, username=JUMP_POD_USER, ssh_private_key=project_ssh_private_key, - # command= in authorized_keys is equivalent to ForceCommand in sshd_config - # By forcing the /bin/false command we only allow proxy jumping, no shell access command=get_add_authorized_keys_script( - [user_ssh_public_key], options='command="/bin/false"' + [user_ssh_public_key], + # The project key, written by _get_jump_pod_commands(), carries no marker, and + # marking only some of our entries defeats the purpose of the marker. It is + # redundant on the jump pod anyway -- every entry there is ours + add_dstack_marker=False, + # command= in authorized_keys is equivalent to ForceCommand in sshd_config + # By forcing the /bin/false command we only allow proxy jumping, no shell access + options='command="/bin/false"', ), ) if ssh_exit_status != 0: diff --git a/src/dstack/_internal/core/backends/slurm/compute.py b/src/dstack/_internal/core/backends/slurm/compute.py index 2083278bc..2fa9ab86b 100644 --- a/src/dstack/_internal/core/backends/slurm/compute.py +++ b/src/dstack/_internal/core/backends/slurm/compute.py @@ -322,11 +322,14 @@ def _run_slurm_job( for job_node in job_nodes ] - # command= in authorized_keys is equivalent to ForceCommand in sshd_config. - # By forcing the /bin/false command we only allow proxy jumping through the login - # node, no shell access res = client.exec( - get_add_authorized_keys_script(authorized_keys, options='command="/bin/false"') + get_add_authorized_keys_script( + authorized_keys, + # command= in authorized_keys is equivalent to ForceCommand in sshd_config. + # By forcing the /bin/false command we only allow proxy jumping through + # the login node, no shell access + options='command="/bin/false"', + ) ) if not res.ok: raise ComputeError(f"Failed to add authorized keys: {res}") diff --git a/src/tests/_internal/core/backends/base/test_authorized_keys.py b/src/tests/_internal/core/backends/base/test_authorized_keys.py index 94e97d673..05db5dc8b 100644 --- a/src/tests/_internal/core/backends/base/test_authorized_keys.py +++ b/src/tests/_internal/core/backends/base/test_authorized_keys.py @@ -40,6 +40,13 @@ def test_marker_is_the_whole_comment_if_the_key_has_none(self): assert entries(script) == [f"{ED25519_KEY} {DSTACK_PUBLIC_KEY_MARKER}"] + def test_omits_marker(self): + script = get_add_authorized_keys_script( + [f"{ED25519_KEY} dev@host"], add_dstack_marker=False + ) + + assert entries(script) == [f"{ED25519_KEY} dev@host"] + def test_prepends_options(self): script = get_add_authorized_keys_script([ED25519_KEY], options='command="/bin/false"')