krypteia-quantica 0.1.0

//! First-order arithmetic masking for ML-DSA polynomials.
//!
//! Same idea as the ML-KEM `masked` module but adapted to the
//! ML-DSA arithmetic (q = 8 380 417, polynomial coefficients held as
//! `i32`, NTT goes all the way down to length-1 components).
//!
//! Each secret polynomial is represented as two additive shares modulo
//! `q`: `s = (s₀ + s₁) mod q`. All operations on secret data manipulate
//! the shares independently, so a first-order side-channel attacker
//! observing one share at a time learns nothing about the unmasked
//! value.
//!
//! ## Sensitive operations protected
//!
//! In `dsa::sign_internal`, the secret polynomials `s1`, `s2`, `t0`
//! are NTT-transformed once before the rejection-sampling loop, then
//! multiplied by the per-iteration challenge polynomial `c` (which is
//! public — verifier recomputes it):
//!
//! ```text
//!   ŝ1, ŝ2, t̂0 ← NTT(s1), NTT(s2), NTT(t0)
//!   for each rejection iteration:
//!       ĉ ← NTT(c)
//!       cs1[i]  ← ĉ · ŝ1[i]      // secret × public
//!       cs2[i]  ← ĉ · ŝ2[i]      // secret × public
//!       ct0[i]  ← ĉ · t̂0[i]      // secret × public
//! ```
//!
//! The masked variants in this module replace `ŝ1`, `ŝ2`, `t̂0` with
//! `MaskedPoly` containers and provide a `pointwise_mul_public`
//! that multiplies each share independently. Because `ĉ` is public,
//! no secret×secret multiplication occurs and first-order masking is
//! sufficient.
//!
//! ## Available operations
//!
//! | Function                             | Description                                                |
//! |--------------------------------------|------------------------------------------------------------|
//! | `MaskedPoly::mask`                 | Split a plaintext polynomial into two shares               |
//! | `MaskedPoly::unmask`               | Reconstruct the polynomial from shares                     |
//! | `MaskedPoly::refresh`              | Re-randomize shares (prevents correlation buildup)         |
//! | `MaskedPoly::zeroize`              | DSE-resistant wipe of both shares                          |
//! | `masked_ntt`                       | Forward NTT applied to each share                          |
//! | `masked_ntt_inv`                   | Inverse NTT applied to each share                          |
//! | `masked_pointwise_mul_public`      | Masked × public pointwise mul (returns a `MaskedPoly`)   |
//!
//! ## Masked `y` pipeline (`sca-masked-y`)
//!
//! `MaskedPoly::expand_mask` samples `y` directly as two shares
//! drawn from SHAKE256. The shares propagate through
//! `masked_ntt` and `masked_mat_vec_mul` / `masked_mat_vec_mul_lazy`
//! so that the intermediate `w = A·y` stays in masked form until the
//! rejection loop commits to emitting it. This closes the DPA
//! recovery of `s1` from `z = y + c·s1` that exists on any unmasked
//! implementation. See *Side-channel analysis of masked y-sampling
//! in ML-DSA* (IACR ePrint 2025/276) and the countermeasure chapter
//! at `doc/sca/countermeasures/ml_dsa.rst`, section *DPA on `y` —
//! the `sca-masked-y` pipeline*.
//!
//! ## References
//!
//! * *Hardware masking of ML-DSA* (IACR ePrint 2024,
//!   `doc/papers/eprint2024_mldsa_hw_masking.pdf`) — reference
//!   construction, we follow the same share topology.
//! * *Side-channel analysis of masked y-sampling in ML-DSA*
//!   (IACR ePrint 2025/276) — basis for
//!   `MaskedPoly::expand_mask` + propagation through the linear
//!   stage.
//! * *Physical security considerations for ML-DSA* (NIST, 2025) —
//!   masking recommendation for high-assurance profiles.
//!
//! ## Where to look next
//!
//! * Countermeasure description and threat analysis:
//!   `doc/sca/countermeasures/ml_dsa.rst`, sections *DPA — first-
//!   order masking of secret polynomials* and *DPA on `y` — the
//!   `sca-masked-y` pipeline*.
//! * Call sites: [`crate::ml_dsa::dsa::sign_internal`] (look for
//!   `#[cfg(feature = "sca-protected")]` and
//!   `#[cfg(feature = "sca-masked-y")]` blocks).
//!
//! ## Scope and residual risk
//!
//! Masking here is **first-order**. The shipped Tier-1 item
//! `T1-A` (A3, refresh shares at the start of every rejection
//! iteration, head-of-loop refresh block in `dsa.rs`) raises the
//! effort required by a higher-order DPA that combines leakage
//! across iterations. Going beyond first-order (full higher-order
//! masking) is tracked as Tier-4 `T4-C`.

