from __future__ import annotations from concurrent.futures import ThreadPoolExecutor from dataclasses import replace from datetime import date, datetime, time as daytime, timezone from decimal import Decimal from io import BytesIO import json import os from pathlib import Path import subprocess import threading import time import unittest from unittest.mock import patch from uuid import UUID from govoplan_core.security import bounded_process from govoplan_core.security.bounded_process import ( ProcessBudgetError, ProcessLimits, bounded_operation_admission, run_bounded_operation, ) from govoplan_core.security.worker_payload import ( WorkerPayloadError, decode_worker_payload, encode_worker_payload, ) from govoplan_core.settings import settings from govoplan_core.security.process_worker import _read_input from tests import bounded_process_fixtures as operations class BoundedProcessTests(unittest.TestCase): def setUp(self) -> None: self.processes = [] original = subprocess.Popen def start(*args, **kwargs): process = original(*args, **kwargs) self.processes.append(process) return process self.patcher = patch.object(bounded_process.subprocess, "Popen", start) self.patcher.start() self.addCleanup(self.patcher.stop) self.addCleanup(self.assert_reaped) def assert_reaped(self) -> None: for process in self.processes: self.assertIsNotNone(process.returncode) self.assertTrue(all(stream.closed for stream in (process.stdin, process.stdout, process.stderr))) with self.assertRaises(ChildProcessError): os.waitpid(process.pid, os.WNOHANG) self.assertEqual(bounded_process._active, 0) def test_roundtrip_empty_and_pipe_sized_input_exact_output_limit(self) -> None: for payload in (b"", b"x" * 200_000): with self.subTest(length=len(payload)): result = run_bounded_operation(operations.echo, payload, limits=ProcessLimits(output_bytes=max(1, len(payload)))) self.assertEqual(result, payload) def test_controls_and_environment_are_applied_before_operation(self) -> None: with patch.dict(os.environ, {"DATABASE_URL": "synthetic-secret", "PYTHONPATH": "/untrusted", "SYNTHETIC_SECRET": "no"}): result = json.loads(run_bounded_operation(operations.observe, b"", limits=ProcessLimits(cpu_seconds=3))) self.assertEqual(result["pid"], result["pgid"]) self.assertEqual(result["pid"], result["sid"]) self.assertNotEqual(result["pid"], os.getpid()) self.assertEqual(result["cpu"], [3, 4]) self.assertEqual(result["memory"], [256 * 1024 * 1024] * 2) self.assertEqual(result["file"], [0, 0]) self.assertEqual(result["core"], [0, 0]) self.assertEqual(result["cwd"], "/") self.assertFalse({"DATABASE_URL", "PYTHONPATH", "SYNTHETIC_SECRET"} & set(result["env"])) def test_real_regex_cpu_is_stopped_before_long_wall_limit(self) -> None: started = time.monotonic() with self.assertRaises(ProcessBudgetError) as raised: run_bounded_operation(operations.regex_stall, b"", limits=ProcessLimits(cpu_seconds=1, wall_seconds=5)) self.assertEqual(raised.exception.code, "cpu_limit") self.assertLess(time.monotonic() - started, 4) def test_memory_failure_never_allocates_the_large_result_in_parent(self) -> None: with self.assertRaises(ProcessBudgetError) as raised: run_bounded_operation(operations.allocate, b"", limits=ProcessLimits(memory_bytes=64 * 1024 * 1024)) self.assertEqual(raised.exception.code, "memory_limit") def test_noisy_stdout_and_stderr_are_bounded_during_execution(self) -> None: for operation in (operations.too_much_stdout, operations.too_much_stderr): with self.subTest(operation=operation.__name__): with self.assertRaises(ProcessBudgetError) as raised: run_bounded_operation(operation, b"", limits=ProcessLimits(output_bytes=1024, wall_seconds=3)) self.assertEqual(raised.exception.code, "output_limit") def test_sleep_and_closed_pipe_hangs_are_timed_out(self) -> None: for operation in (operations.wait, operations.close_pipes_then_wait): with self.subTest(operation=operation.__name__): with self.assertRaises(ProcessBudgetError) as raised: run_bounded_operation(operation, b"30", limits=ProcessLimits(wall_seconds=0.3)) self.assertEqual(raised.exception.code, "timeout") def test_cancellation_kills_and_reaps(self) -> None: started = time.monotonic() with self.assertRaises(ProcessBudgetError) as raised: run_bounded_operation(operations.wait, b"30", cancelled=lambda: time.monotonic() - started > 0.2) self.assertEqual(raised.exception.code, "cancelled") def test_cancellation_callback_exception_also_cleans_up(self) -> None: calls = 0 def cancel(): nonlocal calls calls += 1 if calls > 2: raise KeyboardInterrupt return False with self.assertRaises(KeyboardInterrupt): run_bounded_operation(operations.wait, b"30", cancelled=cancel) def test_failure_does_not_return_child_exception_or_partial_output(self) -> None: with self.assertRaises(ProcessBudgetError) as raised: run_bounded_operation(operations.fail, b"") self.assertEqual(raised.exception.code, "worker_failed") self.assertNotIn("private", str(raised.exception)) def test_descendant_is_stopped_even_after_successful_leader_exit(self) -> None: pid = int(run_bounded_operation(operations.child_with_closed_pipes, b"")) deadline = time.monotonic() + 2 while time.monotonic() < deadline: try: state = Path(f"/proc/{pid}/stat").read_text().split(") ", 1)[1].split()[0] except (FileNotFoundError, ProcessLookupError): return if state == "Z": return # stopped, awaiting the operating system's orphan reaper time.sleep(0.01) self.fail("Descendant is still running after its owned group was cleaned up.") def test_admission_is_bounded_and_releases_capacity(self) -> None: entered = threading.Event() def pending(): return run_bounded_operation(operations.wait, b"0.4", cancelled=lambda: entered.set() and False) with patch.object(settings, "isolated_process_concurrency", 1), ThreadPoolExecutor(max_workers=2) as executor: future = executor.submit(pending) entered.wait(1) deadline = time.monotonic() + 1 while bounded_process._active == 0 and time.monotonic() < deadline: time.sleep(0.001) with self.assertRaises(ProcessBudgetError) as raised: run_bounded_operation(operations.echo, b"denied") self.assertEqual(raised.exception.code, "busy") self.assertEqual(future.result(), b"done") self.assertEqual(run_bounded_operation(operations.echo, b"after"), b"after") def test_invalid_inputs_never_spawn_a_child(self) -> None: with self.assertRaises(ProcessBudgetError): run_bounded_operation(operations.echo, b"too large", limits=ProcessLimits(input_bytes=1)) with self.assertRaises(ValueError): run_bounded_operation(lambda value: value, b"") self.assertEqual(self.processes, []) for changes in ({"wall_seconds": float("nan")}, {"wall_seconds": float("inf")}, {"cpu_seconds": True}, {"memory_bytes": 1}): with self.subTest(changes=changes), self.assertRaises(ValueError): replace(ProcessLimits(), **changes) def test_explicit_admission_covers_preparation_and_can_be_reused_sequentially(self) -> None: with patch.object(settings, "isolated_process_concurrency", 1): with bounded_operation_admission() as admission: self.assertEqual(bounded_process._active, 1) with self.assertRaises(ProcessBudgetError) as busy: with bounded_operation_admission(): self.fail("Preparation should not start without shared capacity.") self.assertEqual(busy.exception.code, "busy") self.assertEqual(run_bounded_operation(operations.echo, b"first", admission=admission), b"first") self.assertEqual(run_bounded_operation(operations.echo, b"second", admission=admission), b"second") self.assertEqual(bounded_process._active, 0) with self.assertRaises(ValueError): run_bounded_operation(operations.echo, b"expired", admission=admission) def test_explicit_admission_rejects_other_threads_and_overlapping_reuse(self) -> None: with bounded_operation_admission() as admission: with ThreadPoolExecutor(max_workers=1) as executor: future = executor.submit(run_bounded_operation, operations.echo, b"wrong thread", admission=admission) with self.assertRaises(ValueError): future.result() nested = False def poll(): nonlocal nested if admission.running and not nested: nested = True with self.assertRaises(ValueError): run_bounded_operation(operations.echo, b"overlap", admission=admission) return False self.assertEqual(run_bounded_operation(operations.wait, b"0.1", admission=admission, cancelled=poll), b"done") self.assertTrue(nested) def test_preparation_failure_releases_capacity_without_spawning(self) -> None: with self.assertRaisesRegex(RuntimeError, "prepare"): with bounded_operation_admission(): raise RuntimeError("prepare") self.assertEqual(bounded_process._active, 0) self.assertEqual(self.processes, []) @unittest.skipUnless(hasattr(os, "fork"), "Fork ownership requires POSIX fork") def test_inherited_admission_expires_without_releasing_child_capacity(self) -> None: read_fd, write_fd = os.pipe() try: with bounded_operation_admission() as admission: child_pid = os.fork() if child_pid == 0: os.close(read_fd) if child_pid == 0: try: os.write(write_fd, json.dumps({ "active": bounded_process._active, "token_active": admission.active, }).encode()) finally: os.close(write_fd) os._exit(0) os.close(write_fd) write_fd = None observation = json.loads(os.read(read_fd, 256)) _, status = os.waitpid(child_pid, 0) self.assertEqual(0, os.waitstatus_to_exitcode(status)) self.assertEqual({"active": 0, "token_active": False}, observation) self.assertEqual(0, bounded_process._active) finally: os.close(read_fd) if write_fd is not None: os.close(write_fd) def test_worker_stdin_reads_are_incremental_instead_of_allocating_the_cap(self) -> None: class ObservedInput(BytesIO): def read(self, size=-1): self_test.assertLessEqual(size, 65536) return super().read(size) self_test = self self.assertEqual(_read_input(ObservedInput(b"tiny"), 64 * 1024 * 1024), b"tiny") self.assertEqual(_read_input(ObservedInput(b"x" * 100000), 100000), b"x" * 100000) self.assertIsNone(_read_input(ObservedInput(b"x" * 100001), 100000)) class WorkerPayloadTests(unittest.TestCase): def test_roundtrip_explicit_types_and_user_keys_cannot_impersonate_tags(self) -> None: value = {"str": ["bytes", "not transport"], "values": ( None, True, 3, 1.25, Decimal("1.2500"), b"\x00\xff", date(2026, 9, 8), datetime(2026, 9, 8, tzinfo=timezone.utc), daytime(12, 30), UUID(int=4), )} self.assertEqual(decode_worker_payload(encode_worker_payload(value)), value) def test_rejects_arbitrary_objects_duplicate_keys_and_invalid_tags(self) -> None: with self.assertRaises(WorkerPayloadError): encode_worker_payload(object()) key = encode_worker_payload("a")[4:] duplicate_keys = b"GWP\x01\x0e\x00\x00\x00\x02" + (key + b"\x00") * 2 for wire in (b'["pickle","payload"]', duplicate_keys, b"GWP\x01\xff", b"GWP\x01\x03\x00\x00\x00\x01\xff"): with self.subTest(wire=wire), self.assertRaises(WorkerPayloadError): decode_worker_payload(wire) def test_transport_depth_and_byte_limits(self) -> None: with self.assertRaises(WorkerPayloadError): encode_worker_payload("x" * 1000, max_bytes=100) with self.assertRaises(WorkerPayloadError): decode_worker_payload(b" " * 1000, max_bytes=100) value = [] for _index in range(66): value = [value] with self.assertRaises(WorkerPayloadError): encode_worker_payload(value) def test_astral_unicode_uses_utf8_bytes_and_exact_byte_caps(self) -> None: value = "\U0001f30d" * 32_769 limit = 4 + 5 + len(value) * 4 wire = encode_worker_payload(value, max_bytes=limit) self.assertEqual(len(wire), limit) self.assertEqual(value, decode_worker_payload(wire, max_bytes=limit)) with self.assertRaises(WorkerPayloadError): encode_worker_payload(value, max_bytes=limit - 1) with self.assertRaises(WorkerPayloadError): decode_worker_payload(wire, max_bytes=limit - 1) def test_unicode_limit_is_checked_without_whole_encoded_temporary(self) -> None: import tracemalloc value = "\U0001f30d" * 900_000 tracemalloc.start() try: with self.assertRaises(WorkerPayloadError): encode_worker_payload(value, max_bytes=1_000_000) _current, peak = tracemalloc.get_traced_memory() finally: tracemalloc.stop() self.assertLess(peak, 2_000_000) def test_malformed_counts_depth_and_trailing_data_fail_before_children(self) -> None: import struct deep = b"GWP\x01" + (b"\x0d" + struct.pack(">I", 1)) * 66 + b"\x00" for wire in ( b"GWP\x01\x0d" + struct.pack(">I", 0xFFFFFFFF), b"GWP\x01\x0e" + struct.pack(">I", 0xFFFFFFFF), b"GWP\x01\x03" + struct.pack(">I", 0xFFFFFFFF), deep, encode_worker_payload(None) + b"\x00", b"GWP\x01\x05\x00\x00\x00\x00", b"GWP\x01\x06\x00\x00\x00\x01x", b"GWP\x01\x08\x00\x00\x00\x01x", ): with self.subTest(wire=wire[:20]), self.assertRaises(WorkerPayloadError): decode_worker_payload(wire) def test_decoder_checks_node_budget_before_allocating_container(self) -> None: from govoplan_core.security import worker_payload wire = encode_worker_payload([None] * 20) with patch.object(worker_payload, "_MAX_NODES", 10): with self.assertRaises(WorkerPayloadError): decode_worker_payload(wire) def test_large_signed_integers_roundtrip_and_decimal_errors_are_normalized(self) -> None: for value in (0, -1, 127, 128, -128, -129, 1 << 20_000, -(1 << 20_000)): with self.subTest(bits=value.bit_length()): self.assertEqual(value, decode_worker_payload(encode_worker_payload(value))) with self.assertRaises(WorkerPayloadError): decode_worker_payload(b"GWP\x01\x07\x00\x00\x00\x07invalid") with self.assertRaises(WorkerPayloadError): encode_worker_payload(1 << 20_000, max_bytes=100)