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
|
Warning Review limit reached
Next review available in: 26 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (8)
📝 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 and ML-KEM-1024. ChangesML-KEM session-key support
Estimated code review effort: 3 (Moderate) | ~20 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 ML-KEM session-key algorithm
SDK->>Wrapper: Export XT_WITH_SESSION_KEY_ALGORITHM
Wrapper->>CLI: Pass rewrap-key option
CLI-->>SDK: Return decrypted plaintext
SDK-->>Test: Return plaintext
Test->>Audit: Match session_key_type in rewrap event
Audit-->>Test: Confirm requested session-key type
Possibly related PRs
Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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
0cefcf0 to
398f300
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). Signed-off-by: Dave Mihalcik <dmihalcik@virtru.com>
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.
| if assert_keys: | ||
| local_env |= {"XT_WITH_ASSERTION_VERIFICATION_KEYS": assert_keys} | ||
| if ecwrap: | ||
| local_env |= {"XT_WITH_ECWRAP": "true"} | ||
| if session_key_algorithm: | ||
| local_env |= {"XT_WITH_SESSION_KEY_ALGORITHM": session_key_algorithm} |
| audit_logs.assert_rewrap_success( | ||
| session_key_type="ec:secp256r1", min_count=1, since_mark=ec_mark | ||
| ) |
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