Implement authenticated custody observations and durable root-grant baseline

Assistant: codex
Assistant-Model: gpt-6.1-sol
Assistant-Session: 01a1029d-2c93-7500-a0e0-e3e089668b29

src/user_engine/custody_runtime.py

@@ -0,0 +1,168 @@
+"""Bounded owner transport and durable baseline for accepted custody observations."""
+from __future__ import annotations
+
+from datetime import datetime, timezone
+import json
+import math
+from pathlib import Path
+import re
+import time
+from typing import Callable
+from urllib.parse import urlsplit
+from urllib.request import HTTPRedirectHandler, ProxyHandler, Request, build_opener
+
+from user_engine.custody import CustodyBaseline, CustodyRead
+
+
+class CustodyReadUnavailable(RuntimeError):
+    def __init__(self):
+        super().__init__("custody read unavailable")
+
+
+class NoRedirect(HTTPRedirectHandler):
+    def redirect_request(self, *args, **kwargs):
+        return None
+
+
+class OpenBaoCustodyReader:
+    """Dedicated Kubernetes login, fresh exact KV read, best-effort revoke.
+
+    Never reuse the flex-auth JWT, ESO AppRole or a cached document. Tokens and
+    owner response/error bodies are never included in raised errors or reprs.
+    """
+    def __init__(self, *, base_url: str, auth_mount: str, role: str,
+                 jwt_file: str, source_id: str,
+                 data_path: str = "platform/data/custody/net-kingdom/platform-root",
+                 timeout_seconds: float = 3, max_bytes: int = 262144,
+                 allow_http: bool = False, opener=None,
+                 clock: Callable = lambda: datetime.now(timezone.utc)):
+        url = urlsplit(base_url)
+        if (url.scheme not in ({"https", "http"} if allow_http else {"https"})
+                or not url.hostname or url.username or url.password
+                or url.path not in ("", "/") or url.query or url.fragment):
+            raise ValueError("trusted OpenBao origin required")
+        if (not re.fullmatch(r"[a-zA-Z0-9_-]+", auth_mount)
+                or not re.fullmatch(r"[a-zA-Z0-9_-]+", role)
+                or data_path != "platform/data/custody/net-kingdom/platform-root"
+                or not source_id or not jwt_file
+                or not math.isfinite(timeout_seconds) or not 0 < timeout_seconds <= 30
+                or not 1024 <= max_bytes <= 262144):
+            raise ValueError("invalid custody reader configuration")
+        self.base_url = base_url.rstrip("/")
+        self.auth_mount, self.role, self.jwt_file = auth_mount, role, jwt_file
+        self.source_id, self.data_path = source_id, data_path
+        self.timeout_seconds, self.max_bytes = timeout_seconds, max_bytes
+        self.opener = opener or build_opener(ProxyHandler({}), NoRedirect())
+        self.clock = clock
+
+    def _request(self, path, *, deadline, token=None, body=None):
+        remaining = deadline - time.monotonic()
+        if remaining <= 0:
+            raise CustodyReadUnavailable()
+        headers = {"Accept": "application/json", "Cache-Control": "no-cache, no-store"}
+        if token:
+            headers["X-Vault-Token"] = token
+        data = None if body is None else json.dumps(body).encode()
+        if data is not None:
+            headers["Content-Type"] = "application/json"
+        request = Request(self.base_url + "/v1/" + path, data=data, headers=headers,
+                          method="GET" if data is None else "POST")
+        with self.opener.open(request, timeout=remaining) as response:
+            if response.status != 200:
+                raise CustodyReadUnavailable()
+            if (response.headers.get("Age") not in (None, "0")
+                    or response.headers.get("X-Cache", "").lower() not in ("", "miss")
+                    or response.headers.get("Content-Encoding", "identity") != "identity"):
+                raise CustodyReadUnavailable()
+            # Read incrementally under one overall deadline, including body I/O.
+            chunks, size = [], 0
+            while True:
+                if time.monotonic() >= deadline:
+                    raise CustodyReadUnavailable()
+                read_chunk = getattr(response, "read1", response.read)
+                chunk = read_chunk(min(65536, self.max_bytes + 1 - size))
+                size += len(chunk)
+                if size > self.max_bytes:
+                    raise CustodyReadUnavailable()
+                if not chunk:
+                    break
+                chunks.append(chunk)
+            return json.loads(b"".join(chunks))
+
+    def read(self):
+        token = None
+        deadline = time.monotonic() + self.timeout_seconds
+        try:
+            with Path(self.jwt_file).open("rb") as handle:
+                raw = handle.read(16385)
+            if len(raw) > 16384:
+                raise CustodyReadUnavailable()
+            jwt = raw.decode().strip()
+            if not jwt or "\n" in jwt or "\r" in jwt:
+                raise CustodyReadUnavailable()
+            login = self._request("auth/" + self.auth_mount + "/login", deadline=deadline,
+                                  body={"role": self.role, "jwt": jwt})
+            token = login["auth"]["client_token"]
+            if not isinstance(token, str) or not token or "\n" in token or "\r" in token:
+                raise CustodyReadUnavailable()
+            envelope = self._request(self.data_path, deadline=deadline, token=token)
+            document = envelope["data"]["data"]["document"]
+            if not isinstance(document, dict):
+                raise CustodyReadUnavailable()
+            return CustodyRead(self.source_id, document, self.clock())
+        except Exception:
+            raise CustodyReadUnavailable() from None
+        finally:
+            # Revoke via same trusted origin. This never renews observation time.
+            if token:
+                try:
+                    request = Request(self.base_url + "/v1/auth/token/revoke-self",
+                                      data=b"{}", headers={"X-Vault-Token": token}, method="POST")
+                    remaining = deadline - time.monotonic()
+                    if remaining > 0:
+                        with self.opener.open(request, timeout=remaining):
+                            pass
+                except Exception:
+                    pass
+
+
+class PostgresCustodyBaselineStore:
+    """One persistent singleton row, with compare-and-swap across workers.
+
+    Use a separate DB connection; no domain-store transaction is committed here.
+    Backups/restores must preserve this security high-water mark. A rolled-back
+    database requires controlled reconciliation before runtime reactivation.
+    """
+    def __init__(self, connection):
+        self.connection = connection
+        with connection.transaction():
+            with connection.cursor() as cursor:
+                cursor.execute("CREATE TABLE IF NOT EXISTS user_engine_custody_baseline "
+                               "(singleton boolean PRIMARY KEY CHECK (singleton), payload jsonb NOT NULL)")
+
+    @staticmethod
+    def _payload(baseline):
+        return json.dumps({"source_id": baseline.source_id, "document": baseline.document}, sort_keys=True)
+
+    def load(self):
+        with self.connection.transaction():
+            with self.connection.cursor() as cursor:
+                cursor.execute("SELECT payload FROM user_engine_custody_baseline WHERE singleton = true")
+                row = cursor.fetchone()
+        if row is None:
+            return None
+        value = json.loads(row[0]) if isinstance(row[0], str) else row[0]
+        return CustodyBaseline(value["source_id"], value["document"])
+
+    def compare_and_store(self, previous, current):
+        with self.connection.transaction():
+            with self.connection.cursor() as cursor:
+                if previous is None:
+                    cursor.execute("INSERT INTO user_engine_custody_baseline (singleton, payload) "
+                                   "VALUES (true, %s::jsonb) ON CONFLICT DO NOTHING RETURNING singleton",
+                                   (self._payload(current),))
+                else:
+                    cursor.execute("UPDATE user_engine_custody_baseline SET payload = %s::jsonb "
+                                   "WHERE singleton = true AND payload = %s::jsonb RETURNING singleton",
+                                   (self._payload(current), self._payload(previous)))
+                return cursor.fetchone() is not None