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

cmdline/src/main/java/io/opentdf/platform/Command.java

@@ -47,6 +47,9 @@
 import io.opentdf.platform.sdk.KeyType;
 import io.opentdf.platform.sdk.SDK;
 import io.opentdf.platform.sdk.SDKBuilder;
+import io.opentdf.platform.sdk.SDKException;
+import io.opentdf.platform.sdk.spi.KemProvider;
+import io.opentdf.platform.sdk.spi.KemProviders;
 import org.apache.logging.log4j.Level;
 import org.apache.logging.log4j.core.config.Configurator;
 import picocli.CommandLine;
@@ -79,11 +82,36 @@ 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",
                 "partial_decrypt", "http_range_source");
 
+        // session-key-mlkem additionally requires a KemProvider that can generate
+        // ML-KEM-768 and ML-KEM-1024 session keys (e.g. sdk-pqc-bc), which the
+        // fips Maven profile excludes. Probe by generating one of each rather
+        // than checking the registry: a provider can be registered for these
+        // types yet only wrap and unwrap, since generateSessionKey defaults to
+        // unsupported.
+        private static List<String> features() {
+            List<String> features = new ArrayList<>(STATIC_FEATURES);
+            if (canGenerateSessionKey(KeyType.MLKEM768Key) && canGenerateSessionKey(KeyType.MLKEM1024Key)) {
+                features.add("session-key-mlkem");
+            }
+            return features;
+        }
+
+        private static boolean canGenerateSessionKey(KeyType keyType) {
+            KemProvider.KemSessionKey probe;
+            try {
+                probe = KemProviders.get(keyType).generateSessionKey(keyType);
+            } catch (SDKException | UnsupportedOperationException e) {
+                return false;
+            }
+            probe.close();
+            return true;
+        }
+
         @CommandLine.Parameters(index = "0", arity = "0..1", description = "Feature to check (e.g., dpop). Omit to list all supported features.")
         private String feature;
 
@@ -92,12 +120,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

cmdline/src/test/java/io/opentdf/platform/CommandTest.java

@@ -14,6 +14,8 @@
 import java.nio.file.FileSystems;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Map;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -38,6 +40,33 @@ private String captureStdout(Runnable action) {
         return captured.toString(StandardCharsets.UTF_8);
     }
 
+    // session-key-mlkem is reported only when an ML-KEM KemProvider is on the
+    // classpath: the default (non-fips) profile puts sdk-pqc-bc on cmdline's
+    // runtime classpath, the fips profile leaves it off, and CI runs these
+    // tests under fips. Probe for the provider class directly rather than via
+    // KemProviders, so the expectation doesn't just mirror Supports' own check
+    // (a provider jar whose ServiceLoader registration broke would still fail).
+    private static final boolean MLKEM_PROVIDER_ON_CLASSPATH =
+            isClassPresent("io.opentdf.platform.sdk.pqc.bc.BouncyCastleKemProvider");
+
+    private static boolean isClassPresent(String className) {
+        try {
+            Class.forName(className, false, CommandTest.class.getClassLoader());
+            return true;
+        } catch (ClassNotFoundException e) {
+            return false;
+        }
+    }
+
+    private static List<String> expectedFeatures() {
+        List<String> features = new ArrayList<>(
+                List.of("dpop", "dpop_nonce_challenge", "partial_decrypt", "http_range_source"));
+        if (MLKEM_PROVIDER_ON_CLASSPATH) {
+            features.add("session-key-mlkem");
+        }
+        return features;
+    }
+
     @Test
     void supports_dpop_exits_0() {
         int code = new CommandLine(new Command()).execute("supports", "dpop");
@@ -56,6 +85,12 @@ void supports_partial_decrypt_and_http_range_source_exit_0() {
         assertThat(new CommandLine(new Command()).execute("supports", "http_range_source")).isEqualTo(0);
     }
 
+    @Test
+    void supports_session_key_mlkem_exits_0_only_with_mlkem_provider() {
+        int code = new CommandLine(new Command()).execute("supports", "session-key-mlkem");
+        assertThat(code).isEqualTo(MLKEM_PROVIDER_ON_CLASSPATH ? 0 : 1);
+    }
+
     @Test
     void supports_unknown_feature_exits_1() {
         int code = new CommandLine(new Command()).execute("supports", "unknown_feature");
@@ -107,8 +142,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",
-                "partial_decrypt", "http_range_source");
+        assertThat(out.lines()).containsExactlyInAnyOrderElementsOf(expectedFeatures());
     }
 
     @Test
