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:
@@ -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
|
||||||
|
|||||||
@@ -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.
|
||||||
|
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
Reference in New Issue
Block a user