feat(xtest): add ML-KEM rewrap session-key roundtrip test - #571
feat(xtest): add ML-KEM rewrap session-key roundtrip test#571dmihalcik-virtru wants to merge 2 commits into
Conversation
pq-sessions
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThe change defines ML-KEM session-key support across the specification, shared SDK test framework, Go, Java, and JavaScript CLI wrappers. It adds audit-log matching and round-trip tests for ML-KEM-768, ML-KEM-1024, and RSA session keys. ChangesML-KEM session-key support
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant Test as xtest/test_pqc.py
participant SDK as SDK.decrypt
participant Wrapper as SDK CLI wrapper
participant CLI as otdfctl
participant Audit as AuditLogAsserter
Test->>SDK: Set session_key_algorithm
SDK->>Wrapper: Export XT_WITH_SESSION_KEY_ALGORITHM
Wrapper->>CLI: Pass the rewrap-key option
CLI-->>SDK: Return decrypted plaintext
SDK-->>Test: Return plaintext
Test->>Audit: Match session_key_type
Audit-->>Test: Confirm the rewrap audit event
Possibly related PRs
Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
xtest/tdfs.py (1)
578-596: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winDefine precedence for
ecwrapvssession_key_algorithm.
SDK.decrypt()can pass both options, but the Go, Java, and JavaScript CLI shims emit them in separate order. This lets the CLI see two rewrap/session-key values and either fail or pick one implicitly. Reject the combination or emit one option with explicit precedence, such assession_key_algorithmoverridingecwrap.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@xtest/tdfs.py` around lines 578 - 596, Update the option construction in the decrypt helper around session_key_algorithm and ecwrap to define deterministic precedence when both are provided. Reject the combination explicitly or emit only one setting, with session_key_algorithm taking precedence over ecwrap, and ensure the generated environment cannot contain conflicting rewrap/session-key values.
🧹 Nitpick comments (2)
spec/DSPX-4221.md (1)
19-19: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse
client-generatedin the session-key description.Replace
client generated key pairwithclient-generated key pair.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@spec/DSPX-4221.md` at line 19, Update the session encryption key description in DSPX-4221 to use the hyphenated phrase “client-generated key pair” instead of “client generated key pair,” without changing the surrounding requirement.Source: Linters/SAST tools
xtest/sdk/go/cli.sh (1)
209-211: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse
[[in both new Bash conditionals.SonarCloud reports the same conditional-style issue at both sites.
xtest/sdk/go/cli.sh#L209-L211: replace[ -n "$XT_WITH_SESSION_KEY_ALGORITHM" ]with[[ -n "$XT_WITH_SESSION_KEY_ALGORITHM" ]].xtest/sdk/java/cli.sh#L198-L200: replace[ -n "$XT_WITH_SESSION_KEY_ALGORITHM" ]with[[ -n "$XT_WITH_SESSION_KEY_ALGORITHM" ]].🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@xtest/sdk/go/cli.sh` around lines 209 - 211, Update the XT_WITH_SESSION_KEY_ALGORITHM condition in xtest/sdk/go/cli.sh at lines 209-211 and xtest/sdk/java/cli.sh at lines 198-200 to use Bash [[ -n ... ]] syntax instead of single brackets, preserving the existing argument-appending behavior.Source: Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@xtest/tdfs.py`:
- Around line 152-156: Make ML-KEM session-key capability detection
algorithm-specific across xtest/tdfs.py:152-156 and 252-254,
xtest/sdk/go/cli.sh:119-123, xtest/sdk/java/cli.sh:129-133,
xtest/sdk/js/cli.sh:104-108, and xtest/test_pqc.py:379-403: define or expose
distinct capabilities for mlkem:768 and mlkem:1024, include mlkem:1024 in
fallback handling while preserving exact algorithm results, have each SDK probe
and report only the requested algorithm, and gate each parameterized PQC test
with its matching capability.
In `@xtest/test_pqc.py`:
- Around line 399-400: Update the pytest.skip call in the focus-filter logic to
provide an actionable reason, including the excluded encrypt_sdk and decrypt_sdk
values or clearly describing the focus rule; preserve the existing skip
condition.
- Line 421: Update the filecmp.cmp assertion in the decrypted-output test to
pass shallow=False, ensuring pt_file and rt_file are compared by byte content
rather than metadata alone.
---
Outside diff comments:
In `@xtest/tdfs.py`:
- Around line 578-596: Update the option construction in the decrypt helper
around session_key_algorithm and ecwrap to define deterministic precedence when
both are provided. Reject the combination explicitly or emit only one setting,
with session_key_algorithm taking precedence over ecwrap, and ensure the
generated environment cannot contain conflicting rewrap/session-key values.
---
Nitpick comments:
In `@spec/DSPX-4221.md`:
- Line 19: Update the session encryption key description in DSPX-4221 to use the
hyphenated phrase “client-generated key pair” instead of “client generated key
pair,” without changing the surrounding requirement.
In `@xtest/sdk/go/cli.sh`:
- Around line 209-211: Update the XT_WITH_SESSION_KEY_ALGORITHM condition in
xtest/sdk/go/cli.sh at lines 209-211 and xtest/sdk/java/cli.sh at lines 198-200
to use Bash [[ -n ... ]] syntax instead of single brackets, preserving the
existing argument-appending behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 05ddc315-d152-4758-ad81-2a05754b1f38
📒 Files selected for processing (6)
spec/DSPX-4221.mdxtest/sdk/go/cli.shxtest/sdk/java/cli.shxtest/sdk/js/cli.shxtest/tdfs.pyxtest/test_pqc.py
baacb54 to
66d032f
Compare
66d032f to
56a9e8f
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@xtest/sdk/go/cli.sh`:
- Line 209: Update the new XT_WITH_SESSION_KEY_ALGORITHM conditionals to use
Bash [[ -n ... ]] syntax in xtest/sdk/go/cli.sh lines 209-209 and
xtest/sdk/java/cli.sh lines 200-200, preserving the existing non-empty check.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 85a7130d-9c13-48e1-b7df-a64fb5ca9b08
📒 Files selected for processing (6)
xtest/audit_logs.pyxtest/sdk/go/cli.shxtest/sdk/java/cli.shxtest/sdk/js/cli.shxtest/tdfs.pyxtest/test_pqc.py
🚧 Files skipped from review as they are similar to previous changes (3)
- xtest/sdk/js/cli.sh
- xtest/tdfs.py
- xtest/test_pqc.py
9ffd4d6 to
1f7c8b2
Compare
X-Test Failure Report |
1f7c8b2 to
0cefcf0
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@spec/DSPX-4221.md`:
- Line 53: Update the session-key algorithm statement in DSPX-4221 to match the
existing test coverage by removing “and vice versa” from the independence claim;
do not add a reverse-combination test.
- Around line 34-36: Update test_session_key_mlkem_roundtrip readiness gating in
tdfs.py to require a session-key-specific ML-KEM capability probe or an explicit
platform build that accepts ML-KEM clientPublicKey, rather than relying solely
on the shared session-key-mlkem flag. Preserve the test’s existing audit-log
sessionKeyType assertion once the capability check passes.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: ee99aa8a-c8d3-462e-85ec-725886e64223
📒 Files selected for processing (7)
spec/DSPX-4221.mdxtest/audit_logs.pyxtest/sdk/go/cli.shxtest/sdk/java/cli.shxtest/sdk/js/cli.shxtest/tdfs.pyxtest/test_pqc.py
🚧 Files skipped from review as they are similar to previous changes (6)
- xtest/sdk/js/cli.sh
- xtest/sdk/go/cli.sh
- xtest/test_pqc.py
- xtest/sdk/java/cli.sh
- xtest/audit_logs.py
- xtest/tdfs.py
398f300 to
17eda9c
Compare
There was a problem hiding this comment.
Pull request overview
Adds end-to-end xtest coverage for ML-KEM as the client-generated rewrap session key, including verification via KAS rewrap audit logs so tests assert the negotiated session-key type (not just successful decrypt).
Changes:
- Add
session_key_algorithmplumbing totdfs.SDK.decrypt()viaXT_WITH_SESSION_KEY_ALGORITHM, and wire it through Go/Java/JS CLI wrappers (plus SDK capability probes). - Extend audit-log parsing/assertion utilities to match on
eventMetaData.sessionKeyType. - Add new PQC tests covering ML-KEM-768/1024 session-key roundtrips and audit-log verification; update existing EC rewrap audit assertion to pin the negotiated session-key type.
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| xtest/test_tdfs.py | Tightens EC rewrap audit assertion to check session_key_type. |
| xtest/test_pqc.py | Adds ML-KEM session-key roundtrip test (768/1024) plus RSA session-key audit-field sanity test. |
| xtest/tdfs.py | Introduces session-key-mlkem feature flag and threads session_key_algorithm into decrypt via env var. |
| xtest/sdk/go/cli.sh | Wires XT_WITH_SESSION_KEY_ALGORITHM to --session-key-algorithm and adds a supports session-key-mlkem probe. |
| xtest/sdk/java/cli.sh | Wires XT_WITH_SESSION_KEY_ALGORITHM to --rewrap-key-type and adds a supports session-key-mlkem probe. |
| xtest/sdk/js/cli.sh | Wires XT_WITH_SESSION_KEY_ALGORITHM to --rewrapKeyType and adds a supports session-key-mlkem probe. |
| xtest/audit_logs.py | Adds ParsedAuditEvent.session_key_type and session_key_type filtering to rewrap assertions. |
| spec/DSPX-4221.md | Documents DSPX-4221 scope/acceptance criteria and links companion PRs. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
17eda9c to
fb0e5a9
Compare
fb0e5a9 to
3bff44e
Compare
X-Test Failure Report |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
xtest/audit_logs.py (1)
1395-1459: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winImprove diagnostics on the final mismatch assertion.
When the fallback path detects a real type mismatch, the final
assert expected in reported, ...at Line 1450-1452 raises a plainAssertionErrorwith only the expected and reported sets. Every other assertion inAuditLogAsserter(for exampleassert_rewrapat Line 1358) routes through_raise_assertion_error, which includes matching/recent log lines and clock-skew context. Losing that context here makes debugging a real cross-service session-key mismatch harder in CI.Include the raw matching events in the failure message so a real mismatch is as easy to diagnose as any other rewrap assertion failure.
♻️ Proposed diagnostic improvement
- assert expected in reported, ( - f"Expected rewrap session_key_type={expected!r}, but platform reported {reported!r}" - ) + if expected not in reported: + event_lines = "\n".join( + f" [{e.timestamp}] {e.raw_entry.raw_line}" for e in events + ) + raise AssertionError( + f"Expected rewrap session_key_type={expected!r}, but platform reported " + f"{reported!r}.\nMatching successful rewrap events:\n{event_lines}" + )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@xtest/audit_logs.py` around lines 1395 - 1459, Update the final mismatch assertion in assert_rewrap_session_key_type to include the raw matching events from the fallback lookup, alongside the expected and reported session-key types. Preserve the existing platform-field absence warning and ensure the failure message exposes enough event detail for diagnosing cross-service mismatches, consistent with the diagnostic context provided by _raise_assertion_error.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@xtest/audit_logs.py`:
- Around line 1395-1459: Update the final mismatch assertion in
assert_rewrap_session_key_type to include the raw matching events from the
fallback lookup, alongside the expected and reported session-key types. Preserve
the existing platform-field absence warning and ensure the failure message
exposes enough event detail for diagnosing cross-service mismatches, consistent
with the diagnostic context provided by _raise_assertion_error.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 797400c8-5442-438c-bd16-881c678e7a8a
📒 Files selected for processing (8)
spec/DSPX-4221.mdxtest/audit_logs.pyxtest/sdk/go/cli.shxtest/sdk/java/cli.shxtest/sdk/js/cli.shxtest/tdfs.pyxtest/test_pqc.pyxtest/test_tdfs.py
🚧 Files skipped from review as they are similar to previous changes (5)
- xtest/sdk/java/cli.sh
- xtest/sdk/js/cli.sh
- xtest/sdk/go/cli.sh
- spec/DSPX-4221.md
- xtest/tdfs.py
3bff44e to
3f3bc36
Compare
Address CodeRabbit review comments on spec/DSPX-4221.md:
- Note the platform-side session-key-mlkem feature-detection imprecision
(reuses the pre-existing Preview.MLKEMTDFEnabled/mechanism probe rather
than a dedicated readiness check; each SDK's own hardcoded capability
flag is what actually gates the test), matching the comment already in
tdfs.py's feature-detection block.
- Fix invalid Go syntax in the SDK contract example (ocrypto.KeyType is a
string type; bitwise-OR between two enum values doesn't compile).
- Remove the "and vice versa" independence claim, which overstated test
coverage (only RSA-wrapped-TDF + ML-KEM-session-key is tested).
Fix a real CI failure this surfaced: the plain assert_rewrap_success(
session_key_type=X) pattern used for the new EC/RSA/ML-KEM session-key-type
assertions has two problems that only show up once you assert on the
*specific* algorithm rather than just "a rewrap succeeded":
1. eventMetaData.sessionKeyType is new (added by the companion platform PR,
unreleased) and not behind any version/preview flag we can check
statically. Asserting a specific value against a baseline/pre-fix
platform build finds zero matches -- indistinguishable from a real
negotiation bug. Broke test_tdf_roundtrip's new ecwrap assertion and
test_session_key_rsa_roundtrip against CI's baseline (unpatched) platform
matrix.
2. Log collection is poll-based. When two rewraps for different session
keys happen back-to-back in the same test (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-based assertion
then reports success using the wrong (earlier) event instead of waiting
for the right one. Reproduced locally: test_tdf_roundtrip[...-ztdf-ecwrap]
failed 3/3 runs with "Expected session_key_type=ec:secp256r1, but
platform reported {'rsa:2048'}" even against the patched platform PR
branch, where the feature itself works correctly (confirmed via a
standalone otdfctl decrypt + raw platform.log inspection).
Added AuditLogAsserter.assert_rewrap_session_key_type(), which waits for an
event matching the expected type before falling back to a short existence
check (present-but-different-value still fails loudly; field-never-present
logs a warning and passes). Used by test_tdf_roundtrip's ecwrap block and
both test_session_key_rsa_roundtrip and test_session_key_mlkem_roundtrip.
Verified locally against the platform PR branch: all three pass
consistently now (previously ecwrap failed 3/3 runs).
Also address a Copilot review finding on SDK.decrypt(): ecwrap=True and
session_key_algorithm could both be set, which the CLI wrappers turn into
duplicate/competing --session-key-algorithm (or --rewrap-key-type) flags,
with the winner depending on the underlying CLI's parsing order rather than
caller intent. No current call site combines them, but rather than leave
that as an implicit landmine, decrypt() now raises ValueError if both are
passed -- ecwrap is shorthand for session_key_algorithm="ec:secp256r1"
anyway.
Remove the now-redundant decrypt-side XT_WITH_ECWRAP shortcut. It predates
session_key_algorithm and was pure sugar for
session_key_algorithm="ec:secp256r1"; XT_WITH_ECWRAP stays for encrypt
(selects the TDF's own KAO wrapping mechanism, an unrelated axis).
Dropped SDK.decrypt()'s ecwrap bool param (one call site, migrated to
session_key_algorithm), removed the dead decrypt-branch XT_WITH_ECWRAP
handling from all three cli.sh wrappers, and updated each script's env-var
doc comment. Verified against go, java, and js (companion PR branch)
decrypt SDKs -- all pass.
Add a TODO in assert_rewrap_session_key_type() documenting the plan for
closing the remaining gap: once opentdf/platform#3814 merges and cuts a
release, gate on a version-checked "audit-session-key-type" feature
(mirroring the existing audit_logging pattern) and hard-fail on mismatch
instead of warning. Not actionable yet since there is no release version
to check against pre-merge -- see the stacked draft PR reproducing the gap
this leaves open (unpatched web-sdk silently sending the wrong session-key
algorithm on EC rewrap, masked when paired with an unpatched platform that
does not emit the field at all).
Signed-off-by: Dave Mihalcik <dmihalcik@virtru.com>
3f3bc36 to
c7bfda5
Compare
X-Test Failure Report |
|



Summary
Implements the test-first step of DSPX-4221 (Session Keys should support ML-KEM): a client-generated ML-KEM ephemeral key pair used as the rewrap session key (the key KAS wraps the response DEK back to), as distinct from the existing
mechanism-mlkem/mechanism-xwing/mechanism-secpmlkemcoverage which only exercises KAS-managed TDF wrapping keys.Changes
xtest/tdfs.py: newsession-key-mlkemfeature flag (detected alongsidemechanism-mlkem, since both are gated server-side by the samePreview.MLKEMTDFEnabledflag).SDK.decrypt()gains asession_key_algorithmparam, threaded through a newXT_WITH_SESSION_KEY_ALGORITHMenv var.xtest/sdk/{go,java,js}/cli.sh: wire the new env var to each CLI's real flag (--session-key-algorithm/--rewrap-key-type/--rewrapKeyType), and addsupports session-key-mlkemprobes. The Java probe joins wrapped help output before grepping — picocli line-wraps the long--rewrap-key-typechoice list at ~80 columns, which can splitmlkem:768across two lines and produce a false negative (caught by running against a real build, not just eyeballing--helpoutput).xtest/test_pqc.py: newtest_session_key_mlkem_roundtrip, parametrized overmlkem:768/mlkem:1024. Encrypts with a plain RSA-wrapped attribute (attribute_default_rsa) so a failure can only be attributed to the session-key transport channel, not to KAS-managed PQC mechanism support.Audit-log assertion (not just a successful decrypt)
A successful decrypt alone doesn't prove KAS actually used ML-KEM for the session key — if an SDK silently ignored the requested algorithm and fell back to RSA, the roundtrip would still pass either way, since both sides would agree on whatever was actually used. To make the test load-bearing, it now asserts against KAS's own rewrap audit log, which is independent of anything the client reports about itself:
xtest/audit_logs.py:ParsedAuditEventgains asession_key_typeproperty (eventMetaData.sessionKeyType);matches_rewrap/assert_rewrap/assert_rewrap_success/assert_rewrap_failuregain a matchingsession_key_typeparameter, mirroring the existingalgorithmparameter.audit_logs.assert_rewrap_success(session_key_type=session_key_algorithm, since_mark=mark)after decrypting.This depends on the companion platform PR (below) adding the
sessionKeyTypefield to the rewrap audit event — without it,session_key_typeis simply absent fromeventMetaDataand the assertion fails.Skips cleanly (with a clear reason) on platforms/SDKs that don't yet support ML-KEM session keys.
Verified end-to-end locally
Brought up a live platform + KAS instances (built from opentdf/platform#3814) and ran this test against Go, Java, and JS SDKs built from their respective companion-PR branches:
All 18 include a real
session_key_typeaudit-log match, not just a successful roundtrip. This run is also what caught a real web-sdk bug (opentdf/web-sdk#975) where the requested session-key algorithm was silently dropped and every decrypt negotiated RSA regardless of what was requested — the audit assertion is what surfaced it; a plain roundtrip-success check would have missed it entirely.Related work
Companion PRs:
rewrap, wires the Go SDK's response decrypt path, and adds thesessionKeyTypeaudit field this test relies on.Test plan
All pass. Plus the live 18/18 run described above.
Ref: DSPX-4221
Summary by CodeRabbit
New Features
Bug Fixes
Tests