Support derandomized key generation for ML-KEM

Signed-off-by: Spencer Wilson <[email protected]>

Author: SWilson4

src/main/c/KeyEncapsulation.c

@@ -85,6 +85,10 @@ JNIEXPORT jobject JNICALL Java_org_openquantumsafe_KeyEncapsulation_get_1KEM_1de
     jfieldID _length_shared_secret = (*env)->GetFieldID(env, cls, "length_shared_secret", "J");
     (*env)->SetLongField(env, _nativeKED, _length_shared_secret, (jlong) kem->length_shared_secret);
 
+    // long length_keypair_seed;
+    jfieldID _length_keypair_seed = (*env)->GetFieldID(env, cls, "length_keypair_seed", "J");
+    (*env)->SetLongField(env, _nativeKED, _length_keypair_seed, (jlong) kem->length_keypair_seed);
+
     return _nativeKED;
 }
 
@@ -110,6 +114,30 @@ JNIEXPORT jint JNICALL Java_org_openquantumsafe_KeyEncapsulation_generate_1keypa
     return (rv_ == OQS_SUCCESS) ? 0 : -1;
 }
 
+/*
+ * Class:     org_openquantumsafe_KeyEncapsulation
+ * Method:    generate_keypair
+ * Signature: ([B[B)I
+ */
+JNIEXPORT jint JNICALL Java_org_openquantumsafe_KeyEncapsulation_generate_1keypair_1derand
+  (JNIEnv *env, jobject obj, jbyteArray jpublic_key, jbyteArray jsecret_key, jbyteArray jseed)
+{
+    jbyte *public_key_native = (*env)->GetByteArrayElements(env, jpublic_key, 0);
+    jbyte *secret_key_native = (*env)->GetByteArrayElements(env, jsecret_key, 0);
+    jbyte *seed_native = (*env)->GetByteArrayElements(env, jseed, 0);
+
+    // Get pointer to KEM
+    OQS_KEM *kem = (OQS_KEM *) getHandle(env, obj, "native_kem_handle_");
+
+    // Invoke liboqs KEM keypair generation function
+    OQS_STATUS rv_ = OQS_KEM_keypair_derand(kem, (uint8_t*) public_key_native, (uint8_t*) secret_key_native, (uint8_t*) seed_native);
+
+    (*env)->ReleaseByteArrayElements(env, jpublic_key, public_key_native, 0);
+    (*env)->ReleaseByteArrayElements(env, jsecret_key, secret_key_native, 0);
+    (*env)->ReleaseByteArrayElements(env, jseed, seed_native, JNI_ABORT);
+    return (rv_ == OQS_SUCCESS) ? 0 : -1;
+}
+
 /*
  * Class:     org_openquantumsafe_KeyEncapsulation
  * Method:    encap_secret

src/main/java/org/openquantumsafe/KeyEncapsulation.java

@@ -20,6 +20,7 @@ class KeyEncapsulationDetails {
         long length_secret_key;
         long length_ciphertext;
         long length_shared_secret;
+        long length_keypair_seed;
 
         /**
          * \brief Print KEM algorithm details
@@ -33,7 +34,9 @@ void printKeyEncapsulation() {
                 "\n  Length public key (bytes): " + this.length_public_key +
                 "\n  Length secret key (bytes): " + this.length_secret_key +
                 "\n  Length ciphertext (bytes): " + this.length_ciphertext +
-                "\n  Length shared secret (bytes): " + this.length_shared_secret
+                "\n  Length shared secret (bytes): " + this.length_shared_secret +
+                "\n  Length keypair seed (bytes): "
+                + ((this.length_keypair_seed > 0) ? this.length_keypair_seed : "N/A")
             );
         }
 
@@ -114,6 +117,18 @@ public KeyEncapsulation(String alg_name, byte[] secret_key)
      */
     private native int generate_keypair(byte[] public_key, byte[] secret_key);
 
