feat(mosaic): add authenticated external lease broker (#836)

2026-07-18 03:12:23 +00:00
parent d801d6c4c8
commit 8ec67a1126
8 changed files with 1667 additions and 0 deletions
--- a/packages/mosaic/framework/tools/lease-broker/daemon.py
+++ b/packages/mosaic/framework/tools/lease-broker/daemon.py
@@ -0,0 +1,544 @@
+#!/usr/bin/env python3
+"""Mosaic external lease broker for Linux SO_PEERCRED authenticated clients."""
+
+from __future__ import annotations
+
+import argparse
+import copy
+import errno
+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
+
+MAX_FRAME: Final = 64 * 1024
+MAX_STATE: Final = 4 * 1024 * 1024
+MAX_PENDING_TOKENS: Final = 256
+MAX_IN_FLIGHT_CONNECTIONS: Final = 16
+STATE_VERSION: Final = 1
+CONNECTION_DEADLINE_SECONDS: Final = 1.0
+HEX_256_LENGTH: Final = 64
+
+
+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 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")
+        if set(token) != {"session_id", "runtime_generation", "binding", "consumed"}:
+            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")
+    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) -> None:
+        self.store = store
+
+    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_tokens(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 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)
+        try:
+            response = self._handle(peer, request)
+            if self.store.value != previous:
+                self.store.commit()
+            return response
+        except StateCommitUncertain:
+            raise
+        except Exception:
+            self.store.value = previous
+            raise
+
+    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_tokens(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")
+            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": request["runtime_generation"],
+                "binding": binding,
+                "consumed": False,
+            }
+            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}
+        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:
+        deadline = time.monotonic() + CONNECTION_DEADLINE_SECONDS
+        try:
+            raw = connection.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, 12)
+            peer = struct.unpack("3i", raw)
+            request = read_frame(connection, deadline)
+        except BrokerFailure as exc:
+            reply = {"ok": False, "code": exc.code}
+        except OSError:
+            return
+        else:
+            try:
+                with broker_lock:
+                    reply = broker.handle(peer, request)
+            except BrokerFailure as exc:
+                reply = {"ok": False, "code": exc.code}
+        try:
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                return
+            connection.settimeout(remaining)
+            connection.sendall((json.dumps(reply, separators=(",", ":")) + "\n").encode())
+        except OSError:
+            return
+
+
+def serve(socket_path: Path, state_path: Path) -> 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)
+    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)
+    arguments = parser.parse_args()
+    serve(arguments.socket, arguments.state)
+
+
+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)