Fix adding authorized keys in Kubernetes and Slurm

src/dstack/_internal/core/backends/base/authorized_keys.py

@@ -0,0 +1,85 @@
+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],
+    *,
+    add_dstack_marker: bool = True,
+    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, 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.
+
+    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
+        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
+        # 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])

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,
@@ -1033,13 +1036,16 @@ 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=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],
+            # 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:
         logger.debug(

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,15 @@ 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
-            """)
+            res = client.exec(
+                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}")
 

src/dstack/_internal/server/services/ssh_fleets/provisioning.py

@@ -1,14 +1,16 @@
 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
 
 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()

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):