ML-DSA: Implemented sampling for matrix A and short vectors. (#290)

libcrux-ml-dsa/Cargo.toml

@@ -11,6 +11,7 @@ readme.workspace = true
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+libcrux-sha3 = { version = "0.0.2-pre.2", path = "../libcrux-sha3" }
 
 [dev-dependencies]
 hex = { version = "0.4.3", features = ["serde"] }

libcrux-ml-dsa/src/arithmetic.rs

@@ -0,0 +1,16 @@
+use crate::constants::COEFFICIENTS_IN_RING_ELEMENT;
+
+/// Values having this type hold a representative 'x' of the ML-DSA field.
+pub(crate) type FieldElement = i32;
+
+#[derive(Clone, Copy)]
+pub struct PolynomialRingElement {
+    pub(crate) coefficients: [FieldElement; COEFFICIENTS_IN_RING_ELEMENT],
+}
+
+impl PolynomialRingElement {
+    pub const ZERO: Self = Self {
+        // FIXME: hax issue, 256 is COEFFICIENTS_IN_RING_ELEMENT
+        coefficients: [0i32; 256],
+    };
+}

libcrux-ml-dsa/src/constants.rs

@@ -1,3 +1,7 @@
+pub(crate) const FIELD_MODULUS: i32 = 8_380_417;
+
+pub(crate) const COEFFICIENTS_IN_RING_ELEMENT: usize = 256;
+
 pub(crate) const FIELD_MODULUS_MINUS_ONE_BIT_LENGTH: usize = 23;
 
 pub(crate) const DROPPED_BITS_FROM_T: usize = 13;

libcrux-ml-dsa/src/hash_functions.rs

@@ -0,0 +1,36 @@
+#![allow(non_snake_case)]
+pub(crate) fn H<const OUTPUT_LENGTH: usize>(input: &[u8]) -> [u8; OUTPUT_LENGTH] {
+    let mut out = [0u8; OUTPUT_LENGTH];
+    libcrux_sha3::portable::shake256(&mut out, input);
+
+    out
+}
+
+pub(crate) mod H_128 {
+    use libcrux_sha3::portable::{incremental, KeccakState1};
+
+    const BLOCK_SIZE: usize = 168;
+    const FIVE_BLOCKS_SIZE: usize = BLOCK_SIZE * 5;
+
+    #[inline(always)]
+    pub(crate) fn new(seed: [u8; 34]) -> KeccakState1 {
+        let mut state = incremental::shake128_init();
+        incremental::shake128_absorb_final(&mut state, &seed);
+
+        state
+    }
+
+    pub(crate) fn squeeze_first_five_blocks(state: &mut KeccakState1) -> [u8; FIVE_BLOCKS_SIZE] {
+        let mut out = [0u8; FIVE_BLOCKS_SIZE];
+        incremental::shake128_squeeze_first_five_blocks(state, &mut out);
+
+        out
+    }
+
+    pub(crate) fn squeeze_next_block(state: &mut KeccakState1) -> [u8; BLOCK_SIZE] {
+        let mut out = [0u8; BLOCK_SIZE];
+        incremental::shake128_squeeze_next_block(state, &mut out);
+
+        out
+    }
+}

libcrux-ml-dsa/src/lib.rs

@@ -1,4 +1,10 @@
+mod arithmetic;
 mod constants;
+mod hash_functions;
+mod matrix;
+mod sample;
+mod utils;
+
 mod ml_dsa_generic;
 
 pub mod ml_dsa_65;

libcrux-ml-dsa/src/matrix.rs

@@ -0,0 +1,27 @@
+use crate::{arithmetic::PolynomialRingElement, sample::sample_ring_element_uniform};
+
+#[allow(non_snake_case)]
+#[inline(always)]
+pub(crate) fn expand_to_A<const ROWS_IN_A: usize, const COLUMNS_IN_A: usize>(
+    mut seed: [u8; 34],
+    transposed: bool,
+) -> [[PolynomialRingElement; COLUMNS_IN_A]; ROWS_IN_A] {
+    let mut A = [[PolynomialRingElement::ZERO; COLUMNS_IN_A]; ROWS_IN_A];
+
+    for i in 0..ROWS_IN_A {
+        for j in 0..COLUMNS_IN_A {
+            seed[32] = i as u8;
+            seed[33] = j as u8;
+
+            let sampled = sample_ring_element_uniform(seed);
+
+            if transposed {
+                A[j][i] = sampled;
+            } else {
+                A[i][j] = sampled;
+            }
+        }
+    }
+
+    A
+}

libcrux-ml-dsa/src/ml_dsa_65.rs

@@ -5,23 +5,27 @@ use crate::constants::*;
 const ROWS_IN_A: usize = 6;
 const COLUMNS_IN_A: usize = 5;
 
-const PUBLIC_KEY_SIZE: usize =
+const VERIFICATION_KEY_SIZE: usize =
     32 + (32 * ROWS_IN_A * (FIELD_MODULUS_MINUS_ONE_BIT_LENGTH - DROPPED_BITS_FROM_T));
-const SECRET_KEY_SIZE: usize =
+const SIGNING_KEY_SIZE: usize =
     (32 + 32 + 64) + 32 * (((ROWS_IN_A + COLUMNS_IN_A) * 4) + (DROPPED_BITS_FROM_T * ROWS_IN_A));
 
 pub struct MLDSA65KeyPair {
-    pub secret_key: [u8; SECRET_KEY_SIZE],
-    pub public_key: [u8; PUBLIC_KEY_SIZE],
+    pub signing_key: [u8; SIGNING_KEY_SIZE],
+    pub verification_key: [u8; VERIFICATION_KEY_SIZE],
 }
 
 /// Generate an ML-DSA-65 Key Pair
 pub fn generate_key_pair(randomness: [u8; 32]) -> MLDSA65KeyPair {
-    let (secret_key, public_key) =
-        crate::ml_dsa_generic::generate_key_pair::<SECRET_KEY_SIZE, PUBLIC_KEY_SIZE>(randomness);
+    let (signing_key, verification_key) = crate::ml_dsa_generic::generate_key_pair::<
+        ROWS_IN_A,
+        COLUMNS_IN_A,
+        SIGNING_KEY_SIZE,
+        VERIFICATION_KEY_SIZE,
+    >(randomness);
 
     MLDSA65KeyPair {
-        secret_key,
-        public_key,
+        signing_key,
+        verification_key,
     }
 }

libcrux-ml-dsa/src/ml_dsa_generic.rs

@@ -1,6 +1,19 @@
-pub(crate) fn generate_key_pair<const SECRET_KEY_SIZE: usize, const PUBLIC_KEY_SIZE: usize>(
+use crate::{hash_functions::H, matrix::expand_to_A, utils::into_padded_array};
+
+#[allow(non_snake_case)]
+pub(crate) fn generate_key_pair<
+    const ROWS_IN_A: usize,
+    const COLUMNS_IN_A: usize,
+    const SIGNING_KEY_SIZE: usize,
+    const VERIFICATION_KEY_SIZE: usize,
+>(
     randomness: [u8; 32],
-) -> ([u8; SECRET_KEY_SIZE], [u8; PUBLIC_KEY_SIZE]) {
-    let _ = randomness;
+) -> ([u8; SIGNING_KEY_SIZE], [u8; VERIFICATION_KEY_SIZE]) {
+    let seed_expanded = H::<1024>(&randomness);
+    let (seed_for_A, seed_expanded) = seed_expanded.split_at(32);
+    let (_seed_for_short_vectors, _random_seed_for_signing) = seed_expanded.split_at(64);
+
+    let _A_hat = expand_to_A::<ROWS_IN_A, COLUMNS_IN_A>(into_padded_array(seed_for_A), false);
+
     todo!();
 }