use super::MlDsaError;
use super::ntt::{self, mod_q};
use super::params::{N, Q};
use super::rng::CryptoRng;

/// A polynomial split into two additive shares modulo `q`.
///
/// Maintains the invariant `unmask()[i] = (share0[i] + share1[i]) mod q`
/// for all `i in 0..N`. Both shares are stored with coefficients in
/// `[0, q-1]`. Neither share alone reveals any information about the
/// underlying polynomial.
pub struct MaskedPoly {
    /// First additive share.
    pub share0: [i32; N],
    /// Second additive share.
    pub share1: [i32; N],
}

impl MaskedPoly {
    /// Build an all-zero `MaskedPoly`. Both shares are zero, so
    /// `unmask()` returns the zero polynomial. Useful as a stack
    /// initializer for fixed-size arrays of masked polynomials.
    pub const fn zero() -> Self {
        Self {
            share0: [0i32; N],
            share1: [0i32; N],
        }
    }

    /// Masked sampling of a masking vector polynomial from a SHAKE256
    /// stream — the DPA-safe replacement for `sample::expand_mask`.
    ///
    /// Implements ExpandMask (FIPS 204 Algorithm 34) but produces a
    /// two-share arithmetic representation `(share0, share1)` directly,
    /// without ever materializing the unmasked y coefficient in a stack
    /// or heap slot.
    ///
    /// ## Threat model
    ///
    /// Boolean-masked y is attackable with ~300 traces per the
    /// Hermelink-Ning-Petri result (ePrint 2025/276). Arithmetic
    /// masking is more robust but still requires careful
    /// implementation: the key invariant is that the unmasked
    /// coefficient value must only exist transiently in a CPU
    /// register, never be written to RAM.
    ///
    /// ## Implementation
    ///
    /// For each coefficient:
    ///   1. Decode the unmasked `y_i` from SHAKE256 output bytes
    ///      into a stack-local `let y_i: i32 = ...` (register-scoped).
    ///   2. Draw a fresh random mask `r_i` from the same SHAKE256
    ///      stream (a separate squeeze block).
    ///   3. Compute `share1_i = r_i mod q`, `share0_i = (y_i - r_i) mod q`.
    ///   4. Write both shares to the output `MaskedPoly`.
    ///   5. `y_i` and `r_i` go out of scope immediately.
    ///
    /// The two SHAKE256 streams (y bits and mask bits) are drawn from
    /// the same state: we first squeeze the packed-y bytes, then
    /// squeeze additional bytes for the mask. This keeps the function
    /// deterministic for a given `rho'' || nonce`, so the signature
    /// remains reproducible (ACVP-compatible).
    ///
    /// # Arguments
    ///
    /// * `rho_double_prime` — 64-byte seed (FIPS 204).
    /// * `nonce` — the per-polynomial nonce (`kappa + r`).
    /// * `gamma1` — the Γ₁ parameter for the current ML-DSA level.
    /// * `bitlen_gamma1_minus1` — bit length used by ExpandMask (17 or 19).
    pub fn sample_expand_mask(
        rho_double_prime: &[u8; 64],
        nonce: u16,
        gamma1: i32,
        bitlen_gamma1_minus1: usize,
    ) -> Self {
        use super::sha3;
        let c = bitlen_gamma1_minus1 + 1; // bits per coefficient
        let poly_bytes = 32 * c; // packed y bytes for 256 coeffs
        let mask_bytes = N * 4; // 1024 bytes of mask randomness

        let mut state = sha3::shake256();
        state.absorb(rho_double_prime);
        state.absorb(&nonce.to_le_bytes());

        // Two-phase squeeze: y packed bytes, then mask bytes.
        let mut y_buf = [0u8; 640]; // 32 * 20 = max poly_bytes
        state.squeeze(&mut y_buf[..poly_bytes]);
        let mut mask_buf = [0u8; N * 4];
        state.squeeze(&mut mask_buf);

        let mut mp = Self::zero();
        // Tight per-coefficient loop — the unmasked y_i only exists
        // in registers between the two writes to mp.share0 and
        // mp.share1 (the compiler optimizer is free to keep it in
        // registers; no intermediate array is used).
        let mut bit_pos = 0usize;
        let b = gamma1 as u32;
        let range = (gamma1 as u32 - 1) + gamma1 as u32; // a + b (bit_unpack semantics)
        let bits = 32 - range.leading_zeros() as usize;
        debug_assert_eq!(bits, c);

        for i in 0..N {
            // 1. Decode y_i (register-scoped, never on the stack as a poly).
            let mut val = 0u32;
            for bit in 0..bits {
                if (y_buf[bit_pos / 8] >> (bit_pos % 8)) & 1 == 1 {
                    val |= 1 << bit;
                }
                bit_pos += 1;
            }
            let y_i = b as i32 - val as i32;

            // 2. Fresh mask r_i from the 4-byte mask chunk.
            let r = u32::from_le_bytes([
                mask_buf[4 * i],
                mask_buf[4 * i + 1],
                mask_buf[4 * i + 2],
                mask_buf[4 * i + 3],
            ]);
            let r_i = (r % (Q as u32)) as i32;

            // 3. share1 = r, share0 = y - r   (mod q).
            mp.share1[i] = r_i;
            mp.share0[i] = mod_q(y_i - r_i);
            // y_i and r_i fall out of scope here.
        }
        zeroize_bytes(&mut y_buf);
        zeroize_bytes(&mut mask_buf);
        mp
    }

