Expose the SHAKE-128 incremental API in libcrux and add 100,000 Kyber KAT tests using this API

.github/workflows/rust.yml

@@ -109,10 +109,15 @@ jobs:
       # Test ...
 
       - name: 🏃🏻‍♀️ Test
-        run: cargo test --verbose $RUST_TARGET_FLAG
+        run: |
+          cargo test --verbose $RUST_TARGET_FLAG -- --skip shake128_rng
+          cargo test --features slow_kat_tests --verbose $RUST_TARGET_FLAG 100_000_known_answer_tests
+
 
       - name: 🏃🏻‍♀️ Test Release
-        run: cargo test --verbose --release $RUST_TARGET_FLAG
+        run: |
+          cargo test --verbose --release $RUST_TARGET_FLAG -- --skip shake128_rng
+          cargo test --features slow_kat_tests --verbose $RUST_TARGET_FLAG 100_000_known_answer_tests
 
       - name: 🔨 Build sys/hacl
         working-directory: sys/hacl

.gitignore

@@ -10,4 +10,5 @@ fuzz/corpus
 fuzz/artifacts
 proofs/fstar/extraction/.cache
 __pycache__
+tests/kyber_kats/__pycache__
 kyber-crate/

Cargo.toml

@@ -74,3 +74,4 @@ rand = []
 wasm = ["wasm-bindgen"]
 log = ["dep:log"]
 tests = []              # Expose functions for testing.
+slow_kat_tests = []     

src/digest.rs

