fix: wire persistence, CRL force-regen, chain registry, mid-session rotation

Five validated blockers addressed:
- PreCommitmentStore class replaces plain dict for pre-rotation commitments
- RotationChainRegistry gains JSON persistence with _rotated_from set
- register_chain() wired into /register and /handshake handlers
- force_regenerate() on CRL after key compromise rotation
- Startup guard blocks multi-replica + key rotation
- Orchestrator resolves reputation through chain_id for mid-session rotation
- First-write-wins check ordering fixed for correct error messages

685 tests passing.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: shivdeep1

airlock/config.py

@@ -180,6 +180,10 @@ class AirlockConfig(BaseSettings):
     pre_rotation_required_tier: int = Field(default=1, ge=0, le=3)  # TrustTier value
     pre_rotation_update_lockout_hours: int = Field(default=72, ge=1, le=720)
 
+    # Persistence paths for rotation stores (empty = in-memory only).
+    precommit_store_path: str = ""
+    rotation_chain_store_path: str = ""
+
     # Event bus drain timeout during shutdown (seconds).
     event_bus_drain_timeout_seconds: float = Field(default=30.0, ge=1.0, le=600.0)
 

airlock/engine/orchestrator.py

@@ -115,9 +115,11 @@ def __init__(
         vc_allowed_issuers: frozenset[str] | None = None,
         revocation_store: RevocationStore | RedisRevocationStore | None = None,
         airlock_keypair: KeyPair | None = None,
+        chain_registry: Any | None = None,
     ) -> None:
         self._reputation = reputation_store
         self._revocation: RevocationStore | RedisRevocationStore | None = revocation_store
+        self._chain_registry = chain_registry
         self._registry = agent_registry
         self._airlock_did = airlock_did
         self._airlock_keypair = airlock_keypair
@@ -532,10 +534,18 @@ async def _deliver_verdict(self, state: OrchestrationState) -> None:
             sealed_at=now,
         )
 
-        # Update reputation for terminal verdicts (unless local_only)
+        # Update reputation for terminal verdicts (unless local_only).
+        # Resolve through rotation chain so mid-session rotation applies
+        # to the current DID rather than the (possibly rotated-out) original.
         if verdict in (TrustVerdict.VERIFIED, TrustVerdict.REJECTED):
             if not state.get("_local_only", False):
-                self._reputation.apply_verdict(session.initiator_did, verdict)
+                reputation_did = session.initiator_did
+                chain_reg = getattr(self, "_chain_registry", None)
+                if chain_reg is not None:
+                    chain = chain_reg.get_chain_by_did(session.initiator_did)
+                    if chain is not None:
+                        reputation_did = chain.current_did
+                self._reputation.apply_verdict(reputation_did, verdict)
 
         if self._session_mgr is not None:
             prev = await self._session_mgr.get(session.session_id)

airlock/gateway/app.py

@@ -30,7 +30,7 @@
 from airlock.registry.agent_store import AgentRegistryStore
 from airlock.reputation.store import ReputationStore
 from airlock.rotation.chain import RotationChainRegistry
-from airlock.rotation.precommit import PreRotationCommitment
+from airlock.rotation.precommit_store import PreCommitmentStore
 
 logger = logging.getLogger(__name__)
 
@@ -165,6 +165,13 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
 
         audit_trail = AuditTrail()
 
+        # Key rotation chain registry (created early so orchestrator can reference it)
+        chain_registry: Any = None
+        if cfg.key_rotation_enabled:
+            from airlock.rotation.chain import RotationChainRegistry
+
+            chain_registry = RotationChainRegistry()
+
         _tok = (cfg.trust_token_secret or "").strip()
         orchestrator = VerificationOrchestrator(
             reputation_store=reputation,
@@ -178,6 +185,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
             vc_allowed_issuers=vc_allowed,
             revocation_store=revocation_store,
             airlock_keypair=airlock_kp,
+            chain_registry=chain_registry,
         )
         event_bus.register(orchestrator.handle_event)
         await event_bus.start()
