security: harden nonce replay guard, constant-time signature comparison

Author: shivdeep1

.gitignore

@@ -86,4 +86,6 @@ REVIEW_BRIEF.md
 AIRLOCK_COMPANY.md
 PROTOCOL_DEBATE.md
 IMPLEMENTATION_PLAN.md
+V03_IMPLEMENTATION_PLAN.md
+SECURITY_AUDIT_POW_ARGON2.md
 .claude/

airlock/crypto/__init__.py

@@ -3,8 +3,10 @@
 from airlock.crypto.keys import KeyPair, resolve_public_key
 from airlock.crypto.signing import (
     canonicalize,
+    sign_attestation,
     sign_message,
     sign_model,
+    verify_attestation,
     verify_model,
     verify_signature,
 )
@@ -15,9 +17,11 @@
     "canonicalize",
     "issue_credential",
     "resolve_public_key",
+    "sign_attestation",
     "sign_message",
     "sign_model",
     "validate_credential",
+    "verify_attestation",
     "verify_model",
     "verify_signature",
 ]

airlock/crypto/signing.py

@@ -1,8 +1,11 @@
 from __future__ import annotations
 
 import json
-from base64 import b64decode, b64encode
+import logging
+import uuid as uuid_mod
+from base64 import b64decode, b64encode, urlsafe_b64encode
 from datetime import UTC, datetime
+from enum import Enum, IntEnum, StrEnum
 from typing import TYPE_CHECKING, Any
 
 from nacl.exceptions import BadSignatureError
@@ -12,6 +15,71 @@
 if TYPE_CHECKING:
     from airlock.schemas.handshake import SignatureEnvelope
 
+logger = logging.getLogger(__name__)
+
+_SIGNATURE_FIELDS = frozenset({"signature", "airlock_signature", "trust_token"})
+
+
+def _prepare_for_json(obj: Any) -> Any:
+    """Recursively convert Python objects to JSON-safe, cross-language types.
+
+    Ensures deterministic serialization that produces identical output in
+    Python, Go, Rust, and JavaScript implementations (C-09 interop fix).
+
+    Conversion rules:
+    - datetime     -> ISO 8601 with timezone (naive datetimes treated as UTC)
+    - IntEnum      -> int value
+    - StrEnum      -> str value
+    - Enum         -> raw .value
+    - UUID         -> lowercase hyphenated string
+    - bytes        -> base64url encoding (no padding)
+    - BaseModel    -> model.model_dump(mode="json")
+    - set          -> sorted list (recursed)
+    - dict         -> recurse into values
+    - list / tuple -> recurse into elements
+    - str, int, float, bool, None -> pass through
+    - other        -> TypeError
+    """
+    # Enums must be checked first: IntEnum is a subclass of int,
+    # StrEnum is a subclass of str, so they'd pass the scalar check below.
+    # Plain Enum (e.g. Enum with int/str value) is NOT a subclass of int/str.
+    if isinstance(obj, IntEnum):
+        return int(obj)
+    if isinstance(obj, StrEnum):
+        return str(obj.value)
+    if isinstance(obj, Enum):
+        return obj.value
+
+    # JSON-native scalars: pass through unchanged
+    if obj is None or isinstance(obj, (bool, int, float, str)):
+        return obj
+
+    if isinstance(obj, datetime):
+        if obj.tzinfo is None or obj.tzinfo.utcoffset(obj) is None:
+            # Naive datetime: treat as UTC
+            obj = obj.replace(tzinfo=UTC)
+        return obj.isoformat()
+
+    if isinstance(obj, uuid_mod.UUID):
+        return str(obj)
+
+    if isinstance(obj, bytes):
+        return urlsafe_b64encode(obj).rstrip(b"=").decode("ascii")
+
+    if isinstance(obj, BaseModel):
+        return obj.model_dump(mode="json")
+
+    if isinstance(obj, set):
+        return [_prepare_for_json(item) for item in sorted(obj, key=str)]
+
+    if isinstance(obj, dict):
+        return {k: _prepare_for_json(v) for k, v in obj.items()}
+
+    if isinstance(obj, (list, tuple)):
+        return [_prepare_for_json(item) for item in obj]
+
+    raise TypeError(f"Cannot canonicalize type: {type(obj)}")
+
 
 def canonicalize(data: dict[str, Any]) -> bytes:
     """Produce deterministic canonical JSON bytes.
@@ -20,10 +88,20 @@ def canonicalize(data: dict[str, Any]) -> bytes:
     - Sort keys
     - No whitespace
     - UTF-8 encoding
-    Strips 'signature' key if present (we sign the unsigned form).
+    - ensure_ascii=False for cross-language consistency
+
+    All values are first normalized via ``_prepare_for_json`` so that
+    datetimes, enums, UUIDs, bytes, etc. are converted to language-agnostic
+    representations *before* JSON encoding.  This guarantees identical
+    canonical bytes across Python, Go, Rust, and JavaScript (C-09 fix).
+
+    Strips known signature/token fields so we sign the unsigned form.
     """
