diff --git a/apodex/_shell_supervisor.py b/apodex/_shell_supervisor.py new file mode 100644 index 0000000..aa040d4 --- /dev/null +++ b/apodex/_shell_supervisor.py @@ -0,0 +1,105 @@ +"""Linux subprocess supervisor for one local shell command. + +The supervisor is a child subreaper. Once the shell's process group is killed, +descendants that created another session are adopted here and can be killed +and reaped before the caller reports a timeout. It stays alive until the +caller has seen both output pipes close, so an exited shell cannot strand a +background writer outside the cleanup path. +""" + +from __future__ import annotations + +import contextlib +import ctypes +import os +import select +import signal +import sys +import time + +_PR_SET_CHILD_SUBREAPER = 36 +_CLEANUP_SECONDS = 5.0 + + +def _children() -> list[int]: + path = f"/proc/self/task/{os.getpid()}/children" + with open(path, encoding="ascii") as stream: + return [int(pid) for pid in stream.read().split()] + + +def _reap() -> dict[int, int]: + exited: dict[int, int] = {} + while True: + try: + pid, status = os.waitpid(-1, os.WNOHANG) + except ChildProcessError: + break + if pid == 0: + break + exited[pid] = os.waitstatus_to_exitcode(status) + return exited + + +def _kill_descendants(shell_pid: int) -> None: + with contextlib.suppress(ProcessLookupError): + os.killpg(shell_pid, signal.SIGKILL) + deadline = time.monotonic() + _CLEANUP_SECONDS + while time.monotonic() < deadline: + _reap() + children = _children() + if not children: + return + for pid in children: + with contextlib.suppress(ProcessLookupError): + os.kill(pid, signal.SIGKILL) + time.sleep(0.01) + _reap() + + +def main(control_fd: int, command: str) -> int: + libc = ctypes.CDLL(None, use_errno=True) + if libc.prctl(_PR_SET_CHILD_SUBREAPER, 1, 0, 0, 0) != 0: + raise OSError(ctypes.get_errno(), "cannot enable child subreaper") + # A missing procfs would make escaped descendants invisible. Do not run + # an unsupervised command when the promised cleanup is unavailable. + _children() + + shell_pid = os.fork() + if shell_pid == 0: + os.close(control_fd) + os.setsid() + os.execl("/bin/sh", "sh", "-c", command) + + # Only the shell and its descendants keep the output pipes open. The + # caller can therefore detect EOF even while this supervisor waits for a + # success or kill instruction on its separate control pipe. + os.close(1) + os.close(2) + success_requested = False + shell_status: int | None = None + try: + while True: + exited = _reap() + if shell_pid in exited: + shell_status = exited[shell_pid] + try: + os.write(control_fd, f"{shell_status}\n".encode("ascii")) + except OSError: + _kill_descendants(shell_pid) + return 124 + if success_requested and shell_status is not None: + return shell_status + ready, _, _ = select.select([control_fd], [], [], 0.02) + if ready: + instruction = os.read(control_fd, 1) + if instruction == b"S": + success_requested = True + else: # K, EOF, or an invalid instruction all stop the tree. + _kill_descendants(shell_pid) + return 124 + finally: + os.close(control_fd) + + +if __name__ == "__main__": + sys.exit(main(int(sys.argv[1]), sys.argv[2])) diff --git a/apodex/sandbox.py b/apodex/sandbox.py index 4b0a933..be7e4a2 100644 --- a/apodex/sandbox.py +++ b/apodex/sandbox.py @@ -41,9 +41,12 @@ from __future__ import annotations import asyncio +import contextlib import logging import os import shlex +import signal +import socket import sys from dataclasses import dataclass from pathlib import Path @@ -253,13 +256,104 @@ async def run_shell( getattr(result, "stderr", "") or "", ) + if sys.platform.startswith("linux"): + return await _run_supervised_shell(command, cwd, timeout) + return await _run_group_shell(command, cwd, timeout) + + +# Run by file path under -P so neither cwd nor this package's directory +# lands on sys.path: a workspace's own apodex/ or ctypes.py must not +# shadow the supervisor or the stdlib modules it imports. +_SUPERVISOR = str(Path(__file__).with_name("_shell_supervisor.py")) + + +async def _run_supervised_shell( + command: str, cwd: str, timeout: int, +) -> tuple[int, str, str]: + """Run under a Linux child subreaper, including escaped descendants.""" + parent_control, child_control = socket.socketpair() + try: + proc = await asyncio.create_subprocess_exec( + sys.executable, "-P", _SUPERVISOR, + str(child_control.fileno()), command, + cwd=cwd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + pass_fds=(child_control.fileno(),), + start_new_session=True, + ) + except BaseException: + parent_control.close() + raise + finally: + child_control.close() + + parent_control.setblocking(False) + + async def collect() -> tuple[bytes, bytes, bytes]: + assert proc.stdout is not None and proc.stderr is not None + loop = asyncio.get_running_loop() + out, err, status = await asyncio.gather( + proc.stdout.read(), proc.stderr.read(), + loop.sock_recv(parent_control, 32), + ) + return out, err, status + + try: + try: + out, err, status = await asyncio.wait_for(collect(), timeout=timeout) + except BaseException: + # The supervisor kills its shell group, then every adopted child. + # Closing the socket is another kill instruction if the write fails. + with contextlib.suppress(OSError): + parent_control.send(b"K") + with contextlib.suppress(TimeoutError): + await asyncio.wait_for(proc.communicate(), timeout=5) + raise + if not status: + # The supervisor exits before running the shell when it cannot + # supervise it (no subreaper, no /proc children list). Its exit + # code is not the command's, so report it as unavailable. + await asyncio.wait_for(proc.wait(), timeout=5) + detail = err.decode("utf-8", "replace").strip().splitlines() + raise RuntimeError( + "native shell supervisor is unavailable on this system " + f"(exit {proc.returncode}); command was not run" + + (f": {detail[-1]}" if detail else "") + ) + with contextlib.suppress(OSError): + parent_control.send(b"S") + await asyncio.wait_for(proc.wait(), timeout=5) + return ( + int(status), + out.decode("utf-8", "replace"), + err.decode("utf-8", "replace"), + ) + finally: + parent_control.close() + + +async def _run_group_shell( + command: str, cwd: str, timeout: int, +) -> tuple[int, str, str]: + """Use process-group cleanup on platforms without Linux subreapers.""" proc = await asyncio.create_subprocess_shell( command, cwd=cwd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, + start_new_session=True, ) - out, err = await asyncio.wait_for(proc.communicate(), timeout=timeout) + completed = False + try: + out, err = await asyncio.wait_for(proc.communicate(), timeout=timeout) + completed = True + finally: + if not completed: + with contextlib.suppress(OSError): + os.killpg(proc.pid, signal.SIGKILL) + with contextlib.suppress(TimeoutError): + await asyncio.wait_for(proc.wait(), timeout=5) return ( proc.returncode or 0, out.decode("utf-8", "replace"), diff --git a/apodex/tests/test_native.py b/apodex/tests/test_native.py index a04c94a..b73be26 100644 --- a/apodex/tests/test_native.py +++ b/apodex/tests/test_native.py @@ -2,8 +2,12 @@ import asyncio import os +import shlex +import sys from pathlib import Path +import pytest + from apodex import cli, docker, sandbox from apodex.native import prepare_native_runtime from apodex.sandbox import BWRAP, CONTAINER, NATIVE, Strategy, resolve_strategy @@ -133,6 +137,102 @@ def test_native_strategy_is_explicitly_not_os_isolated() -> None: assert "not an OS sandbox" in strategy.describe() +@pytest.mark.parametrize("command", [ + "(sleep 2; touch marker) & wait", + "(sleep 2; touch marker) & exit 0", +]) +@pytest.mark.parametrize("interruption", ["timeout", "cancellation"]) +async def test_run_shell_reaps_background_children( + tmp_path, command, interruption, +) -> None: + strategy = Strategy(NATIVE, "test") + if interruption == "timeout": + with pytest.raises(TimeoutError): + await sandbox.run_shell(command, str(tmp_path), 1, strategy) + else: + task = asyncio.create_task(sandbox.run_shell( + command, str(tmp_path), 10, strategy, + )) + await asyncio.sleep(1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + await asyncio.sleep(2) + assert not (tmp_path / "marker").exists() + + +@pytest.mark.skipif(not sys.platform.startswith("linux"), reason="Linux subreaper") +@pytest.mark.parametrize("shell_exits", [False, True]) +@pytest.mark.parametrize("interruption", ["timeout", "cancellation"]) +async def test_run_shell_reaps_descendants_that_create_new_sessions( + tmp_path, shell_exits, interruption, +) -> None: + child = ( + "import os,time; os.setsid(); open('started','w').close(); " + "time.sleep(2.5); open('marker','w').close()" + ) + command = f"{shlex.quote(sys.executable)} -c {shlex.quote(child)} & " + command += "exit 0" if shell_exits else "wait" + strategy = Strategy(NATIVE, "test") + if interruption == "timeout": + with pytest.raises(TimeoutError): + await sandbox.run_shell(command, str(tmp_path), 1, strategy) + else: + task = asyncio.create_task(sandbox.run_shell( + command, str(tmp_path), 10, strategy, + )) + await asyncio.sleep(1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert (tmp_path / "started").exists() + await asyncio.sleep(2) + assert not (tmp_path / "marker").exists() + + +async def test_run_shell_preserves_output_and_exit_code(tmp_path) -> None: + result = await sandbox.run_shell( + "printf out; printf err >&2; exit 3", str(tmp_path), 5, + Strategy(NATIVE, "test"), + ) + assert result == (3, "out", "err") + + +async def test_run_shell_preserves_signal_exit_code(tmp_path) -> None: + result = await sandbox.run_shell( + "kill -TERM $$", str(tmp_path), 5, Strategy(NATIVE, "test"), + ) + assert result[0] == -15 + + +@pytest.mark.skipif(not sys.platform.startswith("linux"), reason="Linux subreaper") +async def test_run_shell_supervisor_ignores_workspace_modules(tmp_path) -> None: + (tmp_path / "apodex").mkdir() + (tmp_path / "apodex" / "__init__.py").write_text("raise SystemExit(7)\n") + for name in ("ctypes", "select"): + (tmp_path / f"{name}.py").write_text("raise SystemExit(7)\n") + result = await sandbox.run_shell( + "printf ok", str(tmp_path), 5, Strategy(NATIVE, "test"), + ) + assert result == (0, "ok", "") + + +@pytest.mark.skipif(not sys.platform.startswith("linux"), reason="Linux subreaper") +async def test_run_shell_reports_unavailable_supervisor( + tmp_path, monkeypatch, +) -> None: + fake = tmp_path / "supervisor.py" + fake.write_text( + "raise OSError(95, 'cannot enable child subreaper')\n" + ) + monkeypatch.setattr(sandbox, "_SUPERVISOR", str(fake)) + with pytest.raises(RuntimeError, match="command was not run.*subreaper"): + await sandbox.run_shell( + "touch marker", str(tmp_path), 5, Strategy(NATIVE, "test"), + ) + assert not (tmp_path / "marker").exists() + + def test_nonroot_native_current_sandbox_skips_tool_user_warning( tmp_path, monkeypatch, caplog, ) -> None: