Add implementation of NIST standard ML-KEM 1024 (#367)

Co-authored-by: Jordan Rose <jrose@signal.org>

Cargo.lock

@@ -1644,7 +1644,8 @@ dependencies = [
  "itertools",
  "log",
  "num_enum",
- "pqcrypto-kyber",
+ "pqcrypto-kyber 0.7.6",
+ "pqcrypto-kyber 0.8.0",
  "pqcrypto-traits",
  "proptest",
  "prost",
@@ -2097,10 +2098,23 @@ dependencies = [
 ]
 
 [[package]]
-name = "pqcrypto-traits"
-version = "0.3.4"
+name = "pqcrypto-kyber"
+version = "0.8.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "97e91cb6af081c6daad5fa705f8adb0634c027662052cb3174bdf2957bf07e25"
+checksum = "2bc5d857fb0a0a0695dbe379f449a185bf73d0173cdcaffa86c015b5d1b11490"
+dependencies = [
+ "cc",
+ "glob",
+ "libc",
+ "pqcrypto-internals",
+ "pqcrypto-traits",
+]
+
+[[package]]
+name = "pqcrypto-traits"
+version = "0.3.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "94e851c7654eed9e68d7d27164c454961a616cf8c203d500607ef22c737b51bb"
 
 [[package]]
 name = "prettyplease"

bin/verify_duplicate_crates

@@ -11,8 +11,10 @@
 # You can use the `cargo tree` command below to see where they come from,
 # and then document them here.
 #
-# (no duplicate crates at this time)
-EXPECTED=""
+# pqcrypto-kyber: v0.7 is what we shipped PQXDH on, v0.8 contains the NIST standard version
+EXPECTED="
+pqcrypto-kyber v0.7.6
+pqcrypto-kyber v0.8.0"
 
 check_cargo_tree() {
     # Only check the mobile targets, where we care most about code size.

rust/protocol/Cargo.toml

@@ -34,10 +34,15 @@ uuid = "1.1.2"
 displaydoc = "0.2"
 thiserror = "1.0.30"
 pqcrypto-kyber = { version = "0.7.6", default-features = false, features = ["std"] }
+pqcrypto-ml-kem = { version = "0.8.0", default-features = false, features = ["std"], package = "pqcrypto-kyber" }
 pqcrypto-traits = "0.3.4"
 
 [features]
 kyber768 = []
+# ML-KEM matches the NIST standard version of Kyber. It may still change
+# incompatibly until the final version of the standard is published and
+# libsignal will update to match.
+mlkem1024 = []
 
 [dev-dependencies]
 criterion = "0.5"

rust/protocol/src/kem.rs

@@ -53,6 +53,8 @@
 mod kyber1024;
 #[cfg(any(feature = "kyber768", test))]
 mod kyber768;
+#[cfg(any(feature = "mlkem1024", test))]
+mod mlkem1024;
 
 use crate::{Result, SignalProtocolError};
 
@@ -147,6 +149,9 @@ pub enum KeyType {
     Kyber768,
     /// Kyber1024 key
     Kyber1024,
+    /// ML-KEM 1024 key
+    #[cfg(any(feature = "mlkem1024", test))]
+    MLKEM1024,
 }
 
 impl KeyType {
@@ -155,6 +160,8 @@ impl KeyType {
             #[cfg(any(feature = "kyber768", test))]
             KeyType::Kyber768 => 0x07,
             KeyType::Kyber1024 => 0x08,
+            #[cfg(any(feature = "mlkem1024", test))]
+            KeyType::MLKEM1024 => 0x0A,
         }
     }
 
@@ -166,6 +173,8 @@ impl KeyType {
             #[cfg(any(feature = "kyber768", test))]
             KeyType::Kyber768 => &kyber768::Parameters,
             KeyType::Kyber1024 => &kyber1024::Parameters,
+            #[cfg(any(feature = "mlkem1024", test))]
+            KeyType::MLKEM1024 => &mlkem1024::Parameters,
         }
     }
 }
@@ -178,6 +187,8 @@ impl TryFrom<u8> for KeyType {
             #[cfg(any(feature = "kyber768", test))]
             0x07 => Ok(KeyType::Kyber768),
             0x08 => Ok(KeyType::Kyber1024),
+            #[cfg(any(feature = "mlkem1024", test))]
+            0x0A => Ok(KeyType::MLKEM1024),
             t => Err(SignalProtocolError::BadKEMKeyType(t)),
         }
     }