    /// Split a plaintext polynomial into two random additive shares.
    ///
    /// Generates a uniformly distributed `share1 ∈ [0, q-1]^N` from the
    /// RNG, then sets `share0 = poly - share1 mod q`. The intermediate
    /// random bytes are zeroized after use.
    ///
    /// # Errors
    ///
    /// Returns [`MlDsaError::RngFailure`] if the RNG fails.
    pub fn mask(poly: &[i32; N], rng: &mut dyn CryptoRng) -> Result<Self, MlDsaError> {
        let mut share0 = [0i32; N];
        let mut share1 = [0i32; N];
        // q < 2^23 so 4 bytes is enough; we keep rejection-bias-low by
        // taking each 32-bit word and reducing mod q.
        let mut rand_bytes = [0u8; N * 4];
        rng.fill_bytes(&mut rand_bytes)?;

        for i in 0..N {
            let r = u32::from_le_bytes([
                rand_bytes[4 * i],
                rand_bytes[4 * i + 1],
                rand_bytes[4 * i + 2],
                rand_bytes[4 * i + 3],
            ]);
            // Reduce a 32-bit word mod q. Bias is 2^32 mod q ≈ 2^9 over
            // 2^32, ~10^-7 — negligible for masking purposes (security
            // properties of masking only require the share to be
            // statistically close to uniform).
            share1[i] = (r % (Q as u32)) as i32;
        }
        zeroize_bytes(&mut rand_bytes);

        for i in 0..N {
            share0[i] = mod_q(poly[i] - share1[i]);
        }

        Ok(MaskedPoly { share0, share1 })
    }

    /// Reconstruct the plaintext polynomial from the two shares.
    ///
    /// Result coefficients are in `[0, q-1]`. The returned polynomial
    /// is unmasked secret data and should be zeroized after use.
    pub fn unmask(&self) -> [i32; N] {
        let mut out = [0i32; N];
        for i in 0..N {
            out[i] = mod_q(self.share0[i] + self.share1[i]);
        }
        out
    }

    /// Securely erase both shares via volatile writes.
    pub fn zeroize(&mut self) {
        zeroize_poly(&mut self.share0);
        zeroize_poly(&mut self.share1);
    }

