fix: harden SDK protocol boundaries

README.md

@@ -59,6 +59,10 @@ The Service module includes stored credential profiles for each built-in Grant T
 Each factory accepts application-owned credential issuance and storage implementations. The included
 memory store is intended for examples and local development.
 
+Concrete Grant Types and extensions can define additional Grant and Revoke request members. Pass
+those members through `GrantOptions.parameters` or `RevokeOptions.parameters`; the Agent retains
+control of the standard Grant Type and Revoke selector fields.
+
 ## Service ASGI integration
 
 `agent_enrollment_protocol.adapters` provides a framework-neutral ASGI integration with no

src/agent_enrollment_protocol/agent/client.py

@@ -16,6 +16,7 @@
     AEP_MEDIA_TYPE,
     AEP_PROBLEM_MEDIA_TYPE,
     AEP_WELL_KNOWN_PATH,
+    AepAssertionError,
     AepValidationError,
     AgentStatus,
     ApiKeyGrantResponse,
@@ -38,6 +39,7 @@
     SigningAlgorithm,
     StatusResponse,
     command_path_from_inspect,
+    decode_jwt_unverified,
     did_web_document_url,
     media_type_essence,
     missing_required_claim_names,
@@ -113,7 +115,8 @@ def __init__(self, options: AgentOptions) -> None:
             raise ValueError("AEP Agent request timeout must be positive and finite")
         self._allow_insecure_loopback = options.allow_insecure_loopback
         self._assertion_lifetime = int(lifetime)
-        self._clock = options.clock
+        self._clock = _validated_clock(options.clock)
+        self._clock()
         owned_transports: list[AsyncHttpTransport] = []
         command_transport: AsyncHttpTransport
         if options.command_transport is None:
@@ -124,7 +127,7 @@ def __init__(self, options: AgentOptions) -> None:
         else:
             command_transport = options.command_transport
         self._command_transport = command_transport
-        self._credential_store = options.credential_store or MemoryCredentialStore(options.clock)
+        self._credential_store = options.credential_store or MemoryCredentialStore(self._clock)
         self._identity_provider = options.identity_provider
         self._identity_store = options.identity_store or MemoryIdentityStore()
         self._idempotency_keys = options.idempotency_keys or RandomIdempotencyKeyProvider()
