krypteia-quantica 0.2.0

// SPDX-License-Identifier: Apache-2.0
// Copyright 2026 Cédric Mesnil <cslashm@pm.me>

// Index-based loops in this module deliberately mirror the reference
// algorithm (parallel-array access `a[i]`/`b[i]`, or a position-indexed
// packing/reduction state) and keep the constant-time data layout explicit;
// the `iter().enumerate()` rewrite would obscure it. Documented, single-lint
// allow per CLAUDE.md §6 (scoped to this module, not a blanket suppression).
#![allow(clippy::needless_range_loop)]

//! Fisher-Yates shuffle for ML-DSA NTT butterfly index randomization
//! (**countermeasure: SPA / SEMA on secret-polynomial NTT**).
//!
//! Same principle as the ML-KEM `shuffle` module, ported to the
//! ML-DSA arithmetic (`q = 8 380 417`, `i32` coefficients,
//! 256-coefficient polynomials, NTT all the way down to length-1).
//!
//! ## Principle
//!
//! Within a given NTT level, all butterfly groups (and all butterflies
//! within a group) are independent — permuting their execution order
//! does not affect correctness, only the pattern of memory accesses
//! and instantaneous power consumption. By drawing a fresh permutation
//! per level and per group from an SCA-dedicated CSPRNG, two successive
//! NTTs on the same input produce different power / EM traces,
//! defeating trace-alignment-based SPA and template attacks. The
//! CSPRNG (`ScaRng`) is seeded with `K ‖ rnd ‖ tr ‖ M'` inside
//! `dsa::sign_internal`.
//!
//! Applied to `s1`, `s2`, `t0` (the three secret vectors). The public
//! matrix `A` keeps the classical NTT for performance — public values
//! need no shuffling.
//!
//! The primary entry point is `ntt_shuffled`, a drop-in replacement
//! for `super::ntt::ntt` that draws from a [`CryptoRng`].
//!
//! ## References
//!
//! * *Hardware NTT shuffling as a lightweight countermeasure for
//!   ML-KEM* (arXiv, 2024) — original shuffling analysis; the
//!   construction transfers directly to ML-DSA.
//! * *Slothy-assisted Cortex-M4/M7 implementations of ML-DSA*
//!   (IACR ePrint 2025) — performance measurements for the shuffled
//!   variant.
//! * *Physical security considerations for ML-DSA* (NIST, 2025) —
//!   recommended posture including shuffling as an SPA mitigation.
//!
//! ## Where to look next
//!
//! * Countermeasure description and threat analysis:
//!   `doc/sca/countermeasures/ml_dsa.rst`, section *SPA / SEMA —
//!   Fisher-Yates shuffled NTT*.
//! * Call site: `dsa::sign_internal`, Step 1 of
//!   the protected branch (three `for` loops applying
//!   `ntt_shuffled` to each polynomial of `s1_hat`, `s2_hat`,
//!   `t0_hat`).

use super::MlDsaError;
use super::ntt::{ZETAS, mod_q};
use super::params::N;
use super::rng::CryptoRng;
use alloc::vec;

/// Generate a uniform random permutation of `0..n` in-place using
/// Fisher-Yates with rejection sampling on 16-bit RNG output.
///
/// `perm` must be a slice of length `n`; on entry its contents are
/// overwritten with the identity permutation, then shuffled in place.
pub fn generate_permutation(perm: &mut [u16], rng: &mut dyn CryptoRng) -> Result<(), MlDsaError> {
    let n = perm.len();
    for i in 0..n {
        perm[i] = i as u16;
    }
    let mut rand_buf = [0u8; 2];
    for i in (1..n).rev() {
        let bound = (i + 1) as u16;
        // Rejection sampling to avoid modular bias.
        let limit = u16::MAX - (u16::MAX % bound);
        loop {
            rng.fill_bytes(&mut rand_buf)?;
            let r = u16::from_le_bytes(rand_buf);
            if r < limit {
                let j = (r % bound) as usize;
                perm.swap(i, j);
                break;
            }
        }
    }
    Ok(())
}

/// Forward NTT with randomized butterfly ordering (SPA countermeasure).
///
/// Functionally equivalent to `super::ntt::ntt` but draws fresh
/// random permutations from `rng` for both the inter-group and
/// intra-group butterfly orderings at each NTT level.
///
/// Uses the non-Montgomery [`ZETAS`] table together with the public
/// `mul_mod_q` helper, so the implementation is fully self-contained
/// — at the cost of being slightly slower than the in-place
/// Montgomery butterflies in `super::ntt::ntt`. Acceptable for the
/// SCA-protected build because the shuffled NTT only runs three times
/// per signature (on `s1`, `s2`, `t0`), once at the start of
/// `sign_internal` and never inside the rejection loop.
///
/// Output coefficients are in `[0, q-1]`, matching the regular NTT.
pub fn ntt_shuffled(f: &mut [i32; N], rng: &mut dyn CryptoRng) -> Result<(), MlDsaError> {
    let mut m = 0usize;
    let mut len = 128;
    while len >= 1 {
        let num_groups = N / (2 * len);

        // Shuffle the order of butterfly groups at this level.
        let mut group_perm = vec![0u16; num_groups];
        generate_permutation(&mut group_perm, rng)?;

        // Within each group, also shuffle the order of butterflies.
        // We must compute the per-group `m`/`zeta` index up front so
        // that re-ordering the groups doesn't desync the zeta lookup.
        for &gi in group_perm.iter() {
            let group_index = gi as usize;
            let start = group_index * 2 * len;
            // ZETAS table is indexed by the linear butterfly group
            // counter (m + 1 .. m + num_groups). The classical NTT
            // increments `m` once per group; here we project the
            // shuffled group index directly.
            let zeta = ZETAS[m + 1 + group_index];

            let mut inner_perm = vec![0u16; len];
            generate_permutation(&mut inner_perm, rng)?;

            for &ji in inner_perm.iter() {
                let j = start + ji as usize;
                // t = zeta * f[j+len]   (mod q)
                let t = mod_q(((zeta as i64 * f[j + len] as i64) % super::params::Q as i64) as i32);
                f[j + len] = mod_q(f[j] - t);
                f[j] = mod_q(f[j] + t);
            }
        }

        m += num_groups;
        len /= 2;
    }
    // Final reduction to [0, q-1] (most coefficients are already in
    // range thanks to the per-butterfly mod_q above, but we normalize
    // for safety / parity with `super::ntt::ntt`).
    for c in f.iter_mut() {
        *c = mod_q(*c);
    }
    Ok(())
}

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

    /// Tiny deterministic xorshift PRNG used only for tests so we
    /// don't depend on `OsRng` (and thus the `std` feature).
    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(())
        }
    }

    #[test]
    fn shuffled_ntt_matches_regular_ntt() {
        let mut a = [0i32; N];
        for i in 0..N {
            a[i] = (i as i32 * 12345 + 7).rem_euclid(super::super::params::Q);
        }
        let mut b = a;

        ntt::ntt(&mut a);

        let mut rng = TestRng(0xDEADBEEFCAFEBABE);
        ntt_shuffled(&mut b, &mut rng).unwrap();

        assert_eq!(a, b, "shuffled NTT must match the regular NTT bit-for-bit");
    }
}