fix(#828): fail closed on uncertain broker commits

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

@@ -32,6 +32,11 @@ class BrokerFailure(Exception):
         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
 
@@ -79,6 +84,8 @@ def validate_state(value: object) -> dict[str, object]:
     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():
@@ -111,9 +118,9 @@ def validate_state(value: object) -> dict[str, object]:
             raise BrokerFailure("STATE_INTEGRITY")
         if not is_non_negative_integer(generation):
             raise BrokerFailure("STATE_INTEGRITY")
-        if generation > session["runtime_generation"]:
+        if generation != session["runtime_generation"]:
             raise BrokerFailure("STATE_INTEGRITY")
-        if not valid_binding(token["binding"]) or type(token["consumed"]) is not bool:
+        if not valid_binding(token["binding"]) or token["consumed"] is not False:
             raise BrokerFailure("STATE_INTEGRITY")
     return value
 
@@ -164,6 +171,7 @@ 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:
@@ -211,6 +219,8 @@ class StateStore:
         return tokens
 
     def commit(self) -> None:
+        if self.poisoned:
+            raise StateCommitUncertain()
         payload = (
             json.dumps(self.value, sort_keys=True, separators=(",", ":")) + "\n"
         ).encode()
@@ -232,11 +242,15 @@ class StateStore:
                 os.close(descriptor)
             os.replace(temporary, self.path)
             replaced = True
-            directory = os.open(self.path.parent, os.O_RDONLY | os.O_DIRECTORY)
             try:
-                os.fsync(directory)
-            finally:
-                os.close(directory)
+                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:
@@ -294,12 +308,16 @@ class Broker:
             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
@@ -467,3 +485,6 @@ if __name__ == "__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)