"""Bounded framed completion protocol for forked Task22 children.""" from __future__ import annotations import ctypes import errno import json import math import os import select import signal import struct import sys import time from collections.abc import Callable from typing import Protocol, cast SCHEMA = "task22-fork-completion-v1" MAX_RECEIPT = 1024 class PidfdSendSignal(Protocol): argtypes: list[object] restype: object def __call__(self, pidfd: int, sig: int, info: None, flags: int, /) -> int: ... def frame(kind: str, value: int, detail: str = "") -> bytes: payload: dict[str, object] = {"schema": SCHEMA, "kind": kind} payload["signal" if kind == "signal" else "code"] = value if detail: payload["detail"] = detail[:256] content = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode() if len(content) > MAX_RECEIPT: raise RuntimeError("completion receipt exceeds bound") return struct.pack(">I", len(content)) + content def install_signal_receipts(descriptor: int) -> None: frames: dict[int, bytes] = {int(number): frame("signal", int(number)) for number in ( signal.SIGHUP, signal.SIGINT, signal.SIGQUIT, signal.SIGTERM, )} def handler(number: int, _stack: object) -> None: receipt = frames.get(number) if receipt is None: os._exit(128 + number) _ = os.write(descriptor, receipt) _ = signal.signal(signal.Signals(number), signal.SIG_DFL) os.kill(os.getpid(), number) for number in frames: _ = signal.signal(signal.Signals(number), handler) def child_exit(descriptor: int, kind: str, code: int, detail: str = "") -> None: try: _ = os.write(descriptor, frame(kind, code, detail)) finally: os.close(descriptor) os._exit(code) def validate_completion(receipt: bytes, status: int) -> int: try: if len(receipt) < 4: raise ValueError unpacked: tuple[int] = struct.unpack(">I", receipt[:4]) length = unpacked[0] if length > MAX_RECEIPT or len(receipt) != length + 4: raise ValueError raw = cast(object, json.loads(receipt[4:])) except (UnicodeDecodeError, json.JSONDecodeError, ValueError, struct.error) as exc: raise RuntimeError("completion receipt is missing or malformed") from exc if not isinstance(raw, dict): raise RuntimeError("completion receipt schema is invalid") value = cast(dict[str, object], raw) if value.get("schema") != SCHEMA: raise RuntimeError("completion receipt schema is invalid") kind = value.get("kind") if os.WIFSIGNALED(status): number = os.WTERMSIG(status) if kind != "signal" or value.get("signal") != number: raise RuntimeError("completion receipt and signal status mismatch") return 128 + number if not os.WIFEXITED(status): raise RuntimeError("completion wait status is not terminal") code = os.WEXITSTATUS(status) if kind not in {"exit", "error"} or value.get("code") != code: raise RuntimeError("completion receipt and exit status mismatch") if kind == "error" and code == 0: raise RuntimeError("completion error receipt reported success") return code def _load_pidfd_sender() -> PidfdSendSignal: if sys.platform != "linux": raise RuntimeError("pidfd signaling requires Linux") library = ctypes.CDLL(None, use_errno=True) try: sender = cast(PidfdSendSignal, cast(object, library.pidfd_send_signal)) except AttributeError as exc: raise RuntimeError("libc pidfd_send_signal is unavailable") from exc sender.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.c_void_p, ctypes.c_uint] sender.restype = ctypes.c_int return sender def ensure_pidfd_support() -> None: _ = _load_pidfd_sender() def _validate_owned_pidfd(pid: int, pidfd: int) -> None: if isinstance(pid, bool) or pid <= 0: raise ValueError("completion child PID is invalid") if isinstance(pidfd, bool) or pidfd < 0: raise ValueError("completion pidfd is invalid") try: _ = os.fstat(pidfd) with open(f"/proc/self/fdinfo/{pidfd}", encoding="ascii") as stream: identities = [ line.removeprefix("Pid:").strip() for line in stream if line.startswith("Pid:") ] if identities != [str(pid)]: raise RuntimeError("completion pidfd does not identify the requested child") _ = os.waitid(3, pidfd, os.WEXITED | os.WNOHANG | os.WNOWAIT) except ChildProcessError as exc: raise RuntimeError("completion pidfd is not owned by this parent") from exc except OSError as exc: raise RuntimeError("completion pidfd is invalid") from exc def _send_pidfd_sigkill(pid: int, pidfd: int) -> None: _validate_owned_pidfd(pid, pidfd) sender = _load_pidfd_sender() _ = ctypes.set_errno(0) result = sender(pidfd, int(signal.SIGKILL), None, 0) if result != 0: error = ctypes.get_errno() or errno.EIO raise OSError(error, os.strerror(error)) def kill_pidfd(pid: int, pidfd: int) -> None: try: _send_pidfd_sigkill(pid, pidfd) except OSError as exc: if exc.errno != errno.ESRCH: raise reaper = select.poll() reaper.register(pidfd, select.POLLIN) if not reaper.poll(5_000): raise RuntimeError("pidfd-signaled child did not become ready for reaping") waited, _status = os.waitpid(pid, 0) if waited != pid: raise RuntimeError("completion kill reaped the wrong child") _kill_pidfd = kill_pidfd def _require_empty_timeout_receipt(receipt_fd: int, chunks: list[bytes]) -> None: while True: try: block = os.read(receipt_fd, MAX_RECEIPT + 5) except BlockingIOError: break if not block: break chunks.append(block) if chunks: raise RuntimeError("forced timeout requires no completion receipt") def await_child(pid: int, pidfd: int, receipt_fd: int, timeout_ms: int = 60_000) -> int: os.set_blocking(receipt_fd, False) poller = select.poll() poller.register(pidfd, select.POLLIN) poller.register(receipt_fd, select.POLLIN | select.POLLHUP) deadline = time.monotonic() + timeout_ms / 1000 chunks: list[bytes] = [] total, exited = 0, False while True: remaining = max(0, math.ceil((deadline - time.monotonic()) * 1000)) events = poller.poll(remaining) if not events: if exited: waited, _status = os.waitpid(pid, 0) if waited != pid: raise RuntimeError("completion timeout reaped the wrong child") else: kill_pidfd(pid, pidfd) _require_empty_timeout_receipt(receipt_fd, chunks) raise RuntimeError("child completion exceeded the bounded execution window") eof = False for descriptor, event in events: if descriptor == pidfd: exited = True elif descriptor == receipt_fd and event & (select.POLLIN | select.POLLHUP): while True: try: block = os.read(receipt_fd, MAX_RECEIPT + 5 - total) except BlockingIOError: break if not block: eof = True break chunks.append(block) total += len(block) if total > MAX_RECEIPT + 4: kill_pidfd(pid, pidfd) raise RuntimeError("completion receipt exceeds bound") if exited and eof: break waited, status = os.waitpid(pid, 0) if waited != pid: raise RuntimeError("completion wait returned the wrong child") return validate_completion(b"".join(chunks), status) def run_child( body: Callable[[], None], receipt_fd: int, stdout_flush: Callable[[], object], stderr_flush: Callable[[], object], ) -> None: install_signal_receipts(receipt_fd) try: body() except SystemExit as exc: code = exc.code if isinstance(exc.code, int) else 1 _ = stdout_flush() _ = stderr_flush() child_exit(receipt_fd, "exit", code) except BaseException as exc: _ = stderr_flush() child_exit(receipt_fd, "error", 1, f"{type(exc).__name__}: {exc}") _ = stdout_flush() _ = stderr_flush() child_exit(receipt_fd, "exit", 0)