@@ -117,8 +151,11 @@ 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,"
-                + "\"partial_decrypt\":true,\"http_range_source\":true}");
+        String expected = "{\"dpop\":true,\"dpop_nonce_challenge\":true,"
+                + "\"partial_decrypt\":true,\"http_range_source\":true"
+                + (MLKEM_PROVIDER_ON_CLASSPATH ? ",\"session-key-mlkem\":true" : "")
+                + "}";
+        assertThat(out.trim()).isEqualTo(expected);
     }
 
     @Test

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 KemSessionKey generateSessionKey(KeyType keyType) {
+        return HybridCrypto.generateSessionKey(keyType);
+    }
 }

sdk-pqc-bc/src/main/java/io/opentdf/platform/sdk/pqc/bc/BouncyCastleKemSessionKey.java

@@ -0,0 +1,55 @@
+package io.opentdf.platform.sdk.pqc.bc;
+
+import io.opentdf.platform.sdk.spi.KemProvider;
+
+import java.util.Arrays;
+
+/**
+ * {@link KemProvider.KemSessionKey} returned by {@link BouncyCastleKemProvider}.
+ * Owns the raw private key bytes (the same encoding the algorithm classes'
+ * {@code unwrapDEK(byte[] rawPriv, ...)} take) and zeroes them on {@link #close()}.
+ * The private key is never PEM-encoded or turned into a {@code String}.
+ *
+ * <p>BouncyCastle's key parameter objects, built inside each unwrap, keep
+ * their own copies that cannot be wiped; they are unreachable once the
+ * unwrap returns.
+ */
+final class BouncyCastleKemSessionKey implements KemProvider.KemSessionKey {
+
+    /** One algorithm's raw-key unwrap, e.g. {@code MLKEMAlgorithm.MLKEM_768::unwrapDEK}. */
+    @FunctionalInterface
+    interface RawUnwrap {
+        byte[] unwrapDEK(byte[] rawPrivateKey, byte[] wrapped);
+    }
+
+    private final String publicKeyPEM;
+    private final byte[] privateKey;
+    private final RawUnwrap unwrap;
+    private boolean closed;
+
+    /** Takes ownership of {@code privateKey}: the caller must not keep or reuse the array. */
+    BouncyCastleKemSessionKey(String publicKeyPEM, byte[] privateKey, RawUnwrap unwrap) {
+        this.publicKeyPEM = publicKeyPEM;
+        this.privateKey = privateKey;
+        this.unwrap = unwrap;
+    }
+
+    @Override
+    public String publicKeyPEM() {
+        return publicKeyPEM;
+    }
+
+    @Override
+    public synchronized byte[] unwrapDEK(byte[] wrapped) {
+        if (closed) {
+            throw new IllegalStateException("KEM session key has been closed");
+        }
+        return unwrap.unwrapDEK(privateKey, wrapped);
+    }
+
+    @Override
+    public synchronized void close() {
+        Arrays.fill(privateKey, (byte) 0);
+        closed = true;
+    }
+}

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,44 @@ static byte[] unwrapDEK(KeyType keyType, String privateKeyPEM, byte[] wrapped) {
         }
     }
 