    /// Re-randomize the shares without changing the unmasked value.
    ///
    /// Draws a fresh random polynomial `r` and updates the shares as
    /// `share0' = share0 - r mod q`, `share1' = share1 + r mod q`.
    /// The sum is preserved: `share0' + share1' ≡ share0 + share1`.
    /// Refreshing prevents higher-order correlation buildup when the
    /// same masked polynomial is reused across multiple operations.
    pub fn refresh(&mut self, rng: &mut dyn CryptoRng) -> Result<(), MlDsaError> {
        let mut rand_bytes = [0u8; N * 4];
        rng.fill_bytes(&mut rand_bytes)?;
        for i in 0..N {
            let r = (u32::from_le_bytes([
                rand_bytes[4 * i],
                rand_bytes[4 * i + 1],
                rand_bytes[4 * i + 2],
                rand_bytes[4 * i + 3],
            ]) % (Q as u32)) as i32;
            self.share0[i] = mod_q(self.share0[i] - r);
            self.share1[i] = mod_q(self.share1[i] + r);
        }
        zeroize_bytes(&mut rand_bytes);
        Ok(())
    }
}

// =====================================================================
// Linear masked operations
// =====================================================================
//
// The NTT and pointwise-multiplication-by-public-data are linear, so
// we can apply them to each share independently and the masking
// invariant is preserved:
//
//   NTT(s₀ + s₁) = NTT(s₀) + NTT(s₁)
//   p · (s₀ + s₁) = p·s₀ + p·s₁

/// Apply the forward NTT to each share independently.
pub fn masked_ntt(m: &mut MaskedPoly) {
    ntt::ntt(&mut m.share0);
    ntt::ntt(&mut m.share1);
}

/// Apply the inverse NTT to each share independently.
pub fn masked_ntt_inv(m: &mut MaskedPoly) {
    ntt::ntt_inv(&mut m.share0);
    ntt::ntt_inv(&mut m.share1);
}

/// Pointwise multiply a masked polynomial by a **public** polynomial
/// in NTT domain. Returns a fresh `MaskedPoly` holding the product.
///
/// `c_hat` (the per-iteration challenge polynomial in NTT form) is
/// public — the verifier recomputes it, and any side-channel
/// observation of it does not help an attacker recover the secret
/// shares.
pub fn masked_pointwise_mul_public(masked: &MaskedPoly, c_hat: &[i32; N]) -> MaskedPoly {
    MaskedPoly {
        share0: ntt::pointwise_mul(&masked.share0, c_hat),
        share1: ntt::pointwise_mul(&masked.share1, c_hat),
    }
}

/// Masked matrix-vector multiplication in the NTT domain: for each
/// output row `i`, compute `sum_j (A_hat[i][j] · y_hat_m[j])` as a
/// masked accumulator.
///
/// Since `A_hat` is public (derived from the public seed `rho`) and
/// the NTT is linear, each share is multiplied independently by the
/// same public matrix and accumulated into the corresponding output
/// share — no secret×secret operation occurs, so first-order shares
/// remain sufficient.
///
/// `y_hat_m` must have length ≥ `l` (the first `l` masked polynomials
/// are consumed). `out` must have length ≥ `k`.
pub fn masked_mat_vec_mul(
    a_hat: &[[[i32; N]; super::params::MAX_L]; super::params::MAX_K],
    y_hat_m: &[MaskedPoly],
    k: usize,
    l: usize,
    out: &mut [MaskedPoly],
) {
    for i in 0..k {
        out[i].share0 = [0i32; N];
        out[i].share1 = [0i32; N];
        for j in 0..l {
            let prod0 = ntt::pointwise_mul(&a_hat[i][j], &y_hat_m[j].share0);
            let prod1 = ntt::pointwise_mul(&a_hat[i][j], &y_hat_m[j].share1);
            for n in 0..N {
                out[i].share0[n] = mod_q(out[i].share0[n] + prod0[n]);
                out[i].share1[n] = mod_q(out[i].share1[n] + prod1[n]);
            }
        }
    }
}

