feat(sdk): support ML-KEM rewrap session keys - #388
Open
dmihalcik-virtru wants to merge 1 commit into
Open
Conversation
Contributor
|
Warning Review limit reached
Next review available in: 39 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 (7)
📝 WalkthroughWalkthroughThe PR adds PEM key-pair generation to the KEM provider API and Bouncy Castle implementation. ChangesKEM key generation and ML-KEM unwrap
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant KASClient
participant KemProviders
participant BouncyCastleKemProvider
participant KAS
KASClient->>KemProviders: resolve ML-KEM provider
KASClient->>BouncyCastleKemProvider: generateKeyPair
BouncyCastleKemProvider-->>KASClient: return PEM key pair
KASClient->>KAS: send public key in rewrap request
KAS-->>KASClient: return wrapped response key
KASClient->>BouncyCastleKemProvider: unwrap with private key
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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
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 `@sdk/src/main/java/io/opentdf/platform/sdk/KASClient.java`:
- Around line 126-133: Update the unwrap flow around the ML-KEM branch to keep
the generated public key in a local request-scoped variable instead of shared
clientPublicKey. Use that local variable when assigning body.clientPublicKey,
while preserving the existing EC behavior and each call’s private-key pairing.
🪄 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: 9be24f35-9e29-4f8c-b225-82b889ac8cc4
📒 Files selected for processing (5)
sdk-pqc-bc/src/main/java/io/opentdf/platform/sdk/pqc/bc/BouncyCastleKemProvider.javasdk-pqc-bc/src/main/java/io/opentdf/platform/sdk/pqc/bc/HybridCrypto.javasdk-pqc-bc/src/test/java/io/opentdf/platform/sdk/pqc/bc/HybridCryptoGenerateKeyPairTest.javasdk/src/main/java/io/opentdf/platform/sdk/KASClient.javasdk/src/main/java/io/opentdf/platform/sdk/spi/KemProvider.java
Contributor
dmihalcik-virtru
force-pushed
the
DSPX-4221-pq-sessions
branch
from
August 1, 2026 13:26
08e8b8d to
c8ece6d
Compare
Contributor
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 `@sdk/src/main/java/io/opentdf/platform/sdk/KASClient.java`:
- Around line 136-153: Make the RSA lazy initialization in
generateSessionKeyMaterial thread-safe using double-checked locking: retain the
fast null check, synchronize the initialization block on the appropriate shared
lock, and recheck decryptor before generating keys. Mark both decryptor and
clientPublicKey volatile so the initialized key pair is safely published and the
public/private keys remain consistent across concurrent unwrap calls.
🪄 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: 40b7a1b3-6d7c-4bbb-bcd3-ddb7462ba145
📒 Files selected for processing (5)
sdk-pqc-bc/src/main/java/io/opentdf/platform/sdk/pqc/bc/BouncyCastleKemProvider.javasdk-pqc-bc/src/main/java/io/opentdf/platform/sdk/pqc/bc/HybridCrypto.javasdk-pqc-bc/src/test/java/io/opentdf/platform/sdk/pqc/bc/HybridCryptoGenerateKeyPairTest.javasdk/src/main/java/io/opentdf/platform/sdk/KASClient.javasdk/src/main/java/io/opentdf/platform/sdk/spi/KemProvider.java
🚧 Files skipped from review as they are similar to previous changes (4)
- sdk/src/main/java/io/opentdf/platform/sdk/spi/KemProvider.java
- sdk-pqc-bc/src/main/java/io/opentdf/platform/sdk/pqc/bc/HybridCrypto.java
- sdk-pqc-bc/src/main/java/io/opentdf/platform/sdk/pqc/bc/BouncyCastleKemProvider.java
- sdk-pqc-bc/src/test/java/io/opentdf/platform/sdk/pqc/bc/HybridCryptoGenerateKeyPairTest.java
Contributor
dmihalcik-virtru
force-pushed
the
DSPX-4221-pq-sessions
branch
from
August 1, 2026 14:30
c8ece6d to
f15299e
Compare
Contributor
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 `@cmdline/src/main/java/io/opentdf/platform/Command.java`:
- Line 80: Update the FEATURES construction in Command so session-key-mlkem is
included only when the KemProvider registry reports an ML-KEM provider is
registered, using the existing KemProviders.registered() mechanism. Keep dpop
and dpop_nonce_challenge unconditional, and avoid relying on sdk-pqc-bc build
profiles or classpath presence.
🪄 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: 72e05d1a-3e79-44ca-b641-7bd957ff59b1
📒 Files selected for processing (7)
cmdline/src/main/java/io/opentdf/platform/Command.javacmdline/src/test/java/io/opentdf/platform/CommandTest.javasdk-pqc-bc/src/main/java/io/opentdf/platform/sdk/pqc/bc/BouncyCastleKemProvider.javasdk-pqc-bc/src/main/java/io/opentdf/platform/sdk/pqc/bc/HybridCrypto.javasdk-pqc-bc/src/test/java/io/opentdf/platform/sdk/pqc/bc/HybridCryptoGenerateKeyPairTest.javasdk/src/main/java/io/opentdf/platform/sdk/KASClient.javasdk/src/main/java/io/opentdf/platform/sdk/spi/KemProvider.java
🚧 Files skipped from review as they are similar to previous changes (5)
- sdk-pqc-bc/src/main/java/io/opentdf/platform/sdk/pqc/bc/BouncyCastleKemProvider.java
- sdk-pqc-bc/src/test/java/io/opentdf/platform/sdk/pqc/bc/HybridCryptoGenerateKeyPairTest.java
- sdk-pqc-bc/src/main/java/io/opentdf/platform/sdk/pqc/bc/HybridCrypto.java
- sdk/src/main/java/io/opentdf/platform/sdk/spi/KemProvider.java
- sdk/src/main/java/io/opentdf/platform/sdk/KASClient.java
|
Part of DSPX-4221 (Session Keys should support ML-KEM). KASClient.unwrap() only branched EC vs RSA for the client's ephemeral rewrap "session key", so requesting an ML-KEM session key would fail. KeyType already had MLKEM768Key/MLKEM1024Key entries and sdk-pqc-bc already had the BouncyCastle KEM primitives (used for KAS-managed mechanism keys), but nothing generated a fresh ML-KEM keypair for the session-key role. - spi/KemProvider: add generateKeyPair(KeyType), returning a new KeyPairPem (PEM-encoded SPKI/PKCS#8 pair). The core sdk module has no compile-time dependency on BouncyCastle, so the actual keypair generation lives in the optional sdk-pqc-bc module and is reached via the existing ServiceLoader-based KemProviders registry. - sdk-pqc-bc/HybridCrypto: add a generateKeyPair dispatch mirroring the existing wrapDEK/unwrapDEK switch, covering all five supported PQC key types (X-Wing, both NIST hybrids, ML-KEM-768/1024). - sdk-pqc-bc/BouncyCastleKemProvider: implement the new SPI method. - KASClient.unwrap(): branch on sessionKeyType.isMLKEM(), generating a keypair via KemProviders.get(...).generateKeyPair(...) and decrypting the response with KemProviders.get(...).unwrapDEK(...). - New unit test (HybridCryptoGenerateKeyPairTest) round-trips generateKeyPair -> wrapDEK -> unwrapDEK for all five PQC key types. - cmdline: add "session-key-mlkem" to the `supports` subcommand's feature list, gated on KemProviders.registered() actually containing an ML-KEM provider rather than being unconditionally true (so a FIPS-profile build, which excludes sdk-pqc-bc, correctly reports it as unsupported instead of claiming a capability that would throw at runtime). Addresses CodeRabbit review feedback: - unwrap() previously wrote the generated session public key to the shared clientPublicKey instance field before serializing the request. A concurrent unwrap() call on the same KASClient could overwrite that field first, causing this call to send the wrong public key while decrypting with its own private key. Refactored so each call's session-key material (public PEM plus whichever private key it needs for the response) is a call-scoped SessionKeyMaterial value, threaded explicitly through generateSessionKeyMaterial()/unwrapResponseKey() rather than shared mutable fields. This split also resolves a SonarCloud cognitive-complexity (java:S3776) failure on unwrap() from combining request-build and response-decrypt branching in one method. - Follow-up on the same class of bug: the RSA branch's lazy keypair-initialization was still a check-then-act race on non-volatile fields (decryptor/clientPublicKey), so two concurrent first callers could generate distinct keypairs and race to overwrite each other's fields. Fixed with double-checked locking and volatile fields. - Command.Supports' FEATURES list unconditionally claimed "session-key-mlkem" regardless of whether an ML-KEM KemProvider is actually registered (e.g. under the fips Maven profile, which excludes sdk-pqc-bc). Now computed dynamically via KemProviders.registered(). Ref: DSPX-4221 Signed-off-by: Dave Mihalcik <[email protected]>
Contributor
dmihalcik-virtru
force-pushed
the
DSPX-4221-pq-sessions
branch
from
August 1, 2026 14:50
f15299e to
52b7d4f
Compare
Contributor
Contributor
There was a problem hiding this comment.
Pull request overview
Adds ML-KEM (FIPS 203) support for the client-generated “session key” path used in KAS rewrap/unwrap, by extending the KEM SPI to allow ephemeral keypair generation and wiring that into KASClient.unwrap(). This keeps the core sdk module provider-agnostic (no BC dependency) while enabling PQC behavior via the existing ServiceLoader registry.
Changes:
- Extend
KemProviderSPI withgenerateKeyPair(KeyType)and implement it insdk-pqc-bc. - Update
KASClient.unwrap()to generate ML-KEM session keypairs viaKemProvidersand unwrap responses viaunwrapDEK. - Add PQC unit test coverage for
generateKeyPair -> wrapDEK -> unwrapDEKround-trips, and updatecmdline supportsfeature reporting.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| sdk/src/main/java/io/opentdf/platform/sdk/spi/KemProvider.java | Adds SPI method for ephemeral KEM keypair generation and a PEM keypair carrier type. |
| sdk/src/main/java/io/opentdf/platform/sdk/KASClient.java | Adds ML-KEM session-key branch and refactors unwrap flow to support EC/RSA/ML-KEM. |
| sdk-pqc-bc/src/main/java/io/opentdf/platform/sdk/pqc/bc/HybridCrypto.java | Implements keypair generation dispatch for all supported PQC key types. |
| sdk-pqc-bc/src/main/java/io/opentdf/platform/sdk/pqc/bc/BouncyCastleKemProvider.java | Implements the new SPI method by delegating to HybridCrypto. |
| sdk-pqc-bc/src/test/java/io/opentdf/platform/sdk/pqc/bc/HybridCryptoGenerateKeyPairTest.java | Adds parameterized test ensuring generated keypairs can wrap/unwrap a DEK for all PQC types. |
| cmdline/src/main/java/io/opentdf/platform/Command.java | Updates supports to advertise session-key-mlkem when an ML-KEM provider is present. |
| cmdline/src/test/java/io/opentdf/platform/CommandTest.java | Updates supports-output expectations to include session-key-mlkem. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| * @param keyType hybrid or pure-PQC algorithm; must be a member of {@link #supportedKeyTypes()} | ||
| * @return the generated keypair, PEM-encoded | ||
| */ | ||
| KeyPairPem generateKeyPair(KeyType keyType); |
Comment on lines
+156
to
+158
| var encryptionKeypair = CryptoUtils.generateRSAKeypair(); | ||
| decryptor = new AsymDecryption(encryptionKeypair.getPrivate()); | ||
| clientPublicKey = CryptoUtils.getRSAPublicKeyPEM(encryptionKeypair.getPublic()); |
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).
KASClient.unwrap()only branched EC vs RSA for the client's ephemeral rewrap "session key" (as opposed to a KAS-managed wrapping key), so requesting an ML-KEM session key would fail.KeyTypealready hadMLKEM768Key/MLKEM1024Keyentries andsdk-pqc-bcalready had the BouncyCastle KEM primitives (used for the KAS-managed mechanism-key path), but nothing generated a fresh ML-KEM keypair for the session-key role.Changes
spi/KemProvider: addgenerateKeyPair(KeyType), returning a newKeyPairPem(PEM-encoded SPKI/PKCS#8 pair). The coresdkmodule has no compile-time dependency on BouncyCastle (per ADR 0001), so the actual keypair generation lives in the optionalsdk-pqc-bcmodule and is reached via the existingServiceLoader-basedKemProvidersregistry — same pattern already used forwrapDEK/unwrapDEK.sdk-pqc-bc/HybridCrypto: add agenerateKeyPairdispatch mirroring the existingwrapDEK/unwrapDEKswitch, covering all five supported PQC key types (X-Wing, both NIST hybrids, ML-KEM-768/1024) so the SPI contract stays complete for every type insupportedKeyTypes().sdk-pqc-bc/BouncyCastleKemProvider: implement the new SPI method.KASClient.unwrap(): branch onsessionKeyType.isMLKEM(), generating a keypair viaKemProviders.get(...).generateKeyPair(...)and decrypting the response viaKemProviders.get(...).unwrapDEK(...).HybridCryptoGenerateKeyPairTestround-tripsgenerateKeyPair -> wrapDEK -> unwrapDEKfor all five PQC key types.cmdline's--rewrap-key-typealready acceptsmlkem:768/mlkem:1024with no code changes needed: picocli's enum candidate list is derived fromKeyType#toString(), which already returns"mlkem:768"/"mlkem:1024". Verified by buildingcmdline.jarand checkinghelp decryptoutput.Test plan
All existing tests plus the new
HybridCryptoGenerateKeyPairTest(5/5) pass. Also builtcmdline.jarand confirmed--rewrap-key-type=mlkem:768parses correctly (fails downstream only on the expected network/gRPC call, not on argument validation).Related work
Companion PRs:
test_session_key_mlkem_roundtrip) gated behind a newsession-key-mlkemfeature flag.rewrap, and adds the equivalent Go SDK branch.Ref: DSPX-4221
Summary by CodeRabbit
New Features
supportscommand to report ML-KEM session-key support.Bug Fixes