Skip to content

Commit

Permalink
fix mlkem
Browse files Browse the repository at this point in the history
  • Loading branch information
franziskuskiefer committed Jan 28, 2025
1 parent 2359a08 commit 13ea297
Show file tree
Hide file tree
Showing 7 changed files with 332 additions and 26 deletions.
1 change: 1 addition & 0 deletions libcrux-ml-kem/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ rand = { version = "0.8", optional = true }
libcrux-platform = { version = "0.0.2-beta.2", path = "../sys/platform" }
libcrux-sha3 = { version = "0.0.2-beta.2", path = "../libcrux-sha3" }
libcrux-intrinsics = { version = "0.0.2-beta.2", path = "../libcrux-intrinsics" }
libcrux-secrets = { version = "0.0.2-beta.2", path = "../secrets" }
hax-lib = { version = "0.1.0", git = "https://github.com/hacspec/hax/" }

[features]
Expand Down
88 changes: 70 additions & 18 deletions libcrux-ml-kem/src/hash_functions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@
// them to be properly abstracted in F*. We would like hax to do this automatically.
// Related Issue: https://github.com/hacspec/hax/issues/616

use libcrux_secrets::{AsSecret, AsSecretRef};

use crate::constants::{G_DIGEST_SIZE, H_DIGEST_SIZE};

