From 31755c3fbbca9893630fc412a48bd1984ade9505 Mon Sep 17 00:00:00 2001 From: Nas Kavian Date: Wed, 2 Sep 2026 23:03:52 -0700 Subject: [PATCH] fix: validate platform extension results --- .../platform/platform.py | 71 ++++++- tests/test_platform_boundaries.py | 188 ++++++++++++++++-- 2 files changed, 242 insertions(+), 17 deletions(-) diff --git a/src/agent_enrollment_protocol/platform/platform.py b/src/agent_enrollment_protocol/platform/platform.py index 259b2e1..71371d0 100644 --- a/src/agent_enrollment_protocol/platform/platform.py +++ b/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: diff --git a/tests/test_platform_boundaries.py b/tests/test_platform_boundaries.py index 0b0ca35..5f7d231 100644 --- a/tests/test_platform_boundaries.py +++ b/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