#!/usr/bin/env python3 """Mosaic external lease broker for Linux SO_PEERCRED authenticated clients.""" from __future__ import annotations import argparse import copy import errno import hmac from concurrent.futures import ThreadPoolExecutor import json import os import secrets import signal import socket import stat import struct import sys import threading import time from pathlib import Path from typing import Final # This script is loaded both as an executable and through isolated stdlib tests. # Keep its co-located receipt implementation importable in both modes. _MODULE_DIRECTORY = str(Path(__file__).resolve().parent) if _MODULE_DIRECTORY not in sys.path: sys.path.insert(0, _MODULE_DIRECTORY) from normative_fragments import build_payload_from_wire from receipt_challenge import is_verbatim_receipt, latest_assistant_digest, receipt_for from receipt_observer import FileTestReceiptObserver, ReceiptObserver, UnavailableReceiptObserver MAX_FRAME: Final = 64 * 1024 MAX_STATE: Final = 4 * 1024 * 1024 MAX_PENDING_TOKENS: Final = 256 MAX_IN_FLIGHT_CONNECTIONS: Final = 16 MAX_LEASE_TTL_SECONDS: Final = 300 STATE_VERSION: Final = 1 READ_DEADLINE_SECONDS: Final = 1.0 HANDLE_QUEUE_TIMEOUT_SECONDS: Final = 1.0 SEND_TIMEOUT_SECONDS: Final = 1.0 HEX_256_LENGTH: Final = 64 LEASE_UNVERIFIED: Final = "UNVERIFIED" LEASE_PENDING: Final = "PENDING_VERIFICATION" LEASE_PENDING_PROMOTION: Final = "PENDING_PROMOTION" LEASE_VERIFIED: Final = "VERIFIED" READ_ONLY_TOOLS: Final = { "claude": frozenset({"Read", "Grep", "Glob", "Ls", "Find"}), "pi": frozenset({"read", "grep", "find", "ls"}), } RECOVERY_TOOL: Final = "mosaic_context_recover" class BrokerFailure(Exception): def __init__(self, code: str) -> None: super().__init__(code) self.code = code class StateCommitUncertain(RuntimeError): def __init__(self) -> None: super().__init__("STATE_COMMIT_UNCERTAIN") def is_non_negative_integer(value: object) -> bool: return type(value) is int and value >= 0 def is_hex_256(value: object) -> bool: return ( isinstance(value, str) and len(value) == HEX_256_LENGTH and all(character in "0123456789abcdef" for character in value) ) def is_positive_decimal(value: object) -> bool: return ( isinstance(value, str) and len(value) > 0 and value[0] in "123456789" and all(character in "0123456789" for character in value) ) def valid_binding(binding: object) -> bool: if not isinstance(binding, dict): return False required = {"compaction_epoch", "request_epoch", "h_source", "h_payload", "schema_version"} if set(binding) != required: return False if not all( is_non_negative_integer(binding[field]) for field in ("compaction_epoch", "request_epoch", "schema_version") ): return False return all( is_hex_256(binding[field]) for field in ("h_source", "h_payload") ) def valid_receipt_evidence(evidence: object) -> bool: return ( isinstance(evidence, dict) and set(evidence) == {"h_latest_assistant"} and is_hex_256(evidence["h_latest_assistant"]) ) def validate_state(value: object) -> dict[str, object]: if not isinstance(value, dict) or set(value) != {"version", "sessions", "tokens"}: raise BrokerFailure("STATE_INTEGRITY") if type(value["version"]) is not int or value["version"] != STATE_VERSION: raise BrokerFailure("STATE_INTEGRITY") sessions = value["sessions"] tokens = value["tokens"] if not isinstance(sessions, dict) or not isinstance(tokens, dict): raise BrokerFailure("STATE_INTEGRITY") if len(tokens) > MAX_PENDING_TOKENS: raise BrokerFailure("STATE_INTEGRITY") anchors: set[tuple[int, str]] = set() for session_id, session in sessions.items(): if not is_hex_256(session_id) or not isinstance(session, dict): raise BrokerFailure("STATE_INTEGRITY") if set(session) != {"anchor_pid", "anchor_starttime", "runtime_generation"}: raise BrokerFailure("STATE_INTEGRITY") anchor_pid = session["anchor_pid"] anchor_starttime = session["anchor_starttime"] if type(anchor_pid) is not int or anchor_pid <= 0: raise BrokerFailure("STATE_INTEGRITY") if not is_positive_decimal(anchor_starttime): raise BrokerFailure("STATE_INTEGRITY") if not is_non_negative_integer(session["runtime_generation"]): raise BrokerFailure("STATE_INTEGRITY") anchor = (anchor_pid, anchor_starttime) if anchor in anchors: raise BrokerFailure("STATE_INTEGRITY") anchors.add(anchor) for token_value, token in tokens.items(): if not is_hex_256(token_value) or not isinstance(token, dict): raise BrokerFailure("STATE_INTEGRITY") token_fields = set(token) if token_fields not in ( {"session_id", "runtime_generation", "binding", "consumed"}, {"session_id", "runtime_generation", "binding", "consumed", "evidence"}, ): raise BrokerFailure("STATE_INTEGRITY") session_id = token["session_id"] generation = token["runtime_generation"] session = sessions.get(session_id) if isinstance(session_id, str) else None if not is_hex_256(session_id) or not isinstance(session, dict): raise BrokerFailure("STATE_INTEGRITY") if not is_non_negative_integer(generation): raise BrokerFailure("STATE_INTEGRITY") if generation != session["runtime_generation"]: raise BrokerFailure("STATE_INTEGRITY") if not valid_binding(token["binding"]) or token["consumed"] is not False: raise BrokerFailure("STATE_INTEGRITY") evidence = token.get("evidence") if evidence is not None and not valid_receipt_evidence(evidence): raise BrokerFailure("STATE_INTEGRITY") return value def proc_node(pid: int) -> dict[str, int | str]: try: text = Path(f"/proc/{pid}/stat").read_text(encoding="utf-8") except (FileNotFoundError, PermissionError, ProcessLookupError) as exc: raise BrokerFailure("PID_UNAVAILABLE") from exc close = text.rfind(")") fields = text[close + 2 :].split() if close < 0 or len(fields) < 20: raise BrokerFailure("PROC_STAT_INVALID") return {"pid": pid, "ppid": int(fields[1]), "starttime": fields[19]} def verified_ancestry(peer_pid: int, anchor_pid: int, anchor_starttime: str) -> bool: chain: list[dict[str, int | str]] = [] seen: set[int] = set() current = peer_pid while current > 0 and current not in seen: seen.add(current) node = proc_node(current) chain.append(node) if current == anchor_pid: if node["starttime"] != anchor_starttime: return False break current = int(node["ppid"]) else: return False if int(chain[-1]["pid"]) != anchor_pid: return False for original in chain: repeated = proc_node(int(original["pid"])) if repeated["starttime"] != original["starttime"]: raise BrokerFailure("PID_STARTTIME_RACE") return True def secure_parent(path: Path) -> None: parent = path.parent if not parent.is_dir() or stat.S_IMODE(parent.stat().st_mode) != 0o700: raise BrokerFailure("INSECURE_PARENT_MODE") class StateStore: def __init__(self, path: Path) -> None: self.path = path secure_parent(path) self.poisoned = False self.value: dict[str, object] = {"version": STATE_VERSION, "sessions": {}, "tokens": {}} flags = os.O_RDONLY | getattr(os, "O_CLOEXEC", 0) | getattr(os, "O_NOFOLLOW", 0) try: descriptor = os.open(path, flags) except FileNotFoundError: return except OSError as exc: raise BrokerFailure("STATE_INTEGRITY") from exc try: metadata = os.fstat(descriptor) if not stat.S_ISREG(metadata.st_mode): raise BrokerFailure("STATE_INTEGRITY") if stat.S_IMODE(metadata.st_mode) != 0o600: raise BrokerFailure("INSECURE_STATE_MODE") if metadata.st_size > MAX_STATE: raise BrokerFailure("STATE_INTEGRITY") chunks = bytearray() while len(chunks) <= MAX_STATE: chunk = os.read(descriptor, min(64 * 1024, MAX_STATE + 1 - len(chunks))) if not chunk: break chunks.extend(chunk) if len(chunks) > MAX_STATE: raise BrokerFailure("STATE_INTEGRITY") try: loaded = json.loads(chunks) except (json.JSONDecodeError, UnicodeDecodeError) as exc: raise BrokerFailure("STATE_INTEGRITY") from exc self.value = validate_state(loaded) except OSError as exc: raise BrokerFailure("STATE_INTEGRITY") from exc finally: os.close(descriptor) def sessions(self) -> dict[str, dict[str, object]]: sessions = self.value.get("sessions") if not isinstance(sessions, dict): raise BrokerFailure("STATE_INTEGRITY") return sessions def tokens(self) -> dict[str, dict[str, object]]: tokens = self.value.get("tokens") if not isinstance(tokens, dict): raise BrokerFailure("STATE_INTEGRITY") return tokens def commit(self) -> None: if self.poisoned: raise StateCommitUncertain() payload = ( json.dumps(self.value, sort_keys=True, separators=(",", ":")) + "\n" ).encode() if len(payload) > MAX_STATE: raise BrokerFailure("STATE_TOO_LARGE") temporary = self.path.with_name(f".{self.path.name}.{os.getpid()}.tmp") descriptor = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) replaced = False try: try: remaining = memoryview(payload) while remaining: written = os.write(descriptor, remaining) if written == 0: raise OSError(errno.EIO, "state write made no progress") remaining = remaining[written:] os.fsync(descriptor) finally: os.close(descriptor) os.replace(temporary, self.path) replaced = True try: directory = os.open(self.path.parent, os.O_RDONLY | os.O_DIRECTORY) try: os.fsync(directory) finally: os.close(directory) except OSError as exc: self.poisoned = True raise StateCommitUncertain() from exc finally: if not replaced: try: temporary.unlink() except FileNotFoundError: pass class Broker: def __init__(self, store: StateStore, observer: ReceiptObserver | None = None) -> None: self.store = store self.observer: ReceiptObserver = observer if observer is not None else UnavailableReceiptObserver() # Set only by begin_verification after its mandatory revoke-first fence. # It is preserved if later cycle admission is refused; all other broker # actions retain the normal snapshot rollback behavior. self._rejected_cycle_fence: tuple[dict[str, object], dict[str, dict[str, object]]] | None = None # VERIFIED authority is deliberately volatile: broker restart revokes all # leases while preserving WI-1 identity and pending-token integrity. self.leases: dict[str, dict[str, object]] = {} def authenticate(self, peer_pid: int, request: dict[str, object]) -> tuple[str, dict[str, object]]: session_id = request.get("session_id") generation = request.get("runtime_generation") if not isinstance(session_id, str) or not is_non_negative_integer(generation): raise BrokerFailure("INVALID_IDENTITY") session = self.store.sessions().get(session_id) if not isinstance(session, dict): raise BrokerFailure("UNKNOWN_SESSION") anchor_pid = session.get("anchor_pid") anchor_starttime = session.get("anchor_starttime") current_generation = session.get("runtime_generation") if ( type(anchor_pid) is not int or not isinstance(anchor_starttime, str) or not is_non_negative_integer(current_generation) ): raise BrokerFailure("STATE_INTEGRITY") if not verified_ancestry(peer_pid, anchor_pid, anchor_starttime): raise BrokerFailure("ANCESTRY_MISMATCH") if generation < current_generation: raise BrokerFailure("STALE_GENERATION") if generation > current_generation: session["runtime_generation"] = generation self.revoke_session_authority(session_id) return session_id, session def session_for_anchor( self, anchor_pid: int, anchor_starttime: str ) -> tuple[str, dict[str, object]] | None: for session_id, session in self.store.sessions().items(): if ( session["anchor_pid"] == anchor_pid and session["anchor_starttime"] == anchor_starttime ): return session_id, session return None def revoke_session_tokens(self, session_id: str) -> None: tokens = self.store.tokens() for token_value in [ value for value, token in tokens.items() if token["session_id"] == session_id ]: del tokens[token_value] def revoke_session_authority(self, session_id: str) -> None: self.revoke_session_tokens(session_id) session = self.store.sessions().get(session_id) generation = session.get("runtime_generation") if isinstance(session, dict) else None self.leases[session_id] = { "state": LEASE_UNVERIFIED, "runtime_generation": generation, } def mint_token( self, session_id: str, generation: int, binding: dict[str, object], ) -> str: if len(self.store.tokens()) >= MAX_PENDING_TOKENS: raise BrokerFailure("TOKEN_CAPACITY") token = secrets.token_hex(32) self.store.tokens()[token] = { "session_id": session_id, "runtime_generation": generation, "binding": copy.deepcopy(binding), "consumed": False, # Receipt evidence is durably committed before this challenge may # be consumed and promotion made externally visible. "evidence": None, } return token def finish_promotion(self, session_id: str) -> None: lease = self.leases.get(session_id) if not isinstance(lease, dict) or lease.get("state") != LEASE_PENDING_PROMOTION: raise BrokerFailure("STATE_INTEGRITY") lease["state"] = LEASE_VERIFIED def handle(self, peer: tuple[int, int, int], request: dict[str, object]) -> dict[str, object]: if self.store.poisoned: raise StateCommitUncertain() previous = copy.deepcopy(self.store.value) previous_leases = copy.deepcopy(self.leases) self._rejected_cycle_fence = None try: response = self._handle(peer, request) if self.store.value != previous: self.store.commit() if request.get("action") == "promote_lease": session_id = request.get("session_id") if not isinstance(session_id, str): raise BrokerFailure("INVALID_IDENTITY") self.finish_promotion(session_id) response["state"] = LEASE_VERIFIED return response except StateCommitUncertain: raise except Exception: fence = self._rejected_cycle_fence if fence is None: self.store.value = previous self.leases = previous_leases else: # A refused re-verification must never resurrect the preceding # VERIFIED authority. Preserve only this post-revoke fence; # every unrelated partial-write failure still rolls back. self.store.value, self.leases = fence raise finally: self._rejected_cycle_fence = None def _handle(self, peer: tuple[int, int, int], request: dict[str, object]) -> dict[str, object]: peer_pid, peer_uid, peer_gid = peer action = request.get("action") if action == "register_anchor": if "session_id" in request: raise BrokerFailure("CALLER_SESSION_ID_REFUSED") generation = request.get("runtime_generation") if not is_non_negative_integer(generation): raise BrokerFailure("INVALID_GENERATION") anchor = proc_node(peer_pid) anchor_starttime = str(anchor["starttime"]) existing = self.session_for_anchor(peer_pid, anchor_starttime) if existing is None: session_id = secrets.token_hex(32) self.store.sessions()[session_id] = { "anchor_pid": peer_pid, "anchor_starttime": anchor_starttime, "runtime_generation": generation, } else: session_id, session = existing current_generation = session["runtime_generation"] if generation < current_generation: raise BrokerFailure("STALE_GENERATION") if generation > current_generation: session["runtime_generation"] = generation self.revoke_session_authority(session_id) return {"ok": True, "session_id": session_id, "peer": {"pid": peer_pid, "uid": peer_uid, "gid": peer_gid, "starttime": anchor["starttime"]}} if action == "authenticate": self.authenticate(peer_pid, request) return {"ok": True} if action == "mint_token": session_id, _ = self.authenticate(peer_pid, request) binding = request.get("binding") if not valid_binding(binding): raise BrokerFailure("INVALID_BINDING") token = self.mint_token(session_id, request["runtime_generation"], binding) return {"ok": True, "token": token} if action == "consume_token": session_id, _ = self.authenticate(peer_pid, request) token_value = request.get("token") token = self.store.tokens().get(token_value) if isinstance(token_value, str) else None if not isinstance(token, dict) or token.get("session_id") != session_id or token.get("runtime_generation") != request.get("runtime_generation") or token.get("consumed") is not False: raise BrokerFailure("TOKEN_REPLAY") del self.store.tokens()[token_value] return {"ok": True} if action == "begin_verification": session_id, _ = self.authenticate(peer_pid, request) runtime = request.get("runtime") binding = request.get("binding") construction = request.get("construction") ttl_seconds = request.get("ttl_seconds", MAX_LEASE_TTL_SECONDS) if runtime not in READ_ONLY_TOOLS: raise BrokerFailure("INVALID_RUNTIME") if not valid_binding(binding): raise BrokerFailure("INVALID_BINDING") if ( type(ttl_seconds) is not int or ttl_seconds <= 0 or ttl_seconds > MAX_LEASE_TTL_SECONDS ): raise BrokerFailure("INVALID_LEASE_TTL") # Revoke-first is a broker operation, not advisory adapter order. self.revoke_session_authority(session_id) # If subsequent construction admission rejects, handle() restores # this fence rather than the pre-cycle VERIFIED snapshot. self._rejected_cycle_fence = ( copy.deepcopy(self.store.value), copy.deepcopy(self.leases) ) try: constructed = build_payload_from_wire(construction) except ValueError as exc: raise BrokerFailure("INVALID_CONSTRUCTION") from exc if ( constructed.injectionDecision != "ACCEPTED" or not constructed.promotion or not isinstance(constructed.h_source, str) or not isinstance(constructed.h_payload, str) ): raise BrokerFailure("PAYLOAD_CONSTRUCTION_REFUSED") if ( not hmac.compare_digest(binding["h_source"], constructed.h_source) or not hmac.compare_digest(binding["h_payload"], constructed.h_payload) ): raise BrokerFailure("PAYLOAD_BINDING_MISMATCH") cycle_binding = copy.deepcopy(binding) cycle_binding["runtime_generation"] = request["runtime_generation"] challenge = self.mint_token(session_id, request["runtime_generation"], binding) self.leases[session_id] = { "state": LEASE_PENDING, "runtime": runtime, "runtime_generation": request["runtime_generation"], "binding": copy.deepcopy(binding), "receipt_challenge": challenge, "ttl_seconds": ttl_seconds, } # The broker constructs both the bound challenge and the delivery # text. The model's role is an exact copy, never hashing its output. return { "ok": True, "state": LEASE_PENDING, "receipt_challenge": challenge, "receipt": receipt_for(challenge, cycle_binding), "binding": cycle_binding, } if action == "observe_receipt": session_id, _ = self.authenticate(peer_pid, request) challenge = request.get("receipt_challenge") if "latest_assistant_message" in request: raise BrokerFailure("INVALID_RECEIPT") if not isinstance(challenge, str): raise BrokerFailure("INVALID_RECEIPT") token = self.store.tokens().get(challenge) if ( not isinstance(token, dict) or token.get("session_id") != session_id or token.get("runtime_generation") != request.get("runtime_generation") or token.get("consumed") is not False ): raise BrokerFailure("RECEIPT_REPLAY") lease = self.leases.get(session_id) if not isinstance(lease, dict) or lease.get("state") != LEASE_PENDING: raise BrokerFailure("RECEIPT_REPLAY") expected_challenge = lease.get("receipt_challenge") binding = lease.get("binding") if ( not isinstance(expected_challenge, str) or not secrets.compare_digest(challenge, expected_challenge) or not isinstance(binding, dict) or token.get("binding") != binding ): raise BrokerFailure("RECEIPT_REPLAY") receipt_binding = copy.deepcopy(binding) receipt_binding["runtime_generation"] = request["runtime_generation"] message = self.observer.observe_latest_assistant_message( session_id, str(lease["runtime"]), request["runtime_generation"], receipt_binding ) if not isinstance(message, str): raise BrokerFailure("RECEIPT_OBSERVATION_UNAVAILABLE") # This accepts one exact current-cycle assistant entry only. It is # deliberately not a transcript search and rejects quoted/extra text. if not is_verbatim_receipt(message, challenge, receipt_binding): raise BrokerFailure("RECEIPT_MISMATCH") evidence = {"h_latest_assistant": latest_assistant_digest(message)} token["evidence"] = evidence lease["state"] = LEASE_PENDING_PROMOTION lease["evidence"] = copy.deepcopy(evidence) return {"ok": True, "state": LEASE_PENDING_PROMOTION} if action == "promote_lease": session_id, _ = self.authenticate(peer_pid, request) challenge = request.get("receipt_challenge") if not isinstance(challenge, str): raise BrokerFailure("INVALID_RECEIPT") lease = self.leases.get(session_id) if not isinstance(lease, dict) or lease.get("state") == LEASE_UNVERIFIED: raise BrokerFailure("INVALID_LEASE_TRANSITION") token = self.store.tokens().get(challenge) if ( not isinstance(token, dict) or token.get("session_id") != session_id or token.get("runtime_generation") != request.get("runtime_generation") or token.get("consumed") is not False ): raise BrokerFailure("RECEIPT_REPLAY") if lease.get("state") != LEASE_PENDING_PROMOTION: raise BrokerFailure("INVALID_LEASE_TRANSITION") expected_challenge = lease.get("receipt_challenge") if not isinstance(expected_challenge, str) or not secrets.compare_digest(challenge, expected_challenge): raise BrokerFailure("RECEIPT_REPLAY") if token.get("binding") != lease.get("binding") or not valid_receipt_evidence(token.get("evidence")): raise BrokerFailure("PROMOTION_TOKEN_INVALID") del self.store.tokens()[challenge] lease["expires_at"] = time.monotonic() + int(lease["ttl_seconds"]) # handle() commits the evidence-backed consumption before # finish_promotion() makes VERIFIED externally visible: promote-last. return {"ok": True, "state": LEASE_PENDING_PROMOTION} if action == "revoke_lease": session_id, _ = self.authenticate(peer_pid, request) self.revoke_session_authority(session_id) return {"ok": True, "state": LEASE_UNVERIFIED} if action == "authorize_tool": session_id, _ = self.authenticate(peer_pid, request) runtime = request.get("runtime") tool_name = request.get("tool_name") if runtime not in READ_ONLY_TOOLS: raise BrokerFailure("INVALID_RUNTIME") if not isinstance(tool_name, str) or not tool_name or len(tool_name) > 256: raise BrokerFailure("INVALID_TOOL") lease = self.leases.get(session_id) state = lease.get("state") if isinstance(lease, dict) else LEASE_UNVERIFIED if tool_name in READ_ONLY_TOOLS[runtime] or tool_name == RECOVERY_TOOL: return {"ok": True, "decision": "allow", "state": state} if not isinstance(lease, dict) or lease.get("state") != LEASE_VERIFIED: return { "ok": False, "code": "MUTATOR_UNVERIFIED", "decision": "deny", "state": LEASE_UNVERIFIED, } if ( lease.get("runtime") != runtime or lease.get("runtime_generation") != request.get("runtime_generation") ): return { "ok": False, "code": "MUTATOR_UNVERIFIED", "decision": "deny", "state": LEASE_UNVERIFIED, } expires_at = lease.get("expires_at") if not isinstance(expires_at, (int, float)) or time.monotonic() >= expires_at: self.revoke_session_authority(session_id) return { "ok": False, "code": "LEASE_EXPIRED", "decision": "deny", "state": LEASE_UNVERIFIED, } return { "ok": True, "decision": "allow", "state": LEASE_VERIFIED, } raise BrokerFailure("UNKNOWN_ACTION") def read_frame(connection: socket.socket, deadline: float) -> dict[str, object]: data = bytearray() while len(data) <= MAX_FRAME: remaining = deadline - time.monotonic() if remaining <= 0: raise BrokerFailure("MALFORMED_REQUEST") connection.settimeout(remaining) chunk = connection.recv(min(4096, MAX_FRAME + 1 - len(data))) if not chunk: break data.extend(chunk) if len(data) > MAX_FRAME: while True: remaining = deadline - time.monotonic() if remaining <= 0: raise BrokerFailure("MALFORMED_REQUEST") connection.settimeout(remaining) if not connection.recv(4096): break raise BrokerFailure("MALFORMED_REQUEST") if not data.endswith(b"\n") or data.count(b"\n") != 1: raise BrokerFailure("MALFORMED_REQUEST") try: value = json.loads(data) except (json.JSONDecodeError, UnicodeDecodeError) as exc: raise BrokerFailure("MALFORMED_REQUEST") from exc if not isinstance(value, dict): raise BrokerFailure("MALFORMED_REQUEST") return value def handle_connection( connection: socket.socket, broker: Broker, broker_lock: threading.Lock, ) -> None: with connection: read_deadline = time.monotonic() + READ_DEADLINE_SECONDS try: raw = connection.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, 12) peer = struct.unpack("3i", raw) request = read_frame(connection, read_deadline) except BrokerFailure as exc: reply = {"ok": False, "code": exc.code} except OSError: return else: # Queueing is bounded and fails closed before broker.handle mutates # state. Once handling starts it must finish atomically; its reply # then receives an independent send budget so a slow fsync cannot # consume the write opportunity and create client/broker ambiguity. acquired = broker_lock.acquire(timeout=HANDLE_QUEUE_TIMEOUT_SECONDS) if not acquired: reply = {"ok": False, "code": "BROKER_BUSY"} else: try: try: reply = broker.handle(peer, request) except BrokerFailure as exc: reply = {"ok": False, "code": exc.code} finally: broker_lock.release() try: connection.settimeout(SEND_TIMEOUT_SECONDS) connection.sendall((json.dumps(reply, separators=(",", ":")) + "\n").encode()) except OSError: return def serve( socket_path: Path, state_path: Path, observer: ReceiptObserver | None = None, ) -> None: secure_parent(socket_path) if socket_path.exists() or socket_path.is_symlink(): raise BrokerFailure("SOCKET_ALREADY_EXISTS") store = StateStore(state_path) broker = Broker(store, observer) broker_lock = threading.Lock() slots = threading.BoundedSemaphore(MAX_IN_FLIGHT_CONNECTIONS) fatal_lock = threading.Lock() fatal_errors: list[Exception] = [] executor = ThreadPoolExecutor( max_workers=MAX_IN_FLIGHT_CONNECTIONS, thread_name_prefix="mosaic-lease-broker", ) server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) server.bind(str(socket_path)) os.chmod(socket_path, 0o600) owned = (socket_path.stat().st_dev, socket_path.stat().st_ino) stopping = False def stop(_signum: int, _frame: object) -> None: nonlocal stopping stopping = True server.close() def process_connection(connection: socket.socket) -> None: try: handle_connection(connection, broker, broker_lock) except Exception as exc: with fatal_lock: if not fatal_errors: fatal_errors.append(exc) finally: slots.release() def fatal_error() -> Exception | None: with fatal_lock: return fatal_errors[0] if fatal_errors else None signal.signal(signal.SIGTERM, stop) signal.signal(signal.SIGINT, stop) server.listen(MAX_IN_FLIGHT_CONNECTIONS) server.settimeout(0.1) print("READY", flush=True) try: while not stopping: failure = fatal_error() if failure is not None: raise failure if not slots.acquire(timeout=0.1): continue try: connection, _ = server.accept() except socket.timeout: slots.release() continue except OSError: slots.release() failure = fatal_error() if failure is not None: raise failure if stopping: break raise try: executor.submit(process_connection, connection) except Exception: slots.release() connection.close() raise finally: server.close() executor.shutdown(wait=True) try: current = socket_path.stat() if (current.st_dev, current.st_ino) == owned: socket_path.unlink() except FileNotFoundError: pass def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--socket", required=True, type=Path) parser.add_argument("--state", required=True, type=Path) parser.add_argument("--test-observer-file", type=Path) arguments = parser.parse_args() observer = ( FileTestReceiptObserver(arguments.test_observer_file) if arguments.test_observer_file is not None else None ) serve(arguments.socket, arguments.state, observer) if __name__ == "__main__": try: main() except BrokerFailure as failure: print(failure.code, file=sys.stderr) raise SystemExit(1) except StateCommitUncertain as failure: print(str(failure), file=sys.stderr) raise SystemExit(1)