/// Low-memory variant of `masked_mat_vec_mul`: recomputes each
/// `a_hat[i][j]` polynomial on-the-fly from the public seed `rho`
/// via SHAKE128 instead of holding the full k×l matrix in memory.
///
/// Trade-off identical to `mat_vec_mul_lazy` (dsa.rs): saves up to
/// 57 KB of stack for the `a_hat` matrix at the cost of repeated
/// SHAKE128 invocations. Called from `sign_internal` when both
/// `sca-masked-y` and `low-mem` are enabled.
pub fn masked_mat_vec_mul_lazy(rho: &[u8; 32], y_hat_m: &[MaskedPoly], k: usize, l: usize, out: &mut [MaskedPoly]) {
    use super::sample;
    for i in 0..k {
        out[i].share0 = [0i32; N];
        out[i].share1 = [0i32; N];
        for j in 0..l {
            let a_ij = sample::rej_ntt_poly(rho, j as u8, i as u8);
            let prod0 = ntt::pointwise_mul(&a_ij, &y_hat_m[j].share0);
            let prod1 = ntt::pointwise_mul(&a_ij, &y_hat_m[j].share1);
            for n in 0..N {
                out[i].share0[n] = mod_q(out[i].share0[n] + prod0[n]);
                out[i].share1[n] = mod_q(out[i].share1[n] + prod1[n]);
            }
        }
    }
}

// =====================================================================
// Local DSE-resistant zeroization helpers
// =====================================================================

/// Securely erase a polynomial buffer (write_volatile + compiler fence).
fn zeroize_poly(p: &mut [i32; N]) {
    for c in p.iter_mut() {
        unsafe { core::ptr::write_volatile(c, 0) };
    }
    core::sync::atomic::compiler_fence(core::sync::atomic::Ordering::SeqCst);
}

/// Securely erase a byte buffer (write_volatile + compiler fence).
fn zeroize_bytes(b: &mut [u8]) {
    for byte in b.iter_mut() {
        unsafe { core::ptr::write_volatile(byte, 0) };
    }
    core::sync::atomic::compiler_fence(core::sync::atomic::Ordering::SeqCst);
}

#[cfg(test)]
mod tests {
    use super::super::ntt;
    use super::*;

