feat: fix SSH key mounting with multi-key support and unique filenames

SSH key mounting was broken because: 1. Each selected key was mounted to a separate source dir but all targeted the same ~/.ssh path in the container, causing Docker Compose's last-mount-wins behavior 2. All keys were named id_ed25519, so they'd overwrite each other Changes: - ssh_keys.py: add key_filename param to prepare_ssh_key_files for unique key names; add write_ssh_config for combined multi-key config - tool_instances.py: collect all selected keys into a single ~/.ssh mount with sanitized unique filenames (id_ed25519_<name>); generate combined SSH config with all IdentityFile entries - tests: add os.makedirs mock for SSH permission tests Quality gates: pytest (19 passed, 1 skipped)
2026-06-02 15:03:26 +02:00
parent c754984df8
commit c6b804bf0a
4 changed files with 217 additions and 87 deletions
@@ -1611,38 +1611,14 @@ async def start_instance(

    # Mount selected SSH keys into container home dir
    if instance.ssh_key_ids:
+        from src.services.ssh_keys import write_ssh_config, _sanitize_filename
+
+        # Collect all valid keys first
+        ssh_keys_to_mount = []
        for key_id in instance.ssh_key_ids:
            ssh_key = await session.get(SSHKey, uuid.UUID(key_id))
            if ssh_key and ssh_key.user_id == user_id:
-                try:
-                    ssh_dir = prepare_ssh_key_files(
-                        instance_dir,
-                        ssh_key,
-                        subdir=f"mounts/ssh/{key_id}/.ssh",
-                        uid=container_uid,
-                        gid=container_gid,
-                    )
-                    ssh_target = os.path.join(home_dir, ".ssh")
-                    extra_volumes.append(
-                        {
-                            "source": ssh_dir,
-                            "target": ssh_target,
-                            "type": "bind",
-                        }
-                    )
-                    logger.debug(
-                        "Mounted SSH key %s for instance %s to %s",
-                        ssh_key.name,
-                        instance.id,
-                        ssh_target,
-                    )
-                except Exception as exc:
-                    logger.error(
-                        "Failed to prepare SSH key %s for instance %s: %s",
-                        key_id,
-                        instance.id,
-                        exc,
-                    )
+                ssh_keys_to_mount.append(ssh_key)
            else:
                logger.warning(
                    "SSH key %s not found or not authorized for user %s",
@@ -1650,6 +1626,79 @@ async def start_instance(
                    user_id,
                )

+        if ssh_keys_to_mount:
+            # Use a single shared .ssh directory so all keys are visible
+            ssh_dir = os.path.join(instance_dir, "mounts", "ssh", ".ssh")
+            os.makedirs(ssh_dir, exist_ok=True)
+
+            key_filenames = []
+            for ssh_key in ssh_keys_to_mount:
+                # Use sanitized key name as filename prefix to avoid collisions
+                key_name = _sanitize_filename(ssh_key.name)
+                # If multiple keys have the same name, append a short hash
+                base_filename = f"id_ed25519_{key_name}"
+                filename = base_filename
+                counter = 1
+                while filename in key_filenames:
+                    filename = f"{base_filename}_{counter}"
+                    counter += 1
+                key_filenames.append(filename)
+
+                try:
+                    prepare_ssh_key_files(
+                        instance_dir,
+                        ssh_key,
+                        subdir="mounts/ssh/.ssh",
+                        uid=container_uid,
+                        gid=container_gid,
+                        key_filename=filename,
+                        write_config=False,
+                    )
+                    logger.debug(
+                        "Prepared SSH key %s as %s for instance %s",
+                        ssh_key.name,
+                        filename,
+                        instance.id,
+                    )
+                except Exception as exc:
+                    logger.error(
+                        "Failed to prepare SSH key %s for instance %s: %s",
+                        ssh_key.id,
+                        instance.id,
+                        exc,
+                    )
+
+            # Write combined SSH config with all keys
+            try:
+                write_ssh_config(
+                    ssh_dir,
+                    key_filenames,
+                    uid=container_uid,
+                    gid=container_gid,
+                )
+            except Exception as exc:
+                logger.error(
+                    "Failed to write SSH config for instance %s: %s",
+                    instance.id,
+                    exc,
+                )
+
+            # Mount the single .ssh directory into container home
+            ssh_target = os.path.join(home_dir, ".ssh")
+            extra_volumes.append(
+                {
+                    "source": ssh_dir,
+                    "target": ssh_target,
+                    "type": "bind",
+                }
+            )
+            logger.debug(
+                "Mounted %d SSH key(s) for instance %s to %s",
+                len(ssh_keys_to_mount),
+                instance.id,
+                ssh_target,
+            )
+
    # ── MANIFEST-BASED FLOW ──────────────────────────────────────
    resolved_manifest = None

@@ -2,6 +2,7 @@

 import logging
 import os
+import re
 from pathlib import Path

 from cryptography.fernet import Fernet
@@ -22,12 +23,28 @@ def _get_fernet() -> Fernet:
    return Fernet(key)


+def _sanitize_filename(name: str) -> str:
+    """Sanitize a string for use as a filename.
+
+    Replaces non-alphanumeric characters with underscores and strips
+    leading/trailing underscores.
+    """
+    sanitized = re.sub(r"[^a-zA-Z0-9_-]", "_", name)
+    sanitized = sanitized.strip("_")
+    # Ensure it's not empty
+    if not sanitized:
+        sanitized = "key"
+    return sanitized
+
+
 def prepare_ssh_key_files(
    instance_dir: str,
    ssh_key,
    subdir: str = ".ssh",
    uid: int | None = None,
    gid: int | None = None,
+    key_filename: str = "id_ed25519",
+    write_config: bool = True,
 ) -> str:
    """Decrypt and write SSH key files to instance directory for container mounting.

@@ -37,6 +54,12 @@ def prepare_ssh_key_files(
        subdir: Subdirectory within instance_dir to write to (default: ".ssh")
        uid: Optional UID to own the files (for bind-mount into non-root container)
        gid: Optional GID to own the files
+        key_filename: Base filename for the key pair (default: "id_ed25519").
+            The private key will be named "{key_filename}" and the public key
+            "{key_filename}.pub".
+        write_config: Whether to write an SSH config file (default: True).
+            Set to False when combining multiple keys into one directory,
+            then call write_ssh_config() separately.

    Returns:
        Path to the .ssh directory
@@ -49,51 +72,101 @@ def prepare_ssh_key_files(
    private_key = fernet.decrypt(ssh_key.private_key_encrypted.encode()).decode()

    # Write private key with restricted permissions
-    private_key_path = ssh_dir / "id_ed25519"
+    private_key_path = ssh_dir / key_filename
    private_key_path.write_text(private_key)
    os.chmod(private_key_path, 0o600)

    # Write public key
-    public_key_path = ssh_dir / "id_ed25519.pub"
+    public_key_path = ssh_dir / f"{key_filename}.pub"
    public_key_path.write_text(ssh_key.public_key)
    os.chmod(public_key_path, 0o644)

-    # Write SSH config
-    config_path = ssh_dir / "config"
-    config_content = """Host *
+    # Write SSH config (only if requested)
+    if write_config:
+        config_path = ssh_dir / "config"
+        config_content = f"""Host *
    StrictHostKeyChecking no
    UserKnownHostsFile /dev/null
-    IdentityFile ~/.ssh/id_ed25519
+    IdentityFile ~/.ssh/{key_filename}
    IdentitiesOnly yes
 """
+        config_path.write_text(config_content)
+        os.chmod(config_path, 0o644)
+
+        # Set ownership to target container user if requested
+        if uid is not None or gid is not None:
+            effective_uid = uid if uid is not None else -1
+            effective_gid = gid if gid is not None else -1
+            try:
+                os.chown(ssh_dir, effective_uid, effective_gid)
+                os.chown(private_key_path, effective_uid, effective_gid)
+                os.chown(public_key_path, effective_uid, effective_gid)
+                os.chown(config_path, effective_uid, effective_gid)
+                logger.debug(
+                    "Set SSH key ownership to uid=%s gid=%s for %s",
+                    effective_uid,
+                    effective_gid,
+                    ssh_dir,
+                )
+            except PermissionError as exc:
+                logger.warning(
+                    "Cannot chown SSH keys to uid=%s gid=%s (running as uid=%s): %s",
+                    effective_uid,
+                    effective_gid,
+                    os.getuid(),
+                    exc,
+                )
+    else:
+        # Still chown the key files even if we didn't write config
+        if uid is not None or gid is not None:
+            effective_uid = uid if uid is not None else -1
+            effective_gid = gid if gid is not None else -1
+            try:
+                os.chown(private_key_path, effective_uid, effective_gid)
+                os.chown(public_key_path, effective_uid, effective_gid)
+            except PermissionError:
+                pass
+
+    return str(ssh_dir)
+
+
+def write_ssh_config(
+    ssh_dir: str,
+    key_filenames: list[str],
+    uid: int | None = None,
+    gid: int | None = None,
+) -> None:
+    """Write an SSH config file that includes multiple IdentityFile entries.
+
+    Args:
+        ssh_dir: Path to the .ssh directory
+        key_filenames: List of key filenames (without .pub extension)
+        uid: Optional UID to own the config file
+        gid: Optional GID to own the config file
+    """
+    ssh_dir_path = Path(ssh_dir)
+    ssh_dir_path.mkdir(parents=True, exist_ok=True)
+
+    config_path = ssh_dir_path / "config"
+    lines = ["Host *"]
+    lines.append("    StrictHostKeyChecking no")
+    lines.append("    UserKnownHostsFile /dev/null")
+    lines.append("    IdentitiesOnly yes")
+    for filename in key_filenames:
+        lines.append(f"    IdentityFile ~/.ssh/{filename}")
+    lines.append("")
+
+    config_content = "\n".join(lines)
    config_path.write_text(config_content)
    os.chmod(config_path, 0o644)

-    # Set ownership to target container user if requested
    if uid is not None or gid is not None:
        effective_uid = uid if uid is not None else -1
        effective_gid = gid if gid is not None else -1
        try:
-            os.chown(ssh_dir, effective_uid, effective_gid)
-            os.chown(private_key_path, effective_uid, effective_gid)
-            os.chown(public_key_path, effective_uid, effective_gid)
            os.chown(config_path, effective_uid, effective_gid)
-            logger.debug(
-                "Set SSH key ownership to uid=%s gid=%s for %s",
-                effective_uid,
-                effective_gid,
-                ssh_dir,
-            )
-        except PermissionError as exc:
-            logger.warning(
-                "Cannot chown SSH keys to uid=%s gid=%s (running as uid=%s): %s",
-                effective_uid,
-                effective_gid,
-                os.getuid(),
-                exc,
-            )
-
-    return str(ssh_dir)
+        except PermissionError:
+            pass


 def cleanup_ssh_key_files(instance_dir: str) -> None:
@@ -164,7 +164,9 @@ class TerminalSession:
            self._read_handler_set = True
            logger.debug("Started event-driven reading for session %s", self.session_id)
        except Exception as exc:
-            logger.error("Failed to start reading for session %s: %s", self.session_id, exc)
+            logger.error(
+                "Failed to start reading for session %s: %s", self.session_id, exc
+            )

    def _stop_reading(self) -> None:
        """Unregister PTY master fd from asyncio event loop."""
@@ -263,9 +265,7 @@ class TerminalSession:
        if self._ack_timeout_handle:
            self._ack_timeout_handle.cancel()
        loop = asyncio.get_event_loop()
-        self._ack_timeout_handle = loop.call_later(
-            5.0, self._ack_timeout_fallback
-        )
+        self._ack_timeout_handle = loop.call_later(5.0, self._ack_timeout_fallback)

    def _ack_timeout_fallback(self) -> None:
        """If no ack received for 5s, assume client is dead and resume."""
@@ -281,7 +281,11 @@ class TerminalSession:
        """Pause reading from PTY due to flow control."""
        self._paused = True
        self._stop_reading()
-        logger.debug("Paused output for session %s (%d unacked)", self.session_id, self._unacknowledged_bytes)
+        logger.debug(
+            "Paused output for session %s (%d unacked)",
+            self.session_id,
+            self._unacknowledged_bytes,
+        )

    def _resume_output(self) -> None:
        """Resume reading from PTY."""
@@ -327,7 +331,9 @@ class TerminalSession:

        self._cols = cols
        self._rows = rows
-        logger.debug("resize() called for session %s: %sx%s", self.session_id, cols, rows)
+        logger.debug(
+            "resize() called for session %s: %sx%s", self.session_id, cols, rows
+        )
        self._set_terminal_size(cols, rows)

        # Send SIGWINCH to docker exec process