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)
This commit is contained in:
Alex Blank
2026-06-02 15:03:26 +02:00
parent c754984df8
commit c6b804bf0a
4 changed files with 217 additions and 87 deletions
+66 -17
View File
@@ -1611,17 +1611,79 @@ async def start_instance(
# Mount selected SSH keys into container home dir # Mount selected SSH keys into container home dir
if instance.ssh_key_ids: 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: for key_id in instance.ssh_key_ids:
ssh_key = await session.get(SSHKey, uuid.UUID(key_id)) ssh_key = await session.get(SSHKey, uuid.UUID(key_id))
if ssh_key and ssh_key.user_id == user_id: if ssh_key and ssh_key.user_id == user_id:
ssh_keys_to_mount.append(ssh_key)
else:
logger.warning(
"SSH key %s not found or not authorized for user %s",
key_id,
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: try:
ssh_dir = prepare_ssh_key_files( prepare_ssh_key_files(
instance_dir, instance_dir,
ssh_key, ssh_key,
subdir=f"mounts/ssh/{key_id}/.ssh", 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, uid=container_uid,
gid=container_gid, 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") ssh_target = os.path.join(home_dir, ".ssh")
extra_volumes.append( extra_volumes.append(
{ {
@@ -1631,24 +1693,11 @@ async def start_instance(
} }
) )
logger.debug( logger.debug(
"Mounted SSH key %s for instance %s to %s", "Mounted %d SSH key(s) for instance %s to %s",
ssh_key.name, len(ssh_keys_to_mount),
instance.id, instance.id,
ssh_target, ssh_target,
) )
except Exception as exc:
logger.error(
"Failed to prepare SSH key %s for instance %s: %s",
key_id,
instance.id,
exc,
)
else:
logger.warning(
"SSH key %s not found or not authorized for user %s",
key_id,
user_id,
)
# ── MANIFEST-BASED FLOW ────────────────────────────────────── # ── MANIFEST-BASED FLOW ──────────────────────────────────────
resolved_manifest = None resolved_manifest = None
+78 -5
View File
@@ -2,6 +2,7 @@
import logging import logging
import os import os
import re
from pathlib import Path from pathlib import Path
from cryptography.fernet import Fernet from cryptography.fernet import Fernet
@@ -22,12 +23,28 @@ def _get_fernet() -> Fernet:
return Fernet(key) 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( def prepare_ssh_key_files(
instance_dir: str, instance_dir: str,
ssh_key, ssh_key,
subdir: str = ".ssh", subdir: str = ".ssh",
uid: int | None = None, uid: int | None = None,
gid: int | None = None, gid: int | None = None,
key_filename: str = "id_ed25519",
write_config: bool = True,
) -> str: ) -> str:
"""Decrypt and write SSH key files to instance directory for container mounting. """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") subdir: Subdirectory within instance_dir to write to (default: ".ssh")
uid: Optional UID to own the files (for bind-mount into non-root container) uid: Optional UID to own the files (for bind-mount into non-root container)
gid: Optional GID to own the files 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: Returns:
Path to the .ssh directory Path to the .ssh directory
@@ -49,21 +72,22 @@ def prepare_ssh_key_files(
private_key = fernet.decrypt(ssh_key.private_key_encrypted.encode()).decode() private_key = fernet.decrypt(ssh_key.private_key_encrypted.encode()).decode()
# Write private key with restricted permissions # 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) private_key_path.write_text(private_key)
os.chmod(private_key_path, 0o600) os.chmod(private_key_path, 0o600)
# Write public key # 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) public_key_path.write_text(ssh_key.public_key)
os.chmod(public_key_path, 0o644) os.chmod(public_key_path, 0o644)
# Write SSH config # Write SSH config (only if requested)
if write_config:
config_path = ssh_dir / "config" config_path = ssh_dir / "config"
config_content = """Host * config_content = f"""Host *
StrictHostKeyChecking no StrictHostKeyChecking no
UserKnownHostsFile /dev/null UserKnownHostsFile /dev/null
IdentityFile ~/.ssh/id_ed25519 IdentityFile ~/.ssh/{key_filename}
IdentitiesOnly yes IdentitiesOnly yes
""" """
config_path.write_text(config_content) config_path.write_text(config_content)
@@ -92,10 +116,59 @@ def prepare_ssh_key_files(
os.getuid(), os.getuid(),
exc, 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) 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)
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(config_path, effective_uid, effective_gid)
except PermissionError:
pass
def cleanup_ssh_key_files(instance_dir: str) -> None: def cleanup_ssh_key_files(instance_dir: str) -> None:
"""Remove temporary SSH key files from instance directory. """Remove temporary SSH key files from instance directory.
+12 -6
View File
@@ -164,7 +164,9 @@ class TerminalSession:
self._read_handler_set = True self._read_handler_set = True
logger.debug("Started event-driven reading for session %s", self.session_id) logger.debug("Started event-driven reading for session %s", self.session_id)
except Exception as exc: 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: def _stop_reading(self) -> None:
"""Unregister PTY master fd from asyncio event loop.""" """Unregister PTY master fd from asyncio event loop."""
@@ -263,9 +265,7 @@ class TerminalSession:
if self._ack_timeout_handle: if self._ack_timeout_handle:
self._ack_timeout_handle.cancel() self._ack_timeout_handle.cancel()
loop = asyncio.get_event_loop() loop = asyncio.get_event_loop()
self._ack_timeout_handle = loop.call_later( self._ack_timeout_handle = loop.call_later(5.0, self._ack_timeout_fallback)
5.0, self._ack_timeout_fallback
)
def _ack_timeout_fallback(self) -> None: def _ack_timeout_fallback(self) -> None:
"""If no ack received for 5s, assume client is dead and resume.""" """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.""" """Pause reading from PTY due to flow control."""
self._paused = True self._paused = True
self._stop_reading() 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: def _resume_output(self) -> None:
"""Resume reading from PTY.""" """Resume reading from PTY."""
@@ -327,7 +331,9 @@ class TerminalSession:
self._cols = cols self._cols = cols
self._rows = rows 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) self._set_terminal_size(cols, rows)
# Send SIGWINCH to docker exec process # Send SIGWINCH to docker exec process
@@ -837,6 +837,7 @@ class TestStartInstanceSshPermissions:
mock_session.get.side_effect = _get mock_session.get.side_effect = _get
with patch("os.path.exists", return_value=True): with patch("os.path.exists", return_value=True):
with patch("os.makedirs"):
with patch( with patch(
"src.api.tool_instances._prepare_manifest_instance" "src.api.tool_instances._prepare_manifest_instance"
) as mock_prepare: ) as mock_prepare:
@@ -960,6 +961,7 @@ class TestStartInstanceSshPermissions:
mock_session.get.side_effect = _get mock_session.get.side_effect = _get
with patch("os.path.exists", return_value=True): with patch("os.path.exists", return_value=True):
with patch("os.makedirs"):
with patch("src.api.tool_instances._modify_compose_file"): with patch("src.api.tool_instances._modify_compose_file"):
result = await start_instance( result = await start_instance(
project_id=fake_project_id, project_id=fake_project_id,