opentdf · dmihalcik-virtru · Jul 31, 2026
@@ -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<Integer> {
-        // 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<String> FEATURES = List.of("dpop", "dpop_nonce_challenge");
+        private static final List<String> 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<String> features() {
+            List<String> 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<Integer> {
 
         @Override
         public Integer call() {
+            List<String> features = features();
             if (feature == null) {
-                printFeatures(FEATURES);
+                printFeatures(features);
                 return 0;
             }
 
-            Optional<String> canonical = FEATURES.stream().filter(f -> f.equalsIgnoreCase(feature)).findFirst();
+            Optional<String> canonical = features.stream().filter(f -> f.equalsIgnoreCase(feature)).findFirst();
             if (json) {
                 Map<String, Boolean> result = new LinkedHashMap<>();
                 // Emit the canonical feature name for recognized features so casing is stable

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

@@ -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);
+    }
 }
@@ -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.
      */

@@ -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");
+    }
+}
@@ -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<OkHttpClient, String, ProtocolClient> 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<String, AccessServiceClient> stubs = new HashMap<>();

@@ -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;
+        }
+    }
 }