feat(kas): support ML-KEM client session keys in rewrap - #3814
Open
dmihalcik-virtru wants to merge 2 commits into
Open
feat(kas): support ML-KEM client session keys in rewrap#3814dmihalcik-virtru wants to merge 2 commits into
dmihalcik-virtru wants to merge 2 commits into
Conversation
pq-sessions
github-actions
Bot
added
comp:sdk
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: ASSERTIVE Plan: Pro Plus Run ID: 📒 Files selected for processing (8)
📝 WalkthroughWalkthroughDecryption now supports ML-KEM-768 and ML-KEM-1024 session keys. The KAS client processes ML-KEM responses. The rewrap service recognizes ML-KEM keys, enforces preview gating, and records the session key type in audit events. Documentation and a draft specification were added. ChangesML-KEM session-key support
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant CLI
participant KASClient
participant RewrapService
participant AuditLogger
CLI->>KASClient: Select ML-KEM session key
KASClient->>RewrapService: Request rewrap
RewrapService->>RewrapService: Validate ML-KEM preview support
RewrapService-->>KASClient: Return KAO results
RewrapService->>AuditLogger: Record sessionKeyType
KASClient->>KASClient: Decrypt permitted wrapped keys
KASClient-->>CLI: Return decrypted keys and obligations
Possibly related PRs
Suggested labels: 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 |
![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
🤖 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 21: Update the wording in the ML-KEM session encryption requirement to
use the hyphenated compound adjective “client-generated” before “key pair.”
- Line 20: Insert one blank line immediately after each Markdown section heading
in spec/DSPX-4221.md, including the headings at lines 20, 23, 26, 29, 32, 35,
and 38, so the document satisfies MD022.
- Around line 20-40: Complete the DSPX-4221 specification by replacing all
template placeholders with concrete requirements covering mlkem:768 and
mlkem:1024 support, pure versus hybrid scope, preview gating, rewrap and
decryption behavior, contracts, constraints, out-of-scope items, and testable
acceptance criteria.
🪄 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: Repository UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: a027b735-ad03-4660-934b-4c85249d7e80
📒 Files selected for processing (5)
otdfctl/cmd/tdf/decrypt.gootdfctl/docs/man/decrypt/_index.mdsdk/kas_client.goservice/kas/access/rewrap.gospec/DSPX-4221.md
Contributor
Benchmark results, click to expandBenchmark authorization.GetDecisions Results:
Benchmark authorization.v2.GetMultiResourceDecision Results:
Benchmark Statistics
Bulk Benchmark Results
TDF3 Benchmark Results:
|
Contributor
dmihalcik-virtru
force-pushed
the
DSPX-4221-pq-sessions
branch
from
July 31, 2026 22:18
33e7fce to
997ec5f
Compare
Contributor
Benchmark results, click to expandBenchmark authorization.GetDecisions Results:
Benchmark authorization.v2.GetMultiResourceDecision Results:
Benchmark Statistics
Bulk Benchmark Results
TDF3 Benchmark Results:
|
Contributor
dmihalcik-virtru
force-pushed
the
DSPX-4221-pq-sessions
branch
from
August 1, 2026 13:22
997ec5f to
e3725bf
Compare
Contributor
Benchmark results, click to expandBenchmark authorization.GetDecisions Results:
Benchmark authorization.v2.GetMultiResourceDecision Results:
Benchmark Statistics
Bulk Benchmark Results
TDF3 Benchmark Results:
|
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 `@spec/DSPX-4221.md`:
- Line 45: Update the Go SDK contract entry for sdk.WithSessionKeyType to remove
the invalid bitwise OR expression on string-based ocrypto.KeyType values. List
separate calls for MLKEM768Key and MLKEM1024Key if either option is supported,
or describe accepting both options in prose outside code formatting.
🪄 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: Repository UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 0ae928d2-8118-4f28-a031-389390aef2db
📒 Files selected for processing (8)
otdfctl/cmd/tdf/decrypt.gootdfctl/docs/man/decrypt/_index.mdsdk/kas_client.goservice/kas/access/rewrap.goservice/logger/audit/logger_test.goservice/logger/audit/rewrap.goservice/logger/audit/rewrap_test.gospec/DSPX-4221.md
Contributor
Address CodeRabbit review comments on spec/DSPX-4221.md: - 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>
dmihalcik-virtru
force-pushed
the
DSPX-4221-pq-sessions
branch
from
August 1, 2026 15:02
e3725bf to
5f9a4de
Compare
Contributor
Benchmark results, click to expandBenchmark authorization.GetDecisions Results:
Benchmark authorization.v2.GetMultiResourceDecision Results:
Benchmark Statistics
Bulk Benchmark Results
TDF3 Benchmark Results:
|
Contributor
|
Contributor
Contributor
There was a problem hiding this comment.
Pull request overview
Adds end-to-end support for pure ML-KEM (768/1024) client session keys in the KAS rewrap flow, so clients can request a PQ session key for wrapping the response DEK and successfully decrypt the rewrap response.
Changes:
- KAS rewrap request parsing now recognizes ML-KEM SPKI client public keys and gates ML-KEM session-key rewrap behind
Preview.MLKEMTDFEnabled; rewrap audit events now recordsessionKeyType. - Go SDK
KASClient.unwrap()adds an ML-KEM branch to decapsulate/decrypt rewrap responses when the session key type is ML-KEM. otdfctl decrypt --session-key-algorithmand docs now accept/documentmlkem:768andmlkem:1024; added spec/ticket write-up.
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 |
|---|---|
| spec/DSPX-4221.md | Spec/ticket documentation for ML-KEM session-key support and acceptance criteria. |
| service/logger/audit/rewrap.go | Adds sessionKeyType to rewrap audit event metadata. |
| service/logger/audit/rewrap_test.go | Updates audit event unit test expectations to include sessionKeyType. |
| service/logger/audit/logger_test.go | Updates logger JSON assertions to include sessionKeyType in rewrap audit logs. |
| service/kas/access/rewrap.go | Accepts ML-KEM SPKI in extractSRTBody, gates ML-KEM session-key rewrap via preview flag, and populates audit SessionKeyType. |
| sdk/kas_client.go | Adds ML-KEM unwrap/decrypt handling for rewrap responses. |
| otdfctl/docs/man/decrypt/_index.md | Documents mlkem:768/mlkem:1024 for --session-key-algorithm and adds an example. |
| otdfctl/cmd/tdf/decrypt.go | Maps mlkem:768/mlkem:1024 CLI values to ocrypto key types. |
Suppressed comments (2)
service/kas/access/rewrap.go:1045
- The new preview gate for ML-KEM session-key rewrap lacks test coverage. There are tests for other rewrap behavior, but none that assert an ML-KEM
clientPublicKeyis rejected whenPreview.MLKEMTDFEnabledis false and accepted when true.
if ocrypto.IsMLKEMKeyType(asymEncrypt.KeyType()) && !p.Preview.MLKEMTDFEnabled {
p.Logger.ErrorContext(ctx, "ml-kem session key rewrap not enabled")
failAllKaos(requests, results, err400("invalid request"))
return "", results, nil
}
sdk/kas_client.go:318
- Error message is misleading: this failure is from
k.sessionKey.PrivateKeyInPemFormat(), notocrypto.PrivateKeyInPemFormat. Consider aligning wording with the EC path ("failed to get private key").
clientPrivateKey, err := k.sessionKey.PrivateKeyInPemFormat()
if err != nil {
return nil, fmt.Errorf("ocrypto.PrivateKeyInPemFormat failed: %w", err)
}
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+450
to
+454
| // Pure ML-KEM client session keys are SPKI-wrapped under the NIST ML-KEM | ||
| // OIDs (FIPS 203), which x509.ParsePKIXPublicKey does not recognize and | ||
| // would otherwise reject as a parse failure. Accept them here; whether | ||
| // ML-KEM rewrap is actually enabled is checked later in tdf3Rewrap. | ||
| if oid, _, kemErr := ocrypto.ParseKEMPublicSPKI(block.Bytes); kemErr == nil && |
Comment on lines
+186
to
+190
| switch { | ||
| case ocrypto.IsECKeyType(k.sessionKey.GetKeyType()): | ||
| return k.handleECKeyResponse(response) | ||
| case ocrypto.IsMLKEMKeyType(k.sessionKey.GetKeyType()): | ||
| return k.handleKEMKeyResponse(response) |
dmihalcik-virtru
added a commit
to opentdf/tests
that referenced
this pull request
Aug 3, 2026
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Part of DSPX-4221 (Session Keys should support ML-KEM).
rewrap'sclientPublicKey— the client-generated ephemeral "session key" used to wrap the response DEK back to the client — only accepted RSA/EC keys.extractSRTBodyparsed it withx509.ParsePKIXPublicKey, which doesn't recognize ML-KEM's SPKI OID and rejected it as a parse failure before the request ever reached the (already KEM-aware) wrap dispatch.Changes
service/kas/access/rewrap.go:extractSRTBodynow recognizes pure ML-KEM-768/1024 SPKI client public keys viaocrypto.ParseKEMPublicSPKI, mirroring the dispatch orderocrypto.FromPublicPEMWithSaltalready uses internally.tdf3Rewrapgates the ML-KEM session-key path onPreview.MLKEMTDFEnabled, matching the existing EC gate (Preview.ECTDFEnabled). The wrap-to-client-key crypto itself needed no changes —ocrypto.FromPublicPEMWithSaltalready dispatches ML-KEM/hybrid keys to the KEM encryptor, since it's shared with the KAS-managed-key (mlkem-wrappedKAO) path.sdk/kas_client.go:unwrap()only branched EC vs RSA; added an ML-KEM branch (handleKEMKeyResponse/processKEMResponse) so the Go SDK can decrypt a rewrap response wrapped to an ML-KEM session key.sdk.WithSessionKeyType(ocrypto.MLKEM768Key / MLKEM1024Key)already existed but previously caused a decrypt-time failure sinceunwrap()would try (and fail) to type-assert the KEM decryptor as RSA.otdfctl:--session-key-algorithmnow acceptsmlkem:768/mlkem:1024(the underlyingsdk.WithSessionKeyTypeplumbing already supported theseocrypto.KeyTypevalues, just not the CLI flag mapping); updateddocs/man/decrypt.New audit field:
sessionKeyTypeA rewrap succeeding doesn't by itself prove which session-key type KAS actually used — nothing observable previously recorded it, so the companion xtest PR could only infer correctness indirectly from a successful decrypt (which would also pass if a client silently fell back to RSA). To make that test load-bearing:
service/logger/audit/rewrap.go:RewrapAuditEventParamsgains aSessionKeyTypefield, surfaced aseventMetaData.sessionKeyTypeon rewrap audit events.service/kas/access/rewrap.go: populates it fromasymEncrypt.KeyType()(already computed for thePreview.MLKEMTDFEnabledgate above) when building each KAO's audit event.service/logger/auditunit tests for the new field.Scope
Deliberately limited to pure ML-KEM (768/1024), matching the ticket title and the existing
mechanism-mlkemprecedent — hybrid PQ/T session keys (X-Wing, secp+ML-KEM) are out of scope here.Test plan
All pass, no regressions. Manually verified
otdfctl help decryptnow listsmlkem:768/mlkem:1024for--session-key-algorithm.Related work
Companion PRs:
test_session_key_mlkem_roundtrip), asserting on the newsessionKeyTypeaudit field added here.Ref: DSPX-4221
Summary by CodeRabbit
New Features
Security & Auditing