327 lines
16 KiB
Python
327 lines
16 KiB
Python
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)
|