From 52b7d4fe67b0f95d52611759a8d17960ea3f7e1c Mon Sep 17 00:00:00 2001 From: Dave Mihalcik Date: Fri, 31 Jul 2026 17:18:07 -0400 Subject: [PATCH] feat(sdk): support ML-KEM rewrap session keys 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 --- .../java/io/opentdf/platform/Command.java | 22 +++- .../java/io/opentdf/platform/CommandTest.java | 5 +- .../sdk/pqc/bc/BouncyCastleKemProvider.java | 5 + .../platform/sdk/pqc/bc/HybridCrypto.java | 32 +++++ .../bc/HybridCryptoGenerateKeyPairTest.java | 43 +++++++ .../io/opentdf/platform/sdk/KASClient.java | 115 ++++++++++++------ .../opentdf/platform/sdk/spi/KemProvider.java | 22 ++++ 7 files changed, 201 insertions(+), 43 deletions(-) create mode 100644 sdk-pqc-bc/src/test/java/io/opentdf/platform/sdk/pqc/bc/HybridCryptoGenerateKeyPairTest.java diff --git a/cmdline/src/main/java/io/opentdf/platform/Command.java b/cmdline/src/main/java/io/opentdf/platform/Command.java index 7def8b2e..185c84f7 100644 --- a/cmdline/src/main/java/io/opentdf/platform/Command.java +++ b/cmdline/src/main/java/io/opentdf/platform/Command.java @@ -43,6 +43,7 @@ import io.opentdf.platform.sdk.KeyType; import io.opentdf.platform.sdk.SDK; import io.opentdf.platform.sdk.SDKBuilder; +import io.opentdf.platform.sdk.spi.KemProviders; import org.apache.logging.log4j.Level; import org.apache.logging.log4j.core.config.Configurator; import picocli.CommandLine; @@ -75,9 +76,21 @@ class Command { @CommandLine.Command(name = "supports", description = "Check if a feature is supported, or list all supported features") static class Supports implements Callable { - // Static, always-on capabilities of this build. There is no "known but + // Always-on capabilities of this build. There is no "known but // unsupported" feature here: anything not in this list is simply unrecognized. - private static final List FEATURES = List.of("dpop", "dpop_nonce_challenge"); + private static final List STATIC_FEATURES = List.of("dpop", "dpop_nonce_challenge"); + + // session-key-mlkem additionally requires an ML-KEM KemProvider on the + // classpath (e.g. sdk-pqc-bc), which the fips Maven profile excludes. + // Check the live registry rather than assuming build-profile wiring, so + // this doesn't claim support on a classpath that would actually throw. + private static List features() { + List features = new ArrayList<>(STATIC_FEATURES); + if (KemProviders.registered().contains(KeyType.MLKEM768Key)) { + features.add("session-key-mlkem"); + } + return features; + } @CommandLine.Parameters(index = "0", arity = "0..1", description = "Feature to check (e.g., dpop). Omit to list all supported features.") private String feature; @@ -87,12 +100,13 @@ static class Supports implements Callable { @Override public Integer call() { + List features = features(); if (feature == null) { - printFeatures(FEATURES); + printFeatures(features); return 0; } - Optional canonical = FEATURES.stream().filter(f -> f.equalsIgnoreCase(feature)).findFirst(); + Optional canonical = features.stream().filter(f -> f.equalsIgnoreCase(feature)).findFirst(); if (json) { Map result = new LinkedHashMap<>(); // Emit the canonical feature name for recognized features so casing is stable diff --git a/cmdline/src/test/java/io/opentdf/platform/CommandTest.java b/cmdline/src/test/java/io/opentdf/platform/CommandTest.java index 8995271d..130c80d3 100644 --- a/cmdline/src/test/java/io/opentdf/platform/CommandTest.java +++ b/cmdline/src/test/java/io/opentdf/platform/CommandTest.java @@ -95,7 +95,7 @@ void supports_noArgs_listsFeatures_exits0() { String out = captureStdout(() -> code[0] = new CommandLine(new Command()).execute("supports")); assertThat(code[0]).isEqualTo(0); - assertThat(out.lines()).containsExactlyInAnyOrder("dpop", "dpop_nonce_challenge"); + assertThat(out.lines()).containsExactlyInAnyOrder("dpop", "dpop_nonce_challenge", "session-key-mlkem"); } @Test @@ -104,7 +104,8 @@ void supports_noArgs_json_listsFeatures() { String out = captureStdout(() -> code[0] = new CommandLine(new Command()).execute("supports", "--json")); assertThat(code[0]).isEqualTo(0); - assertThat(out.trim()).isEqualTo("{\"dpop\":true,\"dpop_nonce_challenge\":true}"); + assertThat(out.trim()).isEqualTo( + "{\"dpop\":true,\"dpop_nonce_challenge\":true,\"session-key-mlkem\":true}"); } @Test diff --git a/sdk-pqc-bc/src/main/java/io/opentdf/platform/sdk/pqc/bc/BouncyCastleKemProvider.java b/sdk-pqc-bc/src/main/java/io/opentdf/platform/sdk/pqc/bc/BouncyCastleKemProvider.java index cccad461..1ca268b9 100644 --- a/sdk-pqc-bc/src/main/java/io/opentdf/platform/sdk/pqc/bc/BouncyCastleKemProvider.java +++ b/sdk-pqc-bc/src/main/java/io/opentdf/platform/sdk/pqc/bc/BouncyCastleKemProvider.java @@ -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); + } } diff --git a/sdk-pqc-bc/src/main/java/io/opentdf/platform/sdk/pqc/bc/HybridCrypto.java b/sdk-pqc-bc/src/main/java/io/opentdf/platform/sdk/pqc/bc/HybridCrypto.java index 361ec2d3..cbdf1c53 100644 --- a/sdk-pqc-bc/src/main/java/io/opentdf/platform/sdk/pqc/bc/HybridCrypto.java +++ b/sdk-pqc-bc/src/main/java/io/opentdf/platform/sdk/pqc/bc/HybridCrypto.java @@ -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; @@ -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. */ diff --git a/sdk-pqc-bc/src/test/java/io/opentdf/platform/sdk/pqc/bc/HybridCryptoGenerateKeyPairTest.java b/sdk-pqc-bc/src/test/java/io/opentdf/platform/sdk/pqc/bc/HybridCryptoGenerateKeyPairTest.java new file mode 100644 index 00000000..52740be4 --- /dev/null +++ b/sdk-pqc-bc/src/test/java/io/opentdf/platform/sdk/pqc/bc/HybridCryptoGenerateKeyPairTest.java @@ -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". + * + *

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"); + } +} diff --git a/sdk/src/main/java/io/opentdf/platform/sdk/KASClient.java b/sdk/src/main/java/io/opentdf/platform/sdk/KASClient.java index dec68a33..3e3b2fa9 100644 --- a/sdk/src/main/java/io/opentdf/platform/sdk/KASClient.java +++ b/sdk/src/main/java/io/opentdf/platform/sdk/KASClient.java @@ -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; @@ -47,8 +49,8 @@ class KASClient implements SDK.KAS { private final BiFunction protocolClientFactory; private final boolean usePlaintext; private final JWSSigner signer; - private AsymDecryption decryptor; - private String clientPublicKey; + private volatile AsymDecryption decryptor; + private volatile String clientPublicKey; private KASKeyCache kasKeyCache; private static final Logger log = LoggerFactory.getLogger(KASClient.class); @@ -118,25 +120,86 @@ 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. Double-checked locking: decryptor/clientPublicKey are + // volatile so the initializing thread's write is visible to others, + // and the synchronized block keeps two concurrent first-callers from + // generating distinct keypairs and racing to overwrite each other's + // fields (which would send one generation's public key while + // decrypting with another's private key). + 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); + } + + 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); @@ -171,29 +234,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 stubs = new HashMap<>(); diff --git a/sdk/src/main/java/io/opentdf/platform/sdk/spi/KemProvider.java b/sdk/src/main/java/io/opentdf/platform/sdk/spi/KemProvider.java index aa90c0da..06772c1c 100644 --- a/sdk/src/main/java/io/opentdf/platform/sdk/spi/KemProvider.java +++ b/sdk/src/main/java/io/opentdf/platform/sdk/spi/KemProvider.java @@ -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; + } + } }