feat(sdk): support ML-KEM rewrap session keys - #388
Open
dmihalcik-virtru wants to merge 1 commit into
Open
Conversation
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. cmdlines --rewrap-key-type already accepts mlkem:768/mlkem:1024 with no changes needed: picoclis enum candidate list is derived from KeyType#toString(), which already returns "mlkem:768"/"mlkem:1024". Verified by building cmdline.jar and checking help decrypt output. Ref: DSPX-4221 Signed-off-by: Dave Mihalcik <dmihalcik@virtru.com>
Contributor
📝 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) | ~20 minutes 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}
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
Comment on lines
+126
to
+133
| KemProvider.KeyPairPem kemKeyPair = null; | ||
| if (sessionKeyType.isEc()) { | ||
| var curve = sessionKeyType.getECCurve(); | ||
| ecKeyPair = new ECKeyPair(curve); | ||
| clientPublicKey = ecKeyPair.publicKeyInPEMFormat(); | ||
| } else if (sessionKeyType.isMLKEM()) { | ||
| kemKeyPair = KemProviders.get(sessionKeyType).generateKeyPair(sessionKeyType); | ||
| clientPublicKey = kemKeyPair.publicKeyPEM; |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Keep the ML-KEM public key local to this unwrap call.
Line 133 writes the generated public PEM to shared clientPublicKey. Another concurrent unwrap call can overwrite this field before line 145 serializes the request. The first call then sends the second call's public key but unwraps with its own private key.
Store the request public key in a local variable and assign body.clientPublicKey from that variable.
Proposed fix
public byte[] unwrap(Manifest.KeyAccess keyAccess, String policy, KeyType sessionKeyType) {
ECKeyPair ecKeyPair = null;
KemProvider.KeyPairPem kemKeyPair = null;
+ String requestClientPublicKey;
if (sessionKeyType.isEc()) {
var curve = sessionKeyType.getECCurve();
ecKeyPair = new ECKeyPair(curve);
- clientPublicKey = ecKeyPair.publicKeyInPEMFormat();
+ requestClientPublicKey = ecKeyPair.publicKeyInPEMFormat();
} else if (sessionKeyType.isMLKEM()) {
kemKeyPair = KemProviders.get(sessionKeyType).generateKeyPair(sessionKeyType);
- clientPublicKey = kemKeyPair.publicKeyPEM;
+ requestClientPublicKey = kemKeyPair.publicKeyPEM;
} else {
// Existing RSA initialization
+ requestClientPublicKey = clientPublicKey;
}
RewrapRequestBody body = new RewrapRequestBody();
- body.clientPublicKey = clientPublicKey;
+ body.clientPublicKey = requestClientPublicKey;📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| KemProvider.KeyPairPem kemKeyPair = null; | |
| if (sessionKeyType.isEc()) { | |
| var curve = sessionKeyType.getECCurve(); | |
| ecKeyPair = new ECKeyPair(curve); | |
| clientPublicKey = ecKeyPair.publicKeyInPEMFormat(); | |
| } else if (sessionKeyType.isMLKEM()) { | |
| kemKeyPair = KemProviders.get(sessionKeyType).generateKeyPair(sessionKeyType); | |
| clientPublicKey = kemKeyPair.publicKeyPEM; | |
| public byte[] unwrap(Manifest.KeyAccess keyAccess, String policy, KeyType sessionKeyType) { | |
| ECKeyPair ecKeyPair = null; | |
| KemProvider.KeyPairPem kemKeyPair = null; | |
| String requestClientPublicKey; | |
| if (sessionKeyType.isEc()) { | |
| var curve = sessionKeyType.getECCurve(); | |
| ecKeyPair = new ECKeyPair(curve); | |
| requestClientPublicKey = ecKeyPair.publicKeyInPEMFormat(); | |
| } else if (sessionKeyType.isMLKEM()) { | |
| kemKeyPair = KemProviders.get(sessionKeyType).generateKeyPair(sessionKeyType); | |
| requestClientPublicKey = kemKeyPair.publicKeyPEM; | |
| } else { | |
| // Existing RSA initialization | |
| requestClientPublicKey = clientPublicKey; | |
| } | |
| RewrapRequestBody body = new RewrapRequestBody(); | |
| body.clientPublicKey = requestClientPublicKey; |
🤖 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 `@sdk/src/main/java/io/opentdf/platform/sdk/KASClient.java` around lines 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.
|
Contributor
This was referenced Jul 31, 2026
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
1 participant
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
Bug Fixes