@@ -202,13 +210,9 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
         app.state.pow_challenges: dict[str, Any] = {}
         app.state.redis_client = redis_client
 
-        # Key rotation (behind feature flag)
-        if cfg.key_rotation_enabled:
-            app.state.chain_registry = RotationChainRegistry()
-            app.state.precommit_store: dict[str, PreRotationCommitment] = {}
-        else:
-            app.state.chain_registry = None
-            app.state.precommit_store = {}
+        # Key rotation — assign chain_registry (created above) and precommit store
+        app.state.chain_registry = chain_registry
+        app.state.precommit_store = PreCommitmentStore()
 
         # Argon2id bounded verification worker pool
         import asyncio as _asyncio

airlock/gateway/crl.py

@@ -127,6 +127,15 @@ async def generate(self) -> SignedCRL:
         )
         return crl
 
+    async def force_regenerate(self) -> SignedCRL:
+        """Immediately regenerate the CRL, bypassing the cache.
+
+        Used when a key compromise requires instant propagation to
+        relying parties polling ``GET /crl``.
+        """
+        self._cached_crl = None
+        return await self.generate()
+
     async def get_or_generate(self) -> SignedCRL:
         """Return the cached CRL if fresh, otherwise regenerate."""
         if self._is_cache_fresh():

airlock/gateway/handlers.py

