Skip to content

feat(sdk): support ML-KEM rewrap session keys #388

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
dmihalcik-virtru wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(sdk): support ML-KEM rewrap session keys #388

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions sdk-pqc-bc/src/main/java/io/opentdf/platform/sdk/pqc/bc/BouncyCastleKemProvider.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,4 +51,9 @@ public byte[] wrapDEK(KeyType keyType, String publicKeyPEM, byte[] dek) {
public byte[] unwrapDEK(KeyType keyType, String privateKeyPEM, byte[] wrapped) {
return HybridCrypto.unwrapDEK(keyType, privateKeyPEM, wrapped);
}

@Override
public KeyPairPem generateKeyPair(KeyType keyType) {
return HybridCrypto.generateKeyPair(keyType);
}
}
32 changes: 32 additions & 0 deletions sdk-pqc-bc/src/main/java/io/opentdf/platform/sdk/pqc/bc/HybridCrypto.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import io.opentdf.platform.sdk.ECKeyPair;
import io.opentdf.platform.sdk.KeyType;
import io.opentdf.platform.sdk.SDKException;
import io.opentdf.platform.sdk.spi.KemProvider;

import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
Expand Down Expand Up @@ -96,6 +97,37 @@ static byte[] unwrapDEK(KeyType keyType, String privateKeyPEM, byte[] wrapped) {
}
}

/**
* Generate a fresh ephemeral keypair for {@code keyType}. Same dispatch
* table as {@link #wrapDEK} / {@link #unwrapDEK}.
*/
static KemProvider.KeyPairPem generateKeyPair(KeyType keyType) {
switch (keyType) {
case HybridXWingKey: {
XWingKeyPair kp = XWingKeyPair.generate();
return new KemProvider.KeyPairPem(kp.publicKeyInPemFormat(), kp.privateKeyInPemFormat());
}
case HybridSecp256r1MLKEM768Key: {
HybridNISTKeyPair kp = HybridNISTAlgorithm.P256_MLKEM768.generate();
return new KemProvider.KeyPairPem(kp.publicKeyInPemFormat(), kp.privateKeyInPemFormat());
}
case HybridSecp384r1MLKEM1024Key: {
HybridNISTKeyPair kp = HybridNISTAlgorithm.P384_MLKEM1024.generate();
return new KemProvider.KeyPairPem(kp.publicKeyInPemFormat(), kp.privateKeyInPemFormat());
}
case MLKEM768Key: {
MLKEMKeyPair kp = MLKEMAlgorithm.MLKEM_768.generate();
return new KemProvider.KeyPairPem(kp.publicKeyInPemFormat(), kp.privateKeyInPemFormat());
}
case MLKEM1024Key: {
MLKEMKeyPair kp = MLKEMAlgorithm.MLKEM_1024.generate();
return new KemProvider.KeyPairPem(kp.publicKeyInPemFormat(), kp.privateKeyInPemFormat());
}
default:
throw new SDKException("unsupported PQC key type: " + keyType);
}
}

/**
* Build the ASN.1 envelope from a hybrid KEM ciphertext and the AES-GCM(iv||ct) encrypted DEK.
*/
Expand Down
43 changes: 43 additions & 0 deletions sdk-pqc-bc/src/test/java/io/opentdf/platform/sdk/pqc/bc/HybridCryptoGenerateKeyPairTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
package io.opentdf.platform.sdk.pqc.bc;

import io.opentdf.platform.sdk.KeyType;
import io.opentdf.platform.sdk.spi.KemProvider;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.EnumSource;

import java.nio.charset.StandardCharsets;

import static org.junit.jupiter.api.Assertions.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertTrue;

/**
* Unit tests for {@link HybridCrypto#generateKeyPair}, the dispatcher backing
* {@link BouncyCastleKemProvider#generateKeyPair} and (in the core {@code sdk}
* module) {@code KASClient}'s KEM branch for the rewrap client "session key".
*
* <p>Each variant is exercised end-to-end through the same PEM-based
* {@link KemProvider} contract KASClient uses: generate a fresh keypair, wrap
* a DEK against its public PEM, unwrap against its private PEM, assert equal.
*/
class HybridCryptoGenerateKeyPairTest {

private static final byte[] DEK = "0123456789abcdef0123456789abcdef".getBytes(StandardCharsets.UTF_8);

@ParameterizedTest
@EnumSource(value = KeyType.class, names = {
"HybridXWingKey",
"HybridSecp256r1MLKEM768Key",
"HybridSecp384r1MLKEM1024Key",
"MLKEM768Key",
"MLKEM1024Key"})
void generatedKeyPairRoundTrips(KeyType keyType) {
KemProvider.KeyPairPem kp = HybridCrypto.generateKeyPair(keyType);

assertTrue(kp.publicKeyPEM.startsWith("-----BEGIN PUBLIC KEY-----"), "public PEM header");
assertTrue(kp.privateKeyPEM.startsWith("-----BEGIN PRIVATE KEY-----"), "private PEM header");

byte[] wrapped = HybridCrypto.wrapDEK(keyType, kp.publicKeyPEM, DEK);
byte[] unwrapped = HybridCrypto.unwrapDEK(keyType, kp.privateKeyPEM, wrapped);
assertArrayEquals(DEK, unwrapped, "DEK round-trip via generated keypair");
}
}
101 changes: 66 additions & 35 deletions sdk/src/main/java/io/opentdf/platform/sdk/KASClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@
import io.opentdf.platform.kas.RewrapRequest;
import io.opentdf.platform.kas.RewrapResponse;
import io.opentdf.platform.sdk.SDK.KasBadRequestException;
import io.opentdf.platform.sdk.spi.KemProvider;
import io.opentdf.platform.sdk.spi.KemProviders;