@@ -24,6 +24,8 @@ use crate::hacl::{
 
 use libcrux_platform::{simd128_support, simd256_support};
 
+pub use sha3::incremental;
+
 #[derive(Debug)]
 pub enum Error {
     InvalidStateFinished,

src/hacl/sha3.rs

@@ -94,6 +94,103 @@ pub(crate) fn shake256<const BYTES: usize>(data: &[u8]) -> [u8; BYTES] {
     out
 }
 
+/// This module groups together functions that can be used to absorb or squeeze
+/// bytes in increments.
+/// TODO: This module should not be public, see: https://github.com/cryspen/libcrux/issues/157
+pub mod incremental {
+    use libcrux_hacl::{
+        Hacl_Hash_SHA3_Scalar_shake128_absorb, Hacl_Hash_SHA3_Scalar_shake128_squeeze_nblocks,
+        Hacl_Hash_SHA3_Scalar_state_free, Hacl_Hash_SHA3_Scalar_state_malloc,
+    };
+
+    /// SHAKE 128
+    ///
+    /// The caller must only call absorb() once, but may call squeeze_nblocks()
+    /// as many times as it wishes, provided the output size passed to
+    /// squeeze_nblocks() is multiple of the SHAKE128 block size (168 bytes).
+    pub struct AbsorbOnceSqueezeManyShake128 {
+        state: *mut u64,
+    }
+
+    impl AbsorbOnceSqueezeManyShake128 {
+        pub fn new() -> Self {
+            let state = Self {
+                state: unsafe { Hacl_Hash_SHA3_Scalar_state_malloc() },
+            };
+
+            state
+        }
+
+        pub fn absorb(&mut self, input: &[u8]) {
+            unsafe {
+                Hacl_Hash_SHA3_Scalar_shake128_absorb(
+                    self.state,
+                    input.as_ptr() as _,
+                    input.len() as u32,
+                )
+            };
+        }
+
+        pub fn squeeze_nblocks<const OUTPUT_BYTES: usize>(&mut self) -> [u8; OUTPUT_BYTES] {
+            debug_assert!(OUTPUT_BYTES % 168 == 0);
+            let mut output = [0u8; OUTPUT_BYTES];
+            unsafe {
+                Hacl_Hash_SHA3_Scalar_shake128_squeeze_nblocks(
+                    self.state,
+                    output.as_mut_ptr(),
+                    OUTPUT_BYTES as u32,
+                )
+            };
+
+            output
+        }
+    }
+    impl Drop for AbsorbOnceSqueezeManyShake128 {
+        fn drop(&mut self) {
+            unsafe { Hacl_Hash_SHA3_Scalar_state_free(self.state) }
+        }
+    }
+
+    use libcrux_hacl::{
+        Hacl_Hash_SHA3_free, Hacl_Hash_SHA3_malloc, Hacl_Hash_SHA3_squeeze,
+        Hacl_Hash_SHA3_state_t_s, Hacl_Hash_SHA3_update, Spec_Hash_Definitions_Shake128,
+    };
+
+    /// SHAKE 128
+    ///
+    /// The caller may call absorb() as many times as it wishes, but must only
+    /// call squeeze() once.
+    pub struct AbsorbManySqueezeOnceShake128 {
+        state: *mut Hacl_Hash_SHA3_state_t_s,
+    }
+
+    impl AbsorbManySqueezeOnceShake128 {
+        pub fn new() -> Self {
+            Self {
+                state: unsafe {
+                    Hacl_Hash_SHA3_malloc(Spec_Hash_Definitions_Shake128.try_into().unwrap())
+                },
+            }
+        }
+
+        pub fn absorb(&mut self, input: &[u8]) {
+            unsafe { Hacl_Hash_SHA3_update(self.state, input.as_ptr() as _, input.len() as u32) };
+        }
+
+        pub fn squeeze<const OUTPUT_BYTES: usize>(&mut self) -> [u8; OUTPUT_BYTES] {
+            let mut output = [0u8; OUTPUT_BYTES];
+            unsafe { Hacl_Hash_SHA3_squeeze(self.state, output.as_mut_ptr(), OUTPUT_BYTES as u32) };
+
+            output
+        }
+    }
+    impl Drop for AbsorbManySqueezeOnceShake128 {
+        fn drop(&mut self) {
+            unsafe { Hacl_Hash_SHA3_free(self.state) }
+        }
+    }
+}
+
 #[cfg(simd256)]
 pub mod x4 {
     use libcrux_hacl::{

src/kem/kyber/hash_functions.rs

@@ -1,4 +1,5 @@
 #![allow(non_snake_case)]
+#![allow(non_camel_case_types)]
 
 use libcrux_platform::simd256_support;
 
@@ -65,3 +66,40 @@ pub(crate) fn XOFx4<const K: usize>(
 
     out
 }
+
+// The following API uses the repeated squeeze API
+// Currently it only supports Scalar SHAKE, adapting it to SIMD SHAKE is a todo
+type XOF_state = crate::digest::incremental::AbsorbOnceSqueezeManyShake128;
+
+#[inline(always)]
+pub(crate) fn XOF_absorb<const K: usize>(input: [[u8; 34]; K]) -> [XOF_state; K] {
+    let mut out =
+        core::array::from_fn(|_| crate::digest::incremental::AbsorbOnceSqueezeManyShake128::new());
+
+    for i in 0..K {
+        out[i].absorb(&input[i]);
+    }
+    out
+}
+
+#[inline(always)]
+pub(crate) fn XOF_squeeze_three_blocks<const K: usize>(
+    state: &mut [XOF_state; K],
+) -> [[u8; 168 * 3]; K] {
+    let mut out = [[0; 168 * 3]; K];
+
+    for i in 0..K {
+        out[i] = state[i].squeeze_nblocks();
+    }
+    out
+}
+
+#[inline(always)]
+pub(crate) fn XOF_squeeze_block<const K: usize>(state: &mut [XOF_state; K]) -> [[u8; 168]; K] {
+    let mut out = [[0; 168]; K];
+
+    for i in 0..K {
+        out[i] = state[i].squeeze_nblocks();
+    }
+    out
+}

src/kem/kyber/matrix.rs

@@ -6,7 +6,7 @@ use super::{
     constants::COEFFICIENTS_IN_RING_ELEMENT,
     hash_functions::XOFx4,
     ntt::{invert_ntt_montgomery, ntt_multiply},
-    sampling::sample_from_uniform_distribution,
+    sampling::{sample_from_uniform_distribution, sample_from_xof},
 };
 use crate::cloop;
 
@@ -24,16 +24,15 @@ pub(in crate::kem::kyber) fn sample_matrix_A<const K: usize>(
             seeds[j][32] = i as u8;
             seeds[j][33] = j as u8;
         }
-        let xof_bytes = XOFx4::<K>(seeds);
 
-        for j in 0..K {
-            let sampled = sample_from_uniform_distribution(xof_bytes[j]);
+        let sampled = sample_from_xof(seeds);
 
+        for j in 0..K {
             // A[i][j] = A_transpose[j][i]
             if transpose {
-                A_transpose[j][i] = sampled;
+                A_transpose[j][i] = sampled[j];
             } else {
-                A_transpose[i][j] = sampled;
+                A_transpose[i][j] = sampled[j];
             }
         }
     }

src/kem/kyber/sampling.rs

@@ -2,7 +2,65 @@ use super::{
     arithmetic::{FieldElement, PolynomialRingElement},
     constants::{COEFFICIENTS_IN_RING_ELEMENT, FIELD_MODULUS, REJECTION_SAMPLING_SEED_SIZE},
 };
-use crate::{cloop, hax_utils::hax_debug_assert};
+use crate::{
+    cloop,
+    hax_utils::hax_debug_assert,
+    kem::kyber::hash_functions::{XOF_absorb, XOF_squeeze_block, XOF_squeeze_three_blocks},
+};
+
+pub fn sample_from_uniform_distribution_next<const K: usize, const N: usize>(
+    randomness: [[u8; N]; K],
+    sampled_coefficients: &mut [usize; K],
+    out: &mut [PolynomialRingElement; K],
+) -> bool {
+    let mut done = true;
+    for i in 0..K {
+        for bytes in randomness[i].chunks(3) {
+            let b1 = bytes[0] as i32;
+            let b2 = bytes[1] as i32;
+            let b3 = bytes[2] as i32;
+
+            let d1 = ((b2 & 0xF) << 8) | b1;
+            let d2 = (b3 << 4) | (b2 >> 4);
+
+            if d1 < FIELD_MODULUS && sampled_coefficients[i] < COEFFICIENTS_IN_RING_ELEMENT {
+                out[i].coefficients[sampled_coefficients[i]] = d1;
+                sampled_coefficients[i] += 1
+            }
+            if d2 < FIELD_MODULUS && sampled_coefficients[i] < COEFFICIENTS_IN_RING_ELEMENT {
+                out[i].coefficients[sampled_coefficients[i]] = d2;
+                sampled_coefficients[i] += 1;
+            }
+        }
+        if sampled_coefficients[i] < COEFFICIENTS_IN_RING_ELEMENT {
+            done = false
+        }
+    }
+    done
+}
+
+pub fn sample_from_xof<const K: usize>(seeds: [[u8; 34]; K]) -> [PolynomialRingElement; K] {
+    let mut sampled_coefficients: [usize; K] = [0; K];
+    let mut out: [PolynomialRingElement; K] = [PolynomialRingElement::ZERO; K];
+
+    let mut xof_states = XOF_absorb::<K>(seeds);
+    let randomness = XOF_squeeze_three_blocks(&mut xof_states);
+
+    let mut done = false;
+    done = sample_from_uniform_distribution_next(randomness, &mut sampled_coefficients, &mut out);
+
+    while !done {
+        let randomness = XOF_squeeze_block(&mut xof_states);
+        done =
+            sample_from_uniform_distribution_next(randomness, &mut sampled_coefficients, &mut out);
+    }
+
+    hax_debug_assert!(out[0]
+        .coefficients
+        .into_iter()
+        .all(|coefficient| coefficient >= 0 && coefficient < FIELD_MODULUS));
+    out
+}
 
 fn rejection_sampling_panic_with_diagnostic() {
     panic!()

sys/hacl/build.rs

@@ -109,9 +109,14 @@ fn create_bindings(platform: &Platform, home_dir: &Path) {
         // Allow function we want to have in
         .allowlist_function("Hacl_AEAD_Chacha20Poly1305_.*")
         .allowlist_function("Hacl_Curve25519.*")
+        .allowlist_function("Hacl_Hash_SHA1.*")
         .allowlist_function("Hacl_Hash_SHA2.*")
-        .allowlist_function("Hacl_Hash_.*")
-        .allowlist_function("Hacl_Blake2.*")
+        .allowlist_function("Hacl_Hash_SHA3.*")
+        .allowlist_function("Hacl_Hash_Blake2.*")
+        .allowlist_function("Hacl_Hash_SHA3_Scalar_state_malloc")
+        .allowlist_function("Hacl_Hash_SHA3_Scalar_state_free")
+        .allowlist_function("Hacl_Hash_SHA3_Scalar_shake128_absorb")
+        .allowlist_function("Hacl_Hash_SHA3_Scalar_shake128_squeeze_nblocks")
         .allowlist_function("Hacl_HMAC_DRBG.*")
         .allowlist_function("Hacl_Ed25519.*")
         .allowlist_function("Hacl_HKDF_.*")
@@ -124,7 +129,10 @@ fn create_bindings(platform: &Platform, home_dir: &Path) {
         .allowlist_type("Spec_.*")
         .allowlist_type("Hacl_HMAC_DRBG.*")
         .allowlist_type("Hacl_Streaming.*")
-        .allowlist_type("Hacl_Hash_.*")
+        .allowlist_type("Hacl_Hash_SHA1.*")
+        .allowlist_type("Hacl_Hash_SHA2.*")
+        .allowlist_type("Hacl_Hash_SHA3.*")
+        .allowlist_type("Hacl_Hash_Blake2.*")
         .allowlist_var("Spec_.*")
         .allowlist_var("Hacl_Streaming.*")
         // XXX: These functions use uint128 in the API, which is not FFI safe
@@ -133,6 +141,10 @@ fn create_bindings(platform: &Platform, home_dir: &Path) {
         .blocklist_function("Hacl_Hash_Blake2b_update_last")
         .blocklist_function("Hacl_Hash_Blake2b_Simd256_update_multi")
         .blocklist_function("Hacl_Hash_Blake2b_Simd256_update_last")
+        // The following functions aren't needed at the moment
+        .blocklist_function("Hacl_Hash_SHA3_Scalar_sha3.*")
+        .blocklist_function("Hacl_Hash_SHA3_Scalar_shake128")
+        .blocklist_function("Hacl_Hash_SHA3_Scalar_shake256")
         // Disable tests to avoid warnings and keep it portable
         .layout_tests(false)
         // Generate bindings

sys/hacl/c/config/hacl.h

@@ -19,5 +19,6 @@
 #include "Hacl_HMAC.h"
 #include "Hacl_P256.h"
 #include "Hacl_RSAPSS.h"
+#include "Hacl_Hash_SHA3_Scalar.h"
 
 void hacl_free(void *ptr);

sys/hacl/c/include/Hacl_Hash_SHA3_Scalar.h

@@ -37,27 +37,41 @@ extern "C" {
 
 void
 Hacl_Hash_SHA3_Scalar_shake128(
-  uint32_t inputByteLen,
-  uint8_t *input,
+  uint8_t *output,
   uint32_t outputByteLen,
-  uint8_t *output
+  uint8_t *input,
+  uint32_t inputByteLen
 );
 
 void
 Hacl_Hash_SHA3_Scalar_shake256(
-  uint32_t inputByteLen,
-  uint8_t *input,
+  uint8_t *output,
   uint32_t outputByteLen,
-  uint8_t *output
+  uint8_t *input,
+  uint32_t inputByteLen
 );
 
-void Hacl_Hash_SHA3_Scalar_sha3_224(uint32_t inputByteLen, uint8_t *input, uint8_t *output);
+void Hacl_Hash_SHA3_Scalar_sha3_224(uint8_t *output, uint8_t *input, uint32_t inputByteLen);
+
+void Hacl_Hash_SHA3_Scalar_sha3_256(uint8_t *output, uint8_t *input, uint32_t inputByteLen);
+
+void Hacl_Hash_SHA3_Scalar_sha3_384(uint8_t *output, uint8_t *input, uint32_t inputByteLen);
 
-void Hacl_Hash_SHA3_Scalar_sha3_256(uint32_t inputByteLen, uint8_t *input, uint8_t *output);
+void Hacl_Hash_SHA3_Scalar_sha3_512(uint8_t *output, uint8_t *input, uint32_t inputByteLen);
 
-void Hacl_Hash_SHA3_Scalar_sha3_384(uint32_t inputByteLen, uint8_t *input, uint8_t *output);
+uint64_t *Hacl_Hash_SHA3_Scalar_state_malloc(void);
 
-void Hacl_Hash_SHA3_Scalar_sha3_512(uint32_t inputByteLen, uint8_t *input, uint8_t *output);
+void Hacl_Hash_SHA3_Scalar_state_free(uint64_t *s);
+
+void
+Hacl_Hash_SHA3_Scalar_shake128_absorb(uint64_t *state, uint8_t *input, uint32_t inputByteLen);
+
+void
+Hacl_Hash_SHA3_Scalar_shake128_squeeze_nblocks(
+  uint64_t *state,
+  uint8_t *output,
+  uint32_t outputByteLen
+);
 
 #if defined(__cplusplus)
 }