@@ -495,6 +506,23 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_mlkem1024_kem() -> Result<()> {
+        // test data for kyber1024
+        let pk_bytes = include_bytes!("kem/test-data/mlkem-pk.dat");
+        let sk_bytes = include_bytes!("kem/test-data/mlkem-sk.dat");
+
+        let pubkey = PublicKey::deserialize(pk_bytes).expect("deserialize pubkey");
+        let secretkey = SecretKey::deserialize(sk_bytes).expect("deserialize secretkey");
+
+        assert_eq!(pubkey.key_type, KeyType::MLKEM1024);
+        let (ss_for_sender, ct) = pubkey.encapsulate();
+        let ss_for_recipient = secretkey.decapsulate(&ct).expect("decapsulation works");
+
+        assert_eq!(ss_for_sender, ss_for_recipient);
+
+        Ok(())
+    }
     #[test]
     fn test_kyber1024_keypair() -> Result<()> {
         let kp = KeyPair::generate(KeyType::Kyber1024);
@@ -512,4 +540,13 @@ mod tests {
         assert_eq!(ss_for_recipient, ss_for_sender);
         Ok(())
     }
+
+    #[test]
+    fn test_mlkem1024_keypair() -> Result<()> {
+        let kp = KeyPair::generate(KeyType::MLKEM1024);
+        let (ss_for_sender, ct) = kp.public_key.encapsulate();
+        let ss_for_recipient = kp.secret_key.decapsulate(&ct).expect("decapsulation works");
+        assert_eq!(ss_for_recipient, ss_for_sender);
+        Ok(())
+    }
 }

rust/protocol/src/kem/mlkem1024.rs

@@ -0,0 +1,51 @@
+//
+// Copyright 2023 Signal Messenger, LLC.
+// SPDX-License-Identifier: AGPL-3.0-only
+//
+
+use crate::Result;
+
+use pqcrypto_traits::kem::{Ciphertext, PublicKey, SecretKey, SharedSecret};
+
+use super::{KeyMaterial, Public, Secret};
+use pqcrypto_ml_kem::ffi::{
+    PQCLEAN_KYBER1024_CLEAN_CRYPTO_BYTES, PQCLEAN_KYBER1024_CLEAN_CRYPTO_CIPHERTEXTBYTES,
+    PQCLEAN_KYBER1024_CLEAN_CRYPTO_PUBLICKEYBYTES, PQCLEAN_KYBER1024_CLEAN_CRYPTO_SECRETKEYBYTES,
+};
+
+pub(crate) struct Parameters;
+
+impl super::Parameters for Parameters {
+    const PUBLIC_KEY_LENGTH: usize = PQCLEAN_KYBER1024_CLEAN_CRYPTO_PUBLICKEYBYTES;
+    const SECRET_KEY_LENGTH: usize = PQCLEAN_KYBER1024_CLEAN_CRYPTO_SECRETKEYBYTES;
+    const CIPHERTEXT_LENGTH: usize = PQCLEAN_KYBER1024_CLEAN_CRYPTO_CIPHERTEXTBYTES;
+    const SHARED_SECRET_LENGTH: usize = PQCLEAN_KYBER1024_CLEAN_CRYPTO_BYTES;
+
+    fn generate() -> (KeyMaterial<Public>, KeyMaterial<Secret>) {
+        let (pk, sk) = pqcrypto_ml_kem::kyber1024::keypair();
+        (
+            KeyMaterial::new(pk.as_bytes().into()),
+            KeyMaterial::new(sk.as_bytes().into()),
+        )
+    }
+
+    fn encapsulate(pub_key: &KeyMaterial<Public>) -> (super::SharedSecret, super::RawCiphertext) {
+        let mlkem_pk = pqcrypto_ml_kem::kyber1024::PublicKey::from_bytes(pub_key)
+            .expect("valid ML-KEM 1024 public key bytes");
+        let (mlkem_ss, mlkem_ct) = pqcrypto_ml_kem::kyber1024::encapsulate(&mlkem_pk);
+        (mlkem_ss.as_bytes().into(), mlkem_ct.as_bytes().into())
+    }
+
+    fn decapsulate(
+        secret_key: &KeyMaterial<Secret>,
+        ciphertext: &[u8],
+    ) -> Result<super::SharedSecret> {
+        let mlkem_sk = pqcrypto_ml_kem::kyber1024::SecretKey::from_bytes(secret_key)
+            .expect("valid ML-KEM 1024 secret key bytes");
+        let mlkem_ct = pqcrypto_ml_kem::kyber1024::Ciphertext::from_bytes(ciphertext)
+            .expect("valid ML-KEM 1024 ciphertext");
+        let mlkem_ss = pqcrypto_ml_kem::kyber1024::decapsulate(&mlkem_ct, &mlkem_sk);
+
+        Ok(mlkem_ss.as_bytes().into())
+    }
+}