Reference implementation of Kyber768 in libcrux. (#33)

This change
* Starts adding the reference implementation of Kyber based on the spec.
* Adds KAT tests for Kyber768 ref implementation.

Cargo.toml

@@ -31,7 +31,8 @@ quickcheck = "1"
 quickcheck_macros = "1"
 serde_json = { version = "1.0" }
 serde = { version = "1.0", features = ["derive"] }
-hex = "0.4"
+hex = { version = "0.4.3", features = ["serde"] }
+pqcrypto-kyber = { version = "0.7.6", default-features = false }
 
 # Benchmarking "RustCrypto"
 chacha20poly1305 = "0.10"
@@ -70,6 +71,10 @@ harness = false
 name = "drbg"
 harness = false
 
+[[bench]]
+name = "kyber768"
+harness = false
+
 [features]
 hacspec = [] # TODO: #7 Use specs instead of efficient implementations
 rand = []

benches/kyber768.rs

@@ -0,0 +1,101 @@
+use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
+
+use libcrux::kem::Algorithm;
+use libcrux::drbg::Drbg;
+use libcrux::digest;
+
+pub fn comparisons_key_generation(c: &mut Criterion) {
+    let mut drbg = Drbg::new(digest::Algorithm::Sha256).unwrap();
+    let mut group = c.benchmark_group("Kyber768 Key Generation");
+
+    group.bench_function("libcrux reference implementation", |b| {
+        b.iter(
+            || {
+                let (_secret_key, _public_key) = libcrux::kem::key_gen(Algorithm::Kyber768, &mut drbg).unwrap();
+            }
+        )
+    });
+
+    group.bench_function("pqclean reference implementation", |b| {
+        b.iter(|| {
+            let (_public_key, _secret_key) = pqcrypto_kyber::kyber768::keypair();
+        })
+    });
+}
+
+pub fn comparisons_encapsulation(c: &mut Criterion) {
+    let mut group = c.benchmark_group("Kyber768 Encapsulation");
+
+    group.bench_function("libcrux reference implementation", |b| {
+        b.iter_batched(
+            || {
+                let mut drbg = Drbg::new(digest::Algorithm::Sha256).unwrap();
+                let (_secret_key, public_key) = libcrux::kem::key_gen(Algorithm::Kyber768, &mut drbg).unwrap();
+
+                (drbg, public_key)
+            },
+            |(mut rng, public_key)| {
+                let (_shared_secret, _ciphertext) =
+                    libcrux::kem::encapsulate(Algorithm::Kyber768, &public_key, &mut rng).unwrap();
+            },
+            BatchSize::SmallInput,
+        )
+    });
+
+    group.bench_function("pqclean reference implementation", |b| {
+        b.iter_batched(
+            || {
+                let (public_key, _secret_key) = pqcrypto_kyber::kyber768::keypair();
+
+                public_key
+            },
+            |public_key| {
+                let (_shared_secret, _ciphertext) = pqcrypto_kyber::kyber768::encapsulate(&public_key);
+            },
+            BatchSize::SmallInput,
+        )
+    });
+}
+
+pub fn comparisons_decapsulation(c: &mut Criterion) {
+    let mut group = c.benchmark_group("Kyber768 Decapsulation");
+
+    group.bench_function("libcrux specification", |b| {
+        b.iter_batched(
+            || {
+                let mut drbg = Drbg::new(digest::Algorithm::Sha256).unwrap();
+                let (secret_key, public_key) = libcrux::kem::key_gen(Algorithm::Kyber768, &mut drbg).unwrap();
+                let (_shared_secret, ciphertext) = libcrux::kem::encapsulate(Algorithm::Kyber768, &public_key, &mut drbg).unwrap();
+                (secret_key, ciphertext)
+            },
+            |(secret_key, ciphertext)| {
+                let _shared_secret = libcrux::kem::decapsulate(Algorithm::Kyber768, &ciphertext, &secret_key);
+            },
+            BatchSize::SmallInput,
+        )
+    });
+
+    group.bench_function("pqclean reference implementation", |b| {
+        b.iter_batched(
+            || {
+                let (public_key, secret_key) = pqcrypto_kyber::kyber768::keypair();
+                let (_shared_secret, ciphertext) = pqcrypto_kyber::kyber768::encapsulate(&public_key);
+
+                (ciphertext, secret_key)
+            },
+            |(ciphertext, secret_key)| {
+                let _shared_secret = pqcrypto_kyber::kyber768::decapsulate(&ciphertext, &secret_key);
+            },
+            BatchSize::SmallInput,
+        )
+    });
+}
+
+pub fn comparisons(c: &mut Criterion) {
+    comparisons_key_generation(c);
+    comparisons_encapsulation(c);
+    comparisons_decapsulation(c);
+}
+
+criterion_group!(benches, comparisons);
+criterion_main!(benches);

specs/kyber/src/lib.rs

@@ -28,9 +28,7 @@ pub type PublicKey = [u8; KYBER768_PUBLIC_KEY_SIZE];
 pub type PrivateKey = [u8; KYBER768_SECRET_KEY_SIZE];
 
 pub type Ciphertext = [u8; KYBER768_CIPHERTEXT_SIZE];
-pub type Enc = [u8; KYBER768_SHARED_SECRET_SIZE];
 pub type SharedSecret = [u8; KYBER768_SHARED_SECRET_SIZE];
-pub type Randomness = [u8; KYBER768_KEY_GENERATION_SEED_SIZE];
 
 #[derive(Debug)]
 pub struct BadRejectionSamplingRandomnessError;
@@ -105,7 +103,7 @@ pub fn generate_keypair(
 pub fn encapsulate(
     public_key: PublicKey,
     randomness: [u8; KYBER768_SHARED_SECRET_SIZE],
-) -> Result<(Ciphertext, Enc), BadRejectionSamplingRandomnessError> {
+) -> Result<(Ciphertext, SharedSecret), BadRejectionSamplingRandomnessError> {
     let randomness_hashed = H(&randomness);
 
     let to_hash: [u8; 2 * H_DIGEST_SIZE] = randomness_hashed.push(&H(&public_key));

src/drbg.rs

@@ -171,7 +171,13 @@ impl Drbg {
 /// Implementation of the [`RngCore`] trait for the [`Drbg`].
 impl RngCore for Drbg {
     fn next_u32(&mut self) -> u32 {
-        todo!()
+        let mut bytes : [u8; 4] = [0; 4];
+        self.generate(&mut bytes).unwrap();
+
+        (bytes[0] as u32) |
+        (bytes[1] as u32) << 8 |
+        (bytes[2] as u32) << 16 |
+        (bytes[3] as u32) << 24
     }
 
     fn next_u64(&mut self) -> u64 {

src/hpke/kem.rs

@@ -243,7 +243,7 @@ fn shared_secret_from_dh(alg: KEM, mut secret: Vec<u8>) -> Vec<u8> {
 /// [`ValidationError`](`HpkeError::ValidationError`) as described in
 /// [validation](#validation-of-inputs-and-outputs).
 pub fn DH(alg: KEM, sk: &PrivateKeyIn, pk: &PublicKeyIn) -> Result<SharedSecret, HpkeError> {
-    match crate::ecdh::derive(kem_to_named_group(alg).into(), pk, sk) {
+    match crate::ecdh::derive(kem_to_named_group(alg).try_into().unwrap(), pk, sk) {
         Ok(secret) => HpkeBytesResult::Ok(shared_secret_from_dh(alg, secret)),
         Err(_) => HpkeBytesResult::Err(HpkeError::ValidationError),
     }
@@ -444,7 +444,7 @@ pub fn DeriveKeyPair(alg: KEM, ikm: &InputKeyMaterial) -> Result<KeyPair, HpkeEr
                 )?;
                 bytes[0] = bytes[0] & bitmask;
                 // This check ensure sk != 0 or sk < order
-                if crate::ecdh::validate_scalar(named_group.into(), &bytes).is_ok() {
+                if crate::ecdh::validate_scalar(named_group.try_into().unwrap(), &bytes).is_ok() {
                     sk = bytes;
                 }
             }

src/jasmin.rs

@@ -1,9 +1,9 @@
 pub mod chacha20;
+pub mod kyber_derand;
 pub mod poly1305;
 pub mod sha2;
 pub mod sha3;
 pub mod x25519;
-pub mod kyber_derand;
 
 #[cfg(test)]
 mod testing {

src/kem.rs

@@ -6,35 +6,50 @@ use rand::{CryptoRng, Rng};
 
 use crate::ecdh;
 
+mod kyber768;
+
+// TODO: These functions are currently exposed simply in order to make NIST KAT
+// testing possible without an implementation of the NIST AES-CTR DRBG. Remove them
+// (and change the visibility of the exported functions to pub(crate)) the
+// moment we have an implementation of one. This is tracked by:
+// https://github.com/cryspen/libcrux/issues/36
+pub use kyber768::generate_keypair as kyber768_generate_keypair_derand;
+pub use kyber768::encapsulate as kyber768_encapsulate_derand;
+pub use kyber768::decapsulate as kyber768_decapsulate_derand;
+
 /// KEM Algorithms
 ///
 /// This includes named elliptic curves or dedicated KEM algorithms like Kyber.
-///
-/// Currently only `Secp256r1` and `X25519` are supported.
 #[derive(Clone, Copy, PartialEq, Debug)]
 pub enum Algorithm {
     X25519,
     X448,
     Secp256r1,
     Secp384r1,
     Secp521r1,
+    Kyber768,
 }
 
 #[derive(Debug, PartialEq, Eq)]
 pub enum Error {
     EcDhError(ecdh::Error),
     KeyGen,
+    Encapsulate,
+    Decapsulate,
     UnsupportedAlgorithm,
 }
 
-impl From<Algorithm> for ecdh::Algorithm {
-    fn from(value: Algorithm) -> Self {
+impl TryFrom<Algorithm> for ecdh::Algorithm {
+    type Error = &'static str;
+
+    fn try_from(value: Algorithm) -> Result<Self, Self::Error> {
         match value {
-            Algorithm::X25519 => ecdh::Algorithm::X25519,
-            Algorithm::X448 => ecdh::Algorithm::X448,
-            Algorithm::Secp256r1 => ecdh::Algorithm::P256,
-            Algorithm::Secp384r1 => ecdh::Algorithm::P384,
-            Algorithm::Secp521r1 => ecdh::Algorithm::P521,
+            Algorithm::X25519 => Ok(ecdh::Algorithm::X25519),
+            Algorithm::X448 => Ok(ecdh::Algorithm::X448),
+            Algorithm::Secp256r1 => Ok(ecdh::Algorithm::P256),
+            Algorithm::Secp384r1 => Ok(ecdh::Algorithm::P384),
+            Algorithm::Secp521r1 => Ok(ecdh::Algorithm::P521),
+            _ => Err("provided algorithm is not an ECDH algorithm"),
         }
     }
 }
@@ -58,10 +73,11 @@ pub enum PublicKey {
 }
 
 /// Compute the public key for a private key of the given [`Algorithm`].
+/// Applicable only to X25519 and secp256r1.
 pub fn secret_to_public(alg: Algorithm, sk: impl AsRef<[u8]>) -> Result<Vec<u8>, Error> {
     match alg {
         Algorithm::X25519 | Algorithm::Secp256r1 => {
-            ecdh::secret_to_public(alg.into(), sk.as_ref()).map_err(|e| e.into())
+            ecdh::secret_to_public(alg.try_into().unwrap(), sk.as_ref()).map_err(|e| e.into())
         }
         _ => Err(Error::UnsupportedAlgorithm),
     }
@@ -78,7 +94,18 @@ pub fn key_gen(
 ) -> Result<(Vec<u8>, Vec<u8>), Error> {
     match alg {
         Algorithm::X25519 | Algorithm::Secp256r1 => {
-            ecdh::key_gen(alg.into(), rng).map_err(|e| e.into())
+            ecdh::key_gen(alg.try_into().unwrap(), rng).map_err(|e| e.into())
+        }
+
+        Algorithm::Kyber768 => {
+            let mut seed = [0; kyber768::KEY_GENERATION_SEED_SIZE];
+            rng.try_fill_bytes(&mut seed).map_err(|_| Error::KeyGen)?;
+
+            if let Ok((pk, sk)) = kyber768::generate_keypair(seed) {
+                Ok((sk.into(), pk.into()))
+            } else {
+                Err(Error::KeyGen)
+            }
         }
         _ => Err(Error::UnsupportedAlgorithm),
     }
@@ -93,9 +120,23 @@ pub fn encapsulate(
     let (new_sk, new_pk) = key_gen(alg, rng)?;
     match alg {
         Algorithm::X25519 | Algorithm::Secp256r1 => {
-            let gxy = ecdh::derive(alg.into(), pk, &new_sk)?;
+            let gxy = ecdh::derive(alg.try_into().unwrap(), pk, &new_sk)?;
             Ok((gxy, new_pk))
         }
+
+        Algorithm::Kyber768 => {
+            let mut seed = [0; kyber768::SHARED_SECRET_SIZE];
+            rng.try_fill_bytes(&mut seed).map_err(|_| Error::KeyGen)?;
+
+            // TODO: Don't unwrap() here. See
+            // https://github.com/cryspen/libcrux/issues/35
+            if let Ok((ct, ss)) = kyber768::encapsulate(pk.try_into().unwrap(), seed) {
+                Ok((ss.into(), ct.into()))
+            } else {
+                Err(Error::Encapsulate)
+            }
+        }
+
         _ => Err(Error::UnsupportedAlgorithm),
     }
 }
@@ -104,9 +145,16 @@ pub fn encapsulate(
 pub fn decapsulate(alg: Algorithm, ct: &[u8], sk: &[u8]) -> Result<Vec<u8>, Error> {
     match alg {
         Algorithm::X25519 | Algorithm::Secp256r1 => {
-            let gxy = ecdh::derive(alg.into(), ct, sk)?;
+            let gxy = ecdh::derive(alg.try_into().unwrap(), ct, sk)?;
             Ok(gxy)
         }
+        Algorithm::Kyber768 => {
+            // TODO: Don't unwrap() here. See
+            // https://github.com/cryspen/libcrux/issues/35
+            let ss = kyber768::decapsulate(sk.try_into().unwrap(), ct.try_into().unwrap());
+
+            Ok(ss.into())
+        }
         _ => Err(Error::UnsupportedAlgorithm),
     }
 }