/// The SHA3 block size.
Expand Down Expand Up @@ -89,7 +91,7 @@ pub(crate) mod portable {
#[inline(always)]
fn G(input: &[u8]) -> [u8; G_DIGEST_SIZE] {
let mut digest = [0u8; G_DIGEST_SIZE];
portable::sha512(&mut digest, input);
portable::sha512(&mut digest, input.as_secret());
digest
}

Expand All @@ -99,7 +101,7 @@ pub(crate) mod portable {
#[inline(always)]
fn H(input: &[u8]) -> [u8; H_DIGEST_SIZE] {
let mut digest = [0u8; H_DIGEST_SIZE];
portable::sha256(&mut digest, input);
portable::sha256(&mut digest, input.as_secret());
digest
}

Expand All @@ -110,7 +112,7 @@ pub(crate) mod portable {
#[inline(always)]
fn PRF<const LEN: usize>(input: &[u8]) -> [u8; LEN] {
let mut digest = [0u8; LEN];
portable::shake256(&mut digest, input);
portable::shake256(&mut digest, input.as_secret());
digest
}

Expand All @@ -124,7 +126,7 @@ pub(crate) mod portable {

let mut out = [[0u8; LEN]; K];
for i in 0..K {
portable::shake256(&mut out[i], &input[i]);
portable::shake256(&mut out[i], (&input[i]).as_secret());
}
out
}
Expand All @@ -135,7 +137,7 @@ pub(crate) mod portable {

let mut shake128_state = [incremental::shake128_init(); K];
for i in 0..K {
incremental::shake128_absorb_final(&mut shake128_state[i], &input[i]);
incremental::shake128_absorb_final(&mut shake128_state[i], (&input[i]).as_secret());
}
PortableHash { shake128_state }
}
Expand Down Expand Up @@ -498,7 +500,7 @@ pub(crate) mod neon {
#[inline(always)]
fn G(input: &[u8]) -> [u8; G_DIGEST_SIZE] {
let mut digest = [0u8; G_DIGEST_SIZE];
libcrux_sha3::neon::sha512(&mut digest, input);
libcrux_sha3::neon::sha512(&mut digest, input.as_secret());
digest
}

Expand All @@ -508,7 +510,7 @@ pub(crate) mod neon {
#[inline(always)]
fn H(input: &[u8]) -> [u8; H_DIGEST_SIZE] {
let mut digest = [0u8; H_DIGEST_SIZE];
libcrux_sha3::neon::sha256(&mut digest, input);
libcrux_sha3::neon::sha256(&mut digest, input.as_secret());
digest
}

Expand All @@ -520,7 +522,12 @@ pub(crate) mod neon {
fn PRF<const LEN: usize>(input: &[u8]) -> [u8; LEN] {
let mut digest = [0u8; LEN];
let mut dummy = [0u8; LEN];
x2::shake256(input, input, &mut digest, &mut dummy);
x2::shake256(
input.as_secret(),
input.as_secret(),
&mut digest,
&mut dummy,
);
digest
}

Expand All @@ -538,20 +545,45 @@ pub(crate) mod neon {
let mut out3 = [0u8; LEN];
match K as u8 {
2 => {
x2::shake256(&input[0], &input[1], &mut out0, &mut out1);
x2::shake256(
(&input[0]).as_secret(),
(&input[1]).as_secret(),
&mut out0,
&mut out1,
);
out[0] = out0;
out[1] = out1;
}
3 => {
x2::shake256(&input[0], &input[1], &mut out0, &mut out1);
x2::shake256(&input[2], &input[2], &mut out2, &mut out3);
x2::shake256(
(&input[0]).as_secret(),
(&input[1]).as_secret(),
&mut out0,
&mut out1,
);
x2::shake256(
(&input[2]).as_secret(),
(&input[2]).as_secret(),
&mut out2,
&mut out3,
);
out[0] = out0;
out[1] = out1;
out[2] = out2;
}
4 => {
x2::shake256(&input[0], &input[1], &mut out0, &mut out1);
x2::shake256(&input[2], &input[3], &mut out2, &mut out3);
x2::shake256(
(&input[0]).as_secret(),
(&input[1]).as_secret(),
&mut out0,
&mut out1,
);
x2::shake256(
(&input[2]).as_secret(),
(&input[3]).as_secret(),
&mut out2,
&mut out3,
);
out[0] = out0;
out[1] = out1;
out[2] = out2;
Expand All @@ -568,15 +600,35 @@ pub(crate) mod neon {
let mut state = [x2::incremental::init(), x2::incremental::init()];
match K as u8 {
2 => {
x2::incremental::shake128_absorb_final(&mut state[0], &input[0], &input[1]);
x2::incremental::shake128_absorb_final(
&mut state[0],
(&input[0]).as_secret(),
(&input[1]).as_secret(),
);
}
3 => {
x2::incremental::shake128_absorb_final(&mut state[0], &input[0], &input[1]);
x2::incremental::shake128_absorb_final(&mut state[1], &input[2], &input[2]);
x2::incremental::shake128_absorb_final(
&mut state[0],
(&input[0]).as_secret(),
(&input[1]).as_secret(),
);
x2::incremental::shake128_absorb_final(
&mut state[1],
(&input[2]).as_secret(),
(&input[2]).as_secret(),
);
}
4 => {
x2::incremental::shake128_absorb_final(&mut state[0], &input[0], &input[1]);
x2::incremental::shake128_absorb_final(&mut state[1], &input[2], &input[3]);
x2::incremental::shake128_absorb_final(
&mut state[0],
(&input[0]).as_secret(),
(&input[1]).as_secret(),
);
x2::incremental::shake128_absorb_final(
&mut state[1],
(&input[2]).as_secret(),
(&input[3]).as_secret(),
);
}
_ => unreachable!("This function can only called be called with N = 2, 3, 4"),
}
Expand Down
11 changes: 6 additions & 5 deletions libcrux-ml-kem/tests/nistkats.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
use libcrux_secrets::{AsSecret, AsSecretRef};
use serde::Deserialize;
use serde_json;
use std::{fs::File, io::BufReader, path::Path};
Expand Down Expand Up @@ -52,23 +53,23 @@ macro_rules! impl_nist_known_answer_tests {
let pk = unpacked::key_pair_serialized_public_key(&unpacked_keys);
let sk = unpacked::key_pair_serialized_private_key(&unpacked_keys);

let public_key_hash = sha256(pk.as_slice());
let secret_key_hash = sha256(sk.as_slice());
let public_key_hash = sha256(pk.as_slice().as_secret());
let secret_key_hash = sha256(sk.as_slice().as_secret());

assert_eq!(public_key_hash, kat.sha3_256_hash_of_public_key, "lhs: computed public key hash, rhs: hash from kat");
assert_eq!(secret_key_hash, kat.sha3_256_hash_of_secret_key, "lhs: computed secret key hash, rhs: hash from kat");
}

let public_key_hash = sha256(key_pair.pk());
let public_key_hash = sha256(key_pair.pk().as_secret());
eprintln!("pk hash: {}", hex::encode(public_key_hash));
let secret_key_hash = sha256(key_pair.sk());
let secret_key_hash = sha256(key_pair.sk().as_secret());

assert_eq!(public_key_hash, kat.sha3_256_hash_of_public_key, "lhs: computed public key hash, rhs: hash from kat");
assert_eq!(secret_key_hash, kat.sha3_256_hash_of_secret_key, "lhs: computed secret key hash, rhs: hash from kat");

let (ciphertext, shared_secret) =
encapsulate(key_pair.public_key(), kat.encapsulation_seed);
let ciphertext_hash = sha256(ciphertext.as_ref());
let ciphertext_hash = sha256(ciphertext.as_ref().as_secret());

assert_eq!(ciphertext_hash, kat.sha3_256_hash_of_ciphertext, "lhs: computed ciphertext hash, rhs: hash from akt");
assert_eq!(shared_secret.as_ref(), kat.shared_secret, "lhs: computed shared secret from encapsulate, rhs: shared secret from kat");
Expand Down
3 changes: 2 additions & 1 deletion libcrux-ml-kem/tests/self.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
use libcrux_ml_kem::{MlKemCiphertext, MlKemPrivateKey};

use libcrux_secrets::AsSecret;
use libcrux_sha3::shake256;
use rand::{rngs::OsRng, thread_rng, RngCore};

Expand Down Expand Up @@ -212,7 +213,7 @@ fn compute_implicit_rejection_shared_secret<const CLEN: usize, const LEN: usize>
let mut to_hash = secret_key[MlKemPrivateKey::<LEN>::len() - SHARED_SECRET_SIZE..].to_vec();
to_hash.extend_from_slice(ciphertext.as_ref());

shake256(&to_hash)
shake256((&to_hash).as_secret())
}

macro_rules! impl_modified_secret_key {
Expand Down
Loading

0 comments on commit 13ea297

Please sign in to comment.