From 9ef0037490ee8748e6958a36c0e3d04d7ea6eca0 Mon Sep 17 00:00:00 2001 From: Aly Raffauf Date: Sat, 27 Jun 2026 23:36:20 -0400 Subject: [PATCH] jail: native-async foreground command execution --- tartarus/agent_loop.py | 17 ++-- tartarus/background.py | 9 +- tartarus/broker.py | 26 +++--- tartarus/cli.py | 2 +- tartarus/jail.py | 192 ++++++++++++++++++--------------------- tests/test_agent_loop.py | 51 +++++------ tests/test_broker.py | 52 ++++++----- tests/test_jail.py | 88 +++++++++++------- 8 files changed, 221 insertions(+), 216 deletions(-) diff --git a/tartarus/agent_loop.py b/tartarus/agent_loop.py index 72ce1dd..2edba09 100644 --- a/tartarus/agent_loop.py +++ b/tartarus/agent_loop.py @@ -12,7 +12,6 @@ mid-stream or mid-tool always leaves a valid transcript. """ import asyncio -import threading from dataclasses import dataclass from tartarus.broker import Broker @@ -102,16 +101,14 @@ class AgentLoop: messages.extend(self._provider.tool_result_messages(results)) async def _run_tool(self, call: ToolCall): - event_loop = asyncio.get_running_loop() output_queue: asyncio.Queue[str] = asyncio.Queue() - cancellation = threading.Event() + # The broker runs on this loop, so streamed output lands on the queue + # directly. def emit_output(text: str) -> None: - event_loop.call_soon_threadsafe(output_queue.put_nowait, text) + output_queue.put_nowait(text) - worker = asyncio.create_task( - asyncio.to_thread(self._broker.handle, call, emit_output, cancellation) - ) + worker = asyncio.create_task(self._broker.handle(call, emit_output)) pending_get: asyncio.Task[str] | None = None try: while True: @@ -134,10 +131,12 @@ class AgentLoop: yield worker.result() return except (asyncio.CancelledError, GeneratorExit): - cancellation.set() + # Cancel the worker and await it shielded so the jail tears its + # process group down before the turn unwinds. + worker.cancel() try: await asyncio.shield(worker) - except Exception: + except (asyncio.CancelledError, Exception): pass raise finally: diff --git a/tartarus/background.py b/tartarus/background.py index e1e3334..c349ff7 100644 --- a/tartarus/background.py +++ b/tartarus/background.py @@ -5,10 +5,8 @@ handed here. The registry assigns it a short id (`bg-1`, `bg-2`, …), tails its combined output, and reports liveness to the control-plane capabilities (`bg_status` / `bg_output` / `bg_stop`). -This object bridges two worlds: the **sync broker** registers a task from a -worker thread (`Broker.handle` runs under `asyncio.to_thread`), while the -**async loop** owns the completion notifications. `register` is therefore -thread-safe and schedules the per-task monitor onto the captured event loop; +The broker registers tasks inline on the event loop thread. The registry +spawns a per-task asyncio monitor for completion notifications. `status`/`output` refresh liveness directly so they work even with no loop attached (e.g. unit tests that poll). """ @@ -84,8 +82,7 @@ class BackgroundRegistry: ) self._tasks[bg_id] = task if self._loop is not None: - # register runs in the broker's worker thread; hop to the loop thread. - self._loop.call_soon_threadsafe(self._start_monitor, task) + self._loop.call_soon(self._start_monitor, task) return bg_id # --- control-plane operations (called by the broker) -------------------- diff --git a/tartarus/broker.py b/tartarus/broker.py index 2f6d2dc..3df1672 100644 --- a/tartarus/broker.py +++ b/tartarus/broker.py @@ -5,10 +5,10 @@ human), then run inside a bwrap jail via the injected JailBuilder. Argument validation and shell-safe interpolation are permanent and shared by all phases. """ +import asyncio import shlex from collections import defaultdict from collections.abc import Callable -from threading import Event from tartarus.audit import AuditEvent, AuditSink, NullAuditLog from tartarus.background import BackgroundError, BackgroundRegistry @@ -91,11 +91,10 @@ class Broker: self._audit = audit if audit is not None else NullAuditLog() self._registry = registry - def handle( + async def handle( self, call: ToolCall, output_callback: Callable[[str], None] | None = None, - cancellation: Event | None = None, ) -> ToolResult: if call.argument_error: return self._audited_error( @@ -121,25 +120,27 @@ class Broker: decision=_broker_decision(f"invalid arguments: {validation_error}"), ) - return self._run( + return await self._run( call.id, capability, call.arguments, output_callback=output_callback, - cancellation=cancellation, ) - def _run( + async def _run( self, call_id: str, capability: Capability, arguments: dict, output_callback: Callable[[str], None] | None = None, - cancellation: Event | None = None, ) -> ToolResult: command = interpolate(capability.runner, arguments) - decision = self._policy.decide(capability, arguments, command) + # The policy may prompt the human (blocking input); keep it off the event + # loop so an ask-* approval never stalls concurrent streaming/cancellation. + decision = await asyncio.to_thread( + self._policy.decide, capability, arguments, command + ) if not decision.allowed: return self._finish( AuditEvent( @@ -160,17 +161,16 @@ class Broker: return self._run_background( call_id, capability, arguments, command, decision ) - return self._run_command( + return await self._run_command( call_id, capability, arguments, command, decision, output_callback=output_callback, - cancellation=cancellation, ) - def _run_command( + async def _run_command( self, call_id: str, capability: Capability, @@ -178,18 +178,16 @@ class Broker: command: str, decision: Decision, output_callback: Callable[[str], None] | None = None, - cancellation: Event | None = None, ) -> ToolResult: # A capability runs unbounded unless it declares its own timeout; # None means "wait forever" in the jail's process wait loop. try: spec = self._jail.build(capability.grants) - result = self._jail.exec( + result = await self._jail.exec( spec, command, timeout=capability.timeout, output_callback=output_callback, - cancellation=cancellation, ) except JailError as error: return self._jail_error( diff --git a/tartarus/cli.py b/tartarus/cli.py index 27a103c..2f633e8 100644 --- a/tartarus/cli.py +++ b/tartarus/cli.py @@ -381,7 +381,7 @@ async def _async_main(argv: list[str]) -> int: shell_closure=manifest.shell_closure, ) policy = PolicyEngine(headless=config.headless) - # Background tasks: the registry bridges the sync broker to this async loop. + # Background tasks: the registry monitors detached runs on this async loop. # Completion notices land on `notices`, which the run drains to inject # follow-up turns; `shutdown_all` reaps every task on exit. notices: asyncio.Queue[Notice] = asyncio.Queue() diff --git a/tartarus/jail.py b/tartarus/jail.py index d03dc5f..4abc9a3 100644 --- a/tartarus/jail.py +++ b/tartarus/jail.py @@ -10,17 +10,15 @@ plain raw-socket containment is a later namespace/firewall step. Unrestricted grants skip bwrap entirely after policy approval, but still use the shell PATH. """ +import asyncio import os import signal import shlex import shutil import subprocess -import threading -import time import uuid from collections.abc import Callable from dataclasses import dataclass, field -from queue import Empty, Queue from typing import Literal from pydantic import ConfigDict, ValidationError, field_validator @@ -128,31 +126,29 @@ class JailBuilder: except ValidationError as error: raise JailError(str(error)) from error - def exec( + async def exec( self, spec: JailSpec, command: str, timeout: int | None = DEFAULT_JAIL_TIMEOUT_SECONDS, output_callback: Callable[[str], None] | None = None, - cancellation: threading.Event | None = None, ) -> ExecResult: if spec.unrestricted: - return self._exec_unrestricted( - spec, command, timeout, output_callback, cancellation + return await self._exec_unrestricted( + spec, command, timeout, output_callback ) if spec.network == "proxy": with self._proxy_factory(spec.allowed_hosts) as proxy: - result = self._exec_argv( + result = await self._exec_argv( self._bwrap_argv(spec, command, proxy.url), timeout, output_callback, - cancellation, ) return _append_stderr(result, proxy.summary()) - return self._exec_argv( - self._bwrap_argv(spec, command), timeout, output_callback, cancellation + return await self._exec_argv( + self._bwrap_argv(spec, command), timeout, output_callback ) def exec_background(self, spec: JailSpec, command: str) -> BackgroundHandle: @@ -211,44 +207,40 @@ class JailBuilder: os.makedirs(bg_dir, exist_ok=True) return os.path.join(bg_dir, f"{uuid.uuid4().hex}.log") - def _exec_argv( + async def _exec_argv( self, argv: list[str], timeout: int | None, output_callback: Callable[[str], None] | None = None, - cancellation: threading.Event | None = None, ) -> ExecResult: try: - proc = subprocess.Popen( - argv, + proc = await asyncio.create_subprocess_exec( + *argv, env={}, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - text=True, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, start_new_session=True, # isolate process group for timeout kills ) except FileNotFoundError as missing: raise JailError(f"jail runtime not found: {missing}") from missing - return _wait_for_process(proc, timeout, output_callback, cancellation) + return await _wait_for_process(proc, timeout, output_callback) - def _exec_unrestricted( + async def _exec_unrestricted( self, spec: JailSpec, command: str, timeout: int | None, output_callback: Callable[[str], None] | None = None, - cancellation: threading.Event | None = None, ) -> ExecResult: env = {"PATH": self._compose_path(spec), **spec.base_env} try: - proc = subprocess.Popen( - shlex.split(command), + proc = await asyncio.create_subprocess_exec( + *shlex.split(command), cwd=spec.work_tree, env=env, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - text=True, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, start_new_session=True, ) except FileNotFoundError as missing: @@ -256,7 +248,7 @@ class JailBuilder: f"unrestricted command runtime not found: {missing}" ) from missing - return _wait_for_process(proc, timeout, output_callback, cancellation) + return await _wait_for_process(proc, timeout, output_callback) def _bwrap_argv( self, @@ -352,112 +344,100 @@ class JailBuilder: return ":".join([spec.shell_path, *spec.extra_path]) -def _wait_for_process( - proc: subprocess.Popen, +async def _wait_for_process( + proc: asyncio.subprocess.Process, timeout: int | None, output_callback: Callable[[str], None] | None, - cancellation: threading.Event | None, ) -> ExecResult: - output_queue: Queue[tuple[str, str | None]] = Queue() stdout_parts: list[str] = [] stderr_parts: list[str] = [] - threads = [ - _start_pipe_reader(proc.stdout, "stdout", output_queue), - _start_pipe_reader(proc.stderr, "stderr", output_queue), + # Drain both pipes concurrently so a chatty command can never deadlock on a + # full pipe buffer; each reader ends at EOF when the process closes the pipe. + pumps = [ + asyncio.create_task(_pump(proc.stdout, stdout_parts, output_callback)), + asyncio.create_task(_pump(proc.stderr, stderr_parts, output_callback)), ] - deadline = None if timeout is None else time.monotonic() + timeout - termination_error: str | None = None - - while proc.poll() is None: - _drain_output_queue(output_queue, stdout_parts, stderr_parts, output_callback) - if cancellation is not None and cancellation.is_set(): - termination_error = "command cancelled" - _terminate_process_group(proc) - break - if deadline is not None and time.monotonic() >= deadline: - termination_error = f"command timed out after {timeout}s" - _terminate_process_group(proc) - break - time.sleep(0.01) - - proc.wait() - for thread in threads: - thread.join() - _drain_output_queue(output_queue, stdout_parts, stderr_parts, output_callback) - - stdout = "".join(stdout_parts) - stderr = "".join(stderr_parts) - if termination_error is not None: - stderr = "\n".join(part for part in (stderr.strip(), termination_error) if part) - if stderr: - stderr += "\n" - return ExecResult(TIMEOUT_EXIT_CODE, stdout, stderr) - return ExecResult(proc.returncode, stdout, stderr) - - -def _start_pipe_reader(pipe, stream_name: str, output_queue: Queue): - def read_lines() -> None: - if pipe is None: - output_queue.put((stream_name, None)) - return - try: - for line in pipe: - output_queue.put((stream_name, line)) - finally: - pipe.close() - output_queue.put((stream_name, None)) - - thread = threading.Thread( - target=read_lines, - name=f"tartarus-{stream_name}-reader", - daemon=True, - ) - thread.start() - return thread + try: + async with asyncio.timeout(timeout): # timeout=None waits indefinitely + await asyncio.gather(*pumps) + await proc.wait() + except TimeoutError: + await _kill_and_drain(proc, pumps) + stderr = _with_note( + "".join(stderr_parts), f"command timed out after {timeout}s" + ) + return ExecResult(TIMEOUT_EXIT_CODE, "".join(stdout_parts), stderr) + except asyncio.CancelledError: + # The turn was aborted (Ctrl-C / task.cancel): tear the process tree down + # before unwinding, then let cancellation propagate to the agent loop. + await _kill_and_drain(proc, pumps) + raise + + code = proc.returncode + assert code is not None # set once proc.wait() has returned + return ExecResult(code, "".join(stdout_parts), "".join(stderr_parts)) -def _drain_output_queue( - output_queue: Queue[tuple[str, str | None]], - stdout_parts: list[str], - stderr_parts: list[str], + +async def _pump( + stream: asyncio.StreamReader | None, + parts: list[str], output_callback: Callable[[str], None] | None, ) -> None: + """Forward one pipe to `parts` line by line, mirroring to the callback.""" + if stream is None: + return while True: - try: - stream_name, text = output_queue.get_nowait() - except Empty: + line = await stream.readline() + if not line: return - if text is None: - continue - if stream_name == "stdout": - stdout_parts.append(text) - else: - stderr_parts.append(text) + text = line.decode(errors="replace") + parts.append(text) if output_callback is not None: output_callback(text) -def _terminate_process_group(proc: subprocess.Popen) -> None: +async def _kill_and_drain( + proc: asyncio.subprocess.Process, + pumps: list[asyncio.Task], +) -> None: + """Tear the process group down and let the reader coroutines finish.""" + await _terminate_process_group(proc) + await asyncio.gather(*pumps, return_exceptions=True) + + +async def _terminate_process_group(proc: asyncio.subprocess.Process) -> None: + """SIGTERM the whole session, escalating to SIGKILL if it lingers. + + `start_new_session=True` made the child its own session leader, so its pid is + the process-group id and one `killpg` reaches the entire tree. + """ try: os.killpg(proc.pid, signal.SIGTERM) except ProcessLookupError: - # The process group may already be gone; fall through so we still reap - # any stragglers below. - pass + return # already gone; nothing left to reap try: - proc.wait(timeout=1) - except subprocess.TimeoutExpired: + await asyncio.wait_for(proc.wait(), timeout=1) + except TimeoutError: try: os.killpg(proc.pid, signal.SIGKILL) except ProcessLookupError: pass + await proc.wait() + + +def _with_note(stderr: str, note: str) -> str: + combined = "\n".join(part for part in (stderr.strip(), note) if part) + return combined + "\n" if combined else "" def _append_stderr(result: ExecResult, message: str) -> ExecResult: - stderr = "\n".join(part for part in (result.stderr.strip(), message) if part) - if stderr: - stderr += "\n" - return ExecResult(result.code, result.stdout, stderr, network_summary=message) + return ExecResult( + result.code, + result.stdout, + _with_note(result.stderr, message), + network_summary=message, + ) def _store_bind_args(bind_paths: list[str]) -> list[str]: diff --git a/tests/test_agent_loop.py b/tests/test_agent_loop.py index 4a64b9f..225ed02 100644 --- a/tests/test_agent_loop.py +++ b/tests/test_agent_loop.py @@ -1,7 +1,6 @@ import asyncio import shlex import subprocess -import threading from collections.abc import AsyncIterator from typing import cast @@ -28,9 +27,7 @@ class LocalJail: def build(self, grant): return grant - def exec( - self, spec, command, timeout=None, output_callback=None, cancellation=None - ): + async def exec(self, spec, command, timeout=None, output_callback=None): completed = subprocess.run(shlex.split(command), capture_output=True, text=True) if output_callback is not None and completed.stdout: output_callback(completed.stdout) @@ -227,25 +224,25 @@ def test_cancel_mid_turn_leaves_a_valid_transcript(): assert messages == [{"role": "user", "content": "hi"}] -def test_aclose_mid_tool_sets_cancellation_and_leaves_transcript(): - """Closing the turn generator while a tool is running must signal the jail - to terminate instead of leaking the worker or committing a result.""" +def test_aclose_mid_tool_cancels_worker_and_leaves_transcript(): + """Closing the turn generator while a tool is running must cancel the worker + so the jail terminates, instead of leaking it or committing a result.""" manifest = echo_manifest() class CancellingBroker: def __init__(self): - self.cancelled = threading.Event() + self.cancelled = False - def handle(self, call, output_callback=None, cancellation=None): - # Emit a chunk so the caller can advance past ToolStarted and then - # close the generator while this worker is still running. + async def handle(self, call, output_callback=None): + # Emit a chunk so the caller can advance past ToolStarted, then block + # until the worker is cancelled by the closing turn. if output_callback is not None: output_callback("partial output") - # Block the worker until the loop has signalled cancellation, proving - # the aclose path propagated the cancellation event. - if cancellation is not None: - cancellation.wait(timeout=5) - self.cancelled.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + self.cancelled = True + raise return ToolResult(call.id, "should-not-appear", is_error=False) provider = ScriptedProvider( @@ -271,28 +268,30 @@ def test_aclose_mid_tool_sets_cancellation_and_leaves_transcript(): asyncio.run(run_and_cancel()) - # The closing handshake signalled cancellation and did not commit anything. - assert broker.cancelled.is_set() + # The closing handshake cancelled the worker and did not commit anything. + assert broker.cancelled assert messages == [{"role": "user", "content": "use echo"}] def test_task_cancel_mid_tool_terminates_worker_synchronously(): """The CLI cancels an in-flight turn with ``task.cancel()`` (not ``aclose``). That CancelledError chains through ``__anext__`` into ``_run_tool``'s own - ``await``, so the jail's cancellation event must be set — and the worker - torn down — before awaiting the cancelled task returns.""" + ``await``, which cancels the worker and awaits its teardown before the + cancelled task returns.""" manifest = echo_manifest() class CancellingBroker: def __init__(self): - self.cancelled = threading.Event() + self.cancelled = False - def handle(self, call, output_callback=None, cancellation=None): + async def handle(self, call, output_callback=None): if output_callback is not None: output_callback("partial output") - if cancellation is not None: - cancellation.wait(timeout=5) - self.cancelled.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + self.cancelled = True + raise return ToolResult(call.id, "should-not-appear", is_error=False) provider = ScriptedProvider( @@ -326,7 +325,7 @@ def test_task_cancel_mid_tool_terminates_worker_synchronously(): except asyncio.CancelledError: pass # Teardown is synchronous on this path: no extra loop pumping needed. - return broker.cancelled.is_set() + return broker.cancelled cancelled_on_return = asyncio.run(run_and_cancel()) diff --git a/tests/test_broker.py b/tests/test_broker.py index b26922b..9edf91f 100644 --- a/tests/test_broker.py +++ b/tests/test_broker.py @@ -1,3 +1,5 @@ +import asyncio + import pytest from tartarus.audit import AuditEvent @@ -34,13 +36,12 @@ class FakeJail(JailBuilder): self.built_grants.append(grant) return JailSpec(work_tree="", shell_path="", base_env={}) - def exec( + async def exec( self, spec: JailSpec, command: str, timeout: int | None = 0, output_callback=None, - cancellation=None, ) -> ExecResult: self.exec_commands.append(command) self.exec_timeouts.append(timeout) @@ -96,6 +97,11 @@ def _call(name, arguments, argument_error=None): ) +def _handle(broker, call): + """Drive the async broker to completion from a sync test.""" + return asyncio.run(broker.handle(call)) + + def _broker(jail=None, policy=None, audit=None): return Broker( echo_manifest(), @@ -109,7 +115,7 @@ def test_echo_round_trips_through_jail(): jail = FakeJail(result=ExecResult(0, "hello world", "")) broker = _broker(jail) - result = broker.handle(_call("echo", {"message": "hello world"})) + result = _handle(broker, _call("echo", {"message": "hello world"})) assert not result.is_error assert result.output == "hello world" @@ -121,7 +127,7 @@ def test_allowed_call_records_one_audit_event(): jail = FakeJail(result=ExecResult(0, "hello world", "")) broker = _broker(jail, audit=audit) - broker.handle(_call("echo", {"message": "hello world"})) + _handle(broker, _call("echo", {"message": "hello world"})) assert len(audit.events) == 1 event = audit.events[0] @@ -140,7 +146,7 @@ def test_nonzero_exit_is_reported_as_error(): jail = FakeJail(result=ExecResult(1, "", "boom")) broker = _broker(jail) - result = broker.handle(_call("echo", {"message": "x"})) + result = _handle(broker, _call("echo", {"message": "x"})) assert result.is_error assert "boom" in result.output @@ -152,7 +158,7 @@ def test_undeclared_timeout_runs_unbounded(): jail = FakeJail(result=ExecResult(0, "hi", "")) broker = _broker(jail) - broker.handle(_call("echo", {"message": "x"})) + _handle(broker, _call("echo", {"message": "x"})) assert jail.exec_timeouts == [None] @@ -170,7 +176,7 @@ def test_declared_timeout_reaches_the_jail(): jail = FakeJail(result=ExecResult(0, "ok", "")) broker = Broker(build_manifest({"run_tests": capability}), jail, PolicyEngine()) - broker.handle(_call("run_tests", {})) + _handle(broker, _call("run_tests", {})) assert jail.exec_timeouts == [300] @@ -179,7 +185,7 @@ def test_output_truncate_is_configurable(): jail = FakeJail(result=ExecResult(0, "abcdef", "")) broker = Broker(echo_manifest(), jail, PolicyEngine(), output_truncate=3) - result = broker.handle(_call("echo", {"message": "x"})) + result = _handle(broker, _call("echo", {"message": "x"})) assert result.output == "abc\n...(truncated)" @@ -189,7 +195,7 @@ def test_jail_error_is_surfaced_not_raised(): jail = FakeJail(error=JailError("network grants not implemented")) broker = _broker(jail, audit=audit) - result = broker.handle(_call("echo", {"message": "x"})) + result = _handle(broker, _call("echo", {"message": "x"})) assert result.is_error assert "jail error" in result.output @@ -203,7 +209,7 @@ def test_unknown_tool_never_reaches_the_jail(): jail = FakeJail() broker = _broker(jail, audit=audit) - result = broker.handle(_call("nope", {})) + result = _handle(broker, _call("nope", {})) assert result.is_error assert "unknown tool" in result.output @@ -219,7 +225,7 @@ def test_malformed_arguments_never_reach_the_jail(): jail = FakeJail() broker = _broker(jail, audit=audit) - result = broker.handle(_call("echo", {}, argument_error="bad JSON")) + result = _handle(broker, _call("echo", {}, argument_error="bad JSON")) assert result.is_error assert "invalid arguments" in result.output @@ -232,7 +238,7 @@ def test_missing_required_argument_never_reaches_the_jail(): jail = FakeJail() broker = _broker(jail) - result = broker.handle(_call("echo", {})) + result = _handle(broker, _call("echo", {})) assert result.is_error assert "missing required parameter 'message'" in result.output @@ -253,7 +259,7 @@ def test_deny_capability_never_reaches_the_jail(): } broker = Broker(build_manifest(capabilities), jail, PolicyEngine()) - result = broker.handle(_call("locked", {})) + result = _handle(broker, _call("locked", {})) assert result.is_error assert "denied" in result.output @@ -294,7 +300,7 @@ def test_ask_always_decline_denies_and_skips_the_jail(): audit=audit, ) - result = broker.handle(_call("run_command", {"command": "ls"})) + result = _handle(broker, _call("run_command", {"command": "ls"})) assert result.is_error assert "denied" in result.output @@ -313,7 +319,7 @@ def test_ask_always_approval_runs_command_in_jail(): PolicyEngine(prompt=lambda *_: True), ) - result = broker.handle(_call("run_command", {"command": "ls -la"})) + result = _handle(broker, _call("run_command", {"command": "ls -la"})) assert not result.is_error # The command is a single shell-escaped argument to bash -c. @@ -328,7 +334,7 @@ def test_unrestricted_decline_denies_and_skips_the_jail(): PolicyEngine(prompt=lambda *_: False), ) - result = broker.handle(_call("shell_escape", {"command": "cat /etc/passwd"})) + result = _handle(broker, _call("shell_escape", {"command": "cat /etc/passwd"})) assert result.is_error assert "denied" in result.output @@ -343,7 +349,7 @@ def test_unrestricted_approval_reaches_the_jail_builder(): PolicyEngine(prompt=lambda *_: True), ) - result = broker.handle(_call("shell_escape", {"command": "cat /etc/passwd"})) + result = _handle(broker, _call("shell_escape", {"command": "cat /etc/passwd"})) assert not result.is_error assert jail.built_grants == [Grant(unrestricted=True)] @@ -387,7 +393,7 @@ def test_background_capability_launches_and_registers(): registry=registry, ) - result = broker.handle(_call("run_bg", {"command": "sleep 1"})) + result = _handle(broker, _call("run_bg", {"command": "sleep 1"})) assert not result.is_error assert result.output == "started bg-1" @@ -400,7 +406,7 @@ def test_background_without_registry_reports_unavailable(): jail = FakeJail() broker = Broker(_background_manifest(), jail, PolicyEngine(prompt=lambda *_: True)) - result = broker.handle(_call("run_bg", {"command": "sleep 1"})) + result = _handle(broker, _call("run_bg", {"command": "sleep 1"})) assert result.is_error assert "not available" in result.output @@ -417,7 +423,7 @@ def test_control_capability_dispatches_to_registry_not_jail(): registry=registry, ) - result = broker.handle(_call("bg_output", {"task": "bg-1"})) + result = _handle(broker, _call("bg_output", {"task": "bg-1"})) assert not result.is_error assert result.output == "output bg-1 from 0" @@ -435,7 +441,7 @@ def test_control_stop_still_honors_policy(): registry=registry, ) - result = broker.handle(_call("bg_stop", {"task": "bg-1"})) + result = _handle(broker, _call("bg_stop", {"task": "bg-1"})) assert result.is_error assert "denied" in result.output @@ -454,7 +460,7 @@ def test_control_unknown_task_is_reported_as_error(): registry=RaisingRegistry(), ) - result = broker.handle(_call("bg_status", {"task": "bg-9"})) + result = _handle(broker, _call("bg_status", {"task": "bg-9"})) assert result.is_error assert "unknown background task" in result.output @@ -508,7 +514,7 @@ def test_shell_injection_is_escaped_before_it_reaches_the_jail(): jail = FakeJail() broker = _broker(jail) - broker.handle(_call("echo", {"message": "hi; echo pwned"})) + _handle(broker, _call("echo", {"message": "hi; echo pwned"})) # The payload is quoted into a single argument; the injected command can't run. assert jail.exec_commands == ["echo 'hi; echo pwned'"] diff --git a/tests/test_jail.py b/tests/test_jail.py index 1e689f7..c836a08 100644 --- a/tests/test_jail.py +++ b/tests/test_jail.py @@ -5,6 +5,7 @@ prove the security invariants: confinement, content purity, and reach isolation. """ +import asyncio import shutil import shlex import subprocess @@ -25,6 +26,11 @@ _NEEDS_SANDBOX = pytest.mark.skipif( ) +def _exec(jail, *args, **kwargs): + """Drive the async JailBuilder.exec to completion from a sync test.""" + return asyncio.run(jail.exec(*args, **kwargs)) + + def _store_root(path: str) -> str: """/nix/store//anything -> /nix/store/.""" return "/".join(path.split("/")[:4]) @@ -60,7 +66,7 @@ def shell_closure(shell_path): @_NEEDS_SANDBOX def test_echo_runs_confined(tmp_path, shell_path, shell_closure): jail = JailBuilder(str(tmp_path), shell_path, shell_closure=shell_closure) - result = jail.exec(jail.build(Grant()), "echo banana") + result = _exec(jail, jail.build(Grant()), "echo banana") assert result.code == 0 assert "banana" in result.stdout @@ -73,7 +79,7 @@ def test_bwrap_parent_environment_does_not_leak_into_proc( monkeypatch.setenv("TARTARUS_TEST_HOST_SECRET", "secret-from-host-env") jail = JailBuilder(str(tmp_path), shell_path, shell_closure=shell_closure) - result = jail.exec(jail.build(Grant()), "cat /proc/1/environ") + result = _exec(jail, jail.build(Grant()), "cat /proc/1/environ") assert "TARTARUS_TEST_HOST_SECRET" not in result.stdout assert "secret-from-host-env" not in result.stdout @@ -87,7 +93,8 @@ def test_proc_file_descriptors_are_not_available_for_output_injection( ): jail = JailBuilder(str(tmp_path), shell_path, shell_closure=shell_closure) - result = jail.exec( + result = _exec( + jail, jail.build(Grant()), "bash -c 'echo injected-output > /proc/1/fd/1; echo normal-output'", ) @@ -102,7 +109,7 @@ def test_host_only_tool_is_absent_inside_jail(tmp_path, shell_path, shell_closur jail = JailBuilder(str(tmp_path), shell_path, shell_closure=shell_closure) # git is not in the shell closure and was not granted, so it cannot resolve # by name even though the baseline shell binaries do. - result = jail.exec(jail.build(Grant()), "git --version") + result = _exec(jail, jail.build(Grant()), "git --version") assert result.code != 0 assert "no such file" in result.stderr.lower() @@ -122,7 +129,7 @@ def test_ungranted_tool_unreachable_by_absolute_store_path( pytest.skip(f"cannot resolve git package: {error}") jail = JailBuilder(str(tmp_path), shell_path, shell_closure=shell_closure) - result = jail.exec(jail.build(Grant()), f"{shlex.quote(git_bin)}/git --version") + result = _exec(jail, jail.build(Grant()), f"{shlex.quote(git_bin)}/git --version") assert result.code != 0 assert "no such file" in result.stderr.lower() @@ -131,7 +138,7 @@ def test_ungranted_tool_unreachable_by_absolute_store_path( @_NEEDS_SANDBOX def test_no_host_filesystem_beyond_work_tree(tmp_path, shell_path, shell_closure): jail = JailBuilder(str(tmp_path), shell_path, shell_closure=shell_closure) - result = jail.exec(jail.build(Grant()), "ls /") + result = _exec(jail, jail.build(Grant()), "ls /") visible = set(result.stdout.split()) assert visible <= {"dev", "nix", "proc", "work"} @@ -143,7 +150,7 @@ def test_no_host_filesystem_beyond_work_tree(tmp_path, shell_path, shell_closure def test_no_network_interfaces_inside_jail(tmp_path, shell_path, shell_closure): jail = JailBuilder(str(tmp_path), shell_path, shell_closure=shell_closure) # --unshare-all removed the network namespace, so /sys/class/net is gone. - result = jail.exec(jail.build(Grant()), "ls /sys/class/net") + result = _exec(jail, jail.build(Grant()), "ls /sys/class/net") assert result.code != 0 @@ -187,7 +194,7 @@ def test_proxy_jail_routes_curl_through_allowed_host( ) ) - result = jail.exec(spec, f"curl -fsS http://{upstream_host}:{upstream_port}/") + result = _exec(jail, spec, f"curl -fsS http://{upstream_host}:{upstream_port}/") assert result.code == 0 assert result.stdout == "hello" @@ -214,7 +221,7 @@ def test_proxy_jail_blocks_unlisted_host(tmp_path, shell_path, shell_closure): ) ) - result = jail.exec(spec, f"curl -fsS http://{upstream_host}:{upstream_port}/") + result = _exec(jail, spec, f"curl -fsS http://{upstream_host}:{upstream_port}/") assert result.code != 0 assert "proxy decisions: 0 allowed, 1 blocked" in result.stderr @@ -232,8 +239,8 @@ def test_writable_grant_allows_only_declared_path(tmp_path, shell_path, shell_cl jail = JailBuilder(str(tmp_path), shell_path, shell_closure=shell_closure) spec = jail.build(Grant(writable=["allowed"])) - allowed = jail.exec(spec, "bash -c 'echo yes > allowed/file.txt'") - denied = jail.exec(spec, "bash -c 'echo no > readonly/file.txt'") + allowed = _exec(jail, spec, "bash -c 'echo yes > allowed/file.txt'") + denied = _exec(jail, spec, "bash -c 'echo no > readonly/file.txt'") assert allowed.code == 0 assert (writable_dir / "file.txt").read_text().strip() == "yes" @@ -308,7 +315,7 @@ def test_unrestricted_grant_bypasses_bwrap_after_approval_path(tmp_path): jail = JailBuilder(str(work_tree), str(bin_dir)) spec = jail.build(Grant(unrestricted=True)) - result = jail.exec(spec, "read-secret") + result = _exec(jail, spec, "read-secret") assert result.code == 0 assert result.stdout == "outside work tree" @@ -320,7 +327,8 @@ def test_exec_streams_unrestricted_output_lines(tmp_path): lines: list[str] = [] code = "import sys; print('one', flush=True); print('two', flush=True)" - result = jail.exec( + result = _exec( + jail, spec, f"{shlex.quote(sys.executable)} -c {shlex.quote(code)}", output_callback=lines.append, @@ -331,32 +339,50 @@ def test_exec_streams_unrestricted_output_lines(tmp_path): assert lines == ["one\n", "two\n"] -def test_exec_cancellation_stops_unrestricted_process(tmp_path): +_SLOW_PROGRAM = ( + "import time; " + "print('started', flush=True); " + "time.sleep(10); " + "print('finished', flush=True)" +) + + +def test_exec_timeout_kills_unrestricted_process(tmp_path): jail = JailBuilder(str(tmp_path), "/unused/bin") spec = jail.build(Grant(unrestricted=True)) - cancellation = threading.Event() - lines: list[str] = [] - code = ( - "import time; " - "print('started', flush=True); " - "time.sleep(10); " - "print('finished', flush=True)" - ) - def capture(line: str) -> None: - lines.append(line) - cancellation.set() - - result = jail.exec( + result = _exec( + jail, spec, - f"{shlex.quote(sys.executable)} -c {shlex.quote(code)}", - output_callback=capture, - cancellation=cancellation, + f"{shlex.quote(sys.executable)} -c {shlex.quote(_SLOW_PROGRAM)}", + timeout=1, ) assert result.code == 124 + assert result.stdout == "started\n" + assert "timed out" in result.stderr + + +def test_exec_cancellation_terminates_unrestricted_process(tmp_path): + jail = JailBuilder(str(tmp_path), "/unused/bin") + spec = jail.build(Grant(unrestricted=True)) + lines: list[str] = [] + command = f"{shlex.quote(sys.executable)} -c {shlex.quote(_SLOW_PROGRAM)}" + + async def run_and_cancel(): + task = asyncio.create_task( + jail.exec(spec, command, output_callback=lines.append) + ) + while not lines: + await asyncio.sleep(0.01) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + asyncio.run(run_and_cancel()) + + # The process was killed during its sleep, so its later output never arrives. assert lines == ["started\n"] - assert "command cancelled" in result.stderr class _HelloHandler(BaseHTTPRequestHandler): -- 2.51.2