+    /**
+     * \brief Wrapper for OQS_API OQS_STATUS OQS_KEM_keypair_derand(const OQS_KEM *kem,
+     *                              uint8_t *public_key, uint8_t *secret_key,
+     *                              const uint8_t *seed);
+     * \param Public key
+     * \param Secret key
+     * \param Seed
+     * \return Status
+     */
+    private native int generate_keypair_derand(byte[] public_key,
+                                               byte[] secret_key, byte[] seed);
+
     /**
      * \brief Wrapper for OQS_API OQS_STATUS OQS_KEM_encaps(const OQS_KEM *kem,
      *                                               uint8_t *ciphertext,
@@ -159,6 +174,27 @@ public byte[] generate_keypair() throws RuntimeException {
         return this.public_key_;
     }
 
+    /**
+     * \brief Invoke native generate_keypair_derand method using the PK and SK lengths
+     * from alg_details_. Check return value and if != 0 throw Exception.
+     */
+    public byte[] generate_keypair(byte[] seed) throws RuntimeException {
+        if (seed.length != alg_details_.length_keypair_seed) {
+            throw new RuntimeException("Incorrect seed length");
+        }
+
+        int rv_ = generate_keypair_derand(this.public_key_, this.secret_key_, seed);
+        if (rv_ != 0) throw new RuntimeException("Cannot generate keypair from seed");
+        return this.public_key_;
+    }
+
+    /**
+     * \brief Return seed length
+     */
+    public long get_keypair_seed_length() {
+        return alg_details_.length_keypair_seed;
+    }
+
     /**
      * \brief Return public key
      */

src/test/java/org/openquantumsafe/KEMTest.java

@@ -8,6 +8,7 @@
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.stream.Stream;
 
 public class KEMTest {
@@ -56,6 +57,43 @@ public void testAllKEMs(String kem_name) {
         System.out.println(sb.toString());
     }
 
+    /**
+     * Test KEMs with derandomized keypair generation.
+     */
+    @ParameterizedTest(name = "Testing {arguments}")
+    @MethodSource("getDerandSupportedKEMsAsStream")
+    public void testKEMsWithDerand(String kem_name) {
+        StringBuilder sb = new StringBuilder();
+        sb.append(kem_name);
+        sb.append(" (derand)");
+        sb.append(String.format("%1$" + (40 - kem_name.length() - 9) + "s", ""));
+
+        // Create client and server
+        KeyEncapsulation client = new KeyEncapsulation(kem_name);
+        KeyEncapsulation server = new KeyEncapsulation(kem_name);
+
+        // Generate seed
+        byte[] seed = Rand.randombytes(client.get_keypair_seed_length());
+
+        // Generate client key pair
+        byte[] client_public_key = client.generate_keypair(seed);
+
+        // Server: encapsulate secret with client's public key
+        Pair<byte[], byte[]> server_pair = server.encap_secret(client_public_key);
+        byte[] ciphertext = server_pair.getLeft();
+        byte[] shared_secret_server = server_pair.getRight();
+
+        // Client: decapsulate
+        byte[] shared_secret_client = client.decap_secret(ciphertext);
+
+        // Check if equal
+        assertArrayEquals(shared_secret_client, shared_secret_server, kem_name);
+
+        // If successful print KEM name, otherwise an exception will be thrown
+        sb.append("\033[0;32m").append("PASSED").append("\033[0m");
+        System.out.println(sb.toString());
+    }
+
     /**
      * Test the MechanismNotSupported Exception
      */
@@ -71,4 +109,12 @@ private static Stream<String> getEnabledKEMsAsStream() {
         return enabled_kems.parallelStream();
     }
 
+    /**
+     * Method to convert the list of derand-supported KEMs to a stream for input to testAllSigs
+     */
+    private static Stream<String> getDerandSupportedKEMsAsStream() {
+        return Arrays.asList(
+                "ML-KEM-512", "ML-KEM-768", "ML-KEM-1024"
+            ).parallelStream();
+    }
 }