feat(xtest): add ML-KEM rewrap session-key roundtrip test

spec/DSPX-4221.md

@@ -0,0 +1,66 @@
+---
+ticket: DSPX-4221
+title: Session Keys should support ML-KEM
+status: in-review
+authors:
+    - dmihalcik@virtru.com
+branches:
+    - opentdf/tests:DSPX-4221-pq-sessions
+    - opentdf/platform:DSPX-4221-pq-sessions
+    - opentdf/java-sdk:DSPX-4221-pq-sessions
+    - opentdf/web-sdk:DSPX-4221-pq-sessions
+prs:
+    - opentdf/tests#571
+    - opentdf/platform#3814
+    - opentdf/java-sdk#388
+    - opentdf/web-sdk#975
+created: 2026-07-31T00:00:00Z
+updated: 2026-08-01T00:00:00Z
+jira_priority: Medium
+---
+
+# Session Keys should support ML-KEM
+
+## Summary
+
+Make sure all clients (and the server) support ML-KEM as the session encryption key (client-generated key pair).
+
+## Problem / Motivation
+
+The rewrap "session key" is the ephemeral key pair a client generates and sends as `clientPublicKey` on a rewrap request, so KAS can wrap the response DEK back to the client. This is a separate concept from the KAS-managed TDF/KAO wrapping key (the `mechanism-mlkem`/`mechanism-xwing`/`mechanism-secpmlkem` features), which already supports post-quantum algorithms. Before this work, the session-key channel only supported RSA and EC, so the rewrap *transport* remained a classical-crypto dependency even for a client and KAS that had otherwise fully adopted PQC-safe wrapping keys — a gap for future-proofing against a cryptographically-relevant quantum computer.
+
+## Proposed Solution
+
+- **Platform (KAS)**: accept pure ML-KEM-768/1024 SPKI client public keys in rewrap, gated behind the same preview flag used for KAS-managed ML-KEM support.
+- **Go, Java, and Web SDKs**: generate an ML-KEM ephemeral session key on request and decapsulate the corresponding rewrap response.
+- **xtest**: a new `session-key-mlkem` feature flag and `test_session_key_mlkem_roundtrip`, which asserts against the KAS rewrap audit log's `sessionKeyType` field rather than just checking that decrypt succeeded — a successful roundtrip alone doesn't prove ML-KEM was actually negotiated, since a client silently falling back to RSA and a server responding in kind would still "work." This test is what caught (and led to a fix for) a real Web SDK bug where the requested session-key algorithm was silently dropped. The platform side of this flag currently reuses the pre-existing `Preview.MLKEMTDFEnabled`/KAS-managed-mechanism probe rather than a dedicated session-key readiness check; each SDK's own hardcoded `session-key-mlkem` capability flag is what actually gates the test on the fix landing (see the comment on `tdfs.py`'s feature-detection block for the known imprecision this leaves on the platform side alone).
+
+## Inputs / Outputs / Contracts
+
+- Client sends `clientPublicKey` as a PEM-encoded SPKI public key on the rewrap request; the server infers the session-key type from the SPKI's algorithm OID (there is no explicit "key type" field on the request).
+- New CLI/API surface accepting `mlkem:768` / `mlkem:1024`:
+  - `otdfctl decrypt --session-key-algorithm mlkem:768`
+  - Go SDK: `sdk.WithSessionKeyType(ocrypto.MLKEM768Key)` (or `ocrypto.MLKEM1024Key`)
+  - Java cmdline: `--rewrap-key-type mlkem:768`
+  - Web SDK CLI: `--rewrapKeyType mlkem:768`
+- New audit field: `eventMetaData.sessionKeyType` on KAS rewrap audit events, recording the negotiated session-key type independently of anything the client reports about itself.
+- xtest: `SDK.decrypt(session_key_algorithm=...)`, threaded through `XT_WITH_SESSION_KEY_ALGORITHM` to each SDK CLI wrapper; `audit_logs.assert_rewrap_success(session_key_type=...)` for verifying the negotiated type.
+
+## Edge Cases & Constraints
+
+- Scope is limited to pure ML-KEM (768/1024); hybrid PQ/T session keys (X-Wing, secp+ML-KEM composites) are explicitly out of scope for this ticket.
+- Gated behind the platform's ML-KEM preview flag; a platform without it enabled rejects an ML-KEM `clientPublicKey` the same way it always rejected any non-RSA/EC key.
+- The session-key algorithm is independent of the TDF's own KAO wrapping mechanism: an RSA-wrapped TDF can be rewrapped over an ML-KEM session key (verified via `test_session_key_mlkem_roundtrip`, which deliberately uses a plain RSA-wrapped attribute).
+
+## Out of Scope
+
+- Hybrid PQ/T session keys (X-Wing, NIST-hybrid EC+ML-KEM composites).
+- Changes to the KAO/TDF wrapping mechanism itself (`mechanism-mlkem`, etc.), which already existed before this work.
+- NanoTDF session keys.
+
+## Acceptance Criteria
+
+- [x] KAS rewrap accepts ML-KEM-768/1024 client session keys, gated by the platform's ML-KEM preview flag.
+- [x] Go SDK, Java SDK, and Web SDK can each generate an ML-KEM session key and successfully decrypt a rewrap response wrapped to it.
+- [x] Cross-SDK interop verified: any encrypt SDK paired with any decrypt SDK, for both `mlkem:768` and `mlkem:1024`.
+- [x] xtest coverage asserts the negotiated session-key type via the KAS audit log, not just roundtrip success.

