krypteia-quantica 0.1.0

Pure-Rust post-quantum cryptography: FIPS 203 ML-KEM, FIPS 204 ML-DSA, and FIPS 205 SLH-DSA. First-order arithmetic masking, shuffled NTT, FORS recompute-and-compare redundancy, constant-time rejection sampling. Targets embedded (no_std), STM32 M0/M4/M33, ESP32-C3 RISC-V. Zero runtime dependencies.
Documentation
/// Polynomial sampling algorithms for ML-KEM (FIPS 203 Section 4.2.2).
///
/// Two sampling methods are provided:
///
/// - [`sample_ntt`] -- Algorithm 7 (SampleNTT): rejection sampling from a SHAKE128
///   XOF to produce a uniformly random polynomial in NTT domain. Operates on
///   **public** data (the seed rho is public), so rejection branching does not
///   leak secret information.
///
/// - [`sample_poly_cbd`] -- Algorithm 8 (SamplePolyCBD): centered binomial
///   distribution sampling from PRF output. Used for secret and error
///   polynomials. Fully constant-time (no secret-dependent branches).
use super::params::{N, Q};
use super::sha3;

/// Sample a uniformly random NTT-domain polynomial (Algorithm 7: SampleNTT).
///
/// Uses SHAKE128 as an XOF (extendable output function) seeded with the
/// 34-byte input `seed = rho || j || i`. Pairs of 12-bit candidates are
/// extracted from each 3-byte block and accepted if less than q = 3329
/// (rejection sampling).
///
/// Since the seed rho is public data, the variable-time rejection loop
/// does not leak secret information through timing.
///
/// # Arguments
///
/// * `seed` - A 34-byte XOF seed: `rho (32 bytes) || column_index (1 byte) || row_index (1 byte)`.
///
/// # Returns
///
/// A 256-coefficient polynomial in NTT domain with coefficients in `[0, q-1]`.
pub fn sample_ntt(seed: &[u8; 34]) -> [i16; N] {
    let mut a_hat = [0i16; N];
    let mut xof = sha3::Xof::new();
    xof.absorb(seed);

    let mut j = 0usize;
    let mut buf = [0u8; 3];

    while j < N {
        xof.squeeze(&mut buf);
        let d1 = (buf[0] as u16) | (((buf[1] as u16) & 0x0F) << 8);
        let d2 = ((buf[1] as u16) >> 4) | ((buf[2] as u16) << 4);

        if d1 < Q {
            a_hat[j] = d1 as i16;
            j += 1;
        }
        if d2 < Q && j < N {
            a_hat[j] = d2 as i16;
            j += 1;
        }
    }
    a_hat
}

/// Sample a polynomial from the centered binomial distribution CBD_eta (Algorithm 8).
///
/// For each of the 256 coefficients, sums `eta` random bits for `x` and
/// `eta` random bits for `y`, then computes `(x - y) mod q`. The result
/// lies in `[-eta, eta]` before reduction, corresponding to the centered
/// binomial distribution.
///
/// Fully constant-time: no branches depend on secret bit values. The
/// branchless modular reduction adds q when the difference is negative
/// using an arithmetic shift mask.
///
/// # Arguments
///
/// * `eta` - The CBD parameter (2 or 3 for ML-KEM).
/// * `bytes` - Exactly `64 * eta` bytes of PRF output.
///
/// # Returns
///
/// A 256-coefficient polynomial with coefficients in `[0, q-1]`.
///
/// # Panics
///
/// Debug-asserts that `bytes.len() == 64 * eta`.
pub fn sample_poly_cbd(eta: usize, bytes: &[u8]) -> [i16; N] {
    debug_assert_eq!(bytes.len(), 64 * eta);
    let mut f = [0i16; N];

    for i in 0..N {
        let mut x = 0i16;
        let mut y = 0i16;
        for j in 0..eta {
            let bit_x = (i * 2 * eta + j) as usize;
            let bit_y = (i * 2 * eta + eta + j) as usize;
            x += ((bytes[bit_x >> 3] >> (bit_x & 7)) & 1) as i16;
            y += ((bytes[bit_y >> 3] >> (bit_y & 7)) & 1) as i16;
        }
        // Branchless mod q: diff is in [-η, η], add q if negative
        let diff = x - y;
        // Constant-time: q & (diff >> 15) adds q iff diff < 0
        f[i] = diff.wrapping_add(super::ntt::Q & (diff >> 15));
    }
    f
}