Add an ML-KEM implementation (#2)

This PR adds an implementation of ML-KEM variant of the Kyber KEM as described
in the initial public draft of FIPS 203. This implementation covers all three
parameter sets described in the specification by making heavy use of Rust
generics and the typenum crate.

In addition to self-compatibility testing, we test correctness by verifying the
test vectors supplied by NIST; see tests/nist.rs.

Cargo.toml

@@ -0,0 +1,5 @@
+[workspace]
+resolver = "2"
+members = [
+    "ml-kem",
+]

ml-kem/Cargo.toml

@@ -0,0 +1,28 @@
+[package]
+name = "ml-kem"
+version = "0.1.0"
+edition = "2021"
+
+[features]
+default = []
+deterministic = [] # Expose deterministic generation and encapsulation functions
+
+[dependencies]
+const-default = "1.0.0"
+crypto-common = { version = "0.1.6", features = ["getrandom"] }
+generic-array = { version = "1.0.0", features = ["const-default"] }
+hybrid-array = { version = "0.2.0-rc.6" }
+sha3 = "0.10.8"
+
+[dev-dependencies]
+criterion = "0.5.1"
+hex = "0.4.3"
+hex-literal = "0.4.1"
+rand = "0.8.5"
+
+[profile.bench]
+debug = true
+
+[[bench]]
+name = "mlkem"
+harness = false

ml-kem/benches/mlkem.rs

@@ -0,0 +1,59 @@
+use criterion::{criterion_group, criterion_main, Criterion};
+use crypto_common::rand_core::CryptoRngCore;
+use hybrid_array::{Array, ArraySize};
+use ml_kem::*;
+
+pub fn rand<L: ArraySize>(rng: &mut impl CryptoRngCore) -> Array<u8, L> {
+    let mut val = Array::<u8, L>::default();
+    rng.fill_bytes(&mut val);
+    val
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let mut rng = rand::thread_rng();
+    let d: B32 = rand(&mut rng);
+    let z: B32 = rand(&mut rng);
+    let m: B32 = rand(&mut rng);
+
+    let (dk, ek) = MlKem768::generate_deterministic(&d, &z);
+    let dk_bytes = dk.as_bytes();
+    let ek_bytes = ek.as_bytes();
+    let (ct, _sk) = ek.encapsulate(&mut rng).unwrap();
+
+    // Key generation
+    c.bench_function("keygen", |b| {
+        b.iter(|| {
+            let (dk, ek) = <MlKem768 as KemCore>::generate_deterministic(&d, &z);
+            let _dk_bytes = dk.as_bytes();
+            let _ek_bytes = ek.as_bytes();
+        })
+    });
+
+    // Encapsulation
+    c.bench_function("encapsulate", |b| {
+        b.iter(|| {
+            let ek = <MlKem768 as KemCore>::EncapsulationKey::from_bytes(&ek_bytes);
+            ek.encapsulate_deterministic(&m).unwrap();
+        })
+    });
+
+    // Decapsulation
+    c.bench_function("decapsulate", |b| {
+        b.iter(|| {
+            let dk = <MlKem768 as KemCore>::DecapsulationKey::from_bytes(&dk_bytes);
+            dk.decapsulate(&ct).unwrap();
+        })
+    });
+
+    // Round trip
+    c.bench_function("round_trip", |b| {
+        b.iter(|| {
+            let (dk, ek) = <MlKem768 as KemCore>::generate_deterministic(&d, &z);
+            let (ct, _sk) = ek.encapsulate(&mut rng).unwrap();
+            dk.decapsulate(&ct).unwrap();
+        })
+    });
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);