Honor sshproxy enforcement in container-based backends

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

@@ -2,6 +2,7 @@
 from typing import Optional
 from uuid import uuid4
 
+from dstack._internal.core.errors import ComputeError
 from dstack._internal.utils.logging import get_logger
 from dstack._internal.utils.ssh import parse_public_key
 
@@ -20,6 +21,57 @@
 DSTACK_PUBLIC_KEY_MARKER = "# added by dstack"
 
 
+def build_authorized_keys(
+    project_ssh_public_key: str,
+    extra_authorized_keys: list[str],
+) -> list[str]:
+    """
+    Builds the list of public keys to authorize on a job container.
+
+    Args:
+        project_ssh_public_key: The project public key, always authorized. Must be valid --
+            the server cannot reach the container without it.
+        extra_authorized_keys: The public keys to authorize in addition to the project key,
+            as passed to `Compute.run_job()`. Untrusted -- keys that cannot be parsed are
+            skipped with a warning, one bad key does not keep the rest out.
+
+    Returns:
+        The normalized keys in OpenSSH disk format, the project key first.
+
+    Raises:
+        ComputeError: The project key is invalid.
+    """
+    project_authorized_keys = normalize_authorized_keys([project_ssh_public_key])
+    if not project_authorized_keys:
+        raise ComputeError("Invalid project SSH key")
+    return project_authorized_keys + normalize_authorized_keys(extra_authorized_keys)
+
+
+def normalize_authorized_keys(authorized_keys: list[str]) -> list[str]:
+    """
+    Rebuilds the given public keys from their parsed fields.
+
+    Every returned entry is a single `type blob [comment]` line with the comment whitespace
+    normalized, so that nothing unvalidated reaches a command or a file.
+
+    Args:
+        authorized_keys: The public keys in OpenSSH disk format.
+
+    Returns:
+        The normalized keys, in the original order. Keys that cannot be parsed are skipped with
+        a warning, therefore the result may be shorter than the input, or empty.
+    """
+    normalized: 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
+        normalized.append(str(key))
+    return normalized
+
+
 def get_add_authorized_keys_script(
     authorized_keys: list[str],
     *,

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

@@ -148,11 +148,21 @@ def run_job(
         volumes: List[Volume],
         placement_group: Optional[PlacementGroup],
         requirements: Requirements,
+        extra_authorized_keys: list[str],
     ) -> JobProvisioningData:
         """
         Launches a new instance for the job. It should return `JobProvisioningData` ASAP.
         If required to wait to get the IP address or SSH port, return partially filled `JobProvisioningData`
         and implement `update_provisioning_data()`.
+
+        `extra_authorized_keys` are the public keys to authorize on the job container in addition
+        to `project_ssh_public_key`, as decided by the caller -- typically the user key, or
+        nothing if the server does not let the user connect to the container directly. The project
+        key is never among them, and the keys are not validated; pass them to
+        `base.authorized_keys.build_authorized_keys()` to get the complete, validated list to
+        authorize. Only Computes that add the keys themselves need this argument; VM-based
+        (shim-based) Computes ignore it, as the server submits the keys to the shim once the
+        instance is up.
         """
         pass
 
@@ -390,10 +400,14 @@ def run_job(
         volumes: List[Volume],
         placement_group: Optional[PlacementGroup],
         requirements: Requirements,
+        extra_authorized_keys: list[str],
     ) -> JobProvisioningData:
         """
         The default `run_job()` implementation for all backends that support `create_instance()`.
         Override only if custom `run_job()` behavior is required.
+
+        `extra_authorized_keys` is ignored -- all such backends are VM-based, and the server
+        submits the keys to the shim later, see `Compute.run_job()`.
         """
         instance_config = InstanceConfiguration(
             project_name=run.project_name,
@@ -442,7 +456,13 @@ def run_jobs(
         project_ssh_private_key: str,
         placement_group: Optional[PlacementGroup],
         requirements: Requirements,
+        extra_authorized_keys: list[str],
     ) -> ComputeGroupProvisioningData:
+        """
+        Launches a compute group -- instances created all at once via the provider API -- running
+        one job per instance. See `Compute.run_job()` for the arguments shared with it, including
+        `extra_authorized_keys`.
+        """
         pass
 
     @abstractmethod

src/dstack/_internal/core/backends/kubernetes/compute.py

@@ -15,7 +15,9 @@
 from typing_extensions import Self
 
 from dstack._internal.core.backends.base.authorized_keys import (
+    build_authorized_keys,
     get_add_authorized_keys_script,
+    normalize_authorized_keys,
 )
 from dstack._internal.core.backends.base.compute import (
     Compute,
@@ -126,7 +128,15 @@ class Operator(str, Enum):
 class KubernetesBackendData(CoreModel):
     jump_pod_name: str
     jump_pod_service_name: str
-    user_ssh_public_key: str
+    # TODO: remove `user_ssh_public_key` once servers before 0.21.4 are no longer supported.
+    user_ssh_public_key: str = ""
+    """Superseded by `extra_authorized_keys`, but still written -- as the first extra key, or an
+    empty string if there are none -- because servers before 0.21.4 require this field and may
+    read this record during a rolling deployment. Read only if `extra_authorized_keys` is empty,
+    that is, if such a server wrote the record.
+    """
+    extra_authorized_keys: list[str] = []
+    """Empty in records written before the field was introduced, see `user_ssh_public_key`."""
 
     @classmethod
     def load(cls, raw: str) -> Self:
@@ -193,6 +203,7 @@ def run_job(
         volumes: list[Volume],
         placement_group: Optional[PlacementGroup],
         requirements: Requirements,
+        extra_authorized_keys: list[str],
     ) -> JobProvisioningData:
         cluster = self.region_cluster_map.get(instance_offer.region)
         if cluster is None:
@@ -202,6 +213,13 @@ def run_job(
         api = client.CoreV1Api(cluster.api_client)
         namespace = cluster.namespace
 
+        # The jump pod, created below only if it does not exist yet, gets the project key alone.
+        # This job's extra keys are added to it later by update_provisioning_data(), which runs
+        # for every job, whether or not the job created the pod. The job pod gets both at once.
+        project_authorized_keys = build_authorized_keys(project_ssh_public_key, [])
+        extra_authorized_keys = normalize_authorized_keys(extra_authorized_keys)
+        authorized_keys = project_authorized_keys + extra_authorized_keys
+
         # There is one jump pod per project that is used as an ssh proxy jump to connect
         # to all job pods of the same project.
         # The service is created here and configured later in update_provisioning_data()
@@ -214,7 +232,7 @@ def run_job(
             jump_pod_name=jump_pod_name,
             jump_pod_service_name=jump_pod_service_name,
             jump_pod_port=cluster.proxy_jump.port,
-            project_ssh_public_key=project_ssh_public_key.strip(),
+            authorized_keys=project_authorized_keys,
         )
 
         pod_name = generate_unique_instance_name_for_job(
@@ -256,8 +274,6 @@ def run_job(
                     should_delete_manually_if_failed=True,
                 )
 
-            assert run.run_spec.ssh_key_pub is not None
-            authorized_keys = [run.run_spec.ssh_key_pub.strip(), project_ssh_public_key.strip()]
             _create_job_pod(
                 api=api,
                 namespace=namespace,
@@ -329,7 +345,9 @@ def run_job(
         backend_data = KubernetesBackendData(
             jump_pod_name=jump_pod_name,
             jump_pod_service_name=jump_pod_service_name,
-            user_ssh_public_key=run.run_spec.ssh_key_pub.strip(),
+            extra_authorized_keys=extra_authorized_keys,
+            # For compatibility with server replicas < 0.21.4
+            user_ssh_public_key=next(iter(extra_authorized_keys), ""),
         )
 
         return JobProvisioningData(
@@ -369,14 +387,19 @@ def update_provisioning_data(
         if provisioning_data.backend_data is not None:
             # Before running a job, ensure the jump pod is running and has user's public SSH key.
             backend_data = KubernetesBackendData.load(provisioning_data.backend_data)
+            extra_authorized_keys = backend_data.extra_authorized_keys
+            # TODO: remove the fallback once servers before 0.21.4 are no longer supported.
+            if not extra_authorized_keys and backend_data.user_ssh_public_key:
+                # Written by a server that predates the `extra_authorized_keys` field.
+                extra_authorized_keys = [backend_data.user_ssh_public_key]
             ssh_proxy = _check_and_configure_jump_pod_service(
                 api=api,
                 namespace=namespace,
                 jump_pod_name=backend_data.jump_pod_name,
                 jump_pod_service_name=backend_data.jump_pod_service_name,
                 jump_pod_hostname=cluster.proxy_jump.hostname,
                 project_ssh_private_key=project_ssh_private_key,
-                user_ssh_public_key=backend_data.user_ssh_public_key,
+                authorized_keys=extra_authorized_keys,
             )
             if ssh_proxy is None:
                 # Jump pod is not ready yet
@@ -837,7 +860,7 @@ def _create_jump_pod_service_if_not_exists(
     jump_pod_name: str,
     jump_pod_service_name: str,
     jump_pod_port: Optional[int],
-    project_ssh_public_key: str,
+    authorized_keys: list[str],
 ) -> None:
     base_labels = build_base_labels(
         component="ssh-proxy",
@@ -914,7 +937,7 @@ def _create_jump_pod_service_if_not_exists(
             )
         if not tolerations:
             logger.warning("No appropriate node found, the jump pod may never be scheduled")
-    commands = _get_jump_pod_commands(authorized_keys=[project_ssh_public_key])
+    commands = _get_jump_pod_commands(authorized_keys)
     pod = client.V1Pod(
         metadata=client.V1ObjectMeta(
             name=jump_pod_name,
@@ -977,7 +1000,7 @@ def _check_and_configure_jump_pod_service(
     jump_pod_service_name: str,
     jump_pod_hostname: Optional[str],
     project_ssh_private_key: str,
-    user_ssh_public_key: str,
+    authorized_keys: list[str],
 ) -> Optional[SSHConnectionParams]:
     jump_pod = api.read_namespaced_pod(
         namespace=namespace,
@@ -1031,21 +1054,26 @@ def _check_and_configure_jump_pod_service(
     if (jump_pod_port := jump_pod_service_ports[0].node_port) is None:
         raise ProvisioningError("Jump pod service %s port is not set", jump_pod_service_name)
 
-    ssh_exit_status, ssh_output = _run_ssh_command(
-        hostname=jump_pod_hostname,
-        port=jump_pod_port,
-        username=JUMP_POD_USER,
-        ssh_private_key=project_ssh_private_key,
-        command=get_add_authorized_keys_script(
-            [user_ssh_public_key],
+    if authorized_keys:
+        command = get_add_authorized_keys_script(
+            authorized_keys,
             # 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"',
-        ),
+        )
+    else:
+        # No keys to add, but the connection itself is still the check that sshd is up
+        command = "true"
+    ssh_exit_status, ssh_output = _run_ssh_command(
+        hostname=jump_pod_hostname,
+        port=jump_pod_port,
+        username=JUMP_POD_USER,
+        ssh_private_key=project_ssh_private_key,
+        command=command,
     )
     if ssh_exit_status != 0:
         logger.debug(
@@ -1072,11 +1100,11 @@ def _check_and_configure_jump_pod_service(
 
 
 def _get_jump_pod_commands(authorized_keys: list[str]) -> list[str]:
-    authorized_keys_content = "\n".join(authorized_keys).strip()
+    authorized_keys_content = "\n".join(authorized_keys)
     commands = [
         "mkdir -p ~/.ssh",
         "chmod 700 ~/.ssh",
-        f"echo '{authorized_keys_content}' > ~/.ssh/authorized_keys",
+        f"echo {shlex.quote(authorized_keys_content)} > ~/.ssh/authorized_keys",
         "chmod 600 ~/.ssh/authorized_keys",
         # regenerate host keys
         "rm -rf /etc/ssh/ssh_host_*",

src/dstack/_internal/core/backends/runpod/compute.py

@@ -4,6 +4,7 @@
 from datetime import timedelta
 from typing import Callable, List, Optional
 
+from dstack._internal.core.backends.base.authorized_keys import build_authorized_keys
 from dstack._internal.core.backends.base.backend import Compute
 from dstack._internal.core.backends.base.compute import (
     ComputeWithAllOffersCached,
@@ -34,7 +35,6 @@
     InstanceAvailability,
     InstanceConfiguration,
     InstanceOfferWithAvailability,
-    SSHKey,
 )
 from dstack._internal.core.models.placement import PlacementGroup
 from dstack._internal.core.models.resources import Memory, Range
@@ -135,20 +135,17 @@ def run_job(
         volumes: List[Volume],
         placement_group: Optional[PlacementGroup],
         requirements: Requirements,
+        extra_authorized_keys: list[str],
     ) -> JobProvisioningData:
-        assert run.run_spec.ssh_key_pub is not None
         instance_config = InstanceConfiguration(
             project_name=run.project_name,
             instance_name=get_job_instance_name(run, job),
-            ssh_keys=[
-                SSHKey(public=run.run_spec.ssh_key_pub.strip()),
-                SSHKey(public=project_ssh_public_key.strip()),
-            ],
+            ssh_keys=[],
             user=run.user,
         )
 
         pod_name = generate_unique_instance_name(instance_config, max_length=MAX_RESOURCE_NAME_LEN)
-        authorized_keys = instance_config.get_public_keys()
+        authorized_keys = build_authorized_keys(project_ssh_public_key, extra_authorized_keys)
         memory_size = round(instance_offer.instance.resources.memory_mib / 1024)
         disk_size = round(instance_offer.instance.resources.disk.size_mib / 1024)
 
@@ -251,6 +248,7 @@ def run_jobs(
         project_ssh_private_key: str,
         placement_group: Optional[PlacementGroup],
         requirements: Requirements,
+        extra_authorized_keys: list[str],
     ) -> ComputeGroupProvisioningData:
         master_job_configuration = job_configurations[0]
         master_job = master_job_configuration.job
@@ -259,15 +257,12 @@ def run_jobs(
         instance_config = InstanceConfiguration(
             project_name=run.project_name,
             instance_name=get_job_instance_name(run, master_job),
-            ssh_keys=[
-                SSHKey(public=get_or_error(run.run_spec.ssh_key_pub).strip()),
-                SSHKey(public=project_ssh_public_key.strip()),
-            ],
+            ssh_keys=[],
             user=run.user,
         )
 
         pod_name = generate_unique_instance_name(instance_config, max_length=MAX_RESOURCE_NAME_LEN)
-        authorized_keys = instance_config.get_public_keys()
+        authorized_keys = build_authorized_keys(project_ssh_public_key, extra_authorized_keys)
         disk_size = round(instance_offer.instance.resources.disk.size_mib / 1024)
 
         network_volume_id = None