From 34f70797edf22f44ec08c638ca823096a77753f5 Mon Sep 17 00:00:00 2001 From: Kai Date: Sat, 24 Oct 2026 16:28:53 +0800 Subject: [PATCH 1/8] feat(environments): execute durable operations in Linux VMs --- opencollab/adapters/_env_vm.py | 308 +++++++++++++++++++ opencollab/adapters/_vm_remote.py | 334 +++++++++++++++++++++ opencollab/adapters/env.py | 4 + opencollab/adapters/safe_anchored_files.py | 4 + opencollab/domain/vm.py | 43 +++ tests/runtime/test_vm_environment.py | 133 ++++++++ 6 files changed, 826 insertions(+) create mode 100644 opencollab/adapters/_env_vm.py create mode 100644 opencollab/adapters/_vm_remote.py create mode 100644 opencollab/domain/vm.py create mode 100644 tests/runtime/test_vm_environment.py diff --git a/opencollab/adapters/_env_vm.py b/opencollab/adapters/_env_vm.py new file mode 100644 index 00000000..e27ae856 --- /dev/null +++ b/opencollab/adapters/_env_vm.py @@ -0,0 +1,308 @@ +from __future__ import annotations + +import asyncio +import base64 +import contextlib +import json +import math +import posixpath +import re +import shlex +import uuid +from collections.abc import AsyncIterator +from typing import Any + +from opencollab import __version__ +from opencollab.adapters._env_base import ENV_FILE_WRITE_LIMIT_BYTES, Environment, ExecResult, TextFileRange +from opencollab.adapters._env_process import ProcessCleanupError, run_process +from opencollab.application.async_timeout import await_owned_operation +from opencollab.domain.vm import VMConnection, validate_vm_path + +_MESSAGE_LIMIT = 16 * 1024 * 1024 +_COMPONENT = re.compile(r"[A-Za-z0-9][A-Za-z0-9_.-]{0,127}") + + +class VMOperationError(RuntimeError): + def __init__(self, message: str, *, operation_id: str, state: str) -> None: + super().__init__(message) + self.operation_id = operation_id + self.state = state + + +class VMTransportError(VMOperationError): + def __init__(self, message: str, *, operation_id: str, returncode: int | None = None) -> None: + super().__init__(message, operation_id=operation_id, state="unknown") + self.transport_returncode = returncode + + +class VMEnvironment(Environment): + """Borrow an existing Linux guest with the same installed OC version. + + Operation records remain in state_directory/run_id for recovery after a + controller or SSH connection exits. cleanup stops owned operations and + removes owned temporary files while preserving the guest and records. + """ + + process_isolated = True + + def __init__(self, connection: VMConnection, workspace: str, *, state_directory: str, run_id: str | None = None): + super().__init__() + self.connection = connection + self.workspace = validate_vm_path(workspace) + self.source_workspace = self.workspace + self.state_directory = validate_vm_path(state_directory) + if self.state_directory == self.workspace or self.state_directory.startswith(self.workspace + "/"): + raise ValueError("VM operation state must be outside the tool workspace") + self.run_id = run_id or uuid.uuid4().hex + if not _COMPONENT.fullmatch(self.run_id): + raise ValueError("VM run_id must be a simple path component") + self._identity: dict[str, Any] | None = None + self._setup_lock = asyncio.Lock() + self._operations: set[str] = set() + self._uncertain: set[str] = set() + self._temporary_files: dict[str, list[int]] = {} + + @property + def revoked(self) -> bool: + return self._aborted or bool(self._uncertain) + + @property + def pending_operations(self) -> tuple[str, ...]: + return tuple(sorted(self._operations)) + + def _ensure_active(self) -> None: + super()._ensure_active() + if self._uncertain: + operation = sorted(self._uncertain)[0] + raise VMOperationError("Recover uncertain VM operations before starting further work", + operation_id=operation, state="unknown") + + def _ssh_command(self) -> list[str]: + connection = self.connection + command = ["ssh", "-T", "-F", "/dev/null", "-o", "BatchMode=yes", + "-o", "StrictHostKeyChecking=yes", "-o", "ControlMaster=no", "-o", "ControlPath=none", + "-o", "ServerAliveInterval=5", "-o", "ServerAliveCountMax=2", + "-o", f"ConnectTimeout={math.ceil(connection.connect_timeout)}", + "-p", str(connection.port)] + if connection.identity_file is not None: + command += ["-i", connection.identity_file, "-o", "IdentitiesOnly=yes"] + if connection.known_hosts_file is not None: + command += ["-o", f"UserKnownHostsFile={connection.known_hosts_file}"] + command += ["-l", connection.user, connection.host, + shlex.join([connection.python, "-m", "opencollab.adapters._vm_remote", "rpc", + self.state_directory, self.run_id])] + return command + + async def _rpc(self, payload: dict[str, Any]) -> Any: + operation_id = payload.get("operation_id", "setup") + try: + result = await run_process(self._ssh_command(), shell=False, + timeout=self.connection.connect_timeout + 15, + input_bytes=json.dumps(payload).encode(), output_limit=_MESSAGE_LIMIT) + except asyncio.TimeoutError as exc: + raise VMTransportError("VM SSH request timed out", operation_id=operation_id) from exc + if result.returncode != 0 or result.stdout_dropped_bytes or result.stderr_dropped_bytes: + detail = result.stderr.decode("utf-8", errors="replace").strip() + raise VMTransportError(f"VM SSH request failed with status {result.returncode}: {detail}", + operation_id=operation_id, returncode=result.returncode) + try: + response = json.loads(result.stdout) + except (ValueError, UnicodeError) as exc: + raise VMTransportError("VM helper response was incomplete", operation_id=operation_id) from exc + if not response["ok"]: + error = {"ValueError": ValueError, "PermissionError": PermissionError, + "FileNotFoundError": FileNotFoundError, + "FileExistsError": FileExistsError, "UnicodeDecodeError": UnicodeError + }.get(response["error"], OSError) + raise error(response["message"]) + return response["value"] + + async def setup(self, mount_dir: str | None = None) -> str: + self._ensure_active() + if mount_dir is not None: + raise ValueError("mount_dir is supported only by container environments") + async with self._setup_lock: + if self._identity is None: + response = await self._rpc({"method": "setup", "workspace": self.workspace}) + if response["version"] != __version__: + raise RuntimeError(f"VM OpenCollab version must match controller {__version__}") + self._identity = response["identity"] + self.workspace = self._identity["workspace"] + self.source_workspace = self.workspace + return self.workspace + + async def _cancel_operation(self, operation_id: str) -> None: + status = await self._rpc({"method": "cancel", "operation_id": operation_id}) + deadline = asyncio.get_running_loop().time() + 12 + while status["status"] in {"running", "preparing", "cleanup_failed"}: + if asyncio.get_running_loop().time() >= deadline: + raise ProcessCleanupError(f"VM operation cleanup did not quiesce, operation {operation_id}") + await asyncio.sleep(0.05) + status = await self._rpc({"method": "query", "operation_id": operation_id}) + if status["status"] in {"unknown", "cleanup_failed"}: + self._uncertain.add(operation_id) + raise ProcessCleanupError(f"VM operation requires recovery, operation {operation_id}") + self._operations.discard(operation_id) + self._uncertain.discard(operation_id) + + async def _wait(self, operation_id: str, status: dict[str, Any]) -> Any: + while status["status"] in {"running", "preparing"}: + await asyncio.sleep(0.05) + status.update(await self._rpc({"method": "query", "operation_id": operation_id})) + state = status["status"] + if state in {"unknown", "cleanup_failed"}: + self._uncertain.add(operation_id) + raise VMOperationError(status.get("message", "VM operation state requires recovery"), + operation_id=operation_id, state=state) + self._operations.discard(operation_id) + self._uncertain.discard(operation_id) + if state == "cancelled": + raise VMOperationError("VM operation was cancelled", operation_id=operation_id, state=state) + if state == "failed": + error = {"ValueError": ValueError, "PermissionError": PermissionError, + "FileNotFoundError": FileNotFoundError, "FileExistsError": FileExistsError, + "UnicodeDecodeError": UnicodeError}.get(status["error"], OSError) + raise error(status["message"]) + return status["value"] + + async def _operate(self, action: str, args: dict[str, Any], operation_id: str | None = None, + *, require_active: bool = True) -> Any: + if require_active: + self._ensure_active() + await self.setup() + operation_id = operation_id or uuid.uuid4().hex + if not isinstance(operation_id, str) or not _COMPONENT.fullmatch(operation_id): + raise ValueError("VM operation_id must be a simple path component") + self._operations.add(operation_id) + payload = {"method": "start", "operation_id": operation_id, + "request": {"action": action, "args": args, "workspace": self.workspace, "identity": self._identity}} + try: + # Retrying submission retains the operation primary key and exact + # request. The guest owns execution independently of the SSH client. + for attempt in range(3): + try: + status = await self._rpc(payload) + return await self._wait(operation_id, status) + except VMTransportError: + if attempt == 2: + raise + await asyncio.sleep(0.1) + except VMTransportError: + self._uncertain.add(operation_id) + raise + except asyncio.CancelledError: + try: + await await_owned_operation(self._cancel_operation(operation_id)) + except BaseException as exc: + self._uncertain.add(operation_id) + raise ProcessCleanupError(f"VM cancellation requires recovery, operation {operation_id}") from exc + raise + + async def exec_operation(self, operation_id: str, cmd: str, timeout: float | None = 120.0, + *, input_bytes: bytes | None = None) -> ExecResult: + if timeout is not None and (isinstance(timeout, bool) or not isinstance(timeout, (int, float)) + or math.isnan(timeout) or timeout <= 0): + raise ValueError("timeout must be positive or None") + if not isinstance(cmd, str) or "\0" in cmd: + raise ValueError("VM command must be text without NUL bytes") + if input_bytes is not None and not isinstance(input_bytes, bytes): + raise TypeError("input_bytes must be bytes") + encoded_input = base64.b64encode(input_bytes).decode() if input_bytes is not None else None + value = await self._operate("exec", {"cmd": cmd, "timeout": timeout if timeout != math.inf else None, + "input": encoded_input}, operation_id) + return self._exec_result(value) + + @staticmethod + def _exec_result(value: dict[str, Any]) -> ExecResult: + value = dict(value) + if value["stdout_bytes"] is not None: + value["stdout_bytes"] = base64.b64decode(value["stdout_bytes"]) + return ExecResult(**value) + + async def exec_cmd(self, cmd: str, timeout: float | None = 120.0, + *, input_bytes: bytes | None = None) -> ExecResult: + return await self.exec_operation(uuid.uuid4().hex, cmd, timeout, input_bytes=input_bytes) + + async def recover_operation(self, operation_id: str) -> Any: + self._operations.add(operation_id) + try: + status = await self._rpc({"method": "query", "operation_id": operation_id}) + value = await self._wait(operation_id, status) + except VMTransportError: + self._uncertain.add(operation_id) + raise + action = status.get("action", "exec") + if action == "temporary": + self._temporary_files[status["args"]["path"]] = value + elif action == "remove": + self._temporary_files.pop(status["args"]["path"], None) + return self._exec_result(value) if action == "exec" else value + + async def read_file(self, path: str) -> str: + return await self._operate("read", {"path": path}) + + async def read_text_range(self, path: str, *, offset: int, limit: int, max_chars: int) -> TextFileRange: + return TextFileRange(**await self._operate("range", {"path": path, "offset": offset, "limit": limit, + "max_chars": max_chars})) + + async def write_file(self, path: str, content: str) -> None: + await self.write_file_with_change(path, content) + + async def write_file_with_change(self, path: str, content: str) -> bool | None: + if len(content.encode()) > ENV_FILE_WRITE_LIMIT_BYTES: + raise ValueError("VM file exceeds write limit") + return await self._operate("write", {"path": path, "content": content}) + + async def write_temp_file(self, content: str, *, prefix: str, suffix: str = ".tmp") -> str: + if not re.fullmatch(r"[A-Za-z0-9._-]+", prefix) or not re.fullmatch(r"[A-Za-z0-9._-]+", suffix): + raise ValueError("temporary file prefix and suffix must be simple path components") + path = f"{prefix}{uuid.uuid4().hex}{suffix}" + identity = await self._operate("temporary", {"path": path, "content": content}) + self._temporary_files[path] = identity + return posixpath.join(self.workspace, path) + + async def remove_file(self, path: str) -> None: + target = posixpath.relpath(path, self.workspace) if path.startswith("/") else posixpath.normpath(path) + if target not in self._temporary_files: + raise OSError(f"refusing to remove unowned VM temporary file {path}") + await self._operate("remove", {"path": target, "identity": self._temporary_files[target]}) + self._temporary_files.pop(target) + + @contextlib.asynccontextmanager + async def file_write_lock(self, path: str) -> AsyncIterator[None]: + self._ensure_active() + await self.setup() + operation_id = uuid.uuid4().hex + self._operations.add(operation_id) + try: + status = await self._rpc({"method": "start", "operation_id": operation_id, + "request": {"action": "lock", "args": {}, "workspace": self.workspace, "identity": self._identity}}) + while status["status"] in {"preparing", "running"} and not status.get("lock_acquired"): + await asyncio.sleep(0.05) + status = await self._rpc({"method": "query", "operation_id": operation_id}) + if not status.get("lock_acquired"): + await self._wait(operation_id, status) + raise RuntimeError("VM edit lock exited before acquisition") + yield None + except VMTransportError: + self._uncertain.add(operation_id) + raise + finally: + if not self._uncertain: + await await_owned_operation(self._cancel_operation(operation_id)) + + async def cleanup(self) -> None: + self.revoke() + for operation_id in tuple(self._operations): + try: + await await_owned_operation(self._cancel_operation(operation_id), propagate_cancellation=True) + except BaseException: + self._uncertain.add(operation_id) + raise + for path in tuple(self._temporary_files): + await self._operate("remove", {"path": path, "identity": self._temporary_files[path]}, + require_active=False) + self._temporary_files.pop(path) + + +__all__ = ["VMEnvironment", "VMOperationError", "VMTransportError"] diff --git a/opencollab/adapters/_vm_remote.py b/opencollab/adapters/_vm_remote.py new file mode 100644 index 00000000..e8d3aa43 --- /dev/null +++ b/opencollab/adapters/_vm_remote.py @@ -0,0 +1,334 @@ +"""Guest entry point for durable SSH operations using the installed package.""" + +from __future__ import annotations + +import asyncio +import base64 +import ctypes +import fcntl +import json +import math +import os +import re +import signal +import subprocess +import sys +import time +from dataclasses import asdict +from pathlib import Path +from typing import Any + +from opencollab import __version__ +from opencollab.adapters._env_base import ENV_FILE_WRITE_LIMIT_BYTES +from opencollab.adapters._env_local import LocalEnvironment +from opencollab.adapters._env_process import ProcessCleanupError +from opencollab.adapters.safe_anchored_files import ( + create_regular_bytes_atomic_at, + open_directory_anchor, + read_regular_bytes_at, + unlink_regular_file_durable_at, + write_regular_bytes_atomic_at, +) +from opencollab.domain.vm import validate_vm_path + +_MAX_MESSAGE = 16 * 1024 * 1024 +_ID = re.compile(r"[A-Za-z0-9][A-Za-z0-9_.-]{0,127}") + + +def _write(directory: Path, name: str, value: Any) -> None: + fd = open_directory_anchor(directory) + try: + write_regular_bytes_atomic_at(fd, directory, name, json.dumps(value).encode(), max_bytes=_MAX_MESSAGE) + finally: + os.close(fd) + + +def _read(directory: Path, name: str) -> Any: + fd = open_directory_anchor(directory) + try: + return json.loads(read_regular_bytes_at(fd, directory, name, max_bytes=_MAX_MESSAGE)) + finally: + os.close(fd) + + +def _directory(value: str) -> Path: + validate_vm_path(value) + current = Path("/") + for component in Path(value).parts[1:]: + current /= component + try: + current.mkdir(mode=0o700) + except FileExistsError: + pass + fd = open_directory_anchor(current) + os.close(fd) + if current.stat().st_uid != os.getuid() or current.stat().st_mode & 0o077: + raise PermissionError("VM state directory must be private and owned by the SSH account") + return current + + +def _identity(workspace: str) -> dict[str, Any]: + fd = open_directory_anchor(workspace) + try: + info = os.fstat(fd) + return {"workspace": os.path.realpath(workspace), "device": info.st_dev, "inode": info.st_ino, + "boot_id": Path("/proc/sys/kernel/random/boot_id").read_text().strip()} + finally: + os.close(fd) + + +def _process(pid: int) -> tuple[int, int, bytes]: + value = Path(f"/proc/{pid}/stat").read_bytes() + fields = value[value.rindex(b")") + 1:].split() + return int(fields[1]), int(fields[19]), fields[0] + + +def _status(directory: Path) -> dict[str, Any]: + try: + return _read(directory, "result.json") + except FileNotFoundError: + pass + try: + started = _read(directory, "started.json") + except FileNotFoundError: + return {"status": "preparing"} + try: + live = (_process(started["pid"])[1] == started["started"] + and Path("/proc/sys/kernel/random/boot_id").read_text().strip() == started["boot_id"]) + except (FileNotFoundError, ProcessLookupError): + live = False + return {"status": "running" if live else "unknown", "owner": started, + "lock_acquired": (directory / "lock.json").exists()} + + +def _descendants() -> dict[int, int]: + processes = {} + for entry in os.listdir("/proc"): + if entry.isdecimal(): + try: + processes[int(entry)] = _process(int(entry)) + except (FileNotFoundError, ProcessLookupError): + continue + owned = {os.getpid()} + while True: + children = {pid for pid, (parent, _started, _state) in processes.items() if parent in owned} + if children <= owned: + break + owned |= children + return {pid: processes[pid][1] for pid in owned if pid != os.getpid()} + + +async def _stop_descendants() -> None: + deadline = time.monotonic() + 4.0 + while True: + children = _descendants() + for pid, started in children.items(): + try: + fd = os.pidfd_open(pid) + try: + if _process(pid)[1] == started: + signal.pidfd_send_signal(fd, signal.SIGKILL) + finally: + os.close(fd) + except (FileNotFoundError, ProcessLookupError): + pass + while True: + try: + pid, _status_code = os.waitpid(-1, os.WNOHANG) + except ChildProcessError: + break + if pid == 0: + break + if not _descendants(): + return + if time.monotonic() >= deadline: + raise ProcessCleanupError("VM operation descendants did not quiesce") + await asyncio.sleep(0.02) + + +async def _perform(request: dict[str, Any], directory: Path) -> Any: + if _identity(request["workspace"]) != request["identity"]: + raise OSError("VM workspace or guest boot identity changed") + env = LocalEnvironment(request["workspace"]) + action = request["action"] + args = request.get("args", {}) + try: + if action == "lock": + lock_root = _directory(f"/tmp/opencollab-vm-locks-{os.getuid()}") + identity = request["identity"] + with (lock_root / f"{identity['device']}-{identity['inode']}.lock").open("a+b") as lock: + deadline = time.monotonic() + 10 + while True: + try: + fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) + break + except BlockingIOError: + if time.monotonic() >= deadline: + raise TimeoutError("timed out acquiring VM workspace edit lock") + await asyncio.sleep(0.02) + _write(directory, "lock.json", True) + await asyncio.Future() + if action == "exec": + result = asdict(await env.exec_cmd(args["cmd"], math.inf if args["timeout"] is None else args["timeout"], + input_bytes=base64.b64decode(args["input"]) if args["input"] is not None else None)) + result["stdout_bytes"] = (base64.b64encode(result["stdout_bytes"]).decode() + if result["stdout_bytes"] is not None else None) + return result + if action == "read": + return await env.read_file(args["path"]) + if action == "range": + return asdict(await env.read_text_range(**args)) + if action == "write": + return await env.write_file_with_change(args["path"], args["content"]) + if action in {"temporary", "remove"}: + fd = open_directory_anchor(env.workspace) + try: + if action == "temporary": + create_regular_bytes_atomic_at(fd, env.workspace, args["path"], args["content"].encode(), + max_bytes=ENV_FILE_WRITE_LIMIT_BYTES) + info = os.stat(args["path"], dir_fd=fd, follow_symlinks=False) + return [info.st_dev, info.st_ino] + unlink_regular_file_durable_at(fd, env.workspace, args["path"], + expected_identity=tuple(args["identity"])) + return None + finally: + os.close(fd) + raise ValueError(f"unknown VM operation {action}") + finally: + await env.cleanup() + + +async def _work(directory: Path) -> None: + request = _read(directory, "request.json") + # Each worker is a Linux subreaper so double-forked children remain owned + # after setsid or process-group changes. Its lock spans command and cleanup. + if ctypes.CDLL(None, use_errno=True).prctl(36, 1, 0, 0, 0) != 0: + raise OSError(ctypes.get_errno(), "cannot establish VM child subreaper") + if not hasattr(os, "pidfd_open") or not hasattr(signal, "pidfd_send_signal"): + raise OSError("VM execution requires Linux pidfd support") + started = {"pid": os.getpid(), "started": _process(os.getpid())[1], + "boot_id": Path("/proc/sys/kernel/random/boot_id").read_text().strip()} + _write(directory, "started.json", started) + task = asyncio.create_task(_perform(request, directory)) + cancelled = False + try: + while not task.done(): + if (directory / "cancel.json").exists(): + task.cancel() + cancelled = True + break + await asyncio.sleep(0.05) + value = await task + result = {"status": "completed", "value": value} + except asyncio.CancelledError: + result = {"status": "cancelled"} + except BaseException as exc: + result = {"status": "failed", "error": type(exc).__name__, "message": str(exc)} + try: + await _stop_descendants() + except BaseException as cleanup_error: + failed_cancel = (directory / "cancel.json").stat().st_mtime_ns if (directory / "cancel.json").exists() else None + while True: + _write(directory, "result.json", {"status": "cleanup_failed", "error": type(cleanup_error).__name__, + "message": str(cleanup_error), "descendants": _descendants()}) + while True: + await asyncio.sleep(0.1) + cancel_path = directory / "cancel.json" + current = cancel_path.stat().st_mtime_ns if cancel_path.exists() else None + if current != failed_cancel: + failed_cancel = current + break + try: + await _stop_descendants() + except BaseException: + continue + result = {"status": "cancelled"} + break + if cancelled and result["status"] == "completed": + result = {"status": "cancelled"} + result["action"] = request["action"] + result["args"] = {"path": request["args"]["path"]} if request["action"] in {"temporary", "remove"} else {} + _write(directory, "result.json", result) + + +def _worker(directory: Path) -> None: + with (directory / "worker.lock").open("a+b") as lock: + try: + fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError: + return + # A missing terminal result after a claimed execution is ambiguous. + # Keep that evidence for recovery instead of repeating side effects. + if (directory / "started.json").exists() or (directory / "result.json").exists(): + return + try: + asyncio.run(_work(directory)) + except BaseException as exc: + _write(directory, "result.json", {"status": "failed", "error": type(exc).__name__, "message": str(exc)}) + + +def _rpc(root: Path, request: dict[str, Any]) -> Any: + action = request["method"] + if action == "setup": + if sys.platform != "linux": + raise OSError("VM environments require a Linux guest") + identity = _identity(request["workspace"]) + try: + previous = _read(root, "environment.json") + except FileNotFoundError: + _write(root, "environment.json", identity) + else: + if previous != identity: + raise OSError("VM workspace or guest boot identity changed") + return {"identity": identity, "version": __version__} + operation = request["operation_id"] + if not isinstance(operation, str) or not _ID.fullmatch(operation): + raise ValueError("invalid VM operation ID") + directory = root / operation + if action == "start": + try: + directory.mkdir(mode=0o700) + except FileExistsError: + pass + fd = open_directory_anchor(directory) + os.close(fd) + with (directory / "submit.lock").open("a+b") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + try: + previous = _read(directory, "request.json") + except FileNotFoundError: + _write(directory, "request.json", request["request"]) + else: + if previous != request["request"]: + raise ValueError("VM operation ID is already bound to another request") + if _status(directory)["status"] == "preparing": + subprocess.Popen([sys.executable, "-m", "opencollab.adapters._vm_remote", "worker", str(directory)], + stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, + start_new_session=True, close_fds=True) + if action == "cancel": + _write(directory, "cancel.json", True) + if action not in {"start", "query", "cancel"}: + raise ValueError("unknown VM RPC method") + return _status(directory) + + +def _main() -> None: + if sys.argv[1] == "worker": + _worker(Path(sys.argv[2])) + return + try: + if len(sys.argv) != 4 or sys.argv[1] != "rpc" or not _ID.fullmatch(sys.argv[3]): + raise ValueError("invalid VM helper arguments") + root = _directory(f"{validate_vm_path(sys.argv[2])}/{sys.argv[3]}") + payload = sys.stdin.buffer.read(_MAX_MESSAGE + 1) + if len(payload) > _MAX_MESSAGE: + raise ValueError("VM request exceeds message limit") + value = _rpc(root, json.loads(payload)) + response = {"ok": True, "value": value} + except Exception as exc: + response = {"ok": False, "error": type(exc).__name__, "message": str(exc)} + sys.stdout.buffer.write(json.dumps(response).encode()) + + +if __name__ == "__main__": + _main() diff --git a/opencollab/adapters/env.py b/opencollab/adapters/env.py index b81ed7a6..615bde6d 100644 --- a/opencollab/adapters/env.py +++ b/opencollab/adapters/env.py @@ -9,6 +9,7 @@ from opencollab.adapters._env_docker import DockerEnvironment from opencollab.adapters._env_local import LocalEnvironment from opencollab.adapters._env_process import PROCESS_OUTPUT_CAPTURE_BYTES +from opencollab.adapters._env_vm import VMEnvironment, VMOperationError, VMTransportError from opencollab.adapters._env_worktree import WorktreeEnvironment @@ -40,4 +41,7 @@ def __init__( "LocalEnvironment", "PROCESS_OUTPUT_CAPTURE_BYTES", "WorktreeEnvironment", + "VMEnvironment", + "VMOperationError", + "VMTransportError", ] diff --git a/opencollab/adapters/safe_anchored_files.py b/opencollab/adapters/safe_anchored_files.py index 68736015..2a1f7e33 100644 --- a/opencollab/adapters/safe_anchored_files.py +++ b/opencollab/adapters/safe_anchored_files.py @@ -429,6 +429,8 @@ def unlink_regular_file_durable_at( root_fd: int, root_path: str | os.PathLike[str], path: str | os.PathLike[str], + *, + expected_identity: tuple[int, int] | None = None, ) -> bool: parent_fd, name = _open_parent_beneath( root_fd, @@ -440,6 +442,8 @@ def unlink_regular_file_durable_at( current = _current_regular_at(parent_fd, name, context="owned file") if current is None: return False + if expected_identity is not None and (current.st_dev, current.st_ino) != expected_identity: + raise OSError(f"owned file was replaced: {path}") os.unlink(name, dir_fd=parent_fd) os.fsync(parent_fd) return True diff --git a/opencollab/domain/vm.py b/opencollab/domain/vm.py new file mode 100644 index 00000000..b41e6020 --- /dev/null +++ b/opencollab/domain/vm.py @@ -0,0 +1,43 @@ +from __future__ import annotations + +import math +import posixpath +import re +from dataclasses import dataclass + + +@dataclass(frozen=True) +class VMConnection: + host: str + user: str + port: int = 22 + identity_file: str | None = None + known_hosts_file: str | None = None + python: str = "python3" + connect_timeout: float = 10.0 + + def __post_init__(self) -> None: + for name, value in (("host", self.host), ("user", self.user)): + if not isinstance(value, str) or not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.:@%-]*", value): + raise ValueError(f"VM {name} must be an SSH hostname or account name") + if isinstance(self.port, bool) or not isinstance(self.port, int) or not 1 <= self.port <= 65535: + raise ValueError("VM SSH port must be between 1 and 65535") + for name, value in (("identity_file", self.identity_file), ("known_hosts_file", self.known_hosts_file)): + if value is not None and (not isinstance(value, str) or not value or any(c in value for c in "\0\n\r")): + raise ValueError(f"VM {name} must be a path without control characters") + if not isinstance(self.python, str) or not self.python or any(c in self.python for c in "\0\n\r"): + raise ValueError("VM python must name the guest Python executable") + if ( + isinstance(self.connect_timeout, bool) or not isinstance(self.connect_timeout, (int, float)) + or not math.isfinite(self.connect_timeout) or self.connect_timeout <= 0 + ): + raise ValueError("VM connect_timeout must be finite and positive") + + +def validate_vm_path(value: str) -> str: + if ( + not isinstance(value, str) or not value.startswith("/") or value == "/" + or posixpath.normpath(value) != value or any(c in value for c in "\0\n\r") + ): + raise ValueError("VM paths must be normalized absolute non-root paths") + return value diff --git a/tests/runtime/test_vm_environment.py b/tests/runtime/test_vm_environment.py new file mode 100644 index 00000000..5d3ee3f0 --- /dev/null +++ b/tests/runtime/test_vm_environment.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +import asyncio +import base64 +import math + +import pytest + +from opencollab.adapters.env import ExecResult, VMEnvironment, VMOperationError, VMTransportError +from opencollab.domain.vm import VMConnection + + +def _env(): + return VMEnvironment(VMConnection("guest", "worker"), "/workspace/repo", + state_directory="/var/tmp/oc-state", run_id="test") + + +@pytest.mark.parametrize("kwargs", [{"host": "-oProxyCommand=evil"}, {"user": "a b"}, {"port": True}, + {"port": 0}, {"connect_timeout": math.inf}, {"python": "a\nb"}]) +def test_connection_rejects_invalid_values(kwargs): + arguments = {"host": "guest", "user": "worker", **kwargs} + with pytest.raises(ValueError): + VMConnection(**arguments) + + +def test_ssh_keeps_host_verification_and_executable_quoting(): + env = _env() + arguments = env._ssh_command() + assert "StrictHostKeyChecking=yes" in arguments + assert "BatchMode=yes" in arguments + assert "ControlPath=none" in arguments + assert arguments[-1] == "python3 -m opencollab.adapters._vm_remote rpc /var/tmp/oc-state test" + + +async def test_reconnect_reuses_operation_and_preserves_binary_exit_255(monkeypatch): + env = _env() + env._identity = {"workspace": env.workspace} + requests = [] + value = {"returncode": 255, "stdout": "�", "stderr": "command error", "stdout_bytes": "/wA="} + + async def rpc(payload): + requests.append(payload) + if len(requests) == 1: + raise VMTransportError("SSH disconnected", operation_id="once", returncode=255) + return {"status": "completed", "value": value} + + monkeypatch.setattr(env, "_rpc", rpc) + result = await env.exec_operation("once", "cat; exit 255", timeout=None, input_bytes=b"\xff\0") + assert result.returncode == 255 + assert result.stdout_bytes == b"\xff\0" + assert requests[0] == requests[1] + assert requests[0]["request"]["args"]["input"] == base64.b64encode(b"\xff\0").decode() + assert not env.pending_operations + + +async def test_uncertain_transport_blocks_further_work_until_recovery(monkeypatch): + env = _env() + env._identity = {"workspace": env.workspace} + + async def fail(payload): + raise VMTransportError("SSH unavailable", operation_id=payload["operation_id"], returncode=255) + + monkeypatch.setattr(env, "_rpc", fail) + with pytest.raises(VMTransportError): + await env.exec_operation("once", "touch side-effect") + assert env.revoked + assert env.pending_operations == ("once",) + with pytest.raises(VMOperationError, match="Recover uncertain"): + await env.exec_cmd("touch another") + + async def recovered(payload): + return {"status": "completed", "action": "exec", "value": { + "returncode": 0, "stdout": "recovered", "stderr": "", "stdout_bytes": None}} + + monkeypatch.setattr(env, "_rpc", recovered) + assert (await env.recover_operation("once")).stdout == "recovered" + assert not env.revoked + + +async def test_cancel_waits_for_remote_cleanup(monkeypatch): + env = _env() + env._identity = {"workspace": env.workspace} + started = asyncio.Event() + cancelled = asyncio.Event() + + async def rpc(payload): + if payload["method"] == "start": + started.set() + return {"status": "running"} + if payload["method"] == "cancel": + cancelled.set() + return {"status": "cancelled"} + await asyncio.Future() + + monkeypatch.setattr(env, "_rpc", rpc) + task = asyncio.create_task(env.exec_operation("once", "sleep 30")) + await started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert cancelled.is_set() + assert not env.pending_operations + + +async def test_failed_cleanup_keeps_operation_and_revocation(monkeypatch): + env = _env() + env._operations.add("owned") + + async def fail(payload): + raise VMTransportError("offline", operation_id="owned") + + monkeypatch.setattr(env, "_rpc", fail) + with pytest.raises(VMTransportError): + await env.cleanup() + assert env.pending_operations == ("owned",) + assert env.revoked + + +async def test_recover_file_operation_retains_temporary_identity(monkeypatch): + env = _env() + + async def rpc(payload): + return {"status": "completed", "action": "temporary", "args": {"path": "owned.tmp"}, "value": [1, 2]} + + monkeypatch.setattr(env, "_rpc", rpc) + assert await env.recover_operation("temporary") == [1, 2] + assert env._temporary_files == {"owned.tmp": [1, 2]} + + +async def test_remote_result_uses_existing_execution_result(): + result = VMEnvironment._exec_result({"returncode": 0, "stdout": "ok", "stderr": "", "stdout_bytes": "b2s="}) + assert isinstance(result, ExecResult) + assert result.stdout_bytes == b"ok" From d68532796d9121e38c643f20ba764ab234a84566 Mon Sep 17 00:00:00 2001 From: Kai Date: Sat, 24 Oct 2026 16:46:25 +0800 Subject: [PATCH 2/8] feat(environments): stream VM workspace transfers over SSH --- docs/virtual-machines.md | 51 ++++++ docs/virtual-machines.zh-CN.md | 51 ++++++ opencollab/adapters/_env_process.py | 7 +- opencollab/adapters/_env_vm.py | 172 ++++++++++++++++-- opencollab/adapters/_vm_archive.py | 172 ++++++++++++++++++ opencollab/adapters/_vm_remote.py | 154 +++++++++++++--- opencollab/adapters/safe_anchored_files.py | 11 +- opencollab/domain/vm.py | 3 +- tests/runtime/test_vm_archive.py | 90 ++++++++++ tests/runtime/test_vm_environment.py | 29 ++++ tests/runtime/test_vm_ssh.py | 193 +++++++++++++++++++++ 11 files changed, 894 insertions(+), 39 deletions(-) create mode 100644 docs/virtual-machines.md create mode 100644 docs/virtual-machines.zh-CN.md create mode 100644 opencollab/adapters/_vm_archive.py create mode 100644 tests/runtime/test_vm_archive.py create mode 100644 tests/runtime/test_vm_ssh.py diff --git a/docs/virtual-machines.md b/docs/virtual-machines.md new file mode 100644 index 00000000..94a2b5e3 --- /dev/null +++ b/docs/virtual-machines.md @@ -0,0 +1,51 @@ +# Existing Linux virtual machines + +`VMEnvironment` attaches to an existing Linux VM over system OpenSSH. The SSH account needs Python 3.10 or later, the same OpenCollab distribution as the controller, Linux child-subreaper and pidfd support, and a complete `/proc` view for its operation descendants. Install the controller's wheel on the guest before connecting. The helper runs from the installed `opencollab.adapters._vm_remote` module and reuses `LocalEnvironment` for commands and file access. + +Prepare an existing workspace and a private state directory owned by the SSH account. Keep operation state outside the workspace. Each environment uses a run ID beneath that directory. Preserve the run ID when reconstructing an environment after a controller restart. + +```python +from opencollab.adapters.env import VMEnvironment +from opencollab.domain.vm import VMConnection + +connection = VMConnection( + host="vm.example.net", + user="worker", + identity_file="/home/controller/.ssh/vm-key", + known_hosts_file="/home/controller/.ssh/known_hosts", +) +env = VMEnvironment( + connection, + "/srv/work/repository", + state_directory="/var/tmp/opencollab-operations", + run_id="experiment-42", +) +await env.setup() +result = await env.exec_operation( + "build-1", + "python3 -c 'import sys; sys.stdout.buffer.write(sys.stdin.buffer.read())'", + input_bytes=b"binary\x00input", +) +assert result.stdout_bytes == b"binary\x00input" +await env.write_file("notes.txt", "A remote edit\n") +page = await env.read_text_range("notes.txt", offset=1, limit=20, max_chars=4000) +await env.cleanup() +``` + +Host-key verification is always strict. OpenSSH uses the supplied known-hosts file or its ordinary known-hosts locations. `identity_file` refers to an existing key. `python` can select a guest virtual environment's Python executable. Transport options remain separate from VM resource identity and machine allocation. + +Commands retain the existing bounded stdout and stderr behavior, truncation counters, timeout notices, and binary `stdout_bytes`. `input_bytes` supplies binary stdin. A command's exit code 255 remains a command result. SSH failures raise `VMTransportError` carrying `transport_returncode` and `operation_id`. A timeout of `None` or `math.inf` gives the remote command an unlimited duration. + +Every operation records its full request, operation ID, supervisor process identity, and terminal result in the private run directory. Submission retries compare the original request and use the same operation ID. A detached supervisor holds the operation lock throughout execution and cleanup, so a disconnected SSH client can reconnect to the same execution. `exec_cmd` allocates an operation ID automatically. `exec_operation` accepts an explicit ID for durable callers. + +Use `recover_operation(error.operation_id)` after restoring SSH connectivity. For command operations it returns the original `ExecResult`. File operations return their original value, and temporary-file recovery restores the environment's ownership record. Pending or uncertain operation IDs remain available through `pending_operations`. An uncertain operation prevents further work until recovery settles its state. A missing supervisor after a claimed execution leaves an unknown operation with its evidence retained. Recovery of that state requires inspecting the preserved guest and operation records before the VM is reused. + +Cancellation stops the command's process group, then the supervisor's descendants. The Linux subreaper retains ownership when children call `setsid` or double-fork, and pidfds bind termination signals to the observed process identities. A cleanup failure preserves the supervisor and descendant records for a later cleanup retry. Cleanup also removes owned temporary files after checking their recorded device and inode. The VM remains running after environment cleanup. + +File operations use the existing anchored regular-file implementation, workspace identity checks, read and write limits, and paginated text reads. `file_write_lock(path)` holds a guest filesystem lock across a complete read-modify-write edit. The lock serializes edits throughout the same physical workspace across controller processes and run IDs. An uncertain file operation retains its edit lock until recovery and cleanup can settle it. Git working-tree probes and diff tools execute through the environment's ordinary command interface. + +The opt-in SSH regression suite reads `OPENCOLLAB_TEST_VM_CONFIG` from a JSON file containing `connection`, `workspace`, and `state_directory`. `connection` contains `VMConnection` keyword arguments. The tests create task directories below the supplied workspace, exercise real commands and files, interrupt an SSH client, stop a daemonized descendant, compare concurrent edits, and leave the borrowed guest running. Run the suite with `pytest -q tests/runtime/test_vm_ssh.py` after configuring a dedicated test workspace. + +Adapter integrations can call `copy_to_host(source, target, directory=False)` and `copy_from_host(source, destination, directory=False)` to stream binary files or directory archives across SSH beyond command-output limits. Guest source and destination paths are explicit absolute paths selected by trusted workspace adapters, including repository metadata outside an exposed subdirectory. Directory transfers copy contents into an empty or absent destination and preserve the inode of an existing workspace. Regular permissions, symlinks, empty files, and filenames supported by the destination filesystem survive the roundtrip. Setuid and setgid bits are removed when crossing the transfer boundary. + +Archive extraction uses an owned staging directory beside the destination and descriptor-relative publication. Parent traversal, repeated paths, hard links, device entries, and writes through symlink parents fail with the staging retained. Upload chunks carry an explicit completion frame before a durable import can execute. An interrupted upload keeps its operation ID and partial archive for inspection or cleanup. Operation records also retain failed import staging paths. Copying a directory into an existing populated destination raises `FileExistsError`. diff --git a/docs/virtual-machines.zh-CN.md b/docs/virtual-machines.zh-CN.md new file mode 100644 index 00000000..b48e0575 --- /dev/null +++ b/docs/virtual-machines.zh-CN.md @@ -0,0 +1,51 @@ +# 现有 Linux 虚拟机 + +`VMEnvironment` 通过系统 OpenSSH 接入现有 Linux VM。SSH 账号需要 Python 3.10 或更新版本、与控制端相同的 OpenCollab 安装包、Linux child-subreaper 与 pidfd 能力，以及能够完整查看操作后代进程的 `/proc`。连接前先在客户机安装控制端使用的 wheel。辅助程序从已安装的 `opencollab.adapters._vm_remote` 模块启动，命令和文件访问复用 `LocalEnvironment`。 + +准备已有工作区，以及由 SSH 账号拥有的私有状态目录。操作状态放在工作区之外。每个环境在状态目录内使用一个运行 ID。控制端重启后重建环境时继续使用原运行 ID。 + +```python +from opencollab.adapters.env import VMEnvironment +from opencollab.domain.vm import VMConnection + +connection = VMConnection( + host="vm.example.net", + user="worker", + identity_file="/home/controller/.ssh/vm-key", + known_hosts_file="/home/controller/.ssh/known_hosts", +) +env = VMEnvironment( + connection, + "/srv/work/repository", + state_directory="/var/tmp/opencollab-operations", + run_id="experiment-42", +) +await env.setup() +result = await env.exec_operation( + "build-1", + "python3 -c 'import sys; sys.stdout.buffer.write(sys.stdin.buffer.read())'", + input_bytes=b"binary\x00input", +) +assert result.stdout_bytes == b"binary\x00input" +await env.write_file("notes.txt", "A remote edit\n") +page = await env.read_text_range("notes.txt", offset=1, limit=20, max_chars=4000) +await env.cleanup() +``` + +主机密钥始终采用严格验证。OpenSSH 使用指定的 known-hosts 文件，省略时使用其常规位置。`identity_file` 引用已有密钥。`python` 可以选择客户机虚拟环境中的 Python 可执行文件。连接参数与 VM 资源标识、机器分配分别配置。 + +命令沿用既有的标准输出与错误输出上限、截断计数、超时提示和二进制 `stdout_bytes`。`input_bytes` 提供二进制标准输入。命令退出码 255 保留为命令结果。SSH 故障抛出 `VMTransportError`，并携带 `transport_returncode` 与 `operation_id`。超时参数为 `None` 或 `math.inf` 时，远端命令可以持续运行。 + +每项操作都在私有运行目录保存完整请求、操作编号、监督进程身份与最终结果。重试提交时比较原始请求，并继续使用同一操作编号。独立运行的监督进程在命令执行与回收期间持有操作锁，因此 SSH 客户端断线后仍能重新连接原执行。`exec_cmd` 自动生成操作编号，`exec_operation` 接收调用者指定的编号。 + +恢复 SSH 后调用 `recover_operation(error.operation_id)`。命令操作返回原始 `ExecResult`，文件操作返回原始结果，临时文件恢复还会重建环境持有的文件身份记录。`pending_operations` 列出尚待处理的操作编号。状态不确定时，环境等待恢复结束再接受后续工作。已经开始执行的操作若失去监督进程，会保存为未知状态并保留现场。再次使用这台 VM 前，需要检查客户机和操作记录来处理该状态。 + +取消首先停止命令进程组，再结束监督进程拥有的后代。Linux subreaper 在子进程调用 `setsid` 或双重 fork 后继续持有其归属，pidfd 将终止信号绑定到已观察到的进程身份。回收失败时保留监督进程与后代记录，以便之后重试。清理还会核对临时文件记录的设备号与 inode，再删除环境拥有的临时文件。环境清理完成后，VM 继续运行。 + +文件操作复用既有的目录锚点、普通文件检查、工作区身份检查、读写上限和分页读取。`file_write_lock(path)` 在一次完整的读取、修改、写入期间持有客户机文件锁。同一物理工作区内，跨控制进程、跨运行 ID 的编辑会依次执行。文件操作状态不确定时，环境保留编辑锁，直到恢复和回收确定操作已结束。Git 工作区探测与 diff 工具使用环境的常规命令接口。 + +可选 SSH 回归测试从环境变量 `OPENCOLLAB_TEST_VM_CONFIG` 指定的 JSON 文件读取 `connection`、`workspace` 和 `state_directory`。`connection` 包含 `VMConnection` 的关键字参数。测试在指定工作区下创建任务目录，实际执行命令和文件操作、中断 SSH 客户端、停止守护化后代进程，并检查并发编辑结果。借用的 VM 在测试后继续运行。配置专用测试工作区后，运行 `pytest -q tests/runtime/test_vm_ssh.py`。 + +工作区适配器可以通过 `copy_to_host(source, target, directory=False)` 与 `copy_from_host(source, destination, directory=False)` 流式传输二进制文件和目录归档，传输大小独立于命令输出上限。客户机源路径和目标路径是工作区适配器明确指定的绝对路径，也可以指向暴露子目录之外的仓库元数据。目录复制将内容传入空目录或新目录，并保留已有工作区的 inode。常规权限、符号链接、空文件，以及目标文件系统支持的文件名会随传输保留。传输时移除 setuid 与 setgid 位。 + +归档先在目标旁边的自有临时目录解包，再通过目录描述符写入目标。父目录跳转、重复路径、硬链接、设备文件，以及经由符号链接父目录写入的情况会触发失败，并保留临时目录。上传分块带有明确结束帧，完整接收后才允许执行持久化导入。上传中断时保留操作编号和部分归档，供检查或清理继续处理。操作记录还保存失败导入的临时目录路径。向已有内容的目录复制时抛出 `FileExistsError`。 diff --git a/opencollab/adapters/_env_process.py b/opencollab/adapters/_env_process.py index 18141dce..c673baec 100644 --- a/opencollab/adapters/_env_process.py +++ b/opencollab/adapters/_env_process.py @@ -54,7 +54,7 @@ def to_exec_result(self) -> ExecResult: stdout_bytes=self.stdout, ) -async def _read_bounded( +async def read_bounded_output( stream: asyncio.StreamReader | None, limit: int, ) -> tuple[bytes, int]: @@ -479,8 +479,8 @@ async def spawn() -> asyncio.subprocess.Process: registry=registry, ) - stdout_task = asyncio.create_task(_read_bounded(process.stdout, output_limit)) - stderr_task = asyncio.create_task(_read_bounded(process.stderr, output_limit)) + stdout_task = asyncio.create_task(read_bounded_output(process.stdout, output_limit)) + stderr_task = asyncio.create_task(read_bounded_output(process.stderr, output_limit)) quiesced = False async def write_input_and_wait() -> None: @@ -556,5 +556,6 @@ async def write_input_and_wait() -> None: "ProcessRegistry", "ProcessResult", "run_process", + "read_bounded_output", "terminate_process", ] diff --git a/opencollab/adapters/_env_vm.py b/opencollab/adapters/_env_vm.py index e27ae856..59ca915c 100644 --- a/opencollab/adapters/_env_vm.py +++ b/opencollab/adapters/_env_vm.py @@ -5,20 +5,33 @@ import contextlib import json import math +import os import posixpath import re import shlex +import shutil +import struct +import tempfile import uuid from collections.abc import AsyncIterator +from pathlib import Path from typing import Any from opencollab import __version__ from opencollab.adapters._env_base import ENV_FILE_WRITE_LIMIT_BYTES, Environment, ExecResult, TextFileRange -from opencollab.adapters._env_process import ProcessCleanupError, run_process +from opencollab.adapters._env_process import ( + ProcessCleanupError, + ProcessRegistry, + read_bounded_output, + run_process, +) +from opencollab.adapters._vm_archive import pack, publish, unpack from opencollab.application.async_timeout import await_owned_operation +from opencollab.application.exception_notes import add_exception_note from opencollab.domain.vm import VMConnection, validate_vm_path -_MESSAGE_LIMIT = 16 * 1024 * 1024 +_MESSAGE_LIMIT = 32 * 1024 * 1024 +_FRAME = b"OCVM1\n" _COMPONENT = re.compile(r"[A-Za-z0-9][A-Za-z0-9_.-]{0,127}") @@ -60,6 +73,7 @@ def __init__(self, connection: VMConnection, workspace: str, *, state_directory: self._setup_lock = asyncio.Lock() self._operations: set[str] = set() self._uncertain: set[str] = set() + self._edit_locks: set[str] = set() self._temporary_files: dict[str, list[int]] = {} @property @@ -77,7 +91,7 @@ def _ensure_active(self) -> None: raise VMOperationError("Recover uncertain VM operations before starting further work", operation_id=operation, state="unknown") - def _ssh_command(self) -> list[str]: + def _ssh_command(self, method: str = "rpc") -> list[str]: connection = self.connection command = ["ssh", "-T", "-F", "/dev/null", "-o", "BatchMode=yes", "-o", "StrictHostKeyChecking=yes", "-o", "ControlMaster=no", "-o", "ControlPath=none", @@ -89,7 +103,7 @@ def _ssh_command(self) -> list[str]: if connection.known_hosts_file is not None: command += ["-o", f"UserKnownHostsFile={connection.known_hosts_file}"] command += ["-l", connection.user, connection.host, - shlex.join([connection.python, "-m", "opencollab.adapters._vm_remote", "rpc", + shlex.join([connection.python, "-m", "opencollab.adapters._vm_remote", method, self.state_directory, self.run_id])] return command @@ -105,17 +119,26 @@ async def _rpc(self, payload: dict[str, Any]) -> Any: detail = result.stderr.decode("utf-8", errors="replace").strip() raise VMTransportError(f"VM SSH request failed with status {result.returncode}: {detail}", operation_id=operation_id, returncode=result.returncode) + return self._decode_response(result.stdout, operation_id) + + @staticmethod + def _decode_response(payload: bytes, operation_id: str) -> Any: + start = payload.rfind(_FRAME) try: - response = json.loads(result.stdout) + if start < 0: + raise ValueError("missing VM helper frame") + response = json.loads(payload[start + len(_FRAME):]) + if not isinstance(response, dict) or not isinstance(response.get("ok"), bool): + raise ValueError("invalid VM helper response") except (ValueError, UnicodeError) as exc: - raise VMTransportError("VM helper response was incomplete", operation_id=operation_id) from exc + raise VMTransportError("VM helper response was incomplete or invalid", operation_id=operation_id) from exc if not response["ok"]: error = {"ValueError": ValueError, "PermissionError": PermissionError, "FileNotFoundError": FileNotFoundError, "FileExistsError": FileExistsError, "UnicodeDecodeError": UnicodeError }.get(response["error"], OSError) raise error(response["message"]) - return response["value"] + return response.get("value", response) async def setup(self, mount_dir: str | None = None) -> str: self._ensure_active() @@ -146,11 +169,11 @@ async def _cancel_operation(self, operation_id: str) -> None: self._uncertain.discard(operation_id) async def _wait(self, operation_id: str, status: dict[str, Any]) -> Any: - while status["status"] in {"running", "preparing"}: + while status["status"] in {"running", "preparing", "uploading"}: await asyncio.sleep(0.05) status.update(await self._rpc({"method": "query", "operation_id": operation_id})) state = status["status"] - if state in {"unknown", "cleanup_failed"}: + if state in {"unknown", "cleanup_failed", "upload_incomplete"}: self._uncertain.add(operation_id) raise VMOperationError(status.get("message", "VM operation state requires recovery"), operation_id=operation_id, state=state) @@ -224,10 +247,17 @@ async def exec_cmd(self, cmd: str, timeout: float | None = 120.0, return await self.exec_operation(uuid.uuid4().hex, cmd, timeout, input_bytes=input_bytes) async def recover_operation(self, operation_id: str) -> Any: + known = operation_id in self._operations self._operations.add(operation_id) try: status = await self._rpc({"method": "query", "operation_id": operation_id}) value = await self._wait(operation_id, status) + except FileNotFoundError: + if not known: + self._operations.discard(operation_id) + else: + self._uncertain.add(operation_id) + raise except VMTransportError: self._uncertain.add(operation_id) raise @@ -274,6 +304,7 @@ async def file_write_lock(self, path: str) -> AsyncIterator[None]: await self.setup() operation_id = uuid.uuid4().hex self._operations.add(operation_id) + self._edit_locks.add(operation_id) try: status = await self._rpc({"method": "start", "operation_id": operation_id, "request": {"action": "lock", "args": {}, "workspace": self.workspace, "identity": self._identity}}) @@ -290,10 +321,131 @@ async def file_write_lock(self, path: str) -> AsyncIterator[None]: finally: if not self._uncertain: await await_owned_operation(self._cancel_operation(operation_id)) + self._edit_locks.discard(operation_id) + + @contextlib.asynccontextmanager + async def _transfer_process(self, method: str): + registry = ProcessRegistry() + process = await registry.spawn(lambda: asyncio.create_subprocess_exec( + *self._ssh_command(method), stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, start_new_session=True)) + stderr = asyncio.create_task(read_bounded_output(process.stderr, 65536)) + try: + yield process, stderr + finally: + await await_owned_operation(registry.abort(), propagate_cancellation=True) + await await_owned_operation(stderr, propagate_cancellation=True) + + async def copy_to_host(self, source: str, target: Path, *, directory: bool = False) -> None: + self._ensure_active() + await self.setup() + target = Path(target).absolute() + target.parent.mkdir(parents=True, exist_ok=True) + staging = Path(tempfile.mkdtemp(prefix=".oc-vm-download-", dir=target.parent)) + archive = staging / "archive.tar" + try: + async with self._transfer_process("download") as (process, stderr): + payload = {"source": source, "directory": directory, + "workspace": self.workspace, "identity": self._identity} + process.stdin.write(json.dumps(payload).encode()) + await process.stdin.drain() + process.stdin.close() + prefix = await process.stdout.readuntil(_FRAME) + header = await process.stdout.readline() + response = self._decode_response(prefix + header, "download") + if not response.get("stream"): + raise VMTransportError("VM download stream is missing", operation_id="download") + with archive.open("wb") as output: + while True: + length = struct.unpack("!I", await process.stdout.readexactly(4))[0] + if not length: + break + if length > 65536: + raise VMTransportError("VM download frame exceeds transfer limit", operation_id="download") + output.write(await process.stdout.readexactly(length)) + output.flush() + os.fsync(output.fileno()) + remainder = await process.stdout.read(1) + await process.wait() + errors, _dropped = await stderr + if process.returncode or remainder: + raise VMTransportError(errors.decode(errors="replace"), operation_id="download", + returncode=process.returncode) + payload = await await_owned_operation( + asyncio.to_thread(unpack, archive, staging / "contents", directory=directory), + propagate_cancellation=True) + await await_owned_operation(asyncio.to_thread(publish, payload, target, directory=directory), + propagate_cancellation=True) + except (asyncio.IncompleteReadError, asyncio.LimitOverrunError) as exc: + raise VMTransportError(f"VM download interrupted, staging retained at {staging}", + operation_id="download") from exc + except BaseException as exc: + add_exception_note(exc, f"VM download staging retained at {staging}") + raise + else: + await asyncio.to_thread(shutil.rmtree, staging) + + async def copy_from_host(self, source: Path, destination: str, *, directory: bool = False) -> None: + self._ensure_active() + await self.setup() + operation_id = uuid.uuid4().hex + self._operations.add(operation_id) + request = {"method": "start", "operation_id": operation_id, + "request": {"action": "import", "workspace": self.workspace, "identity": self._identity, + "args": {"destination": destination, "directory": directory}}} + try: + async with self._transfer_process("upload") as (process, stderr): + stdout = asyncio.create_task(read_bounded_output(process.stdout, _MESSAGE_LIMIT)) + read_fd, write_fd = os.pipe() + def produce() -> None: + with os.fdopen(write_fd, "wb") as output: + pack(source, output, directory=directory) + producer = asyncio.create_task(asyncio.to_thread(produce)) + try: + process.stdin.write(json.dumps(request).encode() + b"\n") + await process.stdin.drain() + while True: + chunk = await asyncio.to_thread(os.read, read_fd, 65536) + if not chunk: + break + process.stdin.write(struct.pack("!I", len(chunk)) + chunk) + await process.stdin.drain() + await producer + process.stdin.write(struct.pack("!I", 0)) + await process.stdin.drain() + process.stdin.close() + output, dropped = await stdout + await process.wait() + errors, _stderr_dropped = await stderr + if process.returncode or dropped: + raise VMTransportError(errors.decode(errors="replace"), operation_id=operation_id, + returncode=process.returncode) + status = self._decode_response(output, operation_id) + finally: + os.close(read_fd) + await await_owned_operation(asyncio.gather(producer, return_exceptions=True)) + if not stdout.done(): + stdout.cancel() + await asyncio.gather(stdout, return_exceptions=True) + await self._wait(operation_id, status) + except VMTransportError: + self._uncertain.add(operation_id) + raise + except (BrokenPipeError, ConnectionResetError) as exc: + self._uncertain.add(operation_id) + raise VMTransportError("VM upload transport closed", operation_id=operation_id) from exc + except BaseException: + try: + await await_owned_operation(self._cancel_operation(operation_id)) + except BaseException as exc: + self._uncertain.add(operation_id) + raise ProcessCleanupError(f"VM upload cleanup requires recovery, operation {operation_id}") from exc + raise async def cleanup(self) -> None: self.revoke() - for operation_id in tuple(self._operations): + operations = (self._operations - self._edit_locks, self._edit_locks) + for operation_id in tuple(operation for group in operations for operation in group): try: await await_owned_operation(self._cancel_operation(operation_id), propagate_cancellation=True) except BaseException: diff --git a/opencollab/adapters/_vm_archive.py b/opencollab/adapters/_vm_archive.py new file mode 100644 index 00000000..8cefa87d --- /dev/null +++ b/opencollab/adapters/_vm_archive.py @@ -0,0 +1,172 @@ +from __future__ import annotations + +import os +import stat +import tarfile +from pathlib import Path, PurePosixPath +from typing import BinaryIO + +from opencollab.adapters.safe_anchored_files import ( + open_directory_anchor, + open_parent_beneath, + write_regular_file_atomic_at, +) +from opencollab.adapters.safe_paths import canonicalize_system_path + + +def _absolute_parent(path: str | Path, *, create: bool = False) -> tuple[int, str]: + value = canonicalize_system_path(path) + root = open_directory_anchor("/") + try: + return open_parent_beneath(root, "/", str(value).lstrip("/"), create=create) + finally: + os.close(root) + + +def _add(archive: tarfile.TarFile, parent: int, name: str, relative: str) -> None: + info = os.stat(name, dir_fd=parent, follow_symlinks=False) + member = tarfile.TarInfo(relative) + member.mode = stat.S_IMODE(info.st_mode) & 0o777 + member.mtime = info.st_mtime + if stat.S_ISLNK(info.st_mode): + member.type = tarfile.SYMTYPE + member.linkname = os.readlink(name, dir_fd=parent) + archive.addfile(member) + return + flags = os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK | getattr(os, "O_CLOEXEC", 0) + fd = os.open(name, flags, dir_fd=parent) + try: + opened = os.fstat(fd) + if (opened.st_dev, opened.st_ino) != (info.st_dev, info.st_ino): + raise OSError("archive source changed while opening it") + if stat.S_ISDIR(opened.st_mode): + member.type = tarfile.DIRTYPE + archive.addfile(member) + for child in sorted(os.listdir(fd)): + _add(archive, fd, child, relative + "/" + child) + elif stat.S_ISREG(opened.st_mode): + member.size = opened.st_size + with os.fdopen(fd, "rb") as handle: + fd = -1 + archive.addfile(member, handle) + else: + raise OSError("VM transfers accept regular files, directories and symlinks") + finally: + if fd >= 0: + os.close(fd) + + +def pack(source: str | Path, output: BinaryIO, *, directory: bool) -> None: + parent, name = _absolute_parent(source) + try: + mode = os.stat(name, dir_fd=parent, follow_symlinks=False).st_mode + if stat.S_ISDIR(mode) != directory: + raise ValueError("VM transfer directory flag must match its source") + with tarfile.open(fileobj=output, mode="w|", format=tarfile.PAX_FORMAT) as archive: + _add(archive, parent, name, "payload") + finally: + os.close(parent) + + +def unpack(archive_path: Path, staging: Path, *, directory: bool) -> Path: + staging.mkdir(mode=0o700) + anchor = open_directory_anchor(staging) + seen = set() + modes = [] + try: + with tarfile.open(archive_path, mode="r|") as archive: + for member in archive: + parts = PurePosixPath(member.name).parts + if (not parts or parts[0] != "payload" or any(p in {"", ".", ".."} for p in parts) + or member.name.startswith("/") or member.name in seen): + raise ValueError("VM archive contains an unsafe or repeated path") + if len(parts) == 1 and member.isdir() != directory: + raise ValueError("VM archive root does not match transfer type") + member.mode &= 0o777 + seen.add(member.name) + parent, name = open_parent_beneath(anchor, staging, member.name, create=True) + try: + if member.isdir(): + os.mkdir(name, 0o700, dir_fd=parent) + modes.append((member.name, member.mode)) + elif member.issym(): + os.symlink(member.linkname, name, dir_fd=parent) + elif member.isfile(): + handle = archive.extractfile(member) + if handle is None: + raise ValueError("VM archive file contents are missing") + with handle: + def write(output: BinaryIO) -> None: + remaining = member.size + while remaining: + chunk = handle.read(min(65536, remaining)) + if not chunk: + raise EOFError("VM archive file was truncated") + output.write(chunk) + remaining -= len(chunk) + write_regular_file_atomic_at(anchor, staging, member.name, write, + max_bytes=member.size, mode=member.mode, create_only=True) + fd = os.open(name, os.O_RDONLY | os.O_NOFOLLOW, dir_fd=parent) + try: + os.fchmod(fd, member.mode) + finally: + os.close(fd) + else: + raise ValueError("VM archive contains an unsupported file type") + finally: + os.close(parent) + if "payload" not in seen: + raise ValueError("VM archive payload is missing") + for path, mode in reversed(modes): + parent, name = open_parent_beneath(anchor, staging, path, create=False) + try: + fd = os.open(name, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW, dir_fd=parent) + try: + os.fchmod(fd, mode) + finally: + os.close(fd) + finally: + os.close(parent) + finally: + os.close(anchor) + return staging / "payload" + + +def publish(payload: Path, destination: str | Path, *, directory: bool) -> None: + source_parent, source_name = _absolute_parent(payload) + parent, name = _absolute_parent(destination, create=True) + try: + if directory: + try: + os.mkdir(name, 0o700, dir_fd=parent) + except FileExistsError: + pass + target = os.open(name, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW, dir_fd=parent) + source_mode = stat.S_IMODE(os.stat(source_name, dir_fd=source_parent, follow_symlinks=False).st_mode) + os.chmod(source_name, source_mode | 0o700, dir_fd=source_parent, follow_symlinks=False) + source = os.open(source_name, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW, dir_fd=source_parent) + try: + if os.listdir(target): + raise FileExistsError("VM directory transfer requires an empty destination") + for child in os.listdir(source): + os.rename(child, child, src_dir_fd=source, dst_dir_fd=target) + os.fchmod(target, source_mode) + os.fsync(target) + finally: + os.close(source) + os.close(target) + else: + try: + existing = os.stat(name, dir_fd=parent, follow_symlinks=False) + except FileNotFoundError: + existing = None + if existing is not None and not stat.S_ISREG(existing.st_mode): + raise OSError("VM transfer destination must be a regular file or absent") + os.rename(source_name, name, src_dir_fd=source_parent, dst_dir_fd=parent) + os.fsync(parent) + finally: + os.close(source_parent) + os.close(parent) + + +__all__ = ["pack", "publish", "unpack"] diff --git a/opencollab/adapters/_vm_remote.py b/opencollab/adapters/_vm_remote.py index e8d3aa43..042179eb 100644 --- a/opencollab/adapters/_vm_remote.py +++ b/opencollab/adapters/_vm_remote.py @@ -10,9 +10,12 @@ import math import os import re +import shutil import signal +import struct import subprocess import sys +import tempfile import time from dataclasses import asdict from pathlib import Path @@ -22,6 +25,7 @@ from opencollab.adapters._env_base import ENV_FILE_WRITE_LIMIT_BYTES from opencollab.adapters._env_local import LocalEnvironment from opencollab.adapters._env_process import ProcessCleanupError +from opencollab.adapters._vm_archive import pack, publish, unpack from opencollab.adapters.safe_anchored_files import ( create_regular_bytes_atomic_at, open_directory_anchor, @@ -31,7 +35,8 @@ ) from opencollab.domain.vm import validate_vm_path -_MAX_MESSAGE = 16 * 1024 * 1024 +_MAX_MESSAGE = 32 * 1024 * 1024 +_FRAME = b"OCVM1\n" _ID = re.compile(r"[A-Za-z0-9][A-Za-z0-9_.-]{0,127}") @@ -91,6 +96,14 @@ def _status(directory: Path) -> dict[str, Any]: try: started = _read(directory, "started.json") except FileNotFoundError: + request = _read(directory, "request.json") + if request["action"] == "import" and not (directory / "upload.json").exists(): + with (directory / "upload.lock").open("a+b") as lock: + try: + fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError: + return {"status": "uploading"} + return {"status": "upload_incomplete"} return {"status": "preparing"} try: live = (_process(started["pid"])[1] == started["started"] @@ -150,6 +163,10 @@ async def _perform(request: dict[str, Any], directory: Path) -> Any: if _identity(request["workspace"]) != request["identity"]: raise OSError("VM workspace or guest boot identity changed") env = LocalEnvironment(request["workspace"]) + info = os.fstat(env._root_fd()) + if (info.st_dev, info.st_ino) != (request["identity"]["device"], request["identity"]["inode"]): + await env.cleanup() + raise OSError("VM workspace identity changed while opening it") action = request["action"] args = request.get("args", {}) try: @@ -168,7 +185,17 @@ async def _perform(request: dict[str, Any], directory: Path) -> Any: await asyncio.sleep(0.02) _write(directory, "lock.json", True) await asyncio.Future() + if action == "import": + destination = Path(args["destination"]) + destination.parent.mkdir(parents=True, exist_ok=True) + staging = Path(tempfile.mkdtemp(prefix=".oc-vm-import-", dir=destination.parent)) + _write(directory, "staging.json", str(staging)) + payload = unpack(directory / "archive.tar", staging / "contents", directory=args["directory"]) + publish(payload, destination, directory=args["directory"]) + shutil.rmtree(staging) + return None if action == "exec": + env.workspace = f"/proc/{os.getpid()}/fd/{env._root_fd()}" result = asdict(await env.exec_cmd(args["cmd"], math.inf if args["timeout"] is None else args["timeout"], input_bytes=base64.b64decode(args["input"]) if args["input"] is not None else None)) result["stdout_bytes"] = (base64.b64encode(result["stdout_bytes"]).decode() @@ -267,6 +294,83 @@ def _worker(directory: Path) -> None: _write(directory, "result.json", {"status": "failed", "error": type(exc).__name__, "message": str(exc)}) +def _prepare(root: Path, request: dict[str, Any]) -> Path: + operation = request["operation_id"] + if not isinstance(operation, str) or not _ID.fullmatch(operation): + raise ValueError("invalid VM operation ID") + directory = root / operation + try: + directory.mkdir(mode=0o700) + except FileExistsError: + pass + fd = open_directory_anchor(directory) + os.close(fd) + with (directory / "submit.lock").open("a+b") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + try: + previous = _read(directory, "request.json") + except FileNotFoundError: + _write(directory, "request.json", request["request"]) + else: + if previous != request["request"]: + raise ValueError("VM operation ID is already bound to another request") + return directory + + +def _read_exact(length: int) -> bytes: + result = bytearray() + while len(result) < length: + chunk = sys.stdin.buffer.read(length - len(result)) + if not chunk: + raise EOFError("VM upload ended before its completion frame") + result.extend(chunk) + return bytes(result) + + +def _upload(root: Path, request: dict[str, Any]) -> Any: + directory = _prepare(root, request) + with (directory / "upload.lock").open("a+b") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + status = _status(directory) + receiving = status["status"] in {"preparing", "uploading", "upload_incomplete"} + output = (directory / "archive.partial").open("wb") if receiving else None + try: + while True: + length = struct.unpack("!I", _read_exact(4))[0] + if length == 0: + break + if length > 65536: + raise ValueError("VM upload frame exceeds transfer limit") + chunk = _read_exact(length) + if output is not None: + output.write(chunk) + if output is not None: + output.flush() + os.fsync(output.fileno()) + output.close() + output = None + os.replace(directory / "archive.partial", directory / "archive.tar") + _write(directory, "upload.json", True) + finally: + if output is not None: + output.close() + return _rpc(root, request) + + +class _ArchiveOutput: + def write(self, value: bytes) -> int: + for offset in range(0, len(value), 65536): + chunk = value[offset:offset + 65536] + sys.stdout.buffer.write(struct.pack("!I", len(chunk))) + sys.stdout.buffer.write(chunk) + return len(value) + + +def _response(value: dict[str, Any]) -> None: + sys.stdout.buffer.write(_FRAME + json.dumps(value).encode() + b"\n") + sys.stdout.buffer.flush() + + def _rpc(root: Path, request: dict[str, Any]) -> Any: action = request["method"] if action == "setup": @@ -286,27 +390,26 @@ def _rpc(root: Path, request: dict[str, Any]) -> Any: raise ValueError("invalid VM operation ID") directory = root / operation if action == "start": - try: - directory.mkdir(mode=0o700) - except FileExistsError: - pass - fd = open_directory_anchor(directory) - os.close(fd) - with (directory / "submit.lock").open("a+b") as lock: - fcntl.flock(lock, fcntl.LOCK_EX) - try: - previous = _read(directory, "request.json") - except FileNotFoundError: - _write(directory, "request.json", request["request"]) - else: - if previous != request["request"]: - raise ValueError("VM operation ID is already bound to another request") + directory = _prepare(root, request) + if (request["request"]["action"] == "import" and not (directory / "upload.json").exists() + and _status(directory)["status"] not in {"cancelled", "completed", "failed"}): + raise OSError("VM upload has no completed archive") if _status(directory)["status"] == "preparing": subprocess.Popen([sys.executable, "-m", "opencollab.adapters._vm_remote", "worker", str(directory)], stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, start_new_session=True, close_fds=True) + fd = open_directory_anchor(directory) + os.close(fd) if action == "cancel": _write(directory, "cancel.json", True) + with (directory / "worker.lock").open("a+b") as lock: + try: + fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError: + pass + else: + if not (directory / "started.json").exists() and not (directory / "result.json").exists(): + _write(directory, "result.json", {"status": "cancelled"}) if action not in {"start", "query", "cancel"}: raise ValueError("unknown VM RPC method") return _status(directory) @@ -317,17 +420,28 @@ def _main() -> None: _worker(Path(sys.argv[2])) return try: - if len(sys.argv) != 4 or sys.argv[1] != "rpc" or not _ID.fullmatch(sys.argv[3]): + if len(sys.argv) != 4 or sys.argv[1] not in {"rpc", "upload", "download"} or not _ID.fullmatch(sys.argv[3]): raise ValueError("invalid VM helper arguments") root = _directory(f"{validate_vm_path(sys.argv[2])}/{sys.argv[3]}") - payload = sys.stdin.buffer.read(_MAX_MESSAGE + 1) + if sys.argv[1] == "upload": + payload = sys.stdin.buffer.readline(_MAX_MESSAGE + 1) + else: + payload = sys.stdin.buffer.read(_MAX_MESSAGE + 1) if len(payload) > _MAX_MESSAGE: raise ValueError("VM request exceeds message limit") - value = _rpc(root, json.loads(payload)) + request = json.loads(payload) + if sys.argv[1] == "download": + if _identity(request["workspace"]) != request["identity"]: + raise OSError("VM workspace or guest boot identity changed") + _response({"ok": True, "stream": True}) + pack(request["source"], _ArchiveOutput(), directory=request["directory"]) + sys.stdout.buffer.write(struct.pack("!I", 0)) + return + value = _upload(root, request) if sys.argv[1] == "upload" else _rpc(root, request) response = {"ok": True, "value": value} except Exception as exc: response = {"ok": False, "error": type(exc).__name__, "message": str(exc)} - sys.stdout.buffer.write(json.dumps(response).encode()) + _response(response) if __name__ == "__main__": diff --git a/opencollab/adapters/safe_anchored_files.py b/opencollab/adapters/safe_anchored_files.py index 2a1f7e33..6625bfeb 100644 --- a/opencollab/adapters/safe_anchored_files.py +++ b/opencollab/adapters/safe_anchored_files.py @@ -78,7 +78,7 @@ def _relative_components( raise PermissionError(f"path escapes anchored directory: {value}") return parts[:-1], parts[-1] -def _open_parent_beneath( +def open_parent_beneath( root_fd: int, root_path: str | os.PathLike[str], path: str | os.PathLike[str], @@ -125,7 +125,7 @@ def read_regular_bytes_at( ) -> bytes: """Read a bounded regular file beneath a pinned directory descriptor.""" limit = _require_limit(max_bytes) - parent_fd, name = _open_parent_beneath( + parent_fd, name = open_parent_beneath( root_fd, root_path, path, @@ -183,7 +183,7 @@ def read_regular_text_range_at( or max_chars < 1 ): raise ValueError("offset, limit, and max_chars must be positive integers") - parent_fd, name = _open_parent_beneath( + parent_fd, name = open_parent_beneath( root_fd, root_path, path, @@ -307,7 +307,7 @@ def write_regular_file_atomic_at( ) -> bool | None: """Atomically publish one file beneath a pinned directory descriptor.""" limit = _require_limit(max_bytes) - parent_fd, name = _open_parent_beneath( + parent_fd, name = open_parent_beneath( root_fd, root_path, path, @@ -432,7 +432,7 @@ def unlink_regular_file_durable_at( *, expected_identity: tuple[int, int] | None = None, ) -> bool: - parent_fd, name = _open_parent_beneath( + parent_fd, name = open_parent_beneath( root_fd, root_path, path, @@ -453,6 +453,7 @@ def unlink_regular_file_durable_at( __all__ = [ "create_regular_bytes_atomic_at", "open_directory_anchor", + "open_parent_beneath", "read_regular_bytes_at", "read_regular_text_range_at", "unlink_regular_file_durable_at", diff --git a/opencollab/domain/vm.py b/opencollab/domain/vm.py index b41e6020..2c80da97 100644 --- a/opencollab/domain/vm.py +++ b/opencollab/domain/vm.py @@ -17,8 +17,9 @@ class VMConnection: connect_timeout: float = 10.0 def __post_init__(self) -> None: + patterns = {"host": r"[A-Za-z0-9\[][A-Za-z0-9_.:\[\]%-]*", "user": r"[A-Za-z0-9_][A-Za-z0-9_.-]*"} for name, value in (("host", self.host), ("user", self.user)): - if not isinstance(value, str) or not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.:@%-]*", value): + if not isinstance(value, str) or not re.fullmatch(patterns[name], value): raise ValueError(f"VM {name} must be an SSH hostname or account name") if isinstance(self.port, bool) or not isinstance(self.port, int) or not 1 <= self.port <= 65535: raise ValueError("VM SSH port must be between 1 and 65535") diff --git a/tests/runtime/test_vm_archive.py b/tests/runtime/test_vm_archive.py new file mode 100644 index 00000000..c838d72d --- /dev/null +++ b/tests/runtime/test_vm_archive.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +import io +import os +import stat +import sys +import tarfile +from pathlib import Path + +import pytest + +from opencollab.adapters._vm_archive import pack, publish, unpack + + +def test_archive_preserves_binary_names_symlinks_modes_and_empty_files(tmp_path): + source = tmp_path / "source" + source.mkdir(mode=0o775) + name = "a space\nand\tunicode-中文-" + (os.fsdecode(b"\xff") if sys.platform == "linux" else "legal") + data = source / name + data.write_bytes(b"\0\xff" * 1000000) + data.chmod(0o775) + (source / "empty").touch() + (source / "link").symlink_to(name) + (source / "subdirectory").mkdir() + archive = tmp_path / "archive.tar" + with archive.open("wb") as output: + pack(source, output, directory=True) + previous = os.umask(0o077) + try: + payload = unpack(archive, tmp_path / "staging", directory=True) + target = tmp_path / "target" + target.mkdir() + original_inode = target.stat().st_ino + publish(payload, target, directory=True) + finally: + os.umask(previous) + assert (target / name).read_bytes() == data.read_bytes() + assert stat.S_IMODE((target / name).stat().st_mode) == 0o775 + assert (target / "empty").read_bytes() == b"" + assert os.readlink(target / "link") == name + assert target.stat().st_ino == original_inode + + +def _archive(path: Path, members): + with tarfile.open(path, "w") as archive: + for name, kind, mode in members: + member = tarfile.TarInfo(name) + member.mode = mode + member.type = kind + if kind == tarfile.SYMTYPE: + member.linkname = "/tmp" + archive.addfile(member, io.BytesIO()) + + +@pytest.mark.parametrize("members", [ + [("../escape", tarfile.REGTYPE, 0o644)], + [("/tmp/escape", tarfile.REGTYPE, 0o644)], + [("payload", tarfile.DIRTYPE, 0o755), ("payload/../escape", tarfile.REGTYPE, 0o644)], + [("payload", tarfile.DIRTYPE, 0o755), ("payload/link", tarfile.SYMTYPE, 0o777), + ("payload/link/escape", tarfile.REGTYPE, 0o644)], + [("payload", tarfile.DIRTYPE, 0o755), ("payload/file", tarfile.LNKTYPE, 0o644)], +]) +def test_archive_rejects_paths_links_and_special_members(tmp_path, members): + archive = tmp_path / "archive.tar" + _archive(archive, members) + with pytest.raises((ValueError, OSError)): + unpack(archive, tmp_path / "staging", directory=True) + assert not (tmp_path / "escape").exists() + + +def test_archive_strips_setuid_and_setgid_without_losing_group_write(tmp_path): + archive = tmp_path / "archive.tar" + _archive(archive, [("payload", tarfile.REGTYPE, 0o6775)]) + previous = os.umask(0o077) + try: + payload = unpack(archive, tmp_path / "staging", directory=False) + finally: + os.umask(previous) + assert stat.S_IMODE(payload.stat().st_mode) == 0o775 + + +def test_directory_publication_rejects_existing_content(tmp_path): + source = tmp_path / "source" + source.mkdir() + target = tmp_path / "target" + target.mkdir() + (target / "existing").write_text("keep") + with pytest.raises(FileExistsError, match="empty"): + publish(source, target, directory=True) + assert (target / "existing").read_text() == "keep" diff --git a/tests/runtime/test_vm_environment.py b/tests/runtime/test_vm_environment.py index 5d3ee3f0..3270ee76 100644 --- a/tests/runtime/test_vm_environment.py +++ b/tests/runtime/test_vm_environment.py @@ -131,3 +131,32 @@ async def test_remote_result_uses_existing_execution_result(): result = VMEnvironment._exec_result({"returncode": 0, "stdout": "ok", "stderr": "", "stdout_bytes": "b2s="}) assert isinstance(result, ExecResult) assert result.stdout_bytes == b"ok" + + +def test_owned_file_removal_rejects_replacement(tmp_path): + import os + + from opencollab.adapters.safe_anchored_files import open_directory_anchor, unlink_regular_file_durable_at + + target = tmp_path / "owned" + target.write_text("owned") + identity = (target.stat().st_dev, target.stat().st_ino) + target.rename(tmp_path / "saved") + target.write_text("replacement") + anchor = open_directory_anchor(tmp_path) + try: + with pytest.raises(OSError, match="replaced"): + unlink_regular_file_durable_at(anchor, tmp_path, "owned", expected_identity=identity) + assert target.read_text() == "replacement" + finally: + os.close(anchor) + + +def test_helper_framing_ignores_banner_and_retains_invalid_payload_error(): + framed = b'login banner\nOCVM1\n{"ok": true, "value": {"status": "completed"}}\n' + assert VMEnvironment._decode_response(framed, "once") == {"status": "completed"} + with pytest.raises(ValueError, match="invalid request"): + VMEnvironment._decode_response(b'OCVM1\n{"ok": false, "error": "ValueError", ' + b'"message": "invalid request"}\n', "once") + with pytest.raises(VMTransportError, match="incomplete"): + VMEnvironment._decode_response(b'OCVM1\n{"ok":', "once") diff --git a/tests/runtime/test_vm_ssh.py b/tests/runtime/test_vm_ssh.py new file mode 100644 index 00000000..abecf5ee --- /dev/null +++ b/tests/runtime/test_vm_ssh.py @@ -0,0 +1,193 @@ +"""Opt-in probes against a borrowed Linux guest with installed OpenCollab.""" + +from __future__ import annotations + +import asyncio +import json +import os +import shlex +import uuid +from pathlib import Path + +import pytest +import pytest_asyncio + +from opencollab.adapters.env import VMEnvironment +from opencollab.adapters.working_tree import EnvWorkingTreeProbe +from opencollab.domain.vm import VMConnection + + +@pytest_asyncio.fixture +async def vm(): + config_path = os.environ.get("OPENCOLLAB_TEST_VM_CONFIG") + if config_path is None: + pytest.skip("OPENCOLLAB_TEST_VM_CONFIG must identify a borrowed Linux test guest") + config = json.loads(Path(config_path).read_text()) + connection = VMConnection(**config["connection"]) + parent = VMEnvironment(connection, config["workspace"], state_directory=config["state_directory"]) + child = parent.workspace + "/probe-" + uuid.uuid4().hex + result = await parent.exec_cmd("mkdir -m 700 -- " + shlex.quote(child)) + assert result.returncode == 0 + await parent.cleanup() + env = VMEnvironment(connection, child, state_directory=config["state_directory"]) + try: + await env.setup() + yield env + finally: + await env.cleanup() + + +async def test_ssh_binary_stdin_exit_255_and_output_bounds(vm): + result = await vm.exec_cmd("cat; printf stderr >&2; exit 255", input_bytes=b"\x00\xffabc", timeout=None) + assert result.returncode == 255 + assert result.stdout_bytes == b"\x00\xffabc" + assert result.stderr == "stderr" + result = await vm.exec_cmd("python3 -c 'import sys; sys.stdout.buffer.write(b\"A\"*1200000); " + "sys.stderr.buffer.write(b\"B\"*1200000)'") + assert len(result.stdout_bytes) == 1024 * 1024 + assert result.stdout_dropped_bytes == 1200000 - 1024 * 1024 + assert result.stderr_dropped_bytes == 1200000 - 1024 * 1024 + + +async def test_ssh_files_ranges_symlinks_and_owned_temporary(vm): + await vm.write_file("file.txt", "first\nsecond\nthird\n") + assert await vm.read_file("file.txt") == "first\nsecond\nthird\n" + page = await vm.read_text_range("file.txt", offset=2, limit=1, max_chars=100) + assert page.lines == ["second"] + assert page.has_more + assert (await vm.exec_cmd("ln -s /etc/passwd link")).returncode == 0 + with pytest.raises(OSError): + await vm.read_file("link") + with pytest.raises(PermissionError): + await vm.write_file("../outside", "escape") + temporary = await vm.write_temp_file("owned", prefix="test-") + with pytest.raises(OSError, match="unowned"): + await vm.remove_file("file.txt") + await vm.exec_cmd("mv " + shlex.quote(temporary) + " " + shlex.quote(temporary + ".saved") + + "; printf replacement > " + shlex.quote(temporary)) + with pytest.raises(OSError, match="replaced"): + await vm.remove_file(temporary) + assert await vm.read_file(temporary) == "replacement" + await vm.exec_cmd("rm " + shlex.quote(temporary) + "; mv " + shlex.quote(temporary + ".saved") + + " " + shlex.quote(temporary)) + await vm.remove_file(temporary) + + +async def test_ssh_operation_recovery_does_not_repeat_side_effect(vm): + operation_id = uuid.uuid4().hex + command = "printf once >> counter; printf before; sleep 0.8; printf after" + request = {"method": "start", "operation_id": operation_id, "request": { + "workspace": vm.workspace, "identity": vm._identity, "action": "exec", + "args": {"cmd": command, "timeout": 20, "input": None}}} + await vm._rpc(request) + # Disconnect a real SSH transport while the durable command is running. + ssh = vm._ssh_command() + ssh[-1] = "sleep 30; " + ssh[-1] + transport = await asyncio.create_subprocess_exec(*ssh, stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE) + transport.stdin.write(json.dumps({"method": "query", "operation_id": operation_id}).encode()) + await transport.stdin.drain() + transport.stdin.close() + await asyncio.sleep(0.1) + transport.kill() + await transport.wait() + result = await vm.exec_operation(operation_id, command, timeout=20) + assert result.stdout == "beforeafter" + assert await vm.read_file("counter") == "once" + recovered = VMEnvironment(vm.connection, vm.workspace, state_directory=vm.state_directory, run_id=vm.run_id) + assert (await recovered.recover_operation(operation_id)).stdout == "beforeafter" + await recovered.cleanup() + + +async def test_ssh_cancel_kills_double_forked_daemon_and_preserves_guest(vm): + code = ("import os,time,signal; p=os.fork(); " + "os._exit(0) if p else None; os.setsid(); p=os.fork(); " + "os._exit(0) if p else None; signal.signal(signal.SIGTERM,signal.SIG_IGN); " + "open('daemon.pid','w').write(str(os.getpid())); " + "[(open('ticks','a').write('x'),time.sleep(.05)) for _ in range(1000)]") + task = asyncio.create_task(vm.exec_cmd("python3 -c " + shlex.quote(code) + "; sleep 30", timeout=None)) + for _ in range(60): + await asyncio.sleep(0.05) + if (await vm.exec_cmd("test -f daemon.pid")).returncode == 0: + break + else: + pytest.fail("daemon did not start") + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + size = (await vm.exec_cmd("wc -c < ticks")).stdout.strip() + await asyncio.sleep(0.2) + assert (await vm.exec_cmd("wc -c < ticks")).stdout.strip() == size + assert (await vm.exec_cmd("kill -0 $(cat daemon.pid) 2>/dev/null")).returncode != 0 + assert (await vm.exec_cmd("test -r /proc/1/stat")).returncode == 0 + + +async def test_ssh_timeout_preserves_output_and_stops_background(vm): + result = await vm.exec_cmd("printf before; printf warning >&2; sleep 30", timeout=0.2) + assert result.timed_out + assert result.stdout == "before" + assert "warning" in result.stderr + + +async def test_ssh_workspace_edits_serialize_across_environment_instances(vm): + other = VMEnvironment(vm.connection, vm.workspace, state_directory=vm.state_directory) + async def edit(env): + async with env.file_write_lock("counter"): + value = await env.read_file("counter") + await asyncio.sleep(0.1) + await env.write_file("counter", str(int(value) + 1)) + await vm.write_file("counter", "0") + try: + await asyncio.gather(edit(vm), edit(other)) + assert await vm.read_file("counter") == "2" + finally: + await other.cleanup() + + +async def test_ssh_working_tree_probe_and_diff(vm): + assert (await vm.exec_cmd("git init -q; git -c user.name=Test -c user.email=test@example.invalid " + "commit -q --allow-empty -m initial")).returncode == 0 + probe = EnvWorkingTreeProbe(vm) + assert not await probe.changed() + await vm.write_file("new.txt", "new\n") + assert await probe.changed() + assert "+new" in await probe.diff() + + +async def test_ssh_reused_operation_id_rejects_changed_request(vm): + operation_id = uuid.uuid4().hex + assert (await vm.exec_operation(operation_id, "printf original")).stdout == "original" + with pytest.raises(ValueError, match="another request"): + await vm.exec_operation(operation_id, "printf changed") + with pytest.raises(FileNotFoundError): + await vm.recover_operation("absent") + + +async def test_ssh_streaming_copy_preserves_tree_and_binary_files(vm, tmp_path): + source = tmp_path / "source" + source.mkdir() + name = "with space\nand 中文" + data = b"\0\xffbinary\n" * 200000 + (source / name).write_bytes(data) + (source / name).chmod(0o775) + (source / "empty").touch() + (source / "link").symlink_to(name) + destination = vm.workspace + "/imported" + await vm.copy_from_host(source, destination, directory=True) + downloaded = tmp_path / "downloaded" + await vm.copy_to_host(destination, downloaded, directory=True) + assert (downloaded / name).read_bytes() == data + assert (downloaded / name).stat().st_mode & 0o777 == 0o775 + assert (downloaded / "empty").read_bytes() == b"" + assert os.readlink(downloaded / "link") == name + await vm.copy_from_host(source / "empty", vm.workspace + "/empty-file") + await vm.copy_to_host(vm.workspace + "/empty-file", tmp_path / "empty-file") + assert (tmp_path / "empty-file").read_bytes() == b"" + # A second archive location exercises state and workspace independence. + separate = VMEnvironment(vm.connection, vm.workspace, state_directory="/tmp/oc-vm-transfer-tests") + try: + await separate.copy_from_host(source / name, vm.workspace + "/separate-binary") + await separate.copy_to_host(vm.workspace + "/separate-binary", tmp_path / "binary") + assert (tmp_path / "binary").read_bytes() == data + finally: + await separate.cleanup() From 320c317edb784c187155f7cda5c886bab9926ec0 Mon Sep 17 00:00:00 2001 From: Kai Date: Sat, 24 Oct 2026 16:50:37 +0800 Subject: [PATCH 3/8] fix(environments): stage VM transfers on the destination filesystem --- docs/virtual-machines.md | 2 +- docs/virtual-machines.zh-CN.md | 2 +- opencollab/adapters/_env_vm.py | 3 ++- opencollab/adapters/_vm_archive.py | 11 ++++++++++- opencollab/adapters/_vm_remote.py | 5 +++-- tests/runtime/test_vm_archive.py | 10 ++++++++++ 6 files changed, 27 insertions(+), 6 deletions(-) diff --git a/docs/virtual-machines.md b/docs/virtual-machines.md index 94a2b5e3..11f2983f 100644 --- a/docs/virtual-machines.md +++ b/docs/virtual-machines.md @@ -48,4 +48,4 @@ The opt-in SSH regression suite reads `OPENCOLLAB_TEST_VM_CONFIG` from a JSON fi Adapter integrations can call `copy_to_host(source, target, directory=False)` and `copy_from_host(source, destination, directory=False)` to stream binary files or directory archives across SSH beyond command-output limits. Guest source and destination paths are explicit absolute paths selected by trusted workspace adapters, including repository metadata outside an exposed subdirectory. Directory transfers copy contents into an empty or absent destination and preserve the inode of an existing workspace. Regular permissions, symlinks, empty files, and filenames supported by the destination filesystem survive the roundtrip. Setuid and setgid bits are removed when crossing the transfer boundary. -Archive extraction uses an owned staging directory beside the destination and descriptor-relative publication. Parent traversal, repeated paths, hard links, device entries, and writes through symlink parents fail with the staging retained. Upload chunks carry an explicit completion frame before a durable import can execute. An interrupted upload keeps its operation ID and partial archive for inspection or cleanup. Operation records also retain failed import staging paths. Copying a directory into an existing populated destination raises `FileExistsError`. +Archive extraction uses an owned staging directory on the destination filesystem and descriptor-relative publication. Parent traversal, repeated paths, hard links, device entries, and writes through symlink parents fail with the staging retained. Upload chunks carry an explicit completion frame before a durable import can execute. An interrupted upload keeps its operation ID and partial archive for inspection or cleanup. Operation records also retain failed import staging paths. Copying a directory into an existing populated destination raises `FileExistsError`. diff --git a/docs/virtual-machines.zh-CN.md b/docs/virtual-machines.zh-CN.md index b48e0575..829d803b 100644 --- a/docs/virtual-machines.zh-CN.md +++ b/docs/virtual-machines.zh-CN.md @@ -48,4 +48,4 @@ await env.cleanup() 工作区适配器可以通过 `copy_to_host(source, target, directory=False)` 与 `copy_from_host(source, destination, directory=False)` 流式传输二进制文件和目录归档，传输大小独立于命令输出上限。客户机源路径和目标路径是工作区适配器明确指定的绝对路径，也可以指向暴露子目录之外的仓库元数据。目录复制将内容传入空目录或新目录，并保留已有工作区的 inode。常规权限、符号链接、空文件，以及目标文件系统支持的文件名会随传输保留。传输时移除 setuid 与 setgid 位。 -归档先在目标旁边的自有临时目录解包，再通过目录描述符写入目标。父目录跳转、重复路径、硬链接、设备文件，以及经由符号链接父目录写入的情况会触发失败，并保留临时目录。上传分块带有明确结束帧，完整接收后才允许执行持久化导入。上传中断时保留操作编号和部分归档，供检查或清理继续处理。操作记录还保存失败导入的临时目录路径。向已有内容的目录复制时抛出 `FileExistsError`。 +归档先在目标文件系统内的自有临时目录解包，再通过目录描述符写入目标。父目录跳转、重复路径、硬链接、设备文件，以及经由符号链接父目录写入的情况会触发失败，并保留临时目录。上传分块带有明确结束帧，完整接收后才允许执行持久化导入。上传中断时保留操作编号和部分归档，供检查或清理继续处理。操作记录还保存失败导入的临时目录路径。向已有内容的目录复制时抛出 `FileExistsError`。 diff --git a/opencollab/adapters/_env_vm.py b/opencollab/adapters/_env_vm.py index 59ca915c..f88e8df2 100644 --- a/opencollab/adapters/_env_vm.py +++ b/opencollab/adapters/_env_vm.py @@ -341,7 +341,8 @@ async def copy_to_host(self, source: str, target: Path, *, directory: bool = Fal await self.setup() target = Path(target).absolute() target.parent.mkdir(parents=True, exist_ok=True) - staging = Path(tempfile.mkdtemp(prefix=".oc-vm-download-", dir=target.parent)) + staging_parent = target if directory and target.is_dir() else target.parent + staging = Path(tempfile.mkdtemp(prefix=".oc-vm-download-", dir=staging_parent)) archive = staging / "archive.tar" try: async with self._transfer_process("download") as (process, stderr): diff --git a/opencollab/adapters/_vm_archive.py b/opencollab/adapters/_vm_archive.py index 8cefa87d..1409f38f 100644 --- a/opencollab/adapters/_vm_archive.py +++ b/opencollab/adapters/_vm_archive.py @@ -146,7 +146,16 @@ def publish(payload: Path, destination: str | Path, *, directory: bool) -> None: os.chmod(source_name, source_mode | 0o700, dir_fd=source_parent, follow_symlinks=False) source = os.open(source_name, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW, dir_fd=source_parent) try: - if os.listdir(target): + entries = set(os.listdir(target)) + owned_staging = payload.parent.parent + staging_parent = os.stat(owned_staging.parent, follow_symlinks=False) + target_info = os.fstat(target) + if (staging_parent.st_dev, staging_parent.st_ino) == (target_info.st_dev, target_info.st_ino): + staged = os.stat(owned_staging.name, dir_fd=target, follow_symlinks=False) + expected = os.stat(owned_staging, follow_symlinks=False) + if (staged.st_dev, staged.st_ino) == (expected.st_dev, expected.st_ino): + entries.discard(owned_staging.name) + if entries: raise FileExistsError("VM directory transfer requires an empty destination") for child in os.listdir(source): os.rename(child, child, src_dir_fd=source, dst_dir_fd=target) diff --git a/opencollab/adapters/_vm_remote.py b/opencollab/adapters/_vm_remote.py index 042179eb..8144b36e 100644 --- a/opencollab/adapters/_vm_remote.py +++ b/opencollab/adapters/_vm_remote.py @@ -186,9 +186,10 @@ async def _perform(request: dict[str, Any], directory: Path) -> Any: _write(directory, "lock.json", True) await asyncio.Future() if action == "import": - destination = Path(args["destination"]) + destination = Path(validate_vm_path(args["destination"])) destination.parent.mkdir(parents=True, exist_ok=True) - staging = Path(tempfile.mkdtemp(prefix=".oc-vm-import-", dir=destination.parent)) + staging_parent = destination if args["directory"] and destination.is_dir() else destination.parent + staging = Path(tempfile.mkdtemp(prefix=".oc-vm-import-", dir=staging_parent)) _write(directory, "staging.json", str(staging)) payload = unpack(directory / "archive.tar", staging / "contents", directory=args["directory"]) publish(payload, destination, directory=args["directory"]) diff --git a/tests/runtime/test_vm_archive.py b/tests/runtime/test_vm_archive.py index c838d72d..f6dfa5c5 100644 --- a/tests/runtime/test_vm_archive.py +++ b/tests/runtime/test_vm_archive.py @@ -88,3 +88,13 @@ def test_directory_publication_rejects_existing_content(tmp_path): with pytest.raises(FileExistsError, match="empty"): publish(source, target, directory=True) assert (target / "existing").read_text() == "keep" + + +def test_publication_accepts_only_its_owned_staging_inside_empty_destination(tmp_path): + target = tmp_path / "destination" + payload = target / ".oc-vm-import-owned" / "contents" / "payload" + payload.mkdir(parents=True) + (payload / "file").write_text("copied") + publish(payload, target, directory=True) + assert (target / "file").read_text() == "copied" + assert (target / ".oc-vm-import-owned").is_dir() From ee77c8a4904dd3b227e60d08257a0c2b4e9de432 Mon Sep 17 00:00:00 2001 From: Kai Date: Sat, 24 Oct 2026 17:00:40 +0800 Subject: [PATCH 4/8] fix(environments): retain VM temporary ownership through recovery Initialize workspace identity during recovery and preserve temporary file ownership when callers cancel after creation. Apply archive permissions through the open write descriptor so unreadable modes survive transfers. --- opencollab/adapters/_env_vm.py | 14 ++++++--- opencollab/adapters/_vm_archive.py | 8 ++--- tests/runtime/test_vm_archive.py | 11 +++++++ tests/runtime/test_vm_environment.py | 14 +++++++++ tests/runtime/test_vm_ssh.py | 46 +++++++++++++++++++++++++++- 5 files changed, 82 insertions(+), 11 deletions(-) diff --git a/opencollab/adapters/_env_vm.py b/opencollab/adapters/_env_vm.py index f88e8df2..d14977b7 100644 --- a/opencollab/adapters/_env_vm.py +++ b/opencollab/adapters/_env_vm.py @@ -144,6 +144,10 @@ async def setup(self, mount_dir: str | None = None) -> str: self._ensure_active() if mount_dir is not None: raise ValueError("mount_dir is supported only by container environments") + await self._setup_identity() + return self.workspace + + async def _setup_identity(self) -> None: async with self._setup_lock: if self._identity is None: response = await self._rpc({"method": "setup", "workspace": self.workspace}) @@ -152,7 +156,6 @@ async def setup(self, mount_dir: str | None = None) -> str: self._identity = response["identity"] self.workspace = self._identity["workspace"] self.source_workspace = self.workspace - return self.workspace async def _cancel_operation(self, operation_id: str) -> None: status = await self._rpc({"method": "cancel", "operation_id": operation_id}) @@ -250,6 +253,7 @@ async def recover_operation(self, operation_id: str) -> Any: known = operation_id in self._operations self._operations.add(operation_id) try: + await self._setup_identity() status = await self._rpc({"method": "query", "operation_id": operation_id}) value = await self._wait(operation_id, status) except FileNotFoundError: @@ -287,9 +291,11 @@ async def write_temp_file(self, content: str, *, prefix: str, suffix: str = ".tm if not re.fullmatch(r"[A-Za-z0-9._-]+", prefix) or not re.fullmatch(r"[A-Za-z0-9._-]+", suffix): raise ValueError("temporary file prefix and suffix must be simple path components") path = f"{prefix}{uuid.uuid4().hex}{suffix}" - identity = await self._operate("temporary", {"path": path, "content": content}) - self._temporary_files[path] = identity - return posixpath.join(self.workspace, path) + async def create() -> str: + identity = await self._operate("temporary", {"path": path, "content": content}) + self._temporary_files[path] = identity + return posixpath.join(self.workspace, path) + return await await_owned_operation(create(), propagate_cancellation=True) async def remove_file(self, path: str) -> None: target = posixpath.relpath(path, self.workspace) if path.startswith("/") else posixpath.normpath(path) diff --git a/opencollab/adapters/_vm_archive.py b/opencollab/adapters/_vm_archive.py index 1409f38f..27f0519c 100644 --- a/opencollab/adapters/_vm_archive.py +++ b/opencollab/adapters/_vm_archive.py @@ -104,13 +104,9 @@ def write(output: BinaryIO) -> None: raise EOFError("VM archive file was truncated") output.write(chunk) remaining -= len(chunk) + os.fchmod(output.fileno(), member.mode) write_regular_file_atomic_at(anchor, staging, member.name, write, - max_bytes=member.size, mode=member.mode, create_only=True) - fd = os.open(name, os.O_RDONLY | os.O_NOFOLLOW, dir_fd=parent) - try: - os.fchmod(fd, member.mode) - finally: - os.close(fd) + max_bytes=member.size, create_only=True) else: raise ValueError("VM archive contains an unsupported file type") finally: diff --git a/tests/runtime/test_vm_archive.py b/tests/runtime/test_vm_archive.py index f6dfa5c5..be6f6d58 100644 --- a/tests/runtime/test_vm_archive.py +++ b/tests/runtime/test_vm_archive.py @@ -79,6 +79,17 @@ def test_archive_strips_setuid_and_setgid_without_losing_group_write(tmp_path): assert stat.S_IMODE(payload.stat().st_mode) == 0o775 +@pytest.mark.parametrize("mode", [0o000, 0o200]) +def test_archive_preserves_files_without_owner_read_permission(tmp_path, mode): + archive = tmp_path / "archive.tar" + _archive(archive, [("payload", tarfile.REGTYPE, mode)]) + payload = unpack(archive, tmp_path / "staging", directory=False) + assert stat.S_IMODE(payload.stat().st_mode) == mode + target = tmp_path / "target" + publish(payload, target, directory=False) + assert stat.S_IMODE(target.stat().st_mode) == mode + + def test_directory_publication_rejects_existing_content(tmp_path): source = tmp_path / "source" source.mkdir() diff --git a/tests/runtime/test_vm_environment.py b/tests/runtime/test_vm_environment.py index 3270ee76..451973b0 100644 --- a/tests/runtime/test_vm_environment.py +++ b/tests/runtime/test_vm_environment.py @@ -6,6 +6,7 @@ import pytest +from opencollab import __version__ from opencollab.adapters.env import ExecResult, VMEnvironment, VMOperationError, VMTransportError from opencollab.domain.vm import VMConnection @@ -118,13 +119,26 @@ async def fail(payload): async def test_recover_file_operation_retains_temporary_identity(monkeypatch): env = _env() + identity = {"workspace": env.workspace, "device": 1, "inode": 2, "boot_id": "boot"} + setup_attempts = 0 async def rpc(payload): + nonlocal setup_attempts + if payload["method"] == "setup": + setup_attempts += 1 + if setup_attempts == 1: + raise VMTransportError("offline", operation_id="setup") + return {"identity": identity, "version": __version__} return {"status": "completed", "action": "temporary", "args": {"path": "owned.tmp"}, "value": [1, 2]} monkeypatch.setattr(env, "_rpc", rpc) + with pytest.raises(VMTransportError): + await env.recover_operation("temporary") + assert env.revoked assert await env.recover_operation("temporary") == [1, 2] + assert env._identity == identity assert env._temporary_files == {"owned.tmp": [1, 2]} + assert not env.revoked async def test_remote_result_uses_existing_execution_result(): diff --git a/tests/runtime/test_vm_ssh.py b/tests/runtime/test_vm_ssh.py index abecf5ee..71feefcd 100644 --- a/tests/runtime/test_vm_ssh.py +++ b/tests/runtime/test_vm_ssh.py @@ -163,6 +163,50 @@ async def test_ssh_reused_operation_id_rejects_changed_request(vm): await vm.recover_operation("absent") +async def test_ssh_recovered_temporary_cleanup_initializes_workspace_identity(vm): + operation_id = uuid.uuid4().hex + await vm._operate("temporary", {"path": "recovered.tmp", "content": "owned"}, operation_id) + recovered = VMEnvironment(vm.connection, vm.workspace, state_directory=vm.state_directory, run_id=vm.run_id) + try: + await recovered.recover_operation(operation_id) + finally: + await recovered.cleanup() + assert (await vm.exec_cmd("test ! -e recovered.tmp")).returncode == 0 + + +async def test_ssh_cancelled_temporary_creation_keeps_cleanup_ownership(vm, monkeypatch): + other = VMEnvironment(vm.connection, vm.workspace, state_directory=vm.state_directory) + original_rpc = other._rpc + completed = asyncio.Event() + release = asyncio.Event() + + async def delay_completed_response(payload): + response = await original_rpc(payload) + if response.get("status") == "completed" and response.get("action") == "temporary": + completed.set() + await release.wait() + return response + + monkeypatch.setattr(other, "_rpc", delay_completed_response) + task = asyncio.create_task(other.write_temp_file("owned", prefix="cancelled-")) + try: + await asyncio.wait_for(completed.wait(), 10) + task.cancel() + release.set() + with pytest.raises(asyncio.CancelledError): + await task + assert len(other._temporary_files) == 1 + temporary = next(iter(other._temporary_files)) + assert (await vm.exec_cmd("test -f " + shlex.quote(temporary))).returncode == 0 + finally: + release.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + await other.cleanup() + assert (await vm.exec_cmd("test ! -e " + shlex.quote(temporary))).returncode == 0 + + async def test_ssh_streaming_copy_preserves_tree_and_binary_files(vm, tmp_path): source = tmp_path / "source" source.mkdir() @@ -184,7 +228,7 @@ async def test_ssh_streaming_copy_preserves_tree_and_binary_files(vm, tmp_path): await vm.copy_to_host(vm.workspace + "/empty-file", tmp_path / "empty-file") assert (tmp_path / "empty-file").read_bytes() == b"" # A second archive location exercises state and workspace independence. - separate = VMEnvironment(vm.connection, vm.workspace, state_directory="/tmp/oc-vm-transfer-tests") + separate = VMEnvironment(vm.connection, vm.workspace, state_directory=vm.state_directory + "-separate") try: await separate.copy_from_host(source / name, vm.workspace + "/separate-binary") await separate.copy_to_host(vm.workspace + "/separate-binary", tmp_path / "binary") From 13bb946c5526672dafa02c4a03fdad180d1411b8 Mon Sep 17 00:00:00 2001 From: Kai Date: Sat, 24 Oct 2026 17:26:27 +0800 Subject: [PATCH 5/8] test(environments): escape Unicode filename fixtures --- tests/runtime/test_vm_archive.py | 2 +- tests/runtime/test_vm_ssh.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/runtime/test_vm_archive.py b/tests/runtime/test_vm_archive.py index be6f6d58..461e731f 100644 --- a/tests/runtime/test_vm_archive.py +++ b/tests/runtime/test_vm_archive.py @@ -15,7 +15,7 @@ def test_archive_preserves_binary_names_symlinks_modes_and_empty_files(tmp_path): source = tmp_path / "source" source.mkdir(mode=0o775) - name = "a space\nand\tunicode-中文-" + (os.fsdecode(b"\xff") if sys.platform == "linux" else "legal") + name = "a space\nand\tunicode-\u4e2d\u6587-" + (os.fsdecode(b"\xff") if sys.platform == "linux" else "legal") data = source / name data.write_bytes(b"\0\xff" * 1000000) data.chmod(0o775) diff --git a/tests/runtime/test_vm_ssh.py b/tests/runtime/test_vm_ssh.py index 71feefcd..a5b42b12 100644 --- a/tests/runtime/test_vm_ssh.py +++ b/tests/runtime/test_vm_ssh.py @@ -210,7 +210,7 @@ async def delay_completed_response(payload): async def test_ssh_streaming_copy_preserves_tree_and_binary_files(vm, tmp_path): source = tmp_path / "source" source.mkdir() - name = "with space\nand 中文" + name = "with space\nand \u4e2d\u6587" data = b"\0\xffbinary\n" * 200000 (source / name).write_bytes(data) (source / name).chmod(0o775) From e89390d4e136ae18f4874ea468b848097c1f3930 Mon Sep 17 00:00:00 2001 From: Kai Date: Sat, 24 Oct 2026 17:32:43 +0800 Subject: [PATCH 6/8] refactor(vm): move connection configuration to adapters --- docs/virtual-machines.md | 2 +- docs/virtual-machines.zh-CN.md | 2 +- opencollab/adapters/_env_vm.py | 2 +- opencollab/adapters/_vm_remote.py | 2 +- opencollab/{domain/vm.py => adapters/vm_config.py} | 0 tests/runtime/test_vm_environment.py | 2 +- tests/runtime/test_vm_ssh.py | 2 +- 7 files changed, 6 insertions(+), 6 deletions(-) rename opencollab/{domain/vm.py => adapters/vm_config.py} (100%) diff --git a/docs/virtual-machines.md b/docs/virtual-machines.md index 11f2983f..22d35a7f 100644 --- a/docs/virtual-machines.md +++ b/docs/virtual-machines.md @@ -6,7 +6,7 @@ Prepare an existing workspace and a private state directory owned by the SSH acc ```python from opencollab.adapters.env import VMEnvironment -from opencollab.domain.vm import VMConnection +from opencollab.adapters.vm_config import VMConnection connection = VMConnection( host="vm.example.net", diff --git a/docs/virtual-machines.zh-CN.md b/docs/virtual-machines.zh-CN.md index 829d803b..bec4f5cf 100644 --- a/docs/virtual-machines.zh-CN.md +++ b/docs/virtual-machines.zh-CN.md @@ -6,7 +6,7 @@ ```python from opencollab.adapters.env import VMEnvironment -from opencollab.domain.vm import VMConnection +from opencollab.adapters.vm_config import VMConnection connection = VMConnection( host="vm.example.net", diff --git a/opencollab/adapters/_env_vm.py b/opencollab/adapters/_env_vm.py index d14977b7..dfc91f3c 100644 --- a/opencollab/adapters/_env_vm.py +++ b/opencollab/adapters/_env_vm.py @@ -26,9 +26,9 @@ run_process, ) from opencollab.adapters._vm_archive import pack, publish, unpack +from opencollab.adapters.vm_config import VMConnection, validate_vm_path from opencollab.application.async_timeout import await_owned_operation from opencollab.application.exception_notes import add_exception_note -from opencollab.domain.vm import VMConnection, validate_vm_path _MESSAGE_LIMIT = 32 * 1024 * 1024 _FRAME = b"OCVM1\n" diff --git a/opencollab/adapters/_vm_remote.py b/opencollab/adapters/_vm_remote.py index 8144b36e..32f59af5 100644 --- a/opencollab/adapters/_vm_remote.py +++ b/opencollab/adapters/_vm_remote.py @@ -33,7 +33,7 @@ unlink_regular_file_durable_at, write_regular_bytes_atomic_at, ) -from opencollab.domain.vm import validate_vm_path +from opencollab.adapters.vm_config import validate_vm_path _MAX_MESSAGE = 32 * 1024 * 1024 _FRAME = b"OCVM1\n" diff --git a/opencollab/domain/vm.py b/opencollab/adapters/vm_config.py similarity index 100% rename from opencollab/domain/vm.py rename to opencollab/adapters/vm_config.py diff --git a/tests/runtime/test_vm_environment.py b/tests/runtime/test_vm_environment.py index 451973b0..1f6f9758 100644 --- a/tests/runtime/test_vm_environment.py +++ b/tests/runtime/test_vm_environment.py @@ -8,7 +8,7 @@ from opencollab import __version__ from opencollab.adapters.env import ExecResult, VMEnvironment, VMOperationError, VMTransportError -from opencollab.domain.vm import VMConnection +from opencollab.adapters.vm_config import VMConnection def _env(): diff --git a/tests/runtime/test_vm_ssh.py b/tests/runtime/test_vm_ssh.py index a5b42b12..051bdf63 100644 --- a/tests/runtime/test_vm_ssh.py +++ b/tests/runtime/test_vm_ssh.py @@ -13,8 +13,8 @@ import pytest_asyncio from opencollab.adapters.env import VMEnvironment +from opencollab.adapters.vm_config import VMConnection from opencollab.adapters.working_tree import EnvWorkingTreeProbe -from opencollab.domain.vm import VMConnection @pytest_asyncio.fixture From 35d4b90c29bedd15311aedcca49ed602e0adf6e7 Mon Sep 17 00:00:00 2001 From: Kai Date: Sat, 24 Oct 2026 21:41:21 +0800 Subject: [PATCH 7/8] fix(runtime): quiesce VM work before cleanup returns --- opencollab/adapters/_env_vm.py | 21 ++++- tests/runtime/test_vm_environment.py | 131 +++++++++++++++++++++++++++ tests/runtime/test_vm_ssh.py | 31 +++++++ 3 files changed, 182 insertions(+), 1 deletion(-) diff --git a/opencollab/adapters/_env_vm.py b/opencollab/adapters/_env_vm.py index dfc91f3c..033110d0 100644 --- a/opencollab/adapters/_env_vm.py +++ b/opencollab/adapters/_env_vm.py @@ -75,6 +75,7 @@ def __init__(self, connection: VMConnection, workspace: str, *, state_directory: self._uncertain: set[str] = set() self._edit_locks: set[str] = set() self._temporary_files: dict[str, list[int]] = {} + self._downloads: set[asyncio.Task[None]] = set() @property def revoked(self) -> bool: @@ -145,6 +146,7 @@ async def setup(self, mount_dir: str | None = None) -> str: if mount_dir is not None: raise ValueError("mount_dir is supported only by container environments") await self._setup_identity() + self._ensure_active() return self.workspace async def _setup_identity(self) -> None: @@ -345,6 +347,14 @@ async def _transfer_process(self, method: str): async def copy_to_host(self, source: str, target: Path, *, directory: bool = False) -> None: self._ensure_active() await self.setup() + owner = asyncio.create_task(self._copy_to_host(source, target, directory=directory)) + self._downloads.add(owner) + try: + await owner + finally: + self._downloads.discard(owner) + + async def _copy_to_host(self, source: str, target: Path, *, directory: bool) -> None: target = Path(target).absolute() target.parent.mkdir(parents=True, exist_ok=True) staging_parent = target if directory and target.is_dir() else target.parent @@ -390,7 +400,7 @@ async def copy_to_host(self, source: str, target: Path, *, directory: bool = Fal add_exception_note(exc, f"VM download staging retained at {staging}") raise else: - await asyncio.to_thread(shutil.rmtree, staging) + await await_owned_operation(asyncio.to_thread(shutil.rmtree, staging), propagate_cancellation=True) async def copy_from_host(self, source: Path, destination: str, *, directory: bool = False) -> None: self._ensure_active() @@ -451,6 +461,15 @@ def produce() -> None: async def cleanup(self) -> None: self.revoke() + downloads = tuple(self._downloads) + for download in downloads: + download.cancel() + if downloads: + results = await await_owned_operation(asyncio.gather(*downloads, return_exceptions=True), + propagate_cancellation=True) + for result in results: + if isinstance(result, ProcessCleanupError): + raise result operations = (self._operations - self._edit_locks, self._edit_locks) for operation_id in tuple(operation for group in operations for operation in group): try: diff --git a/tests/runtime/test_vm_environment.py b/tests/runtime/test_vm_environment.py index 1f6f9758..bd3c5750 100644 --- a/tests/runtime/test_vm_environment.py +++ b/tests/runtime/test_vm_environment.py @@ -3,10 +3,13 @@ import asyncio import base64 import math +import sys +import threading import pytest from opencollab import __version__ +from opencollab.adapters import _env_vm from opencollab.adapters.env import ExecResult, VMEnvironment, VMOperationError, VMTransportError from opencollab.adapters.vm_config import VMConnection @@ -174,3 +177,131 @@ def test_helper_framing_ignores_banner_and_retains_invalid_payload_error(): b'"message": "invalid request"}\n', "once") with pytest.raises(VMTransportError, match="incomplete"): VMEnvironment._decode_response(b'OCVM1\n{"ok":', "once") + + +@pytest.mark.parametrize("action", ["setup", "exec", "lock", "upload", "download"]) +async def test_cleanup_during_setup_prevents_later_work(monkeypatch, tmp_path, action): + env = _env() + entered = asyncio.Event() + release = asyncio.Event() + requests = [] + + async def rpc(payload): + requests.append(payload["method"]) + assert payload["method"] == "setup" + entered.set() + await release.wait() + return {"version": __version__, "identity": {"workspace": env.workspace}} + + async def work(): + if action == "setup": + await env.setup() + elif action == "exec": + await env.exec_cmd("touch late") + elif action == "lock": + async with env.file_write_lock("late"): + pytest.fail("edit lock was acquired after cleanup") + elif action == "upload": + await env.copy_from_host(tmp_path / "source", "late") + else: + await env.copy_to_host("late", tmp_path / "target") + + monkeypatch.setattr(env, "_rpc", rpc) + task = asyncio.create_task(work()) + await asyncio.wait_for(entered.wait(), 2) + await env.cleanup() + release.set() + with pytest.raises(RuntimeError, match="aborted"): + await task + assert requests == ["setup"] + assert not env.pending_operations + assert not (tmp_path / "target").exists() + + +@pytest.mark.parametrize("stage", ["stream", "unpack", "publish"]) +async def test_cleanup_waits_for_download_and_local_publication(monkeypatch, tmp_path, stage): + env = _env() + env._identity = {"workspace": env.workspace} + source = tmp_path / "source" + source.write_bytes(b"downloaded\0\xff") + archive = tmp_path / "source.tar" + with archive.open("wb") as output: + _env_vm.pack(source, output, directory=False) + marker = tmp_path / "stream-ready" + code = ( + "import pathlib,struct,sys,time; sys.stdin.buffer.read(); " + "data=pathlib.Path(sys.argv[1]).read_bytes(); " + "sys.stdout.buffer.write(b'OCVM1\\n{\"ok\":true,\"value\":{\"stream\":true}}\\n'); " + "sys.stdout.buffer.flush(); pathlib.Path(sys.argv[2]).touch(); " + "time.sleep(30) if sys.argv[3]=='stream' else None; " + "sys.stdout.buffer.write(struct.pack('!I',len(data))+data+struct.pack('!I',0)); " + "sys.stdout.buffer.flush()" + ) + monkeypatch.setattr(env, "_ssh_command", lambda _method: [ + sys.executable, "-c", code, str(archive), str(marker), stage]) + entered = asyncio.Event() + release = threading.Event() + loop = asyncio.get_running_loop() + if stage != "stream": + original = getattr(_env_vm, stage) + + def blocked(*args, **kwargs): + loop.call_soon_threadsafe(entered.set) + if not release.wait(5): + raise TimeoutError("test did not release download worker") + return original(*args, **kwargs) + + monkeypatch.setattr(_env_vm, stage, blocked) + + target = tmp_path / "target" + task = asyncio.create_task(env.copy_to_host("source", target)) + cleanup = None + try: + if stage == "stream": + async def ready(): + while not marker.exists(): + await asyncio.sleep(0.01) + await asyncio.wait_for(ready(), 2) + else: + await asyncio.wait_for(entered.wait(), 2) + cleanup = asyncio.create_task(env.cleanup()) + if stage != "stream": + await asyncio.wait({cleanup}, timeout=0.05) + assert not cleanup.done() + release.set() + await asyncio.wait_for(cleanup, 2) + assert task.done() + with pytest.raises(asyncio.CancelledError): + await task + if stage == "publish": + assert target.read_bytes() == source.read_bytes() + else: + assert not target.exists() + finally: + release.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + if cleanup is not None: + await asyncio.gather(cleanup, return_exceptions=True) + + +async def test_download_cleanup_failure_reaches_environment_cleanup(monkeypatch, tmp_path): + env = _env() + env._identity = {"workspace": env.workspace} + entered = asyncio.Event() + + async def download(*args, **kwargs): + try: + entered.set() + await asyncio.Future() + finally: + raise _env_vm.ProcessCleanupError("download process remained alive") + + monkeypatch.setattr(env, "_copy_to_host", download) + task = asyncio.create_task(env.copy_to_host("source", tmp_path / "target")) + await asyncio.wait_for(entered.wait(), 2) + with pytest.raises(_env_vm.ProcessCleanupError, match="remained alive"): + await env.cleanup() + with pytest.raises(_env_vm.ProcessCleanupError, match="remained alive"): + await task diff --git a/tests/runtime/test_vm_ssh.py b/tests/runtime/test_vm_ssh.py index 051bdf63..38f14fde 100644 --- a/tests/runtime/test_vm_ssh.py +++ b/tests/runtime/test_vm_ssh.py @@ -207,6 +207,37 @@ async def delay_completed_response(payload): assert (await vm.exec_cmd("test ! -e " + shlex.quote(temporary))).returncode == 0 +async def test_ssh_cleanup_during_setup_prevents_command_submission(vm, monkeypatch): + other = VMEnvironment(vm.connection, vm.workspace, state_directory=vm.state_directory) + original_rpc = other._rpc + entered = asyncio.Event() + release = asyncio.Event() + + async def delayed_setup(payload): + value = await original_rpc(payload) + if payload["method"] == "setup": + entered.set() + await release.wait() + return value + + monkeypatch.setattr(other, "_rpc", delayed_setup) + task = asyncio.create_task(other.exec_cmd("printf late > after-cleanup")) + try: + await asyncio.wait_for(entered.wait(), 10) + await other.cleanup() + release.set() + with pytest.raises(RuntimeError, match="aborted"): + await task + assert not other.pending_operations + assert (await vm.exec_cmd("test ! -e after-cleanup")).returncode == 0 + finally: + release.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + await other.cleanup() + + async def test_ssh_streaming_copy_preserves_tree_and_binary_files(vm, tmp_path): source = tmp_path / "source" source.mkdir() From 9133b9d6c5101584f7ea68b45f2847cd69299f73 Mon Sep 17 00:00:00 2001 From: Kai Date: Sat, 24 Oct 2026 21:51:47 +0800 Subject: [PATCH 8/8] fix(runtime): retain VM transfers until termination succeeds --- opencollab/adapters/_env_vm.py | 31 +++++++++-- tests/runtime/test_vm_environment.py | 77 +++++++++++++++++++++++++++- 2 files changed, 102 insertions(+), 6 deletions(-) diff --git a/opencollab/adapters/_env_vm.py b/opencollab/adapters/_env_vm.py index 033110d0..21b4c9fb 100644 --- a/opencollab/adapters/_env_vm.py +++ b/opencollab/adapters/_env_vm.py @@ -76,6 +76,9 @@ def __init__(self, connection: VMConnection, workspace: str, *, state_directory: self._edit_locks: set[str] = set() self._temporary_files: dict[str, list[int]] = {} self._downloads: set[asyncio.Task[None]] = set() + self._transfers: dict[ProcessRegistry, tuple[ + asyncio.Task[asyncio.subprocess.Process], asyncio.Task[tuple[bytes, int]] | None, + ]] = {} @property def revoked(self) -> bool: @@ -334,15 +337,31 @@ async def file_write_lock(self, path: str) -> AsyncIterator[None]: @contextlib.asynccontextmanager async def _transfer_process(self, method: str): registry = ProcessRegistry() - process = await registry.spawn(lambda: asyncio.create_subprocess_exec( + spawn = asyncio.create_task(registry.spawn(lambda: asyncio.create_subprocess_exec( *self._ssh_command(method), stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, start_new_session=True)) - stderr = asyncio.create_task(read_bounded_output(process.stderr, 65536)) + stderr=asyncio.subprocess.PIPE, start_new_session=True))) + self._transfers[registry] = (spawn, None) try: + process = await await_owned_operation(spawn, propagate_cancellation=True) + stderr = asyncio.create_task(read_bounded_output(process.stderr, 65536)) + self._transfers[registry] = (spawn, stderr) yield process, stderr finally: - await await_owned_operation(registry.abort(), propagate_cancellation=True) - await await_owned_operation(stderr, propagate_cancellation=True) + await self._cleanup_transfer(registry) + + async def _cleanup_transfer(self, registry: ProcessRegistry) -> None: + transfer = self._transfers.get(registry) + if transfer is None: + return + await await_owned_operation(asyncio.gather(transfer[0], return_exceptions=True), + propagate_cancellation=True) + await await_owned_operation(registry.abort(), propagate_cancellation=True) + transfer = self._transfers.get(registry) + try: + if transfer is not None and transfer[1] is not None and not transfer[1].cancelled(): + await await_owned_operation(transfer[1], propagate_cancellation=True) + finally: + self._transfers.pop(registry, None) async def copy_to_host(self, source: str, target: Path, *, directory: bool = False) -> None: self._ensure_active() @@ -470,6 +489,8 @@ async def cleanup(self) -> None: for result in results: if isinstance(result, ProcessCleanupError): raise result + for registry in tuple(self._transfers): + await self._cleanup_transfer(registry) operations = (self._operations - self._edit_locks, self._edit_locks) for operation_id in tuple(operation for group in operations for operation in group): try: diff --git a/tests/runtime/test_vm_environment.py b/tests/runtime/test_vm_environment.py index bd3c5750..1e3781a6 100644 --- a/tests/runtime/test_vm_environment.py +++ b/tests/runtime/test_vm_environment.py @@ -9,7 +9,7 @@ import pytest from opencollab import __version__ -from opencollab.adapters import _env_vm +from opencollab.adapters import _env_process, _env_vm from opencollab.adapters.env import ExecResult, VMEnvironment, VMOperationError, VMTransportError from opencollab.adapters.vm_config import VMConnection @@ -305,3 +305,78 @@ async def download(*args, **kwargs): await env.cleanup() with pytest.raises(_env_vm.ProcessCleanupError, match="remained alive"): await task + + +@pytest.mark.parametrize("direction,stage", [("download", "spawn"), ("download", "stream"), ("upload", "spawn")]) +async def test_transfer_cleanup_retries_failed_process_termination(monkeypatch, tmp_path, direction, stage): + env = _env() + env._identity = {"workspace": env.workspace} + marker = tmp_path / "stream-ready" + code = ( + "import pathlib,sys,time; sys.stdin.buffer.read(); " + "sys.stdout.buffer.write(b'OCVM1\\n{\"ok\":true,\"value\":{\"stream\":true}}\\n'); " + "sys.stdout.buffer.flush(); pathlib.Path(sys.argv[1]).touch(); time.sleep(30)" + ) + monkeypatch.setattr(env, "_ssh_command", lambda _method: [sys.executable, "-c", code, str(marker)]) + entered = asyncio.Event() + release = asyncio.Event() + spawn = asyncio.create_subprocess_exec + if stage == "spawn": + async def delayed_spawn(*args, **kwargs): + process = await spawn(*args, **kwargs) + entered.set() + await release.wait() + return process + monkeypatch.setattr(asyncio, "create_subprocess_exec", delayed_spawn) + terminate = _env_process.terminate_process + attempts = [] + + async def fail_once(process): + attempts.append(process) + return False if len(attempts) == 1 else await terminate(process) + + monkeypatch.setattr(_env_process, "terminate_process", fail_once) + if direction == "upload": + source = tmp_path / "source" + source.write_text("upload") + + async def rpc(payload): + assert payload["method"] == "cancel" + return {"status": "cancelled"} + + monkeypatch.setattr(env, "_rpc", rpc) + task = asyncio.create_task(env.copy_from_host(source, "target")) + else: + task = asyncio.create_task(env.copy_to_host("source", tmp_path / "target")) + try: + if stage == "spawn": + await asyncio.wait_for(entered.wait(), 2) + else: + async def ready(): + while not marker.exists(): + await asyncio.sleep(0.01) + await asyncio.wait_for(ready(), 2) + cleanup = asyncio.create_task(env.cleanup()) + if stage == "spawn": + await asyncio.wait({cleanup}, timeout=0.05) + release.set() + with pytest.raises(_env_vm.ProcessCleanupError, match="did not quiesce"): + await cleanup + if direction == "download": + with pytest.raises(_env_vm.ProcessCleanupError, match="did not quiesce"): + await task + process = attempts[0] + assert process.returncode is None + await asyncio.wait_for(env.cleanup(), 2) + await asyncio.wait_for(asyncio.gather(task, return_exceptions=True), 2) + assert len(attempts) >= 2 + assert all(attempt is process for attempt in attempts) + assert process.returncode is not None + assert not env._transfers + finally: + release.set() + for process in attempts: + await terminate(process) + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True)