    /// Same deterministic test PRNG as in `shuffle.rs::tests`.
    struct TestRng(u64);
    impl CryptoRng for TestRng {
        fn fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), MlDsaError> {
            for chunk in dest.chunks_mut(8) {
                let mut x = self.0;
                x ^= x << 13;
                x ^= x >> 7;
                x ^= x << 17;
                self.0 = x;
                let bytes = x.to_le_bytes();
                for (i, b) in chunk.iter_mut().enumerate() {
                    *b = bytes[i];
                }
            }
            Ok(())
        }
    }

    fn fixture_poly() -> [i32; N] {
        let mut p = [0i32; N];
        for i in 0..N {
            p[i] = ((i as i32 * 12345 + 7).rem_euclid(Q)) as i32;
        }
        p
    }

    #[test]
    fn masked_mat_vec_mul_matches_unmasked() {
        use super::super::ntt as dsa_ntt;
        use super::super::params::{MAX_K, MAX_L, MlDsa65, Params as ParamsT};
        use super::super::sample;

        let k = MlDsa65::K;
        let l = MlDsa65::L;
        let rho = [0x17u8; 32];

        // Reference: unmasked A·y
        let a_hat = sample::expand_a::<MlDsa65>(&rho);
        let mut y = [[0i32; N]; MAX_L];
        for j in 0..l {
            for n in 0..N {
                y[j][n] = ((j as i32 * 100 + n as i32) % Q).abs();
            }
            dsa_ntt::ntt(&mut y[j]);
        }
        let mut w_ref = [[0i32; N]; MAX_K];
        for i in 0..k {
            for j in 0..l {
                let prod = dsa_ntt::pointwise_mul(&a_hat[i][j], &y[j]);
                for n in 0..N {
                    w_ref[i][n] = mod_q(w_ref[i][n] + prod[n]);
                }
            }
        }

        // Masked path: mask y, call masked_mat_vec_mul, unmask result
        let mut rng = TestRng(0xFEEDC0DEBADCAFEu64);
        let mut y_m: [MaskedPoly; MAX_L] = core::array::from_fn(|_| MaskedPoly::zero());
        for j in 0..l {
            y_m[j] = MaskedPoly::mask(&y[j], &mut rng).unwrap();
        }
        let mut w_m: [MaskedPoly; MAX_K] = core::array::from_fn(|_| MaskedPoly::zero());
        masked_mat_vec_mul(&a_hat, &y_m, k, l, &mut w_m);

        for i in 0..k {
            let w_got = w_m[i].unmask();
            assert_eq!(w_got, w_ref[i], "mismatch at row {}", i);
        }
    }

    #[test]
    fn masked_expand_mask_matches_unmasked_expand_mask() {
        use super::super::params::MlDsa65;
        use super::super::params::Params as _;
        use super::super::sample;

        let rho = [0x42u8; 64];
        let kappa = 0u16;
        let gamma1 = MlDsa65::GAMMA1;
        let bitlen = MlDsa65::BITLEN_GAMMA1_MINUS1;

        // Reference: unmasked ExpandMask for row 0
        let y_ref = sample::expand_mask::<MlDsa65>(&rho, kappa);

        // Masked path: two shares, unmask to compare
        let mp = MaskedPoly::sample_expand_mask(&rho, kappa, gamma1, bitlen);
        let unmasked = mp.unmask();

        // y_ref[0] coefficients are in (-gamma1, gamma1] (centered);
        // unmasked() returns in [0, q-1]. Normalize both to the centered
        // range for comparison.
        for i in 0..N {
            let r = y_ref[0][i];
            let u = {
                let mut v = unmasked[i];
                if v > Q / 2 {
                    v -= Q;
                }
                v
            };
            assert_eq!(r, u, "mismatch at i={}: ref={}, masked={}", i, r, u);
        }
    }

    #[test]
    fn mask_unmask_roundtrip() {
        let p = fixture_poly();
        let mut rng = TestRng(0xCAFEF00DDEADBEEF);
        let mp = MaskedPoly::mask(&p, &mut rng).unwrap();
        let recovered = mp.unmask();
        assert_eq!(recovered, p);
    }

    #[test]
    fn refresh_preserves_unmasked_value() {
        let p = fixture_poly();
        let mut rng = TestRng(0x1234567890ABCDEF);
        let mut mp = MaskedPoly::mask(&p, &mut rng).unwrap();
        mp.refresh(&mut rng).unwrap();
        mp.refresh(&mut rng).unwrap();
        assert_eq!(mp.unmask(), p);
    }

    #[test]
    fn masked_ntt_matches_regular_ntt() {
        let p = fixture_poly();
        let mut expected = p;
        ntt::ntt(&mut expected);

        let mut rng = TestRng(0x0123456789ABCDEF);
        let mut mp = MaskedPoly::mask(&p, &mut rng).unwrap();
        masked_ntt(&mut mp);
        // share0 + share1 mod q must equal NTT(p) coefficient-wise.
        for i in 0..N {
            let got = mod_q(mp.share0[i] + mp.share1[i]);
            assert_eq!(got, expected[i], "mismatch at i={}", i);
        }
    }

    #[test]
    fn masked_pointwise_mul_public_matches_unmasked() {
        let secret = fixture_poly();
        // Build a "public" polynomial — anything will do.
        let mut public = [0i32; N];
        for i in 0..N {
            public[i] = ((i as i32 * 991 + 13).rem_euclid(Q)) as i32;
        }
        // Bring both into NTT domain.
        let mut secret_ntt = secret;
        ntt::ntt(&mut secret_ntt);
        let mut public_ntt = public;
        ntt::ntt(&mut public_ntt);

        // Reference: regular pointwise_mul. Output is in /R Montgomery
        // domain and may be in [-q, q]; normalize for comparison.
        let mut expected = ntt::pointwise_mul(&secret_ntt, &public_ntt);
        for c in expected.iter_mut() {
            *c = mod_q(*c);
        }

        // Masked path: mask the (NTT-domain) secret, multiply, unmask.
        // `unmask()` already normalizes to [0, q-1].
        let mut rng = TestRng(0xFEEDFACE12345678);
        let mp_secret = MaskedPoly::mask(&secret_ntt, &mut rng).unwrap();
        let mp_product = masked_pointwise_mul_public(&mp_secret, &public_ntt);
        let got = mp_product.unmask();

        assert_eq!(got, expected);
    }
}