@@ -379,7 +382,8 @@ async def grant(self, options: GrantOptions | None = None) -> CommandResult[Gran
         key = await self._idempotency_key(
             inspection, Command.GRANT.value, options.idempotency_key, grant_type=grant_type
         )
-        request_data: dict[str, object] = {"grant_type": grant_type}
+        request_data: dict[str, object] = dict(options.parameters)
+        request_data["grant_type"] = grant_type
         if options.requested_scopes:
             request_data["requested_scopes"] = options.requested_scopes
         request = GrantRequest.model_validate(request_data)
@@ -395,6 +399,16 @@ async def grant(self, options: GrantOptions | None = None) -> CommandResult[Gran
         )
         raw = result.body
         credential = _parse_credential(grant_type, raw)
+        if isinstance(credential, ApiKeyGrantResponse):
+            config = (inspection.document.commands.grant_types_config or {}).get(grant_type)
+            configured_headers = None if config is None else config.to_wire().get("header_names")
+            if configured_headers is not None:
+                if not isinstance(configured_headers, list) or any(
+                    not isinstance(value, str) for value in configured_headers
+                ):
+                    raise ValueError("AEP API-key header_names configuration is invalid")
+                if credential.header.lower() not in {value.lower() for value in configured_headers}:
+                    raise ValueError("AEP API-key Grant response header was not advertised")
         grant_result = GrantResult(credential=credential, grant_type=grant_type, raw=raw)
         if credential is not None:
             record = CredentialRecord(
@@ -411,9 +425,10 @@ async def grant(self, options: GrantOptions | None = None) -> CommandResult[Gran
 
     async def revoke(self, options: RevokeOptions) -> CommandResult[RevokeResponse]:
         if options.all_grant_types:
-            request_data: dict[str, object] = {"all_grant_types": "true"}
+            request_data: dict[str, object] = dict(options.parameters)
+            request_data["all_grant_types"] = "true"
         else:
-            request_data = {}
+            request_data = dict(options.parameters)
             if options.grant_type is not None:
                 request_data["grant_type"] = options.grant_type
             if options.credential_id is not None:
@@ -457,7 +472,7 @@ async def revoke(self, options: RevokeOptions) -> CommandResult[RevokeResponse]:
             )
             if matches:
                 await self._agent._credential_store.delete_credential(
-                    record.service_did, record.credential_id
+                    inspection.document.service.did, record.credential_id
                 )
         return result
 
@@ -656,6 +671,9 @@ async def _sign_assertion(
         assertion = await signer(claims, tuple(SigningAlgorithm(value) for value in algorithms))
         if not assertion:
             raise ValueError("AEP assertion signer returned an empty assertion")
+        _validate_signed_assertion(
+            assertion, claims, algorithms, self._agent._allow_insecure_loopback
+        )
         return assertion
 
     async def _find_credential(
@@ -829,6 +847,43 @@ def _validate_credential_record(record: CredentialRecord, service_did: str, now:
         raise ValueError("stored AEP credential metadata does not match its payload")
 
 
+def _validated_clock(clock: Clock) -> Clock:
+    def current_time() -> datetime:
+        value = clock()
+        if value.utcoffset() is None:
+            raise ValueError("AEP Agent clock must return a time with a UTC offset")
+        return value
+
+    return current_time
+
+
+def _validate_signed_assertion(
+    assertion: str,
+    claims: ClientAssertionClaims,
+    algorithms: tuple[str, ...],
+    allow_insecure_loopback: bool,
+) -> None:
+    try:
+        header, payload = decode_jwt_unverified(assertion)
+        parsed_claims = ClientAssertionClaims.model_validate_json(
+            json.dumps(payload),
+            context={"allow_insecure_loopback": allow_insecure_loopback},
+        )
+    except (AepAssertionError, ValueError) as error:
+        raise ValueError("AEP assertion signer returned an invalid assertion") from error
+    key_id = header.get("kid")
+    if (
+        header.get("typ") != "JWT"
+        or header.get("alg") not in algorithms
+        or not isinstance(key_id, str)
+        or key_id.partition("#")[0] != claims.iss
+        or parsed_claims != claims
+    ):
+        raise ValueError(
+            "AEP assertion signer returned an assertion that does not match the request"
+        )
+
+
 def _inspection_from_cache(
     entry: InspectCacheEntry, inspect_url: str, service_url: str
 ) -> Inspection:

src/agent_enrollment_protocol/agent/types.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Awaitable, Callable, Mapping, Sequence
+from copy import deepcopy
 from dataclasses import dataclass, field
 from datetime import datetime
 from types import MappingProxyType
@@ -74,7 +75,7 @@ class CredentialRecord:
     expires_at: datetime
     grant_type: str
     issued_at: datetime
-    payload: bytes
+    payload: bytes = field(repr=False)
     service_did: str
     service_url: str
 
@@ -145,19 +146,21 @@ class EnrollOptions:
 class GrantOptions:
     grant_type: str | None = None
     idempotency_key: str | None = None
+    parameters: Mapping[str, object] = field(default_factory=lambda: MappingProxyType({}))
     preferred_grant_types: tuple[str, ...] = ()
     requested_scopes: tuple[str, ...] = ()
 
     def __post_init__(self) -> None:
         object.__setattr__(self, "preferred_grant_types", tuple(self.preferred_grant_types))
+        object.__setattr__(self, "parameters", _copy_parameters(self.parameters))
         object.__setattr__(self, "requested_scopes", tuple(self.requested_scopes))
 
 
 @dataclass(frozen=True, slots=True)
 class GrantResult:
-    credential: BuiltInCredential | None
+    credential: BuiltInCredential | None = field(repr=False)
     grant_type: str
-    raw: bytes
+    raw: bytes = field(repr=False)
 
 
 @dataclass(frozen=True, slots=True)
@@ -166,6 +169,10 @@ class RevokeOptions:
     credential_id: str | None = None
     grant_type: str | None = None
     idempotency_key: str | None = None
+    parameters: Mapping[str, object] = field(default_factory=lambda: MappingProxyType({}))
+
+    def __post_init__(self) -> None:
+        object.__setattr__(self, "parameters", _copy_parameters(self.parameters))
 
 
 @dataclass(frozen=True, slots=True)
@@ -203,3 +210,9 @@ def __init__(self, missing: Sequence[str]) -> None:
             "AEP Agent cannot satisfy the Service's required Claim Names: " + ", ".join(values)
         )
         self.missing = values
+
+
+def _copy_parameters(parameters: Mapping[str, object]) -> Mapping[str, object]:
+    if any(not isinstance(name, str) or not name for name in parameters):
+        raise ValueError("AEP extension parameter names must be non-empty strings")
+    return MappingProxyType(deepcopy(dict(parameters)))

src/agent_enrollment_protocol/core/assertions.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
 import json
+import re
 import time
-from base64 import urlsafe_b64decode
+from base64 import b64decode
 from dataclasses import dataclass
 from typing import Any
 
@@ -11,6 +12,8 @@
 from .errors import AepAssertionError
 from .models import AssertionOperation, ClientAssertionClaims, SigningAlgorithm
 
+_BASE64URL_PATTERN = re.compile(r"\A[A-Za-z0-9_-]+\Z")
+
 
 @dataclass(frozen=True, slots=True)
 class VerifyClientAssertionOptions:
@@ -103,7 +106,7 @@ def verify_client_assertion(
 
 def decode_jwt_unverified(assertion: str) -> tuple[dict[str, Any], dict[str, Any]]:
     parts = assertion.split(".")
-    if len(parts) != 3:
+    if len(parts) != 3 or any(_BASE64URL_PATTERN.fullmatch(part) is None for part in parts):
         raise AepAssertionError("Invalid JWT.")
     try:
         header = _decode_part(parts[0])
@@ -122,8 +125,17 @@ def _require_key_binding(key_id: str, claims: ClientAssertionClaims) -> None:
 
 
 def _decode_part(value: str) -> dict[str, Any]:
-    decoded = urlsafe_b64decode(value + "=" * (-len(value) % 4)).decode()
-    parsed = json.loads(decoded)
+    decoded = b64decode(value + "=" * (-len(value) % 4), altchars=b"-_", validate=True).decode()
+
+    def pairs(items: list[tuple[str, Any]]) -> dict[str, Any]:
+        result: dict[str, Any] = {}
+        for key, item in items:
+            if key in result:
+                raise ValueError(f"duplicate object member: {key}")
+            result[key] = item
+        return result
+
+    parsed = json.loads(decoded, object_pairs_hook=pairs)
     if not isinstance(parsed, dict):
         raise ValueError("JWT part must be an object")
     return parsed

src/agent_enrollment_protocol/core/did_web.py

@@ -1,38 +1,26 @@
 from __future__ import annotations
 
+import re
 from collections.abc import Mapping
+from copy import deepcopy
 from typing import Any
 from urllib.parse import unquote, urlsplit
 
+_INVALID_PERCENT_ENCODING = re.compile(r"%(?![0-9A-Fa-f]{2})")
+
 
 def did_web_document_url(did: str, *, allow_insecure_loopback: bool = False) -> str:
-    prefix = "did:web:"
-    if not did.startswith(prefix):
-        raise ValueError(f"Unsupported DID method: {did}")
-    parts = did[len(prefix) :].split(":")
-    if not parts[0]:
-        raise ValueError(f"Invalid did:web identifier: {did}")
-    host = unquote(parts[0])
-    path = (
-        "/.well-known/did.json"
-        if len(parts) == 1
-        else f"/{'/'.join(unquote(part) for part in parts[1:])}/did.json"
-    )
+    host, path_parts = _did_web_parts(did)
+    path = "/.well-known/did.json" if not path_parts else f"/{'/'.join(path_parts)}/did.json"
     scheme = "http" if allow_insecure_loopback and _is_loopback(host) else "https"
-    url = f"{scheme}://{host}{path}"
-    parsed = urlsplit(url)
-    try:
-        _ = parsed.port
-    except ValueError as error:
-        raise ValueError(f"Invalid did:web identifier: {did}") from error
-    if not parsed.hostname or parsed.username or parsed.password or parsed.query or parsed.fragment:
-        raise ValueError(f"Invalid did:web identifier: {did}")
-    return url
+    return f"{scheme}://{host}{path}"
 
 
 def select_did_web_public_jwk(
     document: Mapping[str, Any], *, did: str, key_id: str
 ) -> dict[str, Any]:
+    if document.get("id") != did:
+        raise ValueError("AEP did:web document ID does not identify the assertion issuer")
     key_did = key_id.partition("#")[0]
     if key_did != did:
         raise ValueError("AEP did:web key ID does not identify the assertion issuer")
@@ -44,10 +32,46 @@ def select_did_web_public_jwk(
             continue
         jwk = method.get("publicKeyJwk")
         if isinstance(jwk, dict):
-            return dict(jwk)
+            return deepcopy(jwk)
     raise ValueError(f"No public JWK found for {key_id}")
 
 
 def _is_loopback(host: str) -> bool:
     hostname = urlsplit(f"//{host}").hostname
     return hostname in {"localhost", "127.0.0.1", "::1"}
+
+
+def _did_web_parts(did: str) -> tuple[str, tuple[str, ...]]:
+    prefix = "did:web:"
+    if not did.startswith(prefix):
+        raise ValueError(f"Unsupported DID method: {did}")
+    encoded_parts = did[len(prefix) :].split(":")
+    if (
+        not did.isascii()
+        or not encoded_parts[0]
+        or any(_INVALID_PERCENT_ENCODING.search(part) for part in encoded_parts)
+    ):
+        raise ValueError(f"Invalid did:web identifier: {did}")
+    host = unquote(encoded_parts[0])
+    if not host.isascii():
+        raise ValueError(f"Invalid did:web identifier: {did}")
+    parsed = urlsplit(f"//{host}")
+    try:
+        _ = parsed.port
+    except ValueError as error:
+        raise ValueError(f"Invalid did:web identifier: {did}") from error
+    if (
+        not parsed.hostname
+        or parsed.username
+        or parsed.password
+        or parsed.path
+        or parsed.query
+        or parsed.fragment
+    ):
+        raise ValueError(f"Invalid did:web identifier: {did}")
+    path_parts = tuple(encoded_parts[1:])
+    for part in path_parts:
+        decoded = unquote(part)
+        if not part or decoded in {".", ".."} or any(value in decoded for value in "/\\?#"):
+            raise ValueError(f"Invalid did:web identifier: {did}")
+    return host, path_parts

src/agent_enrollment_protocol/core/http.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import re
 from collections.abc import Mapping
 from dataclasses import dataclass, field
 from types import MappingProxyType
@@ -14,6 +15,13 @@
     ProtectedResourceAuthorization,
 )
 
+_TOKEN = r"[!#$%&'*+.^_`|~0-9A-Za-z-]+"
+_QUOTED_STRING = r'"(?:[\t !#-\[\]-~\x80-\xff]|\\[\t !-~\x80-\xff])*"'
+_MEDIA_TYPE_PATTERN = re.compile(
+    rf"\A[ \t]*(?P<type>{_TOKEN})/(?P<subtype>{_TOKEN})[ \t]*"
+    rf"(?:;[ \t]*{_TOKEN}[ \t]*=[ \t]*(?:{_TOKEN}|{_QUOTED_STRING})[ \t]*)*\Z"
+)
+
 
 @dataclass(frozen=True, slots=True)
 class HttpRequest:
@@ -37,7 +45,10 @@ def __post_init__(self) -> None:
 
 
 def media_type_essence(value: str) -> str:
-    return value.partition(";")[0].strip().lower()
+    match = _MEDIA_TYPE_PATTERN.fullmatch(value)
+    if match is None:
+        return ""
+    return f"{match.group('type')}/{match.group('subtype')}".lower()
 
 
 def normalize_endpoint_base(endpoint_base: str = DEFAULT_HTTP_ENDPOINT_BASE) -> str: