From c733b66555b6b821489553a7245be6b485214535 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Wed, 8 Jul 2026 17:37:59 +0200 Subject: [PATCH] Shorten SSH attach control socket path --- .../_internal/core/services/ssh/attach.py | 92 ++++++++++- .../core/services/ssh/test_attach.py | 149 ++++++++++++++++++ 2 files changed, 240 insertions(+), 1 deletion(-) create mode 100644 src/tests/_internal/core/services/ssh/test_attach.py diff --git a/src/dstack/_internal/core/services/ssh/attach.py b/src/dstack/_internal/core/services/ssh/attach.py index 7727aef5df..27b494b43c 100644 --- a/src/dstack/_internal/core/services/ssh/attach.py +++ b/src/dstack/_internal/core/services/ssh/attach.py @@ -1,5 +1,8 @@ import atexit +import hashlib +import os import re +import stat import time from pathlib import Path from typing import Optional, Union @@ -26,6 +29,12 @@ # ssh -L option format: [bind_address:]port:host:hostport _SSH_TUNNEL_REGEX = re.compile(r"(?:[\w.-]+:)?(?P\d+):localhost:(?P\d+)") +# Use a fixed short POSIX temp root for OpenSSH ControlPath. tempfile.gettempdir() +# can point to long per-user paths on macOS and reintroduce the socket length failure. +_CONTROL_SOCKET_BASE_DIR = Path("/tmp") +_CONTROL_SOCKET_DIR_MODE = 0o700 +_CONTROL_SOCKET_HASH_LENGTH = 24 +_CONTROL_SOCKET_MAX_PATH_BYTES = 103 class BaseSSHAttach: @@ -38,7 +47,10 @@ class BaseSSHAttach: @classmethod def get_control_sock_path(cls, run_name: str) -> Path: - return ConfigManager().dstack_ssh_dir / f"{run_name}.control.sock" + config_manager = ConfigManager() + if os.name != "posix" or not hasattr(os, "getuid"): + return config_manager.dstack_ssh_dir / f"{run_name}.control.sock" + return _get_control_sock_path(config_manager.dstack_ssh_dir, run_name) @classmethod def reuse_ports_lock(cls, run_name: str) -> Optional[PortsLock]: @@ -149,6 +161,84 @@ def _remove_hosts_from_ssh_config(self): self._hosts_added_to_ssh_config = False +def _get_control_sock_path(dstack_ssh_dir: Path, run_name: str) -> Path: + key = f"{dstack_ssh_dir.resolve(strict=False)}\0{run_name}".encode() + digest = hashlib.sha256(key).hexdigest()[:_CONTROL_SOCKET_HASH_LENGTH] + sock_name = f"{digest}.sock" + short_control_sock_path = _CONTROL_SOCKET_BASE_DIR / f"dstack-ssh-{os.getuid()}" / sock_name + try: + _ensure_control_socket_dir(short_control_sock_path.parent) + short_path_error = _get_control_sock_path_length_error(short_control_sock_path) + except SSHError as e: + short_path_error = str(e) + if short_path_error is None: + return short_control_sock_path + + # Preserve the old location as a fallback for environments where /tmp cannot be used. + # This fallback is only valid if the resulting Unix socket path is still short enough + # for OpenSSH. + legacy_control_sock_path = dstack_ssh_dir / f"{run_name}.control.sock" + legacy_path_error = _get_control_sock_path_length_error(legacy_control_sock_path) + if legacy_path_error is None: + return legacy_control_sock_path + + raise SSHError( + "Cannot create SSH control socket: " + f"short path failed ({short_path_error}); " + f"fallback path failed ({legacy_path_error})" + ) + + +def _ensure_control_socket_dir(runtime_dir: Path) -> None: + base_dir = runtime_dir.parent + # Follow /tmp if it is a symlink (macOS), but require unsafe shared dirs to be sticky. + try: + st = base_dir.stat() + except OSError as e: + raise SSHError(f"Cannot access SSH control socket base directory {base_dir}: {e}") from e + if not stat.S_ISDIR(st.st_mode): + raise SSHError(f"Unsafe SSH control socket base directory {base_dir}: not a directory") + mode = stat.S_IMODE(st.st_mode) + if mode & 0o022 and not (st.st_mode & stat.S_ISVTX): + raise SSHError(f"Unsafe SSH control socket base directory {base_dir}: unsafe permissions") + + created = False + try: + runtime_dir.mkdir(mode=_CONTROL_SOCKET_DIR_MODE) + created = True + except FileExistsError: + pass + except OSError as e: + raise SSHError(f"Cannot create SSH control socket directory {runtime_dir}: {e}") from e + if created: + try: + runtime_dir.chmod(_CONTROL_SOCKET_DIR_MODE) + except OSError as e: + raise SSHError(f"Cannot secure SSH control socket directory {runtime_dir}: {e}") from e + try: + st = runtime_dir.lstat() + except OSError as e: + raise SSHError(f"Cannot access SSH control socket directory {runtime_dir}: {e}") from e + if not stat.S_ISDIR(st.st_mode): + raise SSHError(f"Unsafe SSH control socket directory {runtime_dir}: not a directory") + if st.st_uid != os.getuid(): + raise SSHError(f"Unsafe SSH control socket directory {runtime_dir}: wrong owner") + if stat.S_IMODE(st.st_mode) & 0o077: + raise SSHError(f"Unsafe SSH control socket directory {runtime_dir}: unsafe permissions") + + +def _get_control_sock_path_length_error(path: Path) -> Optional[str]: + # OpenSSH uses Unix-domain sockets for ControlPath. macOS rejects paths at 104 bytes, + # so use 103 as the conservative cross-platform ceiling and count encoded bytes. + path_length = len(os.fsencode(path)) + if path_length > _CONTROL_SOCKET_MAX_PATH_BYTES: + return ( + f"{path} is too long for an SSH control socket " + f"({path_length} bytes > {_CONTROL_SOCKET_MAX_PATH_BYTES})" + ) + return None + + class SSHAttach(BaseSSHAttach): """ `SSHAttach` attaches to a job directly, via a backend-specific chain of hosts. diff --git a/src/tests/_internal/core/services/ssh/test_attach.py b/src/tests/_internal/core/services/ssh/test_attach.py new file mode 100644 index 0000000000..bdb4a9ed19 --- /dev/null +++ b/src/tests/_internal/core/services/ssh/test_attach.py @@ -0,0 +1,149 @@ +import os +import shutil +import stat +import uuid +from pathlib import Path +from typing import Generator + +import pytest + +from dstack._internal.core.errors import SSHError +from dstack._internal.core.services.ssh import attach + +pytestmark = pytest.mark.skipif( + os.name != "posix" or not hasattr(os, "getuid"), + reason="OpenSSH control sockets are only used on POSIX", +) + + +@pytest.fixture +def short_tmp_path() -> Generator[Path, None, None]: + path = Path("/tmp") / f"dstack-test-{uuid.uuid4().hex[:8]}" + path.mkdir(mode=0o700) + try: + yield path + finally: + shutil.rmtree(path, ignore_errors=True) + + +def test_base_attach_control_sock_path_uses_short_runtime_dir( + monkeypatch: pytest.MonkeyPatch, short_tmp_path: Path +) -> None: + monkeypatch.setattr(attach, "_CONTROL_SOCKET_BASE_DIR", short_tmp_path) + dstack_ssh_dir = short_tmp_path / ("nested-" * 20) / ".dstack" / "ssh" + + class FakeConfigManager: + def __init__(self) -> None: + self.dstack_ssh_dir = dstack_ssh_dir + + monkeypatch.setattr(attach, "ConfigManager", FakeConfigManager) + + path = attach.BaseSSHAttach.get_control_sock_path("endpoint-with-a-very-long-name-" * 4) + + assert path.parent == short_tmp_path / f"dstack-ssh-{os.getuid()}" + assert path.suffix == ".sock" + assert not path.is_relative_to(dstack_ssh_dir) + + +def test_control_sock_path_uses_short_private_runtime_dir( + monkeypatch: pytest.MonkeyPatch, short_tmp_path: Path +) -> None: + monkeypatch.setattr(attach, "_CONTROL_SOCKET_BASE_DIR", short_tmp_path) + dstack_ssh_dir = short_tmp_path / ("nested-" * 20) / ".dstack" / "ssh" + run_name = "endpoint-with-a-very-long-name-" * 4 + + path = attach._get_control_sock_path(dstack_ssh_dir, run_name) + + assert path == attach._get_control_sock_path(dstack_ssh_dir, run_name) + assert path.parent == short_tmp_path / f"dstack-ssh-{os.getuid()}" + assert path.suffix == ".sock" + assert len(path.name) == attach._CONTROL_SOCKET_HASH_LENGTH + len(".sock") + assert not path.is_relative_to(dstack_ssh_dir) + assert stat.S_IMODE(path.parent.stat().st_mode) == 0o700 + + +def test_control_sock_path_separates_dstack_ssh_dirs( + monkeypatch: pytest.MonkeyPatch, short_tmp_path: Path +) -> None: + monkeypatch.setattr(attach, "_CONTROL_SOCKET_BASE_DIR", short_tmp_path) + + path1 = attach._get_control_sock_path(short_tmp_path / "home1" / ".dstack" / "ssh", "run") + path2 = attach._get_control_sock_path(short_tmp_path / "home2" / ".dstack" / "ssh", "run") + + assert path1 != path2 + + +def test_control_sock_path_falls_back_from_unsafe_runtime_dir( + monkeypatch: pytest.MonkeyPatch, short_tmp_path: Path +) -> None: + monkeypatch.setattr(attach, "_CONTROL_SOCKET_BASE_DIR", short_tmp_path) + runtime_dir = short_tmp_path / f"dstack-ssh-{os.getuid()}" + runtime_dir.mkdir() + runtime_dir.chmod(0o755) + + legacy_path = short_tmp_path / ".dstack" / "ssh" / "run.control.sock" + + assert attach._get_control_sock_path(short_tmp_path / ".dstack" / "ssh", "run") == legacy_path + + +def test_control_sock_path_rejects_runtime_dir_symlink_and_falls_back( + monkeypatch: pytest.MonkeyPatch, short_tmp_path: Path +) -> None: + monkeypatch.setattr(attach, "_CONTROL_SOCKET_BASE_DIR", short_tmp_path) + runtime_dir = short_tmp_path / f"dstack-ssh-{os.getuid()}" + target_dir = short_tmp_path / "target" + target_dir.mkdir() + runtime_dir.symlink_to(target_dir) + + legacy_path = short_tmp_path / ".dstack" / "ssh" / "run.control.sock" + + assert attach._get_control_sock_path(short_tmp_path / ".dstack" / "ssh", "run") == legacy_path + + +def test_control_sock_path_falls_back_when_short_base_dir_is_unavailable( + monkeypatch: pytest.MonkeyPatch, short_tmp_path: Path +) -> None: + monkeypatch.setattr(attach, "_CONTROL_SOCKET_BASE_DIR", short_tmp_path / "missing") + legacy_path = short_tmp_path / ".dstack" / "ssh" / "run.control.sock" + + assert attach._get_control_sock_path(short_tmp_path / ".dstack" / "ssh", "run") == legacy_path + + +def test_control_sock_path_falls_back_when_short_path_is_too_long( + monkeypatch: pytest.MonkeyPatch, short_tmp_path: Path +) -> None: + long_base_dir = short_tmp_path / ("nested-" * 20) + long_base_dir.mkdir() + monkeypatch.setattr(attach, "_CONTROL_SOCKET_BASE_DIR", long_base_dir) + legacy_path = short_tmp_path / ".dstack" / "ssh" / "run.control.sock" + + assert attach._get_control_sock_path(short_tmp_path / ".dstack" / "ssh", "run") == legacy_path + + +def test_control_sock_path_reports_short_and_fallback_failures( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + monkeypatch.setattr(attach, "_CONTROL_SOCKET_BASE_DIR", tmp_path / "missing") + dstack_ssh_dir = tmp_path / ("nested-" * 20) / ".dstack" / "ssh" + + with pytest.raises(SSHError) as exc_info: + attach._get_control_sock_path(dstack_ssh_dir, "endpoint-with-a-very-long-name-" * 4) + error = str(exc_info.value) + assert "short path failed" in error + assert "fallback path failed" in error + assert "Cannot access SSH control socket base directory" in error + assert "is too long for an SSH control socket" in error + + +def test_control_sock_path_rejects_runtime_dir_symlink( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, short_tmp_path: Path +) -> None: + monkeypatch.setattr(attach, "_CONTROL_SOCKET_BASE_DIR", short_tmp_path) + runtime_dir = short_tmp_path / f"dstack-ssh-{os.getuid()}" + target_dir = short_tmp_path / "target" + target_dir.mkdir() + runtime_dir.symlink_to(target_dir) + dstack_ssh_dir = tmp_path / ("nested-" * 20) / ".dstack" / "ssh" + + with pytest.raises(SSHError, match="not a directory"): + attach._get_control_sock_path(dstack_ssh_dir, "endpoint-with-a-very-long-name-" * 4)