xtest/audit_logs.py

@@ -436,6 +436,16 @@ def algorithm(self) -> str | None:
         """Get the algorithm from rewrap event metadata."""
         return self.event_metadata.get("algorithm")
 
+    @property
+    def session_key_type(self) -> str | None:
+        """Get the client's rewrap session-key type from rewrap event metadata.
+
+        This is the ephemeral key the client generated and sent as
+        clientPublicKey (e.g. "rsa:2048", "ec:secp256r1", "mlkem:768") --
+        distinct from `algorithm`, which is the KAO/TDF wrapping algorithm.
+        """
+        return self.event_metadata.get("sessionKeyType")
+
     @property
     def tdf_format(self) -> str | None:
         """Get the TDF format from rewrap event metadata."""
@@ -467,6 +477,7 @@ def matches_rewrap(
         policy_uuid: str | None = None,
         key_id: str | None = None,
         algorithm: str | None = None,
+        session_key_type: str | None = None,
         attr_fqns: list[str] | None = None,
     ) -> bool:
         """Check if this event matches rewrap criteria.
@@ -476,6 +487,7 @@ def matches_rewrap(
             policy_uuid: Expected policy UUID (object ID)
             key_id: Expected key ID from metadata
             algorithm: Expected algorithm from metadata
+            session_key_type: Expected client session-key type from metadata
             attr_fqns: Expected attribute FQNs (all must be present)
 
         Returns:
@@ -491,6 +503,8 @@ def matches_rewrap(
             return False
         if algorithm is not None and self.algorithm != algorithm:
             return False
+        if session_key_type is not None and self.session_key_type != session_key_type:
+            return False
         if attr_fqns is not None:
             event_attrs = set(self.object_attrs)
             if not all(fqn in event_attrs for fqn in attr_fqns):
@@ -1254,6 +1268,7 @@ def assert_rewrap(
         policy_uuid: str | None = None,
         key_id: str | None = None,
         algorithm: str | None = None,
+        session_key_type: str | None = None,
         attr_fqns: list[str] | None = None,
         min_count: int = 1,
         since_mark: str | None = None,
@@ -1264,13 +1279,16 @@ def assert_rewrap(
         Looks for audit log entries with:
         - msg='rewrap'
         - action.result=<result>
-        - Optionally matching policy_uuid, key_id, algorithm, attr_fqns
+        - Optionally matching policy_uuid, key_id, algorithm, session_key_type, attr_fqns
 
         Args:
             result: Expected action result ('success', 'failure', 'error', 'cancel')
             policy_uuid: Expected policy UUID (object.id)
             key_id: Expected key ID from eventMetaData.keyID
             algorithm: Expected algorithm from eventMetaData.algorithm
+            session_key_type: Expected client session-key type from
+                eventMetaData.sessionKeyType (e.g. "mlkem:768") -- the
+                client's ephemeral rewrap key, not the KAO wrap algorithm
             attr_fqns: Expected attribute FQNs (all must be present)
             min_count: Minimum number of matching entries (default: 1)
             since_mark: Only check logs since marked timestamp
@@ -1306,6 +1324,7 @@ def assert_rewrap(
                     policy_uuid=policy_uuid,
                     key_id=key_id,
                     algorithm=algorithm,
+                    session_key_type=session_key_type,
                     attr_fqns=attr_fqns,
                 ):
                     matching.append(event)
@@ -1331,6 +1350,8 @@ def assert_rewrap(
             criteria.append(f"key_id={key_id}")
         if algorithm:
             criteria.append(f"algorithm={algorithm}")
+        if session_key_type:
+            criteria.append(f"session_key_type={session_key_type}")
         if attr_fqns:
             criteria.append(f"attr_fqns={attr_fqns}")
 
@@ -1349,6 +1370,7 @@ def assert_rewrap_success(
         policy_uuid: str | None = None,
         key_id: str | None = None,
         algorithm: str | None = None,
+        session_key_type: str | None = None,
         attr_fqns: list[str] | None = None,
         min_count: int = 1,
         since_mark: str | None = None,
@@ -1363,17 +1385,78 @@ def assert_rewrap_success(
             policy_uuid=policy_uuid,
             key_id=key_id,
             algorithm=algorithm,
+            session_key_type=session_key_type,
             attr_fqns=attr_fqns,
             min_count=min_count,
             since_mark=since_mark,
             timeout=timeout,
         )
 
+    def assert_rewrap_session_key_type(
+        self,
+        expected: str,
+        since_mark: str | None = None,
+        min_count: int = 1,
+        timeout: float = 20.0,
+    ) -> None:
+        """Assert a successful rewrap negotiated the expected client session-key type.
+
+        Tolerant of platform builds that don't emit eventMetaData.sessionKeyType
+        at all: that field was added by DSPX-4221 and isn't behind any version
+        or preview flag we can check statically, so a plain
+        assert_rewrap_success(session_key_type=expected) call would spuriously
+        find zero matches (not "wrong type") against a baseline/pre-fix
+        platform build -- indistinguishable, from the caller's perspective,
+        from a real negotiation bug.
+
+        Waits (up to `timeout`) for an event matching the expected type first,
+        rather than accepting the first bare "result=success" match: log
+        collection is poll-based, so when two rewraps for different session
+        keys happen back-to-back in the same test (e.g. a plain decrypt
+        immediately followed by one requesting a specific algorithm), the
+        earlier rewrap's log line can still be un-tailed at the time of the
+        later mark and get folded into the same collection batch -- a bare
+        count check would then report success using the wrong (earlier)
+        event instead of waiting for the right one. Only after that wait
+        fails do we check whether the field is present at all, to tell
+        "platform doesn't support this field" apart from a real negotiation
+        bug.
+        """
+        try:
+            self.assert_rewrap(
+                result="success",
+                session_key_type=expected,
+                min_count=min_count,
+                since_mark=since_mark,
+                timeout=timeout,
+            )
+            return
+        except AssertionError:
+            pass
+
+        events = self.assert_rewrap_success(
+            min_count=min_count, since_mark=since_mark, timeout=1.0
+        )
+        reported = {
+            e.session_key_type for e in events if e.session_key_type is not None
+        }
+        if not reported:
+            logger.warning(
+                "Platform build doesn't emit eventMetaData.sessionKeyType; "
+                f"skipping the check that the session key type was {expected!r} "
+                "(the rewrap itself succeeded)."
+            )
+            return
+        assert expected in reported, (
+            f"Expected rewrap session_key_type={expected!r}, but platform reported {reported!r}"
+        )
+
     def assert_rewrap_failure(
         self,
         policy_uuid: str | None = None,
         key_id: str | None = None,
         algorithm: str | None = None,
+        session_key_type: str | None = None,
         attr_fqns: list[str] | None = None,
         min_count: int = 1,
         since_mark: str | None = None,
@@ -1389,6 +1472,7 @@ def assert_rewrap_failure(
             policy_uuid=policy_uuid,
             key_id=key_id,
             algorithm=algorithm,
+            session_key_type=session_key_type,
             attr_fqns=attr_fqns,
             min_count=min_count,
             since_mark=since_mark,

xtest/sdk/go/cli.sh

@@ -17,6 +17,7 @@
 #  XT_WITH_ASSERTION_VERIFICATION_KEYS [string] - Path to assertion verification private key file
 #  XT_WITH_ATTRIBUTES [string] - Attributes to be used for encryption
 #  XT_WITH_MIME_TYPE [string] - MIME type for the encrypted file
+#  XT_WITH_SESSION_KEY_ALGORITHM [string] - Rewrap session key algorithm for decryption (e.g. mlkem:768)
 #
 SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)
 
@@ -115,6 +116,11 @@ if [ "$1" == "supports" ]; then
       "${cmd[@]}" help policy kas-registry key create | grep -iE 'mlkem:768|mlkem:1024'
       exit $?
       ;;
+    session-key-mlkem)
+      set -o pipefail
+      "${cmd[@]}" help decrypt | grep -iE 'mlkem:768|mlkem:1024'
+      exit $?
+      ;;
     dpop | dpop_nonce_challenge)
       set -o pipefail
       "${cmd[@]}" --version --json | jq -e --arg f "$2" '.supported_features |
@@ -200,6 +206,9 @@ elif [ "$1" == "decrypt" ]; then
   if [ "$XT_WITH_ECWRAP" == 'true' ]; then
     args+=(--session-key-algorithm "ec:secp256r1")
   fi
+  if [[ -n "$XT_WITH_SESSION_KEY_ALGORITHM" ]]; then
+    args+=(--session-key-algorithm "$XT_WITH_SESSION_KEY_ALGORITHM")
+  fi
   if [ "$XT_WITH_VERIFY_ASSERTIONS" == 'false' ]; then
     args+=(--no-verify-assertions)
   fi

xtest/sdk/java/cli.sh

@@ -18,6 +18,7 @@
 #  XT_WITH_ATTRIBUTES [string] - Attributes to be used for encryption
 #  XT_WITH_MIME_TYPE [string] - MIME type for the encrypted file
 #  XT_WITH_TARGET_MODE [string] - Target spec mode for the encrypted file
+#  XT_WITH_SESSION_KEY_ALGORITHM [string] - Rewrap session key algorithm for decryption (e.g. mlkem:768)
 #
 SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)
 
@@ -125,6 +126,16 @@ if [ "$1" == "supports" ]; then
       java -jar "$SCRIPT_DIR"/cmdline.jar help encrypt | grep -i "mlkem:768"
       exit $?
       ;;
+    session-key-mlkem)
+      # --rewrap-key-type has long accepted "mlkem:768" as a value (its choices
+      # come from the same KeyType enum used for KAS-managed-key mechanisms),
+      # so grepping --help for it would false-positive on builds that predate
+      # KASClient actually decapsulating an ML-KEM rewrap response. Use the
+      # explicit `supports` subcommand instead, which is a hardcoded,
+      # source-controlled feature list scoped to this exact capability.
+      java -jar "$SCRIPT_DIR"/cmdline.jar supports session-key-mlkem
+      exit $?
+      ;;
     mechanism-rsa-4096 | mechanism-ec-curves-384-521)
       # rsa4096 support in >= 0.13.0
       set -o pipefail
@@ -189,6 +200,9 @@ else
   if [ "$XT_WITH_ECWRAP" == 'true' ]; then
     args+=(--rewrap-key-type="ec:secp256r1")
   fi
+  if [[ -n "$XT_WITH_SESSION_KEY_ALGORITHM" ]]; then
+    args+=(--rewrap-key-type="$XT_WITH_SESSION_KEY_ALGORITHM")
+  fi
 fi
 
 if [ "$1" == "decrypt" ]; then

xtest/sdk/js/cli.sh

@@ -20,6 +20,7 @@
 #  XT_WITH_TARGET_MODE [string] - Target spec mode for the encrypted file
 #  XT_WITH_DPOP [string] - Enable DPoP token binding; value selects algorithm (e.g. ES256)
 #  XT_WITH_DPOP_KEY [string] - Path to PEM-encoded PKCS8 private key for DPoP signing
+#  XT_WITH_SESSION_KEY_ALGORITHM [string] - Rewrap session key algorithm for decryption (e.g. mlkem:768)
 #  CLIENTID [string] - Override OIDC client ID (default: opentdf)
 #  CLIENTSECRET [string] - Override OIDC client secret (default: secret)
 #
@@ -100,6 +101,17 @@ if [[ "$1" == "supports" ]]; then
       npx $CTL encrypt --help | grep -i 'mlkem:768'
       exit $?
       ;;
+    session-key-mlkem)
+      # --rewrapKeyType has long accepted "mlkem:768" as a choice (it shares
+      # PUBLIC_KEY_ALGORITHMS with --encapKeyType, used for KAS-managed-key
+      # mechanisms), so grepping --help for it would false-positive on builds
+      # that predate decryptStreamFrom() actually forwarding the requested
+      # algorithm to unwrapKey(). Use the explicit supportedFeatures list from
+      # --version instead, which is hardcoded and scoped to this capability.
+      set -o pipefail
+      npx $CTL --version | jq -e '.supportedFeatures | index("session-key-mlkem")' >/dev/null
+      exit $?
+      ;;
     mechanism-xwing)
       set -o pipefail
       npx $CTL help | grep -i xwing
@@ -255,6 +267,9 @@ elif [[ "$1" == "decrypt" ]]; then
   if [[ "$XT_WITH_ECWRAP" == 'true' ]]; then
     args+=(--rewrapKeyType "ec:secp256r1")
   fi
+  if [[ -n "$XT_WITH_SESSION_KEY_ALGORITHM" ]]; then
+    args+=(--rewrapKeyType "$XT_WITH_SESSION_KEY_ALGORITHM")
+  fi
   if [[ -n "$XT_WITH_KAS_ALLOW_LIST" ]]; then
     args+=(--allowList "$XT_WITH_KAS_ALLOW_LIST")
   fi