import okhttp3.OkHttpClient;
import org.slf4j.Logger;
Expand Down Expand Up @@ -118,25 +120,76 @@ static class RewrapRequestBody {

private static final Gson gson = new Gson();

@Override
public byte[] unwrap(Manifest.KeyAccess keyAccess, String policy, KeyType sessionKeyType) {
ECKeyPair ecKeyPair = null;
/** Session-key material generated for a single unwrap() call: the PEM to send to KAS, plus whichever private key is needed to process the response. */
private static final class SessionKeyMaterial {
final String publicKeyPem;
final ECKeyPair ecKeyPair;
final KemProvider.KeyPairPem kemKeyPair;

SessionKeyMaterial(String publicKeyPem, ECKeyPair ecKeyPair, KemProvider.KeyPairPem kemKeyPair) {
this.publicKeyPem = publicKeyPem;
this.ecKeyPair = ecKeyPair;
this.kemKeyPair = kemKeyPair;
}
}

private SessionKeyMaterial generateSessionKeyMaterial(KeyType sessionKeyType) {
if (sessionKeyType.isEc()) {
var curve = sessionKeyType.getECCurve();
ecKeyPair = new ECKeyPair(curve);
clientPublicKey = ecKeyPair.publicKeyInPEMFormat();
} else {
// Initialize the RSA key pair only once and reuse it for future unwrap operations
if (decryptor == null) {
var encryptionKeypair = CryptoUtils.generateRSAKeypair();
decryptor = new AsymDecryption(encryptionKeypair.getPrivate());
clientPublicKey = CryptoUtils.getRSAPublicKeyPEM(encryptionKeypair.getPublic());
ECKeyPair ecKeyPair = new ECKeyPair(curve);
return new SessionKeyMaterial(ecKeyPair.publicKeyInPEMFormat(), ecKeyPair, null);
}
if (sessionKeyType.isMLKEM()) {
KemProvider.KeyPairPem kemKeyPair = KemProviders.get(sessionKeyType).generateKeyPair(sessionKeyType);
return new SessionKeyMaterial(kemKeyPair.publicKeyPEM, null, kemKeyPair);
}
// Initialize the RSA key pair only once and reuse it for future unwrap operations
if (decryptor == null) {
var encryptionKeypair = CryptoUtils.generateRSAKeypair();
decryptor = new AsymDecryption(encryptionKeypair.getPrivate());
clientPublicKey = CryptoUtils.getRSAPublicKeyPEM(encryptionKeypair.getPublic());
}
return new SessionKeyMaterial(clientPublicKey, null, null);
}
Comment on lines +136 to +153

@coderabbitai coderabbitai Bot Aug 1, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

Guard the lazy RSA key-pair initialization against concurrent unwrap() calls.

The EC and ML-KEM branches now correctly return per-call local key material, which fixes the previously flagged race on the shared clientPublicKey field. The RSA branch still has a check-then-act race: if (decryptor == null) at Line 147 reads and later writes the shared, non-volatile fields decryptor and clientPublicKey without synchronization.

If two threads call unwrap() with an RSA session key type before decryptor is first set, both can generate distinct RSA key pairs and race to overwrite each other's decryptor/clientPublicKey. One call can then send a public key from one generation while decrypting the response with a different generation's private key, causing decryption failures. This class already synchronizes close() (Line 110) and getStub() (Line 233), so concurrent use of the same KASClient instance is an expected usage pattern, not a theoretical edge case.

Synchronize the lazy-init block with double-checked locking, and mark decryptor and clientPublicKey volatile so the initializing thread's write is visible to other threads.

🔒 Proposed fix: synchronize the lazy RSA key-pair initialization 
-    private AsymDecryption decryptor;
-    private String clientPublicKey;
+    private volatile AsymDecryption decryptor;
+    private volatile String clientPublicKey;
-        // Initialize the RSA key pair only once and reuse it for future unwrap operations
-        if (decryptor == null) {
-            var encryptionKeypair = CryptoUtils.generateRSAKeypair();
-            decryptor = new AsymDecryption(encryptionKeypair.getPrivate());
-            clientPublicKey = CryptoUtils.getRSAPublicKeyPEM(encryptionKeypair.getPublic());
-        }
-        return new SessionKeyMaterial(clientPublicKey, null, null);
+        // Initialize the RSA key pair only once and reuse it for future unwrap operations.
+        if (decryptor == null) {
+            synchronized (this) {
+                if (decryptor == null) {
+                    var encryptionKeypair = CryptoUtils.generateRSAKeypair();
+                    decryptor = new AsymDecryption(encryptionKeypair.getPrivate());
+                    clientPublicKey = CryptoUtils.getRSAPublicKeyPEM(encryptionKeypair.getPublic());
+                }
+            }
+        }
+        return new SessionKeyMaterial(clientPublicKey, null, null);
🤖 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 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.


private byte[] unwrapResponseKey(
KeyType sessionKeyType, SessionKeyMaterial keyMaterial, RewrapResponse response, byte[] wrappedKey) {
if (sessionKeyType.isEc()) {
if (keyMaterial.ecKeyPair == null) {
throw new SDKException("ECKeyPair is null. Unable to proceed with the unwrap operation.");
}

var kasEphemeralPublicKey = response.getSessionPublicKey();
ECPublicKey publicKey;
try {
publicKey = ECKeyPair.publicKeyFromPem(kasEphemeralPublicKey);
} catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
throw new SDKException("error decoding KAS session public key", e);
}
byte[] symKey = ECKeyPair.computeECDHKey(publicKey, keyMaterial.ecKeyPair.getPrivateKey());

var sessionKey = ECKeyPair.calculateHKDF(GLOBAL_KEY_SALT, symKey);

AesGcm gcm = new AesGcm(sessionKey);
AesGcm.Encrypted encrypted = new AesGcm.Encrypted(wrappedKey);
return gcm.decrypt(encrypted);
}
if (sessionKeyType.isMLKEM()) {
if (keyMaterial.kemKeyPair == null) {
throw new SDKException("KEM keypair is null. Unable to proceed with the unwrap operation.");
}
return KemProviders.get(sessionKeyType).unwrapDEK(sessionKeyType, keyMaterial.kemKeyPair.privateKeyPEM, wrappedKey);
}
return decryptor.decrypt(wrappedKey);
}

@Override
public byte[] unwrap(Manifest.KeyAccess keyAccess, String policy, KeyType sessionKeyType) {
SessionKeyMaterial keyMaterial = generateSessionKeyMaterial(sessionKeyType);

RewrapRequestBody body = new RewrapRequestBody();
body.policy = policy;
body.clientPublicKey = clientPublicKey;
body.clientPublicKey = keyMaterial.publicKeyPem;
body.keyAccess = keyAccess;

var requestBody = gson.toJson(body);
Expand Down Expand Up @@ -171,29 +224,7 @@ public byte[] unwrap(Manifest.KeyAccess keyAccess, String policy, KeyType sessi
}

var wrappedKey = response.getEntityWrappedKey().toByteArray();
if (sessionKeyType.isEc()) {

if (ecKeyPair == null) {
throw new SDKException("ECKeyPair is null. Unable to proceed with the unwrap operation.");
}

var kasEphemeralPublicKey = response.getSessionPublicKey();
ECPublicKey publicKey;
try {
publicKey = ECKeyPair.publicKeyFromPem(kasEphemeralPublicKey);
} catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
throw new SDKException("error decoding KAS session public key", e);
}
byte[] symKey = ECKeyPair.computeECDHKey(publicKey, ecKeyPair.getPrivateKey());

var sessionKey = ECKeyPair.calculateHKDF(GLOBAL_KEY_SALT, symKey);

AesGcm gcm = new AesGcm(sessionKey);
AesGcm.Encrypted encrypted = new AesGcm.Encrypted(wrappedKey);
return gcm.decrypt(encrypted);
} else {
return decryptor.decrypt(wrappedKey);
}
return unwrapResponseKey(sessionKeyType, keyMaterial, response, wrappedKey);
}

private final HashMap<String, AccessServiceClient> stubs = new HashMap<>();
Expand Down
22 changes: 22 additions & 0 deletions sdk/src/main/java/io/opentdf/platform/sdk/spi/KemProvider.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,4 +47,26 @@ public interface KemProvider {
* decap path; the production decrypt flow defers unwrap to the KAS.
*/
byte[] unwrapDEK(KeyType keyType, String privateKeyPEM, byte[] wrapped);

/**
* Generate a fresh ephemeral keypair for {@code keyType}. Used for the
* client's rewrap "session key" — the ephemeral key a client generates
* and sends with a rewrap request so KAS can wrap the response DEK back
* to it — as opposed to a KAS-managed wrapping key (see {@link #wrapDEK}).
*
* @param keyType hybrid or pure-PQC algorithm; must be a member of {@link #supportedKeyTypes()}
* @return the generated keypair, PEM-encoded
*/
KeyPairPem generateKeyPair(KeyType keyType);

/** A generated KEM keypair: SPKI public key PEM and PKCS#8 private key PEM. */
final class KeyPairPem {
public final String publicKeyPEM;
public final String privateKeyPEM;

public KeyPairPem(String publicKeyPEM, String privateKeyPEM) {
this.publicKeyPEM = publicKeyPEM;
this.privateKeyPEM = privateKeyPEM;
}
}
}
Loading