fix(#828): bound persisted broker token state

Prune consumed and revoked tokens, cap pending token state, reject oversized serialization before replacement, and roll back in-memory mutations when persistence fails.\n\ncloses #828

packages/mosaic/framework/tools/lease-broker/daemon.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import argparse
+import copy
 import errno
 import json
 import os
@@ -19,6 +20,7 @@ from typing import Final
 
 MAX_FRAME: Final = 64 * 1024
 MAX_STATE: Final = 4 * 1024 * 1024
+MAX_PENDING_TOKENS: Final = 256
 STATE_VERSION: Final = 1
 CONNECTION_DEADLINE_SECONDS: Final = 1.0
 HEX_256_LENGTH: Final = 64
@@ -209,14 +211,16 @@ class StateStore:
         return tokens
 
     def commit(self) -> None:
+        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:
-                payload = (
-                    json.dumps(self.value, sort_keys=True, separators=(",", ":")) + "\n"
-                ).encode()
                 remaining = memoryview(payload)
                 while remaining:
                     written = os.write(descriptor, remaining)
@@ -269,7 +273,6 @@ class Broker:
         if generation > current_generation:
             session["runtime_generation"] = generation
             self.revoke_session_tokens(session_id)
-            self.store.commit()
         return session_id, session
 
     def session_for_anchor(
@@ -284,11 +287,24 @@ class Broker:
         return None
 
     def revoke_session_tokens(self, session_id: str) -> None:
-        for token in self.store.tokens().values():
-            if token["session_id"] == session_id:
-                token["consumed"] = True
+        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]:
+        previous = copy.deepcopy(self.store.value)
+        try:
+            response = self._handle(peer, request)
+            if self.store.value != previous:
+                self.store.commit()
+            return response
+        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":
@@ -307,7 +323,6 @@ class Broker:
                     "anchor_starttime": anchor_starttime,
                     "runtime_generation": generation,
                 }
-                self.store.commit()
             else:
                 session_id, session = existing
                 current_generation = session["runtime_generation"]
@@ -316,7 +331,6 @@ class Broker:
                 if generation > current_generation:
                     session["runtime_generation"] = generation
                     self.revoke_session_tokens(session_id)
-                    self.store.commit()
             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)
@@ -326,6 +340,8 @@ class Broker:
             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,
@@ -333,7 +349,6 @@ class Broker:
                 "binding": binding,
                 "consumed": False,
             }
-            self.store.commit()
             return {"ok": True, "token": token}
         if action == "consume_token":
             session_id, _ = self.authenticate(peer_pid, request)
@@ -341,8 +356,7 @@ class Broker:
             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")
-            token["consumed"] = True
-            self.store.commit()
+            del self.store.tokens()[token_value]
             return {"ok": True}
         raise BrokerFailure("UNKNOWN_ACTION")