@@ -210,6 +210,15 @@ async def handle_handshake(
     if nack is not None:
         return nack
 
+    # Ensure rotation chain exists for this initiator (idempotent)
+    chain_registry = getattr(request.app.state, "chain_registry", None)
+    if chain_registry is not None:
+        try:
+            verify_key = resolve_public_key(body.initiator.did)
+            chain_registry.register_chain(body.initiator.did, bytes(verify_key))
+        except Exception:
+            pass  # Best-effort; chain registration may already exist
+
     # --- Proof-of-Work verification (anti-Sybil, v0.2) ---
     from airlock.config import get_config
 
@@ -405,6 +414,16 @@ async def handle_register(profile: AgentProfile, request: Request) -> dict[str,
     registry: dict[str, AgentProfile] = request.app.state.agent_registry
     registry[profile.did.did] = profile
     request.app.state.agent_store.upsert(profile)
+
+    # Auto-register rotation chain so key rotation works for this DID
+    chain_registry = getattr(request.app.state, "chain_registry", None)
+    if chain_registry is not None:
+        try:
+            verify_key = resolve_public_key(profile.did.did)
+            chain_registry.register_chain(profile.did.did, bytes(verify_key))
+        except Exception as exc:
+            logger.debug("Chain registration skipped for %s: %s", profile.did.did, exc)
+
     _audit_bg(
         request,
         event_type="agent_registered",
@@ -811,9 +830,9 @@ async def handle_rotate_key(
     new_public_key_bytes = bytes(new_verify_key)
 
     # Check pre-rotation commitment BEFORE first-write-wins
-    commitment_store: dict[str, PreRotationCommitment] = getattr(
-        request.app.state, "precommit_store", {}
-    )
+    from airlock.rotation.precommit_store import PreCommitmentStore
+
+    commitment_store: PreCommitmentStore = request.app.state.precommit_store
     existing_commitment = commitment_store.get(rotation_req.rotation_chain_id)
 
     # Mandatory pre-commitment check for higher tiers
@@ -863,6 +882,16 @@ async def handle_rotate_key(
         grace_seconds = cfg.key_rotation_grace_seconds
     await revocation_store.rotate_out(rotation_req.old_did, grace_seconds=grace_seconds)
 
+    # On compromise: force immediate CRL regeneration so relying parties see it
+    if reason == "compromise":
+        crl_gen = getattr(request.app.state, "crl_generator", None)
+        if crl_gen is not None:
+            try:
+                await crl_gen.force_regenerate()
+                logger.info("CRL force-regenerated after key compromise: %s", rotation_req.old_did)
+            except Exception as exc:
+                logger.error("CRL force-regeneration failed: %s", exc)
+
     # Transfer trust score via chain_id with penalty
     if trust_score is not None:
         penalty = cfg.key_rotation_trust_penalty
@@ -893,16 +922,19 @@ async def handle_rotate_key(
 
     # Handle chained commitment (N+2)
     if rotation_req.next_key_commitment:
-        commitment_store[rotation_req.rotation_chain_id] = PreRotationCommitment(
-            alg="sha256",
-            digest=rotation_req.next_key_commitment,
-            committed_at=datetime.now(UTC),
-            committed_by_did=rotation_req.new_did,
-            signature="",  # Commitment is embedded in the signed rotation request
+        commitment_store.put(
+            rotation_req.rotation_chain_id,
+            PreRotationCommitment(
+                alg="sha256",
+                digest=rotation_req.next_key_commitment,
+                committed_at=datetime.now(UTC),
+                committed_by_did=rotation_req.new_did,
+                signature="",  # Commitment is embedded in the signed rotation request
+            ),
         )
     else:
         # Clear the used commitment
-        commitment_store.pop(rotation_req.rotation_chain_id, None)
+        commitment_store.delete(rotation_req.rotation_chain_id)
 
     _audit_bg(
         request,
@@ -998,9 +1030,9 @@ async def handle_pre_commit_key(
         raise HTTPException(status_code=404, detail="DID not registered in any chain")
 
     # Check existing commitment and lockout
-    commitment_store: dict[str, PreRotationCommitment] = getattr(
-        request.app.state, "precommit_store", {}
-    )
+    from airlock.rotation.precommit_store import PreCommitmentStore
+
+    commitment_store: PreCommitmentStore = request.app.state.precommit_store
     existing = commitment_store.get(chain_id)
     if existing is not None:
         if not can_update_commitment(existing, lockout_hours=cfg.pre_rotation_update_lockout_hours):
@@ -1017,7 +1049,7 @@ async def handle_pre_commit_key(
         committed_by_did=commit_req.did,
         signature=commit_req.signature,
     )
-    commitment_store[chain_id] = commitment
+    commitment_store.put(chain_id, commitment)
 
     _audit_bg(
         request,

airlock/gateway/startup_validate.py

@@ -56,6 +56,14 @@ def validate_startup_config(cfg: AirlockConfig) -> None:
             "Production with AIRLOCK_EXPECT_REPLICAS > 1 requires AIRLOCK_REDIS_URL for shared replay and rate limits."
         )
 
+    if getattr(cfg, "key_rotation_enabled", False) and cfg.expect_replicas > 1:
+        raise AirlockStartupError(
+            "Key rotation requires single-replica deployment in the current release. "
+            "Multi-replica key rotation with a Redis-backed chain registry is planned. "
+            "Set AIRLOCK_EXPECT_REPLICAS=1 or disable key rotation with "
+            "AIRLOCK_KEY_ROTATION_ENABLED=false."
+        )
+
     reg = (cfg.default_registry_url or "").strip()
     if reg:
         parsed = urlparse(reg)

airlock/rotation/chain.py

@@ -12,10 +12,13 @@
 from __future__ import annotations
 
 import hashlib
+import json
 import logging
+import os
 import threading
 import time
 from datetime import UTC, datetime
+from pathlib import Path
 
 from pydantic import BaseModel
 
@@ -50,18 +53,28 @@ def compute_chain_id(public_key_bytes: bytes) -> str:
 
 
 class RotationChainRegistry:
-    """In-memory registry mapping DIDs to rotation chains.
+    """Registry mapping DIDs to rotation chains with optional JSON persistence.
 
-    Thread-safe via a reentrant lock.  All mutations are atomic with
+    Thread-safe via a threading lock.  All mutations are atomic with
     respect to the two index dicts (``_by_chain_id`` and ``_by_did``).
+
+    Parameters
+    ----------
+    path:
+        Filesystem path for the backing JSON file.  When *None*, the
+        registry is purely in-memory (suitable for tests and development).
     """
 
-    def __init__(self) -> None:
+    def __init__(self, path: str | None = None) -> None:
+        self._path: str | None = path
         self._by_chain_id: dict[str, RotationChainRecord] = {}
         self._by_did: dict[str, str] = {}  # did -> chain_id
         self._rotated_from: set[str] = set()  # DIDs that have been rotated away
         self._lock = threading.Lock()
 
+        if path and os.path.exists(path):
+            self._load()
+
     def register_chain(
         self,
         did: str,
@@ -89,6 +102,7 @@ def register_chain(
             )
             self._by_chain_id[chain_id] = record
             self._by_did[did] = chain_id
+            self._persist()
             logger.info("Rotation chain registered: chain=%s did=%s", chain_id[:16], did)
             return record
 
@@ -114,18 +128,19 @@ def rotate(
             if record is None:
                 raise ValueError(f"Unknown rotation chain: {chain_id}")
 
-            if record.current_did != old_did:
-                raise ValueError(
-                    f"Chain {chain_id[:16]} current DID is {record.current_did}, "
-                    f"not {old_did}"
-                )
-
             # First-write-wins: if old_did was already rotated, reject
+            # (checked before current_did comparison for clearer errors)
             if old_did in self._rotated_from:
                 raise ValueError(
                     f"DID {old_did} has already been rotated out (first-write-wins)"
                 )
 
+            if record.current_did != old_did:
+                raise ValueError(
+                    f"Chain {chain_id[:16]} current DID is {record.current_did}, "
+                    f"not {old_did}"
+                )
+
             self._rotated_from.add(old_did)
 
             # Update the record
@@ -145,6 +160,7 @@ def rotate(
             self._by_did[new_did] = chain_id
             # Keep old DID in the index for reverse lookups
             # but it is now in _rotated_from
+            self._persist()
 
             logger.info(
                 "Key rotated: chain=%s old=%s new=%s count=%d",
@@ -206,3 +222,64 @@ def check_rotation_rate(
             cutoff = time.time() - 86400.0
             recent = sum(1 for ts in record.rotation_timestamps if ts > cutoff)
             return recent >= max_per_24h
+
+    # ------------------------------------------------------------------
+    # Serialisation helpers
+    # ------------------------------------------------------------------
+
+    def _load(self) -> None:
+        """Deserialise chain records from the backing JSON file."""
+        if self._path is None:
+            return
+        try:
+            raw = Path(self._path).read_text(encoding="utf-8")
+            data: dict[str, object] = json.loads(raw)
+
+            for chain_id, blob in (data.get("chains") or {}).items():
+                record = RotationChainRecord.model_validate(blob)
+                self._by_chain_id[chain_id] = record
+                # Rebuild the DID index from the record
+                self._by_did[record.current_did] = chain_id
+                for prev_did in record.previous_dids:
+                    self._by_did[prev_did] = chain_id
+
+            rotated_list = data.get("rotated_from") or []
+            if isinstance(rotated_list, list):
+                self._rotated_from = set(rotated_list)
+
+            logger.info(
+                "Loaded %d rotation chains from %s",
+                len(self._by_chain_id),
+                self._path,
+            )
+        except Exception:
+            logger.exception("Failed to load chain registry from %s", self._path)
+
+    def _persist(self) -> None:
+        """Atomically write the current state to disk.
+
+        Writes to a temporary sibling file first, then renames it into
+        place.  On POSIX this is atomic; on Windows ``os.replace`` is
+        used which is atomic on NTFS.
+        """
+        if self._path is None:
+            return
+
+        chains_serialised: dict[str, object] = {}
+        for chain_id, record in self._by_chain_id.items():
+            chains_serialised[chain_id] = record.model_dump(mode="json")
+
+        payload = {
+            "chains": chains_serialised,
+            "rotated_from": sorted(self._rotated_from),
+        }
+
+        tmp_path = self._path + ".tmp"
+        try:
+            Path(tmp_path).write_text(
+                json.dumps(payload, indent=2),
+                encoding="utf-8",
+            )
+            os.replace(tmp_path, self._path)
+        except Exception:
+            logger.exception("Failed to persist chain registry to %s", self._path)