-    cleaned = {k: v for k, v in data.items() if k != "signature"}
-    return json.dumps(cleaned, sort_keys=True, separators=(",", ":"), default=str).encode("utf-8")
+    cleaned = {k: v for k, v in data.items() if k not in _SIGNATURE_FIELDS}
+    prepared = _prepare_for_json(cleaned)
+    return json.dumps(prepared, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode(
+        "utf-8"
+    )
 
 
 def sign_message(message_dict: dict[str, Any], signing_key: SigningKey) -> str:
@@ -59,6 +137,13 @@ def sign_model(model: BaseModel, signing_key: SigningKey) -> SignatureEnvelope:
     """Sign a Pydantic model and return a SignatureEnvelope.
 
     Canonical form excludes the 'signature' field.
+
+    NOTE: ``model.model_dump(mode="json")`` already converts datetimes,
+    enums, UUIDs etc. to JSON-safe primitives via Pydantic's serializer,
+    so the dict passed to ``sign_message`` → ``canonicalize`` contains only
+    str/int/float/bool/None/list/dict.  ``_prepare_for_json`` will simply
+    pass these through.  This path is therefore safe for cross-language
+    signature verification without additional pre-conversion.
     """
     from airlock.schemas.handshake import SignatureEnvelope
 
@@ -83,3 +168,40 @@ def verify_model(model: BaseModel, verify_key: VerifyKey) -> bool:
     data = model.model_dump(mode="json")
     data.pop("signature", None)
     return verify_signature(data, sig.value, verify_key)
+
+
+def sign_attestation(attestation: BaseModel, signing_key: SigningKey) -> str:
+    """Sign an AirlockAttestation and return a base64-encoded Ed25519 signature.
+
+    Canonical form excludes ``airlock_signature`` and ``trust_token`` fields
+    (handled by :func:`canonicalize`).  The returned string is suitable for
+    setting on ``AirlockAttestation.airlock_signature``.
+    """
+    data = attestation.model_dump(mode="json")
+    return sign_message(data, signing_key)
+
+
+def verify_attestation(attestation: BaseModel, public_key: VerifyKey | bytes) -> bool:
+    """Verify the ``airlock_signature`` on an :class:`AirlockAttestation`.
+
+    Parameters
+    ----------
+    attestation:
+        The attestation model instance.  Must have an ``airlock_signature``
+        field containing a base64-encoded Ed25519 signature string.
+    public_key:
+        Either a :class:`~nacl.signing.VerifyKey` or raw 32-byte public key.
+
+    Returns
+    -------
+    bool
+        ``True`` if the signature is valid, ``False`` otherwise (including
+        when ``airlock_signature`` is ``None``).
+    """
+    sig_b64 = getattr(attestation, "airlock_signature", None)
+    if sig_b64 is None:
+        return False
+    if isinstance(public_key, bytes):
+        public_key = VerifyKey(public_key)
+    data = attestation.model_dump(mode="json")
+    return verify_signature(data, sig_b64, public_key)

airlock/engine/orchestrator.py

@@ -21,8 +21,8 @@
 
 from langgraph.graph import END, StateGraph
 
-from airlock.crypto.keys import resolve_public_key
-from airlock.crypto.signing import verify_model
+from airlock.crypto.keys import KeyPair, resolve_public_key
+from airlock.crypto.signing import sign_attestation, verify_model
 from airlock.crypto.vc import validate_credential
 from airlock.engine.state import SessionManager
 from airlock.gateway.revocation import RedisRevocationStore, RevocationStore
@@ -114,11 +114,13 @@ def __init__(
         session_mgr: SessionManager | None = None,
         vc_allowed_issuers: frozenset[str] | None = None,
         revocation_store: RevocationStore | RedisRevocationStore | None = None,
+        airlock_keypair: KeyPair | None = None,
     ) -> None:
         self._reputation = reputation_store
         self._revocation: RevocationStore | RedisRevocationStore | None = revocation_store
         self._registry = agent_registry
         self._airlock_did = airlock_did
+        self._airlock_keypair = airlock_keypair
         self._model = litellm_model
         self._api_base = litellm_api_base
         self._on_challenge = on_challenge
@@ -137,6 +139,17 @@ def __init__(
 
         self._graph = self._build_graph()
 
+    def _sign_attestation_if_keypair(self, attestation: AirlockAttestation) -> AirlockAttestation:
+        """Sign the attestation with the gateway keypair if available.
+
+        Must be called BEFORE adding the trust_token so the signature
+        covers the attestation content without ephemeral fields.
+        """
+        if self._airlock_keypair is None:
+            return attestation
+        sig = sign_attestation(attestation, self._airlock_keypair.signing_key)
+        return attestation.model_copy(update={"airlock_signature": sig})
+
     async def _persist_graph_snapshot(self, final_state: OrchestrationState) -> None:
         """Mirror LangGraph session + checks into ``SessionManager`` for HTTP polling."""
         if self._session_mgr is None:
@@ -407,6 +420,8 @@ async def _handle_challenge_response(self, event: ChallengeResponseReceived) ->
             issued_at=now,
             privacy_mode=privacy_mode_str,
         )
+        # Sign BEFORE adding trust_token so the signature covers core content
+        attestation = self._sign_attestation_if_keypair(attestation)
         if verdict == TrustVerdict.VERIFIED and self._trust_token_secret:
             attestation = attestation.model_copy(
                 update={
@@ -493,6 +508,8 @@ async def _deliver_verdict(self, state: OrchestrationState) -> None:
             issued_at=now,
             privacy_mode=privacy_mode_str,
         )
+        # Sign BEFORE adding trust_token so the signature covers core content
+        attestation = self._sign_attestation_if_keypair(attestation)
         if verdict == TrustVerdict.VERIFIED and self._trust_token_secret:
             attestation = attestation.model_copy(
                 update={

airlock/gateway/__init__.py

@@ -1,5 +1,11 @@
 from __future__ import annotations
 
-from airlock.gateway.app import create_app
+
+def create_app(*args, **kwargs):  # type: ignore[no-untyped-def]
+    """Lazy wrapper to avoid circular import between engine and gateway."""
+    from airlock.gateway.app import create_app as _create_app
+
+    return _create_app(*args, **kwargs)
+
 
 __all__ = ["create_app"]

airlock/gateway/admin_routes.py

@@ -10,6 +10,7 @@
 from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
 from pydantic import BaseModel, Field
 
+from airlock.gateway.revocation import RevocationReason
 from airlock.reputation.scoring import INITIAL_SCORE
 from airlock.schemas.identity import AgentProfile
 from airlock.schemas.reputation import TrustScore
@@ -107,26 +108,51 @@ async def delete_agent(
     return {"deleted": True, "did": did}
 
 
+class RevokeBody(BaseModel):
+    reason: RevocationReason = RevocationReason.KEY_COMPROMISE
+
+
 @router.post("/revoke/{did:path}", response_model=dict[str, Any])
 async def revoke_agent(
     did: str,
     request: Request,
     _: Annotated[None, Depends(require_admin_token)],
+    body: RevokeBody | None = None,
+) -> dict[str, Any]:
+    reason = body.reason if body is not None else RevocationReason.KEY_COMPROMISE
+    store = request.app.state.revocation_store
+    changed = await store.revoke(did, reason=reason)
+    return {"revoked": True, "did": did, "changed": changed, "reason": reason.value}
+
+
+@router.post("/suspend/{did:path}", response_model=dict[str, Any])
+async def suspend_agent(
+    did: str,
+    request: Request,
+    _: Annotated[None, Depends(require_admin_token)],
 ) -> dict[str, Any]:
     store = request.app.state.revocation_store
-    changed = await store.revoke(did)
-    return {"revoked": True, "did": did, "changed": changed}
+    changed = await store.suspend(did)
+    return {"suspended": True, "did": did, "changed": changed}
 
 
-@router.post("/unrevoke/{did:path}", response_model=dict[str, Any])
-async def unrevoke_agent(
+@router.post("/reinstate/{did:path}", response_model=dict[str, Any])
+async def reinstate_agent(
     did: str,
     request: Request,
     _: Annotated[None, Depends(require_admin_token)],
 ) -> dict[str, Any]:
     store = request.app.state.revocation_store
-    changed = await store.unrevoke(did)
-    return {"unrevoked": True, "did": did, "changed": changed}
+    changed = await store.reinstate(did)
+    if not changed:
+        # Could be because the DID is permanently revoked or not suspended
+        reason = store.get_revocation_reason(did)
+        if reason is not None:
+            raise HTTPException(
+                status_code=409,
+                detail=f"Cannot reinstate permanently revoked DID (reason: {reason.value})",
+            )
+    return {"reinstated": changed, "did": did}
 
 
 @router.get("/revoked", response_model=dict[str, Any])

airlock/gateway/app.py

@@ -143,6 +143,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
             session_mgr=session_mgr,
             vc_allowed_issuers=vc_allowed,
             revocation_store=revocation_store,
+            airlock_keypair=airlock_kp,
         )
         event_bus.register(orchestrator.handle_event)
         await event_bus.start()

airlock/gateway/handlers.py

@@ -207,13 +207,31 @@ async def handle_handshake(
             status_code=400,
         )
     if body.pow is not None:
-        from airlock.pow import verify_pow
-
-        if not verify_pow(body.pow):
-            return JSONResponse(
-                {"error": "pow_invalid", "detail": "Proof-of-work verification failed"},
-                status_code=400,
-            )
+        from airlock.pow import verify_pow_with_store
+
+        pow_store = getattr(request.app.state, "pow_challenges", None)
+        if pow_store is not None:
+            ok, reason = verify_pow_with_store(body.pow, pow_store)
+            if not ok:
+                error_map: dict[str, tuple[str, int]] = {
+                    "unknown_challenge": ("Challenge ID not recognised or already used", 400),
+                    "expired_challenge": ("PoW challenge has expired", 400),
+                    "invalid_proof": ("Proof-of-work verification failed", 400),
+                }
+                detail, status = error_map.get(reason or "", ("PoW verification failed", 400))
+                return JSONResponse(
+                    {"error": f"pow_{reason}", "detail": detail, "status_code": status},
+                    status_code=status,
+                )
+        else:
+            # Fallback: no challenge store available — hash-only check
+            from airlock.pow import verify_pow
+
+            if not verify_pow(body.pow):
+                return JSONResponse(
+                    {"error": "pow_invalid", "detail": "Proof-of-work verification failed"},
+                    status_code=400,
+                )
 
     session_mgr = request.app.state.session_mgr
     now = datetime.now(UTC)