Merge pull request #56 from Nichokas/new-mlkem-lib

New mlkem lib

Author: Nichokas

Cargo.lock

@@ -23,12 +23,11 @@ chacha20poly1305 = { version = "0.10.1", features = ["std"] }
 hkdf = "0.12.4"
 serde = { version = "1.0.219", features = ["derive"] }
 sha2 = "0.10.9"
-kyberlib = "0.0.6"
 zerocopy = "0.8.25"
 bincode = { version = "2.0.1", features = ["serde"] }
 serde_bytes = "0.11.17"
-rand_chacha = "0.3.1"
-rand = "0.8.5"
+rand_chacha = "0.9.0"
+libcrux-ml-kem = "0.0.2"
 
 
 [dev-dependencies]

Cargo.toml

@@ -1,15 +1,17 @@
-// benches/chacha_bench.rs
+use chacha20poly1305::aead::rand_core::OsRng as ChaChaOsRng;
 use chacha20poly1305::aead::Aead;
 use chacha20poly1305::{aead::AeadCore, ChaCha20Poly1305, KeyInit};
 use criterion::{black_box, criterion_group, criterion_main, BatchSize, Criterion};
-use rand::RngCore;
+use rand_chacha::rand_core::{RngCore, SeedableRng};
+use rand_chacha::ChaCha20Rng;
 
 fn key_derivation_benchmark(c: &mut Criterion) {
     c.bench_function("chacha_key_derivation", |b| {
         b.iter_batched(
             || {
                 let mut key = [0u8; 32];
-                rand::thread_rng().fill_bytes(&mut key);
+                let mut rng = ChaCha20Rng::from_os_rng();
+                rng.fill_bytes(&mut key);
                 key
             },
             |key| {
@@ -28,13 +30,15 @@ fn encryption_benchmark(c: &mut Criterion) {
         group.bench_with_input(format!("encrypt_{}B", size), size, |b, &size| {
             b.iter_batched(
                 || {
-                    let key = ChaCha20Poly1305::generate_key(&mut rand::thread_rng());
+                    let mut rng = ChaChaOsRng;
+                    let key = ChaCha20Poly1305::generate_key(&mut rng);
                     let msg = vec![0u8; size];
                     (key, msg)
                 },
                 |(key, msg)| {
                     let cipher = ChaCha20Poly1305::new(&key);
-                    let nonce = ChaCha20Poly1305::generate_nonce(&mut rand::thread_rng());
+                    let mut rng = ChaChaOsRng;
+                    let nonce = ChaCha20Poly1305::generate_nonce(&mut rng);
                     cipher.encrypt(&nonce, &*msg).unwrap()
                 },
                 BatchSize::SmallInput,
@@ -50,9 +54,11 @@ fn decryption_benchmark(c: &mut Criterion) {
         group.bench_with_input(format!("decrypt_{}B", size), size, |b, &size| {
             b.iter_batched(
                 || {
-                    let key = ChaCha20Poly1305::generate_key(&mut rand::thread_rng());
+                    let mut rng = ChaChaOsRng;
+                    let key = ChaCha20Poly1305::generate_key(&mut rng);
                     let cipher = ChaCha20Poly1305::new(&key);
-                    let nonce = ChaCha20Poly1305::generate_nonce(&mut rand::thread_rng());
+                    let mut rng = ChaChaOsRng;
+                    let nonce = ChaCha20Poly1305::generate_nonce(&mut rng);
                     let msg = vec![0u8; size];
                     let ct = cipher.encrypt(&nonce, &*msg).unwrap();
                     (key, nonce, ct)
@@ -76,4 +82,4 @@ criterion_group! {
         decryption_benchmark
 }
 
-criterion_main!(benches);
+criterion_main!(benches);

benches/chacha_bench.rs

@@ -1,41 +1,47 @@
 // benches/kyber_bench.rs
 use criterion::{black_box, criterion_group, criterion_main, BatchSize, Criterion};
 use kychacha_crypto::{decrypt, encrypt, generate_keypair};
+use libcrux_ml_kem::mlkem768::{self};
+use rand_chacha::rand_core::{RngCore, SeedableRng};
+use rand_chacha::ChaCha20Rng;
 
 fn keygen_benchmark(c: &mut Criterion) {
-    c.bench_function("kyber_keypair_generation", |b| {
+    c.bench_function("mlkem_keypair_generation", |b| {
         b.iter(|| {
-            black_box(generate_keypair().unwrap());
+            black_box(generate_keypair());
         });
     });
 }
 
 fn encapsulation_benchmark(c: &mut Criterion) {
-    let server_kp = generate_keypair().unwrap();
+    let server_kp = generate_keypair();
 
-    c.bench_function("kyber_encapsulation", |b| {
+    c.bench_function("mlkem_encapsulation", |b| {
         b.iter(|| {
-            let (ct, _ss) =
-                kyberlib::encapsulate(&server_kp.public, &mut rand::thread_rng()).unwrap();
-            black_box(ct);
+            let mut rng = ChaCha20Rng::from_os_rng();
+            let mut randomness = [0u8; 32];
+            rng.fill_bytes(&mut randomness);
+            let (_ct, _ss) = black_box(mlkem768::encapsulate(&server_kp.public_key(), randomness));
         });
     });
 }
 
 fn decapsulation_benchmark(c: &mut Criterion) {
-    let server_kp = generate_keypair().unwrap();
-    let (ct, _) = kyberlib::encapsulate(&server_kp.public, &mut rand::thread_rng()).unwrap();
+    let server_kp = generate_keypair();
+    let mut rng = ChaCha20Rng::from_os_rng();
+    let mut randomness = [0u8; 32];
+    rng.fill_bytes(&mut randomness);
+    let (ct, _) = mlkem768::encapsulate(&server_kp.public_key(), randomness);
 
-    c.bench_function("kyber_decapsulation", |b| {
+    c.bench_function("mlkem_decapsulation", |b| {
         b.iter(|| {
-            let ss = kyberlib::decapsulate(&ct, &server_kp.secret).unwrap();
-            black_box(ss);
+            black_box(mlkem768::decapsulate(&server_kp.private_key(), &ct));
         });
     });
 }
 
 fn full_encryption_benchmark(c: &mut Criterion) {
-    let server_kp = generate_keypair().unwrap();
+    let server_kp = generate_keypair();
     let messages = vec![
         ("short", b"test".to_vec()),
         ("medium", vec![0u8; 1024]),
@@ -46,24 +52,24 @@ fn full_encryption_benchmark(c: &mut Criterion) {
         c.bench_function(&format!("full_encryption_{}", name), |b| {
             b.iter_batched(
                 || message.clone(),
-                |msg| black_box(encrypt(&server_kp.public, &msg).unwrap()),
+                |msg| black_box(encrypt(&server_kp.public_key(), &msg)),
                 BatchSize::SmallInput,
             )
         });
     }
 }
 
 fn full_decryption_benchmark(c: &mut Criterion) {
-    let server_kp = generate_keypair().unwrap();
+    let server_kp = generate_keypair();
     let messages = vec![
-        ("short", encrypt(&server_kp.public, b"test").unwrap()),
+        ("short", encrypt(&server_kp.public_key(), b"test").unwrap()),
         (
             "medium",
-            encrypt(&server_kp.public, &vec![0u8; 1024]).unwrap(),
+            encrypt(&server_kp.public_key(), &vec![0u8; 1024]).unwrap(),
         ),
         (
             "long",
-            encrypt(&server_kp.public, &vec![0u8; 4096]).unwrap(),
+            encrypt(&server_kp.public_key(), &vec![0u8; 4096]).unwrap(),
         ),
     ];
 
@@ -87,4 +93,4 @@ criterion_group! {
         full_decryption_benchmark
 }
 
-criterion_main!(benches);
+criterion_main!(benches);

benches/kyber_bench.rs

@@ -14,10 +14,10 @@ use chacha20poly1305::{
 ///
 /// // do not use this on your code, instead use encrypt fn
 /// let key = [0u8; 32];
-/// let (nonce, cifrado) = encrypt_with_key(&key, b"mensaje").unwrap();
+/// let (nonce, encrypted) = encrypt_with_key(&key, b"message").unwrap();
 /// ```
 ///
-/// # Errores
+/// # Errors
 /// - key ≠ 32 bytes
 /// - Error while encrypting
 pub fn encrypt_with_key(key: &[u8; 32], plaintext: &[u8]) -> Result<(Vec<u8>, Vec<u8>)> {
@@ -42,12 +42,12 @@ pub fn encrypt_with_key(key: &[u8; 32], plaintext: &[u8]) -> Result<(Vec<u8>, Ve
 ///
 /// let key = [0u8; 32];
 /// // do not use this on your code, instead use encrypt fn
-/// let (nonce, encrypted) = encrypt_with_key(&key, b"mensaje").unwrap();
+/// let (nonce, encrypted) = encrypt_with_key(&key, b"message").unwrap();
 /// // do not use this on your code, instead use decrypt fn
 /// let decrypted_text = decrypt_with_key(&key, &nonce, &encrypted).unwrap();
 /// ```
 ///
-/// # Errores
+/// # Errors
 /// - Key ≠ 32 bytes or nonce ≠ 12 bytes
 /// - Failed auth or corruption on the data
 pub fn decrypt_with_key(key: &[u8; 32], nonce: &[u8], ciphertext: &[u8]) -> Result<Vec<u8>> {

src/encryption.rs

@@ -1,46 +1,42 @@
 //! Kyber-1024 key exchange implementation (NIST PQC Round 3)
 
-use anyhow::{anyhow, Error};
 use hkdf::Hkdf;
-use kyberlib::{keypair, SharedSecret};
-use rand::SeedableRng;
 use rand_chacha::ChaCha20Rng;
 use sha2::Sha256;
 use zerocopy::IntoBytes;
-use kyberlib::{RngCore, CryptoRng};
-
-/// Kyber-768 key sizes
-pub const KYBER_PUBLIC_KEY_BYTES: usize = 1184;
-/// Kyber-768 key sizes
-pub const KYBER_SECRET_KEY_BYTES: usize = 2400;
+use libcrux_ml_kem::*;
+use libcrux_ml_kem::mlkem768::MlKem768KeyPair;
+use rand_chacha::rand_core::{RngCore, SeedableRng};
 
 /// Derives 256-bit ChaCha20 key from Kyber shared secret
 ///
 /// Uses HKDF-SHA256 with protocol-specific context
 ///
 /// # Security
 /// Context string prevents key reuse in different protocol components
-pub fn derive_chacha_key(shared_secret: &SharedSecret) -> [u8; 32] {
+pub fn derive_chacha_key(shared_secret: &MlKemSharedSecret) -> [u8; 32] {
     let hk = Hkdf::<Sha256>::new(None, shared_secret.as_bytes());
     let mut okm = [0u8; 32];
     hk.expand(b"chacha-encryption-v1", &mut okm)
         .expect("HKDF failed");
     okm
 }
 
-/// Generates CPA-secure Kyber-1024 keypair
+/// Generates ML-KEM keypair
 ///
 /// # Example
 /// ```
 /// # use std::error::Error;
 /// # fn main() -> Result<(), Box<dyn Error>> {
 /// use kychacha_crypto::generate_keypair;
 ///
-/// let keypair = generate_keypair()?;
+/// let keypair = generate_keypair();
 /// Ok(())
 /// # }
 /// ```
-pub fn generate_keypair() -> std::result::Result<kyberlib::Keypair, Error> {
-    let mut rng = ChaCha20Rng::from_entropy();
-    keypair(&mut rng).map_err(|e| anyhow!("Key generation failed: {}", e))
+pub fn generate_keypair() -> MlKem768KeyPair {
+    let mut rng = ChaCha20Rng::from_os_rng();
+    let mut randomness = [0u8; 64];
+    rng.fill_bytes(&mut randomness);
+    mlkem768::generate_key_pair(randomness)
 }