+    /**
+     * Generate a fresh ephemeral rewrap session key for {@code keyType}. Same
+     * dispatch table as {@link #wrapDEK} / {@link #unwrapDEK}. The raw private
+     * key goes straight into the returned {@link BouncyCastleKemSessionKey};
+     * it is never PEM-encoded.
+     */
+    static KemProvider.KemSessionKey generateSessionKey(KeyType keyType) {
+        switch (keyType) {
+            case HybridXWingKey: {
+                XWingKeyPair kp = XWingKeyPair.generate();
+                return new BouncyCastleKemSessionKey(kp.publicKeyInPemFormat(), kp.getPrivateKey(),
+                        XWingKeyPair::unwrapDEK);
+            }
+            case HybridSecp256r1MLKEM768Key: {
+                HybridNISTKeyPair kp = HybridNISTAlgorithm.P256_MLKEM768.generate();
+                return new BouncyCastleKemSessionKey(kp.publicKeyInPemFormat(), kp.getPrivateKey(),
+                        HybridNISTAlgorithm.P256_MLKEM768::unwrapDEK);
+            }
+            case HybridSecp384r1MLKEM1024Key: {
+                HybridNISTKeyPair kp = HybridNISTAlgorithm.P384_MLKEM1024.generate();
+                return new BouncyCastleKemSessionKey(kp.publicKeyInPemFormat(), kp.getPrivateKey(),
+                        HybridNISTAlgorithm.P384_MLKEM1024::unwrapDEK);
+            }
+            case MLKEM768Key: {
+                MLKEMKeyPair kp = MLKEMAlgorithm.MLKEM_768.generate();
+                return new BouncyCastleKemSessionKey(kp.publicKeyInPemFormat(), kp.getPrivateKey(),
+                        MLKEMAlgorithm.MLKEM_768::unwrapDEK);
+            }
+            case MLKEM1024Key: {
+                MLKEMKeyPair kp = MLKEMAlgorithm.MLKEM_1024.generate();
+                return new BouncyCastleKemSessionKey(kp.publicKeyInPemFormat(), kp.getPrivateKey(),
+                        MLKEMAlgorithm.MLKEM_1024::unwrapDEK);
+            }
+            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.
      */

sdk-pqc-bc/src/main/java/io/opentdf/platform/sdk/pqc/bc/XWingKeyPair.java

@@ -62,6 +62,11 @@ public String privateKeyInPemFormat() {
         return HybridSpki.encodePkcs8Pem(HybridSpki.OID_XWING, privateKey);
     }
 
+    /** A copy of the raw 32-byte private key seed, as {@link #unwrapDEK(byte[], byte[])} takes it. */
+    byte[] getPrivateKey() {
+        return privateKey.clone();
+    }
+
     public static byte[] pubKeyFromPem(String pem) {
         byte[] raw = HybridSpki.decodeSpkiPem(pem, HybridSpki.OID_XWING);
         if (raw.length != PUBLIC_KEY_SIZE) {

sdk-pqc-bc/src/test/java/io/opentdf/platform/sdk/pqc/bc/HybridCryptoSessionKeyTest.java

@@ -0,0 +1,116 @@
+package io.opentdf.platform.sdk.pqc.bc;
+
+import io.opentdf.platform.sdk.KeyType;
+import io.opentdf.platform.sdk.SDKException;
+import io.opentdf.platform.sdk.spi.KemProvider;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.MethodSource;
+
+import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
+import java.util.stream.Stream;
+
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+/**
+ * Unit tests for {@link HybridCrypto#generateSessionKey}, the dispatcher backing
+ * {@link BouncyCastleKemProvider#generateSessionKey} 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 the way KAS and KASClient use it:
+ * generate a session key, wrap a DEK against its public PEM, unwrap through
+ * the session key, assert equal. The negative cases check that the session
+ * key and envelope fail closed: unwrapping must throw rather than return a DEK.
+ */
+class HybridCryptoSessionKeyTest {
+
+    private static final byte[] DEK = "0123456789abcdef0123456789abcdef".getBytes(StandardCharsets.UTF_8);
+
+    /** Every hybrid and pure ML-KEM key type, so a newly added PQC type can't skip these tests. */
+    static Stream<KeyType> pqcKeyTypes() {
+        return Arrays.stream(KeyType.values()).filter(kt -> kt.isHybrid() || kt.isMLKEM());
+    }
+
+    @ParameterizedTest
+    @MethodSource("pqcKeyTypes")
+    void generatedSessionKeyRoundTrips(KeyType keyType) {
+        try (KemProvider.KemSessionKey sessionKey = HybridCrypto.generateSessionKey(keyType)) {
+            assertTrue(sessionKey.publicKeyPEM().startsWith("-----BEGIN PUBLIC KEY-----"), "public PEM header");
+
+            byte[] wrapped = HybridCrypto.wrapDEK(keyType, sessionKey.publicKeyPEM(), DEK);
+            assertArrayEquals(DEK, sessionKey.unwrapDEK(wrapped), "DEK round-trip via generated session key");
+        }
+    }
+
+    @ParameterizedTest
+    @MethodSource("pqcKeyTypes")
+    void otherSessionKeyOfSameTypeCannotUnwrap(KeyType keyType) {
+        try (KemProvider.KemSessionKey recipient = HybridCrypto.generateSessionKey(keyType);
+             KemProvider.KemSessionKey other = HybridCrypto.generateSessionKey(keyType)) {
+            byte[] wrapped = HybridCrypto.wrapDEK(keyType, recipient.publicKeyPEM(), DEK);
+            assertThrows(SDKException.class, () -> other.unwrapDEK(wrapped),
+                    "a different session key of the same type must not unwrap the DEK");
+        }
+    }
+
+    @ParameterizedTest
+    @MethodSource("pqcKeyTypes")
+    void tamperedKemCiphertextCannotUnwrap(KeyType keyType) {
+        try (KemProvider.KemSessionKey sessionKey = HybridCrypto.generateSessionKey(keyType)) {
+            byte[][] parts = HybridCrypto.unmarshalEnvelope(
+                    HybridCrypto.wrapDEK(keyType, sessionKey.publicKeyPEM(), DEK));
+
+            // The midpoint of the KEM ciphertext is inside its ML-KEM part for every type
+            // (the EC / X25519 share, where present, is appended after it). Decapsulation
+            // implicitly rejects and yields a wrong wrap key, so AES-GCM must fail.
+            parts[0][parts[0].length / 2] ^= 0x01;
+            byte[] tampered = HybridCrypto.marshalEnvelope(parts[0], parts[1]);
+
+            assertThrows(SDKException.class, () -> sessionKey.unwrapDEK(tampered),
+                    "a flipped byte in the KEM ciphertext must not unwrap the DEK");
+        }
+    }
+
+    @ParameterizedTest
+    @MethodSource("pqcKeyTypes")
+    void tamperedEncryptedDekCannotUnwrap(KeyType keyType) {
+        try (KemProvider.KemSessionKey sessionKey = HybridCrypto.generateSessionKey(keyType)) {
+            byte[][] parts = HybridCrypto.unmarshalEnvelope(
+                    HybridCrypto.wrapDEK(keyType, sessionKey.publicKeyPEM(), DEK));
+
+            parts[1][parts[1].length / 2] ^= 0x01;
+            byte[] tampered = HybridCrypto.marshalEnvelope(parts[0], parts[1]);
+
+            assertThrows(SDKException.class, () -> sessionKey.unwrapDEK(tampered),
+                    "a flipped byte in the AES-GCM sealed DEK must not unwrap");
+        }
+    }
+
+    @ParameterizedTest
+    @MethodSource("pqcKeyTypes")
+    void closedSessionKeyCannotUnwrap(KeyType keyType) {
+        KemProvider.KemSessionKey sessionKey = HybridCrypto.generateSessionKey(keyType);
+        byte[] wrapped = HybridCrypto.wrapDEK(keyType, sessionKey.publicKeyPEM(), DEK);
+
+        sessionKey.close();
+
+        assertThrows(IllegalStateException.class, () -> sessionKey.unwrapDEK(wrapped),
+                "a closed session key must not unwrap");
+    }
+
+    @Test
+    void closeZeroesThePrivateKey() {
+        byte[] privateKey = {1, 2, 3, 4};
+        BouncyCastleKemSessionKey sessionKey =
+                new BouncyCastleKemSessionKey("unused", privateKey, (priv, wrapped) -> priv.clone());
+        assertArrayEquals(new byte[] {1, 2, 3, 4}, sessionKey.unwrapDEK(new byte[0]),
+                "unwrap is handed the private key");
+
+        sessionKey.close();
+
+        assertArrayEquals(new byte[4], privateKey, "close must zero the private key");
+    }
+}