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) |
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