fix(apodex): reap detached shell descendants after timeout

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]))

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"),

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: