From 9f483fdd2e6699596a0f1441e296a6fd5d84e47d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=85=B1=E7=89=9B=E8=82=89?= Date: Mon, 19 Oct 2026 22:25:01 +0800 Subject: [PATCH] fix(execution): bound output and clean up cancelled sandbox startup --- plugins/home/README.md | 4 + plugins/home/src/reme_home/auto_bash.py | 15 +- plugins/home/tests/test_auto_bash.py | 130 +++++++++ reme/steps/common/python_execute.py | 76 ++---- .../file_sandbox/file_sandbox_backend.py | 66 +---- .../utils/file_sandbox/path_bubble_backend.py | 7 +- .../file_sandbox/path_seatbelt_backend.py | 160 ++++++++---- reme/utils/local_backend.py | 247 ++++++++++++++++++ tests/unit/test_file_sandbox_backend.py | 136 +++++++++- tests/unit/test_process_output_limits.py | 156 +++++++++++ .../unit/test_sandbox_startup_cancellation.py | 164 ++++++++++++ 11 files changed, 972 insertions(+), 189 deletions(-) create mode 100644 reme/utils/local_backend.py create mode 100644 tests/unit/test_process_output_limits.py create mode 100644 tests/unit/test_sandbox_startup_cancellation.py diff --git a/plugins/home/README.md b/plugins/home/README.md index a8f61ebb..c3470af3 100644 --- a/plugins/home/README.md +++ b/plugins/home/README.md @@ -19,3 +19,7 @@ macOS), with sandbox paths identical to their host paths. On Linux, the `memory_limit_mib` cgroup limit defaults to 4096 MiB; 0 disables it. Supplying the limit requires a path policy. macOS has no memory limit and warns when a nonzero value is supplied. + +Standard output and standard error each have a 1 MiB (1,048,576-byte) limit on +all backends. Exceeding either limit terminates the command and returns a +failure with bounded partial output and a diagnostic in standard error. diff --git a/plugins/home/src/reme_home/auto_bash.py b/plugins/home/src/reme_home/auto_bash.py index 2970f84e..a464103d 100644 --- a/plugins/home/src/reme_home/auto_bash.py +++ b/plugins/home/src/reme_home/auto_bash.py @@ -1,6 +1,8 @@ """Run allowlisted shell commands without an approval prompt.""" import asyncio +import os +import shlex from functools import cache from agentscope.message import TextBlock, ToolResultState @@ -10,6 +12,8 @@ from agentscope.tool import Bash, ToolChunk from reme.steps.base_step import BaseStep from reme.steps.common.file_sandbox import file_sandbox_for_step +from reme.utils.local_backend import LocalProcessBackend + _READ_COMMANDS = { ":": {}, "cat": {}, @@ -54,8 +58,7 @@ _SAFE_DESTINATIONS = {"/dev/null"} # Trusted shell setup, applied only after validating the user's command. # Git's configuration can launch helpers even for otherwise read-only commands. _READ_ONLY_SETUP = r""" -unset GIT_EXTERNAL_DIFF GIT_TRACE GIT_TRACE_SETUP GIT_TRACE_PERFORMANCE -unset GIT_TRACE_PACK_ACCESS GIT_TRACE_PACKET GIT_TRACE_CURL GIT_TRACE2 GIT_TRACE2_EVENT GIT_TRACE2_PERF +unset GIT_EXTERNAL_DIFF export GIT_OPTIONAL_LOCKS=0 GIT_PAGER=/bin/cat GIT_NO_LAZY_FETCH=1 GIT_TERMINAL_PROMPT=0 sort() { # A large sort otherwise writes temporary files, even without -o. @@ -163,6 +166,8 @@ class AutoBash(Bash): } def __init__(self, *, is_read_only: bool = True, **kwargs): + if kwargs.get("backend") is None: + kwargs["backend"] = LocalProcessBackend() super().__init__(**kwargs) self.is_read_only = is_read_only @@ -227,7 +232,11 @@ class AutoBash(Bash): ) return if self.is_read_only: - command = _READ_ONLY_SETUP + command + # Trace destinations can write files even for metadata-only Git + # commands. Clear the entire inherited family on every invocation, + # including new Git trace variables, without changing os.environ. + trace_variables = " ".join(shlex.quote(name) for name in os.environ if name.startswith("GIT_TRACE")) + command = f"unset {trace_variables}\n" + _READ_ONLY_SETUP + command async for chunk in super().call(command=command, description=description, timeout=timeout): yield chunk diff --git a/plugins/home/tests/test_auto_bash.py b/plugins/home/tests/test_auto_bash.py index 616d498a..9ad9f466 100644 --- a/plugins/home/tests/test_auto_bash.py +++ b/plugins/home/tests/test_auto_bash.py @@ -2,18 +2,21 @@ import asyncio import inspect +import os import subprocess import threading from importlib.metadata import EntryPoint from pathlib import Path import pytest +import psutil import yaml from agentscope.message import ToolResultState from agentscope.permission import PermissionBehavior from agentscope.tool import ExecResult from reme_home.auto_bash import AutoBash, AutoBashStep +from reme.utils.local_backend import LocalProcessBackend from reme.utils.file_sandbox.path_bubble_backend import PathBubbleBackend from reme.components.application_context import ApplicationContext @@ -47,6 +50,75 @@ async def test_auto_bash_runs_in_workspace(tmp_path): assert response.answer == "hello" +@pytest.mark.asyncio +@pytest.mark.skipif(os.name != "posix", reason="POSIX shell pipelines") +@pytest.mark.parametrize("background", [False, True]) +@pytest.mark.parametrize("cancel", [False, True]) +async def test_local_pipeline_timeout_and_cancellation(tmp_path, monkeypatch, background, cancel): + """Timeouts and cancellation clean up pipelines even after the shell exits.""" + (tmp_path / "note.txt").write_text("hello\n", encoding="utf-8") + step = AutoBashStep(app_context=ApplicationContext(workspace_dir=str(tmp_path))) + processes = [] + descendants = [] + original_start = asyncio.create_subprocess_exec + original_stop = LocalProcessBackend._stop_process # pylint: disable=protected-access + started = asyncio.Event() + + def group_children(pid): + children = [] + for child in psutil.process_iter(): + try: + if child.pid != pid and os.getpgid(child.pid) == pid: + children.append(child) + except (ProcessLookupError, PermissionError): + pass + return children + + async def start(*args, **kwargs): + process = await original_start(*args, **kwargs) + processes.append(process) + started.set() + return process + + async def stop(self, process, communication): + descendants.extend(group_children(process.pid)) + await original_stop(self, process, communication) + + monkeypatch.setattr(asyncio, "create_subprocess_exec", start) + monkeypatch.setattr(LocalProcessBackend, "_stop_process", stop) + command = "tail -f note.txt | cat" + (" &" if background else "") + assert AutoBash._is_read_only(command) # pylint: disable=protected-access + task = asyncio.create_task(step(command=command, timeout=60000 if cancel else 100)) + try: + await asyncio.wait_for(started.wait(), timeout=2) + if cancel: + # Wait until pipeline children exist before cancelling the invocation. + async def pipeline_started(): + while len(group_children(processes[0].pid)) < 2: + await asyncio.sleep(0.01) + + await asyncio.wait_for(pipeline_started(), timeout=2) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=0.8) + else: + response = await asyncio.wait_for(task, timeout=0.8) + assert response.success is False + assert "timed out after 100ms" in response.answer + assert descendants + assert processes[0].returncode is not None + assert all(not child.is_running() or child.status() == psutil.STATUS_ZOMBIE for child in descendants) + finally: + if not task.done(): + task.cancel() + for process in processes: + try: + os.killpg(process.pid, 9) + except ProcessLookupError: + pass + await asyncio.gather(task, return_exceptions=True) + + @pytest.mark.asyncio @pytest.mark.parametrize( "command", @@ -345,6 +417,64 @@ async def test_sort_does_not_create_temporary_files(tmp_path): assert "sort:" in "".join(block.text for chunk in large for block in chunk.content) +@pytest.mark.asyncio +@pytest.mark.parametrize("trace_variable", ["GIT_TRACE_REFS", "GIT_TRACE", "GIT_TRACE2_EVENT"]) +async def test_read_only_git_cannot_write_inherited_traces(tmp_path, monkeypatch, trace_variable): + """Allowed metadata queries must not write Git traces from the parent environment.""" + + def git(*args): + return subprocess.run(["git", *args], cwd=tmp_path, check=True, capture_output=True) + + git("init", "-q") + git( + "-c", + "user.name=Test", + "-c", + "user.email=test@example.org", + "-c", + "commit.gpgsign=false", + "commit", + "--allow-empty", + "-qm", + "test", + ) + expected = git("rev-parse", "HEAD").stdout.decode().strip() + marker = tmp_path / "trace output.txt" + tool = AutoBash(cwd=tmp_path) + # Set after construction to check that each invocation uses the current environment. + monkeypatch.setenv(trace_variable, str(marker)) + + chunks = [chunk async for chunk in tool.call("git rev-parse HEAD")] + + assert chunks + assert all(chunk.state != ToolResultState.ERROR for chunk in chunks) + assert "".join(block.text for chunk in chunks for block in chunk.content).strip() == expected + assert not marker.exists() + + # Control: the same environment still enables tracing for unrestricted execution. + unrestricted = AutoBash(cwd=tmp_path, is_read_only=False) + chunks = [chunk async for chunk in unrestricted.call("git rev-parse HEAD")] + assert all(chunk.state != ToolResultState.ERROR for chunk in chunks) + assert marker.stat().st_size > 0 + + +@pytest.mark.asyncio +async def test_read_only_git_clears_future_trace_variables(tmp_path, monkeypatch): + """Trace isolation uses a prefix rule instead of an enumeration of known names.""" + helper = tmp_path / "git" + helper.write_text('#!/bin/sh\nprintf "%s" "${GIT_TRACE_FUTURE-unset}"\n', encoding="utf-8") + helper.chmod(0o700) + monkeypatch.setenv("PATH", str(tmp_path) + os.pathsep + os.environ["PATH"]) + monkeypatch.setenv("GIT_TRACE_FUTURE", "inherited") + tool = AutoBash(cwd=tmp_path) + + chunks = [chunk async for chunk in tool.call("git rev-parse HEAD")] + + assert all(chunk.state != ToolResultState.ERROR for chunk in chunks) + assert "".join(block.text for chunk in chunks for block in chunk.content) == "unset" + assert os.environ["GIT_TRACE_FUTURE"] == "inherited" + + @pytest.mark.asyncio async def test_git_pretty_config_cannot_launch_gpg(tmp_path): """Git format aliases and defaults cannot smuggle signature verification.""" diff --git a/reme/steps/common/python_execute.py b/reme/steps/common/python_execute.py index 5b227ab7..14ac974c 100644 --- a/reme/steps/common/python_execute.py +++ b/reme/steps/common/python_execute.py @@ -1,11 +1,9 @@ """Execute Python code and return printed stdout.""" -import asyncio import sys -from dataclasses import dataclass from typing import Any -from ...utils.file_sandbox.file_sandbox_backend import FileSandboxBackend +from ...utils.local_backend import LocalProcessBackend, ProcessTimeout from ..base_step import BaseStep from ...components import R from .file_sandbox import file_sandbox_for_step @@ -13,17 +11,9 @@ from .file_sandbox import file_sandbox_for_step DEFAULT_TIMEOUT = 60.0 -@dataclass(frozen=True) -class _PythonResult: - stdout: str - stderr: str - returncode: int | None - timed_out: bool = False - - @R.register("python_execute_step") class PythonExecuteStep(BaseStep): - """Run Python code in a subprocess and return stdout as the response answer.""" + """Run Python code and return stdout, limiting each output stream to 1 MiB.""" async def execute(self): assert self.context is not None @@ -41,69 +31,39 @@ class PythonExecuteStep(BaseStep): self.context.response.answer = timeout_error return self.context.response + timed_out = False try: backend = file_sandbox_for_step(self) - result = await self._run_python(code, timeout, backend) + result = await (backend or LocalProcessBackend()).exec_shell( + [sys.executable, "-c", code], + cwd=str(self.workspace_path), + timeout=timeout, + **({"raise_errors": True} if backend is None else {}), + ) + except ProcessTimeout as exc: + result, timed_out = exc.result, True except (TypeError, ValueError, OSError, RuntimeError) as exc: self.context.response.success = False self.context.response.answer = str(exc) return self.context.response - if result.timed_out: + timed_out = timed_out or (result.exit_code == -1 and result.stderr == b"timed out") + stdout = result.stdout.decode(errors="replace") + stderr = result.stderr.decode(errors="replace") + if timed_out: self.context.response.success = False self.context.response.answer = f"Python execution timed out after {timeout:g}s" - stderr = result.stderr else: - stdout = result.stdout or "" - stderr = result.stderr or "" - self.context.response.success = result.returncode == 0 - self.context.response.answer = stdout if stdout or result.returncode == 0 else stderr + self.context.response.success = result.exit_code == 0 + self.context.response.answer = stdout if stdout or result.exit_code == 0 else stderr self.context.response.metadata.update( { - "returncode": result.returncode, + "returncode": result.exit_code, "stderr": stderr, "python_timeout": timeout, }, ) return self.context.response - async def _run_python(self, code: str, timeout: float, backend: FileSandboxBackend | None = None) -> _PythonResult: - if backend is not None: - result = await backend.exec_shell( - [sys.executable, "-c", code], - cwd=str(self.workspace_path), - timeout=timeout, - ) - return _PythonResult( - stdout=result.stdout.decode(errors="replace"), - stderr=result.stderr.decode(errors="replace"), - returncode=result.exit_code, - timed_out=result.exit_code == -1 and result.stderr == b"timed out", - ) - process = await asyncio.create_subprocess_exec( - sys.executable, - "-c", - code, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - cwd=self.workspace_path, - ) - try: - stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=timeout) - return _PythonResult( - stdout=stdout.decode(), - stderr=stderr.decode(), - returncode=process.returncode, - ) - except TimeoutError: - process.kill() - stdout, stderr = await process.communicate() - return _PythonResult( - stdout=stdout.decode(), - stderr=stderr.decode(), - returncode=process.returncode, - timed_out=True, - ) - @staticmethod def _parse_timeout(raw: Any) -> tuple[float, str]: try: diff --git a/reme/utils/file_sandbox/file_sandbox_backend.py b/reme/utils/file_sandbox/file_sandbox_backend.py index 34802c85..220d0e4c 100644 --- a/reme/utils/file_sandbox/file_sandbox_backend.py +++ b/reme/utils/file_sandbox/file_sandbox_backend.py @@ -2,14 +2,12 @@ import asyncio import os -import signal import sys from abc import abstractmethod -import psutil -from agentscope.tool import BackendBase, ExecResult +from agentscope.tool import ExecResult -_CLEANUP_TIMEOUT = 2.0 +from ..local_backend import LocalProcessBackend def contains(root: str, path: str) -> bool: @@ -17,7 +15,7 @@ def contains(root: str, path: str) -> bool: return os.path.commonpath((root, path)) == root -class FileSandboxBackend(BackendBase): +class FileSandboxBackend(LocalProcessBackend): """Execute subprocesses with explicit host-path read and write access.""" def __init__(self, *, cwd: str | os.PathLike[str], read_only_paths: list[str], writable_paths: list[str]) -> None: @@ -76,9 +74,6 @@ class FileSandboxBackend(BackendBase): async def _start_process(self, command: list[str], cwd: str, *, stdin=asyncio.subprocess.DEVNULL): """Start one sandboxed process, never falling back to host execution.""" - def _release_process(self, process: asyncio.subprocess.Process) -> None: - """Release platform resources after the child has exited.""" - async def exec_shell( self, command: list[str], @@ -90,61 +85,6 @@ class FileSandboxBackend(BackendBase): process = await self._start_process(command, cwd or self._cwd) return await self._communicate(process, timeout=timeout) - async def _communicate( - self, - process: asyncio.subprocess.Process, - *, - data: bytes | None = None, - timeout: float | None = None, - ) -> ExecResult: - """Drain output and release process resources on success, timeout, or cancellation.""" - communication = asyncio.create_task(process.communicate(data)) - try: - stdout, stderr = await asyncio.wait_for(asyncio.shield(communication), timeout=timeout) - except TimeoutError: - await self._stop_process(process, communication) - return ExecResult(exit_code=-1, stdout=b"", stderr=b"timed out") - except BaseException: - await self._stop_process(process, communication) - raise - finally: - self._release_process(process) - return ExecResult(exit_code=process.returncode, stdout=stdout, stderr=stderr) - - @staticmethod - def _kill_process(process: asyncio.subprocess.Process) -> None: - """Stop the launcher and known descendants, including detached children.""" - if sys.platform == "darwin": - try: - os.killpg(process.pid, signal.SIGSTOP) - except ProcessLookupError: - pass - try: - descendants = psutil.Process(process.pid).children(recursive=True) - except psutil.Error: - descendants = [] - for child in reversed(descendants): - try: - child.kill() - except psutil.Error: - pass - try: - os.killpg(process.pid, signal.SIGKILL) - except ProcessLookupError: - pass - - async def _stop_process(self, process: asyncio.subprocess.Process, communication: asyncio.Task) -> None: - """Bound pipe draining even if an untracked child kept a pipe open.""" - self._kill_process(process) - try: - await asyncio.wait_for(asyncio.shield(communication), timeout=_CLEANUP_TIMEOUT) - except TimeoutError: - communication.cancel() - try: - await asyncio.wait_for(process.wait(), timeout=_CLEANUP_TIMEOUT) - except TimeoutError: - pass - async def read_file(self, path: str) -> bytes: """Read through a sandboxed command so file rules still apply.""" result = await self.exec_shell(["cat", path]) diff --git a/reme/utils/file_sandbox/path_bubble_backend.py b/reme/utils/file_sandbox/path_bubble_backend.py index 418b99c7..cbca8a79 100644 --- a/reme/utils/file_sandbox/path_bubble_backend.py +++ b/reme/utils/file_sandbox/path_bubble_backend.py @@ -77,12 +77,9 @@ class PathBubbleBackend(FileSandboxBackend): try: if group is not None: argv = [sys.executable, "-I", "-S", "-c", _CGROUP_EXEC, str(group / "cgroup.procs"), *argv] - process = await asyncio.create_subprocess_exec( - *argv, + process = await self._spawn_process( + argv, stdin=stdin, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - start_new_session=True, ) except BaseException: if group is not None: diff --git a/reme/utils/file_sandbox/path_seatbelt_backend.py b/reme/utils/file_sandbox/path_seatbelt_backend.py index 63320e32..ee6a5c67 100644 --- a/reme/utils/file_sandbox/path_seatbelt_backend.py +++ b/reme/utils/file_sandbox/path_seatbelt_backend.py @@ -1,6 +1,7 @@ """macOS Seatbelt file sandbox with host-identical paths.""" import asyncio +import errno import json import os import plistlib @@ -16,6 +17,21 @@ import psutil from agentscope.tool import ExecResult from .file_sandbox_backend import FileSandboxBackend, contains +from ..local_backend import ( + LocalProcessBackend, + OutputLimitExceeded, + capture_output, + finish_cleanup, + output_limit_result, +) + + +class LaunchctlError(RuntimeError): + """Preserve launchctl's error code so missing services can be handled narrowly.""" + + def __init__(self, message: str, exit_code: int): + super().__init__(message) + self.exit_code = exit_code class PathSeatbeltBackend(FileSandboxBackend): @@ -102,17 +118,11 @@ class PathSeatbeltBackend(FileSandboxBackend): """Start Seatbelt directly when the GUI launchd domain is unavailable.""" tempdir = tempfile.mkdtemp(prefix="reme-seatbelt-", dir=os.path.realpath(tempfile.gettempdir())) try: - process = await asyncio.create_subprocess_exec( - "/usr/bin/sandbox-exec", - "-p", - self._profile(tempdir), - *command, + process = await self._spawn_process( + ["/usr/bin/sandbox-exec", "-p", self._profile(tempdir), *command], cwd=cwd, env={**os.environ, "HOME": self._cwd, "TMPDIR": tempdir}, stdin=stdin, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - start_new_session=True, ) except BaseException: shutil.rmtree(tempdir) @@ -127,27 +137,38 @@ class PathSeatbeltBackend(FileSandboxBackend): @staticmethod async def _has_gui_launchd(domain: str) -> bool: - process = await asyncio.create_subprocess_exec( - "/bin/launchctl", - "print", - domain, - stdout=asyncio.subprocess.DEVNULL, - stderr=asyncio.subprocess.DEVNULL, - ) - return await process.wait() == 0 + result = await LocalProcessBackend().exec_shell(["/bin/launchctl", "print", domain], timeout=5) + return result.ok() @staticmethod async def _launchctl(*arguments: str) -> str: - process = await asyncio.create_subprocess_exec( - "/bin/launchctl", - *arguments, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - stdout, stderr = await process.communicate() - if process.returncode: - raise RuntimeError(f"launchctl {' '.join(arguments[:1])} failed: {stderr.decode(errors='replace').strip()}") - return stdout.decode(errors="replace").strip() + result = await LocalProcessBackend().exec_shell(["/bin/launchctl", *arguments], timeout=5) + if not result.ok(): + raise LaunchctlError( + f"launchctl {arguments[0]} failed: {result.stderr.decode(errors='replace').strip()}", + result.exit_code, + ) + return result.stdout.decode(errors="replace").strip() + + async def _bootstrap(self, domain: str, plist: Path, service: str) -> None: + """Resolve registration before rolling it back, even when its reply is interrupted.""" + registration = asyncio.create_task(self._launchctl("bootstrap", domain, str(plist))) + try: + await asyncio.shield(registration) + except BaseException: + + async def rollback(): + # Killing the client cannot undo a registration accepted by launchd. + # Wait for the bounded launchctl operation before issuing bootout. + await asyncio.gather(registration, return_exceptions=True) + try: + await self._launchctl("bootout", service) + except LaunchctlError as exc: + if exc.exit_code != errno.ESRCH: + raise + + await finish_cleanup(rollback()) + raise @staticmethod async def _wait_status(status: Path, pid: int) -> int: @@ -158,6 +179,20 @@ class PathSeatbeltBackend(FileSandboxBackend): raise RuntimeError("launchd sandbox job exited without reporting a status") await asyncio.sleep(0.02) + @staticmethod + async def _read_fifo(fd: int, buffer: bytearray, name: str, exceeded: asyncio.Future, finished: asyncio.Event): + """Drain launchd output through a bounded pipe instead of an unbounded disk file.""" + while True: + try: + chunk = os.read(fd, 65536) + except BlockingIOError: + if finished.is_set(): + return + await asyncio.sleep(0.005) + continue + capture_output(buffer, chunk, name, exceeded) + await asyncio.sleep(0) + async def _run_job(self, command: list[str], cwd: str, *, data: bytes | None, timeout: float | None) -> ExecResult: if sys.platform != "darwin": raise RuntimeError("PathSeatbeltBackend requires macOS") @@ -189,12 +224,34 @@ class PathSeatbeltBackend(FileSandboxBackend): label = f"ai.agentscope.reme.sandbox.{uuid.uuid4().hex}" service = f"{domain}/{label}" registered = False - timed_out = False - completed = False pid = None - stdout = b"" - stderr = b"" + stdout, stderr = bytearray(), bytearray() + exceeded = asyncio.get_running_loop().create_future() + finished = asyncio.Event() + output_fds = [] + readers = [] + + async def cleanup(completed): + try: + if registered: + if not completed and pid is not None: + self._kill_process(SimpleNamespace(pid=pid)) + await self._launchctl("bootout", service) + finally: + finished.set() + try: + await asyncio.wait_for(asyncio.gather(*readers), timeout=0.25) + except TimeoutError: + pass + try: + for path, buffer, name in ((stdout_path, stdout, "stdout"), (stderr_path, stderr, "stderr")): + os.mkfifo(path, 0o600) + # Holding both ends open prevents a transient EOF before launchd + # opens its writer. Kernel pipe capacity bounds pending output. + fd = os.open(path, os.O_RDWR | os.O_NONBLOCK) + output_fds.append(fd) + readers.append(asyncio.create_task(self._read_fifo(fd, buffer, name, exceeded, finished))) if data is not None: stdin_path.write_bytes(data) # The shell reports the command's status in a control file. Its @@ -225,29 +282,34 @@ class PathSeatbeltBackend(FileSandboxBackend): plist_path = control / "job.plist" with os.fdopen(os.open(plist_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600), "wb") as stream: plistlib.dump(job, stream) - await self._launchctl("bootstrap", domain, str(plist_path)) + await self._bootstrap(domain, plist_path, service) registered = True - pid = int(await self._launchctl("kickstart", "-p", service)) + + async def run(): + nonlocal pid + pid = int(await self._launchctl("kickstart", "-p", service)) + return await self._watch_output(self._wait_status(status_path, pid), exceeded) + try: - exit_code = await asyncio.wait_for(self._wait_status(status_path, pid), timeout=timeout) + exit_code = await self._execute(run(), timeout=timeout, cleanup=cleanup) except TimeoutError: - timed_out = True - else: - completed = True + return ExecResult(exit_code=-1, stdout=b"", stderr=b"timed out") + except OutputLimitExceeded as exc: + return output_limit_result(bytes(stdout), bytes(stderr), str(exc)) + if exceeded.done(): + return output_limit_result(bytes(stdout), bytes(stderr), exceeded.result()) + return ExecResult( + exit_code=exit_code, + stdout=bytes(stdout), + stderr=bytes(stderr), + ) finally: - try: - if registered: - if not completed and pid is not None: - self._kill_process(SimpleNamespace(pid=pid)) - await self._launchctl("bootout", service) - finally: - if not timed_out: - stdout = stdout_path.read_bytes() if stdout_path.exists() else b"" - stderr = stderr_path.read_bytes() if stderr_path.exists() else b"" - shutil.rmtree(control) - if timed_out: - return ExecResult(exit_code=-1, stdout=b"", stderr=b"timed out") - return ExecResult(exit_code=exit_code, stdout=stdout, stderr=stderr) + for reader in readers: + reader.cancel() + await asyncio.gather(*readers, return_exceptions=True) + for fd in output_fds: + os.close(fd) + shutil.rmtree(control) async def exec_shell( self, diff --git a/reme/utils/local_backend.py b/reme/utils/local_backend.py new file mode 100644 index 00000000..d8b48cb3 --- /dev/null +++ b/reme/utils/local_backend.py @@ -0,0 +1,247 @@ +"""Shared subprocess execution with bounded timeout and cancellation cleanup.""" + +import asyncio +import os +import signal +import sys +from collections.abc import Awaitable, Callable + +import psutil +from agentscope.tool import ExecResult, LocalBackend + +_CLEANUP_TIMEOUT = 0.25 +OUTPUT_LIMIT_BYTES = 1024 * 1024 + + +async def finish_cleanup(operation: Awaitable): + """Finish rollback after cancellation, deferring additional cancellation requests.""" + task = asyncio.ensure_future(operation) + while True: + try: + return await asyncio.shield(task) + except asyncio.CancelledError: + if task.done(): + return task.result() + + +class OutputLimitExceeded(RuntimeError): + """Identify the output stream that exceeded its byte budget.""" + + +def capture_output(buffer: bytearray, chunk: bytes, stream: str, exceeded: asyncio.Future) -> None: + """Retain at most one stream's budget, signalling overflow without blocking draining.""" + remaining = OUTPUT_LIMIT_BYTES - len(buffer) + buffer.extend(chunk[:remaining]) + if len(chunk) > remaining and not exceeded.done(): + exceeded.set_result(stream) + + +def output_limit_result(stdout: bytes, stderr: bytes, stream: str) -> ExecResult: + """Return bounded partial output and an explicit failure diagnostic.""" + message = f"\nOutput limit exceeded: {stream} exceeds {OUTPUT_LIMIT_BYTES} bytes; command terminated.\n".encode() + diagnostic = (stderr[: max(0, OUTPUT_LIMIT_BYTES - len(message))] + message)[:OUTPUT_LIMIT_BYTES] + return ExecResult(exit_code=1, stdout=stdout[:OUTPUT_LIMIT_BYTES], stderr=diagnostic) + + +class ProcessTimeout(TimeoutError): + """Carry the reaped process output for callers with their own timeout protocol.""" + + def __init__(self, result: ExecResult): + super().__init__("timed out") + self.result = result + + +class LocalProcessBackend(LocalBackend): + """Keep local file helpers; terminate commands exceeding 1 MiB on either output stream.""" + + async def exec_shell( + self, + command: list[str], + *, + cwd: str | None = None, + timeout: float | None = None, + raise_errors: bool = False, + ) -> ExecResult: + """Run an argv in its own session and collect output within the deadline.""" + try: + process = await self._start_process(command, cwd) + except OSError as exc: + if raise_errors: + raise + return ExecResult(exit_code=127, stdout=b"", stderr=str(exc).encode("utf-8")) + return await self._communicate(process, timeout=timeout, raise_errors=raise_errors) + + async def _start_process(self, command: list[str], cwd: str | None, *, stdin=None): + """Start a host process; sandboxes override platform-specific setup.""" + return await self._spawn_process(command, cwd=cwd, stdin=stdin) + + async def _spawn_process(self, command: list[str], **kwargs): + """Capture output in a separate process group for every execution backend.""" + spawning = asyncio.create_task( + asyncio.create_subprocess_exec( + *command, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + start_new_session=os.name == "posix", + **kwargs, + ), + ) + try: + return await asyncio.shield(spawning) + except asyncio.CancelledError: + + async def rollback(): + outcome = (await asyncio.gather(spawning, return_exceptions=True))[0] + if isinstance(outcome, BaseException): + return + exceeded = asyncio.get_running_loop().create_future() + communication = asyncio.create_task( + self._collect_output(outcome, None, bytearray(), bytearray(), exceeded), + ) + await self._stop_process(outcome, communication) + + await finish_cleanup(rollback()) + raise + + def _release_process(self, process: asyncio.subprocess.Process) -> None: + """Release any platform resources owned by a completed execution.""" + + @staticmethod + async def _watch_output(operation: Awaitable, exceeded: asyncio.Future): + """Stop waiting as soon as either captured stream exceeds its budget.""" + task = asyncio.ensure_future(operation) + try: + await asyncio.wait((task, exceeded), return_when=asyncio.FIRST_COMPLETED) + if exceeded.done(): + raise OutputLimitExceeded(exceeded.result()) + return await task + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + @staticmethod + async def _collect_output(process, data, stdout, stderr, exceeded): + """Drain both pipes concurrently while retaining bounded byte buffers.""" + + async def read(reader, buffer, name): + while chunk := await reader.read(65536): + capture_output(buffer, chunk, name, exceeded) + + async def write(): + if process.stdin is not None: + try: + if data is not None: + process.stdin.write(data) + await process.stdin.drain() + except (BrokenPipeError, ConnectionResetError): + pass + finally: + process.stdin.close() + + tasks = [ + asyncio.create_task(read(process.stdout, stdout, "stdout")), + asyncio.create_task(read(process.stderr, stderr, "stderr")), + asyncio.create_task(write()), + ] + try: + await asyncio.gather(*tasks) + await process.wait() + return bytes(stdout), bytes(stderr) + finally: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + @staticmethod + async def _execute(operation: Awaitable, *, timeout: float | None, cleanup: Callable[[bool], Awaitable[None]]): + """Apply one deadline and always clean up success, failure, or cancellation.""" + task = asyncio.ensure_future(operation) + completed = False + try: + result = await asyncio.wait_for(asyncio.shield(task), timeout=timeout) + completed = True + return result + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + await cleanup(completed) + + async def _communicate( + self, + process: asyncio.subprocess.Process, + *, + data: bytes | None = None, + timeout: float | None = None, + raise_errors: bool = False, + ) -> ExecResult: + """Collect pipes and release the process, preserving each caller's timeout protocol.""" + stdout_buffer, stderr_buffer = bytearray(), bytearray() + exceeded = asyncio.get_running_loop().create_future() + communication = asyncio.create_task( + self._collect_output(process, data, stdout_buffer, stderr_buffer, exceeded), + ) + + async def cleanup(completed): + try: + if not completed: + await self._stop_process(process, communication) + finally: + self._release_process(process) + + try: + stdout, stderr = await self._execute( + self._watch_output(asyncio.shield(communication), exceeded), + timeout=timeout, + cleanup=cleanup, + ) + except OutputLimitExceeded as exc: + return output_limit_result(bytes(stdout_buffer), bytes(stderr_buffer), str(exc)) + except TimeoutError: + if raise_errors: + stdout, stderr = communication.result() if not communication.cancelled() else (b"", b"") + raise ProcessTimeout(ExecResult(exit_code=process.returncode, stdout=stdout, stderr=stderr)) from None + return ExecResult(exit_code=-1, stdout=b"", stderr=b"timed out") + return ExecResult(exit_code=process.returncode, stdout=stdout, stderr=stderr) + + @staticmethod + def _kill_process(process: asyncio.subprocess.Process) -> None: + """Kill the group and known detached children on macOS and Windows.""" + if os.name != "posix" or sys.platform == "darwin": + if os.name == "posix": + try: + os.killpg(process.pid, signal.SIGSTOP) + except ProcessLookupError: + pass + try: + parent = psutil.Process(process.pid) + for child in reversed(parent.children(recursive=True)): + try: + child.kill() + except (psutil.Error, OSError): + pass + if os.name != "posix": + parent.kill() + except (psutil.Error, OSError): + pass + if os.name == "posix": + # The launcher may have exited while a pipeline still holds stdout. + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + + async def _stop_process(self, process: asyncio.subprocess.Process, communication: asyncio.Task) -> None: + """Kill descendants and bound draining even if a child escaped the group.""" + self._kill_process(process) + try: + await asyncio.wait_for(asyncio.shield(communication), timeout=_CLEANUP_TIMEOUT) + except TimeoutError: + communication.cancel() + await asyncio.gather(communication, return_exceptions=True) + # asyncio has no public API to close a Process's captured pipes. + # An escaped descendant must not keep these transports alive. + process._transport.close() # pylint: disable=protected-access + try: + await asyncio.wait_for(process.wait(), timeout=_CLEANUP_TIMEOUT) + except TimeoutError: + pass diff --git a/tests/unit/test_file_sandbox_backend.py b/tests/unit/test_file_sandbox_backend.py index 12113e58..ad09c0ad 100644 --- a/tests/unit/test_file_sandbox_backend.py +++ b/tests/unit/test_file_sandbox_backend.py @@ -1,6 +1,7 @@ """Shared file sandbox selection and macOS policy construction.""" import asyncio +import os import plistlib import sys from pathlib import Path @@ -10,6 +11,7 @@ from agentscope.tool import ExecResult from reme.components.application_context import ApplicationContext from reme.steps.common.python_execute import PythonExecuteStep +from reme.utils.local_backend import LocalProcessBackend from reme.utils.file_sandbox import create_file_sandbox_backend from reme.utils.file_sandbox.path_bubble_backend import PathBubbleBackend from reme.utils.file_sandbox.path_seatbelt_backend import PathSeatbeltBackend @@ -109,32 +111,62 @@ def test_seatbelt_excludes_read_only_child_from_writable_parent(tmp_path): @pytest.mark.asyncio -async def test_sandbox_timeout_does_not_wait_forever_for_inherited_pipe(tmp_path, monkeypatch): +@pytest.mark.parametrize("sandbox", [False, True]) +@pytest.mark.parametrize("cancel", [False, True]) +async def test_inherited_pipe_cleanup_is_bounded(tmp_path, monkeypatch, sandbox, cancel): """A surviving child with an output pipe must not hold the timeout response open.""" - backend = PathBubbleBackend(cwd=tmp_path, read_only_paths=[], writable_paths=[], memory_limit_mib=0) + from unittest.mock import Mock + + backend = ( + PathBubbleBackend(cwd=tmp_path, read_only_paths=[], writable_paths=[], memory_limit_mib=0) + if sandbox + else LocalProcessBackend() + ) + started = asyncio.Event() + cancelled = asyncio.Event() + + class StuckPipe: + """Model an inherited pipe held open after termination.""" + + async def read(self, _size): + """Wait forever until the caller cancels pipe draining.""" + started.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() class StuckProcess: """Model a child that keeps its output pipe open after termination.""" pid = 12345 - - async def communicate(self, _data=None): - """Wait forever until the caller cancels pipe draining.""" - await asyncio.Event().wait() + _transport = Mock() + stdout = StuckPipe() + stderr = StuckPipe() + stdin = None async def wait(self): - """Model an already reaped launcher.""" - return 0 + """A stale process transport must also have a bounded wait.""" + await asyncio.Event().wait() async def start(*_args, **_kwargs): return StuckProcess() monkeypatch.setattr(backend, "_start_process", start) monkeypatch.setattr(backend, "_kill_process", lambda _process: None) - monkeypatch.setattr("reme.utils.file_sandbox.file_sandbox_backend._CLEANUP_TIMEOUT", 0.01) + monkeypatch.setattr("reme.utils.local_backend._CLEANUP_TIMEOUT", 0.01) - result = await asyncio.wait_for(backend.exec_shell(["true"], timeout=0.01), timeout=1) - assert result.exit_code == -1 and result.stderr == b"timed out" + task = asyncio.create_task(backend.exec_shell(["true"], timeout=None if cancel else 0.01)) + await asyncio.wait_for(started.wait(), timeout=1) + if cancel: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=1) + else: + result = await asyncio.wait_for(task, timeout=1) + assert result.exit_code == -1 and result.stderr == b"timed out" + assert cancelled.is_set() + StuckProcess._transport.close.assert_called_once() # pylint: disable=protected-access @pytest.mark.asyncio @@ -145,6 +177,11 @@ async def test_seatbelt_launch_failure_cleans_up_without_host_fallback(tmp_path, seen = {} async def fail(*arguments): + if arguments[0] == "bootout": + import errno + from reme.utils.file_sandbox.path_seatbelt_backend import LaunchctlError + + raise LaunchctlError("service not found", errno.ESRCH) seen["arguments"] = arguments with open(arguments[-1], "rb") as stream: seen["job"] = plistlib.load(stream) @@ -181,3 +218,80 @@ async def test_python_job_selects_seatbelt_on_macos(tmp_path, monkeypatch): assert response.success is True assert isinstance(seen["backend"], PathSeatbeltBackend) assert seen["command"] == [sys.executable, "-c", "print('ok')"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["success", "failure", "timeout", "cancel", "error", "kickstart_error"]) +async def test_seatbelt_launchd_lifecycle(tmp_path, monkeypatch, mode): + """Every launchd outcome unregisters the job and removes its control files.""" + from unittest.mock import AsyncMock, Mock + + monkeypatch.setattr(sys, "platform", "darwin") + monkeypatch.setattr("reme.utils.file_sandbox.path_seatbelt_backend.shutil.which", lambda _: "sandbox-exec") + backend = PathSeatbeltBackend(cwd=tmp_path, read_only_paths=[], writable_paths=[]) + monkeypatch.setattr(backend, "_has_gui_launchd", AsyncMock(return_value=True)) + killed = Mock() + monkeypatch.setattr(backend, "_kill_process", killed) + calls = [] + control = None + started = asyncio.Event() + + async def launchctl(*args): + nonlocal control + calls.append(args[0]) + if args[0] == "bootstrap": + control = Path(args[-1]).parent + if args[0] == "kickstart" and mode == "kickstart_error": + raise RuntimeError("kickstart failed") + if args[0] == "bootout": + # Output is collected after launchd has reaped the process group. + (control / "stdout").write_bytes(b"output") + (control / "stderr").write_bytes(b"diagnostic") + return "12345" + + async def wait_status(*_args): + started.set() + if mode == "error": + raise RuntimeError("status failed") + if mode in ("timeout", "cancel"): + await asyncio.Event().wait() + return 7 if mode == "failure" else 0 + + monkeypatch.setattr(backend, "_launchctl", launchctl) + monkeypatch.setattr(backend, "_wait_status", wait_status) + task = asyncio.create_task(backend.exec_shell(["true"], timeout=0.01 if mode == "timeout" else None)) + if mode == "cancel": + await asyncio.wait_for(started.wait(), timeout=1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=1) + elif mode in ("error", "kickstart_error"): + with pytest.raises(RuntimeError, match="failed"): + await asyncio.wait_for(task, timeout=1) + else: + result = await asyncio.wait_for(task, timeout=1) + if mode == "timeout": + assert result == ExecResult(exit_code=-1, stdout=b"", stderr=b"timed out") + else: + assert result.exit_code == (7 if mode == "failure" else 0) + assert result.stdout == b"output" and result.stderr == b"diagnostic" + assert calls == ["bootstrap", "kickstart", "bootout"] + assert not control.exists() + assert killed.call_count == int(mode in ("timeout", "cancel", "error")) + + +@pytest.mark.asyncio +async def test_python_timeout_preserves_exit_status_and_stderr(tmp_path): + """The shared executor retains the existing unsandboxed Python metadata.""" + step = PythonExecuteStep(app_context=ApplicationContext(workspace_dir=str(tmp_path))) + response = await step( + code="import sys, time; print('diagnostic', file=sys.stderr, flush=True); time.sleep(60)", + python_timeout=0.3, + ) + assert response.success is False + assert response.answer == "Python execution timed out after 0.3s" + assert response.metadata == { + "returncode": -9 if os.name == "posix" else 1, + "stderr": "diagnostic\n", + "python_timeout": 0.3, + } diff --git a/tests/unit/test_process_output_limits.py b/tests/unit/test_process_output_limits.py new file mode 100644 index 00000000..55acb81c --- /dev/null +++ b/tests/unit/test_process_output_limits.py @@ -0,0 +1,156 @@ +"""Execution output stays bounded across pipe and launchd backends.""" + +# pylint: disable=protected-access + +import asyncio +import os +import plistlib +import stat +import sys +from pathlib import Path +from unittest.mock import AsyncMock + +import pytest + +from reme.components.application_context import ApplicationContext +from reme.steps.common.python_execute import PythonExecuteStep +from reme.utils.file_sandbox.path_bubble_backend import PathBubbleBackend +from reme.utils.file_sandbox.path_seatbelt_backend import PathSeatbeltBackend +from reme.utils.local_backend import LocalProcessBackend, OUTPUT_LIMIT_BYTES + + +@pytest.mark.asyncio +@pytest.mark.parametrize("backend_name", ["local", "bubble", "seatbelt"]) +@pytest.mark.parametrize("stream", ["stdout", "stderr", "both"]) +@pytest.mark.parametrize("extra", [0, 1]) +async def test_pipe_output_boundary(tmp_path, monkeypatch, backend_name, stream, extra): + """Each stream has its own budget; crossing either budget fails execution.""" + local = LocalProcessBackend() + backend = local + if backend_name != "local": + backend_class = PathBubbleBackend if backend_name == "bubble" else PathSeatbeltBackend + backend = backend_class(cwd=tmp_path, read_only_paths=[], writable_paths=[]) + # Exercise the shared collector with a real process, without requiring + # kernel sandbox support or changing the host launchd domain. + monkeypatch.setattr(backend, "_start_process", local._start_process) + if backend_name == "seatbelt": + monkeypatch.setattr(backend, "_has_gui_launchd", AsyncMock(return_value=False)) + size = OUTPUT_LIMIT_BYTES + extra + code = "import os;" + ";".join( + f"os.write({fd}, b'x' * {size})" for fd in (1, 2) if stream in ("both", "stdout" if fd == 1 else "stderr") + ) + if backend_name == "seatbelt": + with pytest.warns(RuntimeWarning, match="No macOS GUI launchd domain"): + result = await backend.exec_shell([sys.executable, "-c", code], timeout=5) + else: + result = await backend.exec_shell([sys.executable, "-c", code], timeout=5) + assert len(result.stdout) <= OUTPUT_LIMIT_BYTES + assert len(result.stderr) <= OUTPUT_LIMIT_BYTES + if extra: + assert not result.ok() + assert b"Output limit exceeded:" in result.stderr + else: + assert result.ok(), result.stderr + assert result.stdout == (b"x" * size if stream in ("stdout", "both") else b"") + assert result.stderr == (b"x" * size if stream in ("stderr", "both") else b"") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fd", [1, 2]) +async def test_continuous_output_stops_process(tmp_path, fd): + """Overflow terminates a live writer before it can perform a later write.""" + marker = tmp_path / "late.txt" + code = ( + "import os,time;from pathlib import Path;" + f"os.write({fd}, b'x' * {OUTPUT_LIMIT_BYTES + 1});" + f"time.sleep(1);Path({str(marker)!r}).write_text('late')" + ) + backend = LocalProcessBackend() + process = await backend._start_process([sys.executable, "-c", code], str(tmp_path)) + result = await asyncio.wait_for(backend._communicate(process), timeout=3) + assert not result.ok() + assert b"Output limit exceeded:" in result.stderr + assert process.returncode is not None + assert not marker.exists() + + +@pytest.mark.asyncio +async def test_python_step_reports_output_overflow(tmp_path): + """Python callers receive bounded partial output and failure metadata.""" + step = PythonExecuteStep(app_context=ApplicationContext(workspace_dir=str(tmp_path))) + response = await step(code=f"import os;os.write(1, b'x' * {OUTPUT_LIMIT_BYTES + 1})") + assert response.success is False + assert len(response.answer) == OUTPUT_LIMIT_BYTES + assert "Output limit exceeded: stdout" in response.metadata["stderr"] + assert response.metadata["returncode"] != 0 + + +@pytest.mark.asyncio +@pytest.mark.skipif(os.name != "posix", reason="launchd uses POSIX FIFOs") +@pytest.mark.parametrize("stream", ["stdout", "stderr", "both"]) +@pytest.mark.parametrize("extra", [0, 1]) +async def test_launchd_output_boundary(tmp_path, monkeypatch, stream, extra): + """Real FIFO writers obey each stream's budget and are reaped on overflow.""" + backend = PathSeatbeltBackend(cwd=tmp_path, read_only_paths=[], writable_paths=[]) + monkeypatch.setattr("reme.utils.file_sandbox.path_seatbelt_backend.sys.platform", "darwin") + monkeypatch.setattr("reme.utils.file_sandbox.path_seatbelt_backend.shutil.which", lambda _: "sandbox-exec") + monkeypatch.setattr(backend, "_has_gui_launchd", AsyncMock(return_value=True)) + size = OUTPUT_LIMIT_BYTES + extra + code = "import os,time;" + ";".join( + f"os.write({fd}, b'x' * {size})" for fd in (1, 2) if stream in ("both", "stdout" if fd == 1 else "stderr") + ) + if extra: + code += ";time.sleep(60)" + job = {} + process = None + control = None + calls = [] + + async def launchctl(*args): + nonlocal job, process, control + calls.append(args[0]) + if args[0] == "bootstrap": + control = Path(args[-1]).parent + with open(args[-1], "rb") as source: + job = plistlib.load(source) + for key in ("StandardOutPath", "StandardErrorPath"): + assert stat.S_ISFIFO(os.stat(job[key]).st_mode) + elif args[0] == "kickstart": + with open(job["StandardOutPath"], "wb") as stdout, open(job["StandardErrorPath"], "wb") as stderr: + process = await asyncio.create_subprocess_exec( + sys.executable, + "-c", + code, + stdout=stdout, + stderr=stderr, + start_new_session=True, + ) + return str(process.pid) + elif args[0] == "bootout": + if process.returncode is None: + process.kill() + await process.wait() + return "" + + async def wait_status(*_args): + return await process.wait() + + monkeypatch.setattr(backend, "_launchctl", launchctl) + monkeypatch.setattr(backend, "_wait_status", wait_status) + try: + result = await asyncio.wait_for(backend.exec_shell(["test-command"], timeout=5), timeout=7) + assert len(result.stdout) <= OUTPUT_LIMIT_BYTES + assert len(result.stderr) <= OUTPUT_LIMIT_BYTES + assert result.ok() is (extra == 0) + if extra: + assert b"Output limit exceeded:" in result.stderr + else: + assert result.stdout == (b"x" * size if stream in ("stdout", "both") else b"") + assert result.stderr == (b"x" * size if stream in ("stderr", "both") else b"") + assert calls == ["bootstrap", "kickstart", "bootout"] + assert process.returncode is not None + assert not control.exists() + finally: + if process is not None and process.returncode is None: + process.kill() + await process.wait() diff --git a/tests/unit/test_sandbox_startup_cancellation.py b/tests/unit/test_sandbox_startup_cancellation.py new file mode 100644 index 00000000..88783383 --- /dev/null +++ b/tests/unit/test_sandbox_startup_cancellation.py @@ -0,0 +1,164 @@ +"""Cancellation during startup must not abandon OS-owned sandbox resources.""" + +# pylint: disable=protected-access + +import asyncio +import errno +import sys +from pathlib import Path +from unittest.mock import AsyncMock + +import pytest + +from reme.utils.file_sandbox.path_bubble_backend import PathBubbleBackend +from reme.utils.file_sandbox.path_seatbelt_backend import LaunchctlError, PathSeatbeltBackend + + +@pytest.mark.asyncio +@pytest.mark.parametrize("late_registration", [False, True]) +@pytest.mark.parametrize("cancel_twice", [False, True]) +async def test_cancelled_bootstrap_waits_and_unregisters(tmp_path, monkeypatch, late_registration, cancel_twice): + """Unregister after a pending bootstrap settles, including a late successful reply.""" + backend = PathSeatbeltBackend(cwd=tmp_path, read_only_paths=[], writable_paths=[]) + monkeypatch.setattr("reme.utils.file_sandbox.path_seatbelt_backend.sys.platform", "darwin") + monkeypatch.setattr("reme.utils.file_sandbox.path_seatbelt_backend.shutil.which", lambda _: "sandbox-exec") + monkeypatch.setattr(backend, "_has_gui_launchd", AsyncMock(return_value=True)) + entered = asyncio.Event() + reply = asyncio.Event() + bootout_entered = asyncio.Event() + bootout_reply = asyncio.Event() + services = set() + calls = [] + control = None + + async def launchctl(*args): + nonlocal control + calls.append(args[0]) + if args[0] == "bootstrap": + control = Path(args[-1]).parent + if not late_registration: + services.add("sandbox") + entered.set() + await reply.wait() + services.add("sandbox") + elif args[0] == "bootout": + assert reply.is_set(), "bootout must not race an unfinished registration" + assert control.exists(), "keep control files until the service is removed" + bootout_entered.set() + await bootout_reply.wait() + services.clear() + else: + pytest.fail(f"Unexpected launchctl call: {args[0]}") + return "" + + monkeypatch.setattr(backend, "_launchctl", launchctl) + task = asyncio.create_task(backend.exec_shell(["true"])) + await asyncio.wait_for(entered.wait(), timeout=1) + task.cancel() + await asyncio.sleep(0) + assert not task.done() + reply.set() + if cancel_twice: + await asyncio.wait_for(bootout_entered.wait(), timeout=1) + task.cancel() + await asyncio.sleep(0) + assert not task.done() + bootout_reply.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=1) + assert calls == ["bootstrap", "bootout"] + assert not services + assert not control.exists() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("accepted", [False, True]) +async def test_failed_bootstrap_attempts_unregister(tmp_path, monkeypatch, accepted): + """A failed reply may still hide an accepted registration; absence alone is benign.""" + backend = PathSeatbeltBackend(cwd=tmp_path, read_only_paths=[], writable_paths=[]) + plist = tmp_path / "job.plist" + calls = [] + services = set() + + async def launchctl(*args): + calls.append(args[0]) + if args[0] == "bootstrap": + if accepted: + services.add("sandbox") + raise LaunchctlError("registration reply failed", 5) + if not services: + raise LaunchctlError("service not found", errno.ESRCH) + services.clear() + return "" + + monkeypatch.setattr(backend, "_launchctl", launchctl) + with pytest.raises(LaunchctlError, match="registration reply failed"): + await backend._bootstrap("gui/501", plist, "gui/501/sandbox") + assert calls == ["bootstrap", "bootout"] + assert not services + + +@pytest.mark.asyncio +async def test_bootstrap_rollback_reports_unregistration_failure(tmp_path, monkeypatch): + """Permission and other cleanup failures must not be mistaken for a missing service.""" + backend = PathSeatbeltBackend(cwd=tmp_path, read_only_paths=[], writable_paths=[]) + launchctl = AsyncMock(side_effect=[LaunchctlError("bootstrap failed", 5), LaunchctlError("bootout denied", 1)]) + monkeypatch.setattr(backend, "_launchctl", launchctl) + with pytest.raises(LaunchctlError, match="bootout denied"): + await backend._bootstrap("gui/501", tmp_path / "job.plist", "gui/501/sandbox") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("memory_limit", [0, 1]) +@pytest.mark.parametrize("spawn_fails", [False, True]) +async def test_cancelled_bubble_start_reaps_before_removing_cgroup(tmp_path, monkeypatch, memory_limit, spawn_fails): + """A delayed process handle is reclaimed before its cgroup is removed.""" + backend = PathBubbleBackend(cwd=tmp_path, read_only_paths=[], writable_paths=[], memory_limit_mib=memory_limit) + monkeypatch.setattr("reme.utils.file_sandbox.path_bubble_backend.sys.platform", "linux") + monkeypatch.setattr("reme.utils.file_sandbox.path_bubble_backend.shutil.which", lambda _: "bwrap") + group = tmp_path / "cgroup" + + def create_group(): + group.mkdir() + return group + + monkeypatch.setattr(backend, "_create_cgroup", create_group) + entered = asyncio.Event() + reply = asyncio.Event() + processes = [] + create_process = asyncio.create_subprocess_exec + + async def delayed_spawn(*_args, **kwargs): + if not spawn_fails: + processes.append(await create_process(sys.executable, "-c", "import time;time.sleep(60)", **kwargs)) + entered.set() + await reply.wait() + if spawn_fails: + raise FileNotFoundError("spawn failed after cancellation") + return processes[0] + + monkeypatch.setattr("reme.utils.local_backend.asyncio.create_subprocess_exec", delayed_spawn) + task = asyncio.create_task(backend._start_process(["true"], str(tmp_path))) + try: + await asyncio.wait_for(entered.wait(), timeout=2) + task.cancel() + await asyncio.sleep(0) + task.cancel() + await asyncio.sleep(0) + assert not task.done() + assert group.exists() is bool(memory_limit) + reply.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=2) + assert all(process.returncode is not None for process in processes) + assert not group.exists() + assert not backend._process_cgroups + finally: + reply.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + for process in processes: + if process.returncode is None: + process.kill() + await process.communicate()