fix: validate platform extension results

src/agent_enrollment_protocol/platform/platform.py

@@ -54,6 +54,7 @@
     IdempotentOperation,
     IdentityListQuery,
     IdentityRecord,
+    KeyStore,
     PlatformIdempotencyInput,
     PlatformIdempotencyState,
     PlatformOptions,
@@ -132,6 +133,8 @@ async def did_document(self, agent_did_id: str) -> PlatformResult[dict[str, Any]
         if identity is None or identity.status is not ManagedAgentStatus.ACTIVE:
             return _problem(404, "not_recognized", "Identity not recognized")
         validate_identity_record(identity)
+        if identity.agent_did_id != agent_did_id:
+            raise ValueError("AEP Platform identity store returned a mismatched record")
         method = await self._key_store.did_verification_method(identity)
         return PlatformResult(
             status=200,
@@ -173,12 +176,13 @@ async def list(
             return _problem(404, "not_recognized", "Identity not recognized")
         listed = await self._identity_store.list(context.principal, effective)
         if any(
-            identity.principal != context.principal
-            or (effective.service_did is not None and identity.service_did != effective.service_did)
-            or (effective.status is not None and identity.status is not effective.status)
+            not _valid_listed_identity(identity, context.principal, effective)
             for identity in listed.identities
         ):
             raise ValueError("AEP Platform identity store returned an unauthorized record")
+        identifiers = [identity.agent_identity_id for identity in listed.identities]
+        if listed.total < len(identifiers) or len(set(identifiers)) != len(identifiers):
+            raise ValueError("AEP Platform identity store returned an invalid list result")
         return _success(200, list_response(listed))
 
     async def provision(
@@ -258,7 +262,7 @@ async def _sign(
                 SignHandlerInput(identity, _clone_model(request)), context
             )
             if handled is not None:
-                _validate_sign_result(handled, identity, request, lifetime)
+                await _validate_sign_result(handled, identity, request, lifetime, self._key_store)
                 return handled
         now = _aware(self._request_time(context))
         claim_data: dict[str, Any] = {
@@ -286,6 +290,9 @@ async def _sign(
         if request.platform_context is not None:
             response_data["platform_context"] = request.platform_context
         body = PlatformSignCompleted.model_validate(response_data)
+        await _validate_sign_result(
+            _success(200, body), identity, request, lifetime, self._key_store
+        )
         return _success(200, body)
 
     async def update_identity(
@@ -311,6 +318,9 @@ async def update_identity(
         )
         if updated is None:
             return _problem(404, "not_recognized", "Identity not recognized")
+        validate_identity_record(updated)
+        if not _is_lifecycle_update(identity, updated, request.status):
+            raise ValueError("AEP Platform identity store returned a mismatched record")
         return _success(200, identity_response(updated))
 
     async def verify(
@@ -347,6 +357,10 @@ async def _verify(
         ):
             return unrecognized
         identity = await self._identity_store.find_by_agent_did(agent_did)
+        if identity is not None:
+            validate_identity_record(identity)
+            if identity.agent_did != agent_did:
+                raise ValueError("AEP Platform identity store returned a mismatched record")
         if (
             identity is None
             or identity.service_did != request.service_did
@@ -443,6 +457,8 @@ async def _authorized_identity(
         if identity is None:
             return None
         validate_identity_record(identity)
+        if identity.agent_identity_id != agent_identity_id:
+            raise ValueError("AEP Platform identity store returned a mismatched record")
         request = replace(request, identity=identity)
         if (
             not context.principal
@@ -553,11 +569,12 @@ def _problem(status: int, code: str, title: str) -> PlatformResult[BodyT]:
     )
 
 
-def _validate_sign_result(
+async def _validate_sign_result(
     result: PlatformResult[PlatformSignResponse],
     identity: IdentityRecord,
     request: PlatformSignRequest,
     requested_lifetime: int,
+    key_store: KeyStore,
 ) -> None:
     body = result.body
     if isinstance(body, PlatformSignPending):
@@ -575,6 +592,50 @@ def _validate_sign_result(
         or expires - issued != timedelta(seconds=requested_lifetime)
     ):
         raise ValueError("AEP Platform signing handler response does not match the request")
+    try:
+        header, _ = decode_jwt_unverified(body.client_assertion)
+        claims = verify_client_assertion(
+            body.client_assertion,
+            key=await key_store.verification_key(identity),
+            options=VerifyClientAssertionOptions(
+                algorithms=identity.signing_algorithms,
+                audience=request.service_did,
+                current_time=int(issued.timestamp()),
+                issuer=identity.agent_did,
+                operation=request.op,
+                resource=request.resource,
+                subject=identity.agent_did,
+            ),
+        )
+    except AepAssertionError as error:
+        raise ValueError("AEP Platform signer returned an invalid client assertion") from error
+    if (
+        header.get("kid") != identity.key_id
+        or claims.iat != int(issued.timestamp())
+        or claims.exp != int(expires.timestamp())
+        or claims.jti != request.jti
+    ):
+        raise ValueError("AEP Platform signer returned mismatched assertion claims")
+
+
+def _valid_listed_identity(
+    identity: IdentityRecord, principal: str, query: IdentityListQuery
+) -> bool:
+    try:
+        validate_identity_record(identity)
+    except ValueError:
+        return False
+    return (
+        identity.principal == principal
+        and (query.service_did is None or identity.service_did == query.service_did)
+        and (query.status is None or identity.status is query.status)
+    )
+
+
+def _is_lifecycle_update(
+    before: IdentityRecord, after: IdentityRecord, status: ManagedAgentStatus
+) -> bool:
+    return after == replace(before, status=status, updated_at=after.updated_at)
 
 
 def _seconds(value: timedelta, name: str) -> int:

tests/test_platform_boundaries.py

@@ -9,6 +9,7 @@
 
 from agent_enrollment_protocol.core import (
     AssertionOperation,
+    ClientAssertionClaims,
     ManagedAgentStatus,
     PlatformAgentIdentity,
     PlatformLifecycleRequest,
@@ -20,6 +21,7 @@
     PlatformVerificationResponse,
     ProblemDetails,
     SigningAlgorithm,
+    sign_client_assertion,
 )
 from agent_enrollment_protocol.platform import (
     DidVerificationMethod,
@@ -434,6 +436,60 @@ async def list(self, principal: str, query: IdentityListQuery) -> IdentityListRe
     with pytest.raises(ValueError, match="unauthorized"):
         await unauthorized.list(IdentityListQuery(), context(key=None))
 
+    class MismatchedLookupStore(MemoryIdentityStore):
+        async def get(self, agent_identity_id: str) -> IdentityRecord | None:
+            del agent_identity_id
+            return replace(record(), agent_identity_id="pai_other", principal="owner-one")
+
+        async def find_by_agent_did_id(self, agent_did_id: str) -> IdentityRecord | None:
+            del agent_did_id
+            return replace(record(), agent_did_id="other")
+
+    mismatched_lookup = Platform(options(identity_store=MismatchedLookupStore()))
+    with pytest.raises(ValueError, match="mismatched"):
+        await mismatched_lookup.get_identity("pai_one", context(key=None))
+    with pytest.raises(ValueError, match="mismatched"):
+        await mismatched_lookup.did_document("one")
+
+    class MismatchedUpdateStore(MemoryIdentityStore):
+        async def update_status(
+            self,
+            agent_identity_id: str,
+            status: ManagedAgentStatus,
+            updated_at: datetime,
+        ) -> IdentityRecord | None:
+            updated = await super().update_status(agent_identity_id, status, updated_at)
+            return None if updated is None else replace(updated, service_did="did:web:other")
+
+    mismatched_update = Platform(options(identity_store=MismatchedUpdateStore()))
+    update_id, _ = await provisioned(mismatched_update)
+    with pytest.raises(ValueError, match="mismatched"):
+        await mismatched_update.update_identity(
+            update_id,
+            PlatformLifecycleRequest(status=ManagedAgentStatus.SUSPENDED),
+            context(key=None),
+        )
+
+    class InvalidListStore(MemoryIdentityStore):
+        async def list(self, principal: str, query: IdentityListQuery) -> IdentityListResult:
+            del principal, query
+            item = replace(record(), principal="owner-one")
+            return IdentityListResult((item, item), 1)
+
+    invalid_list = Platform(options(identity_store=InvalidListStore()))
+    with pytest.raises(ValueError, match="invalid list"):
+        await invalid_list.list(IdentityListQuery(), context(key=None))
+
+    class MalformedListStore(MemoryIdentityStore):
+        async def list(self, principal: str, query: IdentityListQuery) -> IdentityListResult:
+            del principal, query
+            item = replace(record(), created_at=NOW.replace(tzinfo=None), principal="owner-one")
+            return IdentityListResult((item,), 1)
+
+    malformed_list = Platform(options(identity_store=MalformedListStore()))
+    with pytest.raises(ValueError, match="unauthorized"):
+        await malformed_list.list(IdentityListQuery(), context(key=None))
+
 
 @pytest.mark.asyncio
 async def test_sign_handler_validation() -> None:
@@ -521,20 +577,34 @@ async def unhandled(*_: Any) -> None:
     assert authenticated.body.platform_context == {"handle": "opaque"}
     assert received[0].platform_context == {"changed": True, "handle": "opaque"}
 
-    good = PlatformSignCompleted(
-        status="completed",
-        agent_did=authenticated.body.agent_did,
-        client_assertion="jwt",
-        expires_at="2026-01-02T03:04:35Z",
-        issued_at="2026-01-02T03:04:05Z",
-        jti="custom",
-        service_did=SERVICE_DID,
-    )
+    keys = KeyStore()
 
-    async def completed(*_: Any) -> PlatformResult[Any]:
-        return PlatformResult(200, good, "application/aep+json")
+    async def completed(value: Any, request_context: Any) -> PlatformResult[Any]:
+        del request_context
+        claims = ClientAssertionClaims(
+            aud=SERVICE_DID,
+            exp=int(NOW.timestamp()) + 30,
+            iat=int(NOW.timestamp()),
+            iss=value.identity.agent_did,
+            jti="custom",
+            op=AssertionOperation.STATUS,
+            sub=value.identity.agent_did,
+        )
+        return PlatformResult(
+            200,
+            PlatformSignCompleted(
+                status="completed",
+                agent_did=value.identity.agent_did,
+                client_assertion=await keys.sign(value.identity, claims),
+                expires_at="2026-01-02T03:04:35Z",
+                issued_at="2026-01-02T03:04:05Z",
+                jti="custom",
+                service_did=SERVICE_DID,
+            ),
+            "application/aep+json",
+        )
 
-    custom = Platform(options(sign_handler=completed))
+    custom = Platform(options(key_store=keys, sign_handler=completed))
     custom_id, _ = await provisioned(custom)
     result = await custom.sign(
         custom_id,
@@ -549,6 +619,37 @@ async def completed(*_: Any) -> PlatformResult[Any]:
     assert isinstance(result.body, PlatformSignCompleted)
     assert result.body.jti == "custom"
 
+    class InvalidSigner(KeyStore):
+        async def sign(self, identity: IdentityRecord, claims: ClientAssertionClaims) -> str:
+            del identity, claims
+            return "not-a-jwt"
+
+    invalid_signer = Platform(options(key_store=InvalidSigner()))
+    invalid_id, _ = await provisioned(invalid_signer)
+    with pytest.raises(ValueError, match="invalid client assertion"):
+        await invalid_signer.sign(
+            invalid_id,
+            PlatformSignRequest(
+                jti="invalid", op=AssertionOperation.STATUS, service_did=SERVICE_DID
+            ),
+            context("invalid-signer"),
+        )
+
+    class MismatchedSigner(KeyStore):
+        async def sign(self, identity: IdentityRecord, claims: ClientAssertionClaims) -> str:
+            return await super().sign(identity, claims.model_copy(update={"jti": "other"}))
+
+    mismatched_signer = Platform(options(key_store=MismatchedSigner()))
+    signer_id, _ = await provisioned(mismatched_signer)
+    with pytest.raises(ValueError, match="mismatched assertion claims"):
+        await mismatched_signer.sign(
+            signer_id,
+            PlatformSignRequest(
+                jti="mismatched", op=AssertionOperation.STATUS, service_did=SERVICE_DID
+            ),
+            context("mismatched-signer"),
+        )
+
 
 @pytest.mark.asyncio
 async def test_idempotency_store_must_supply_response() -> None:
@@ -619,6 +720,69 @@ async def test_hosted_verification_rejects_unrecognized_assertions() -> None:
     assert isinstance(invalid_signature.body, PlatformVerificationResponse)
     assert invalid_signature.body.verified is False
 
+    unknown_did = "did:web:platform.example:agents:unknown"
+    unknown_assertion = sign_client_assertion(
+        ClientAssertionClaims(
+            aud=SERVICE_DID,
+            exp=int(NOW.timestamp()) + 30,
+            iat=int(NOW.timestamp()),
+            iss=unknown_did,
+            jti="unknown",
+            op=AssertionOperation.STATUS,
+            sub=unknown_did,
+        ),
+        key=keys.key,
+        algorithm=SigningAlgorithm.ES256,
+    )
+    unknown = await platform.verify(
+        PlatformVerificationRequest(
+            client_assertion=unknown_assertion,
+            op=AssertionOperation.STATUS,
+            service_did=SERVICE_DID,
+        ),
+        context("verify-unknown"),
+    )
+    assert isinstance(unknown.body, PlatformVerificationResponse)
+    assert unknown.body.verified is False
+
+    class MismatchedVerificationStore(MemoryIdentityStore):
+        async def find_by_agent_did(self, agent_did: str) -> IdentityRecord | None:
+            identity = await super().find_by_agent_did(agent_did)
+            return (
+                None
+                if identity is None
+                else replace(identity, agent_did="did:web:other", key_id="did:web:other")
+            )
+
+    mismatched_store = MismatchedVerificationStore()
+    mismatched = Platform(
+        options(
+            discovery=discovery,
+            hosted_verification=True,
+            identity_store=mismatched_store,
+            key_store=keys,
+            replay_store=MemoryReplayStore(),
+        )
+    )
+    mismatched_id, _ = await provisioned(mismatched)
+    mismatched_signed = await mismatched.sign(
+        mismatched_id,
+        PlatformSignRequest(
+            jti="mismatched", op=AssertionOperation.STATUS, service_did=SERVICE_DID
+        ),
+        context("mismatched-sign"),
+    )
+    assert isinstance(mismatched_signed.body, PlatformSignCompleted)
+    with pytest.raises(ValueError, match="mismatched"):
+        await mismatched.verify(
+            PlatformVerificationRequest(
+                client_assertion=mismatched_signed.body.client_assertion,
+                op=AssertionOperation.STATUS,
+                service_did=SERVICE_DID,
+            ),
+            context("mismatched-verify"),
+        )
+
     header, _, signature = signed.body.client_assertion.split(".")
     import base64
     import json