krypteia-quantica 0.1.0

//! First-order arithmetic masking for ML-KEM polynomials
//! (**countermeasure: DPA / DEMA / CPA on the K-PKE secret ``s``**).
//!
//! ## Principle
//!
//! Every secret polynomial is kept as two additive shares modulo `q`:
//! `s = s_0 + s_1 (mod q)`. Shares are drawn uniformly at random and
//! each operation consuming the secret (NTT, pointwise multiplication
//! with a public polynomial, additions / subtractions) is rewritten
//! to operate on the shares without ever materialising the unmasked
//! `s`. Fresh-randomness refreshes (`MaskedPoly::refresh`) break
//! cross-operation correlations that a high-order DPA could otherwise
//! exploit.
//!
//! Because the unmasked `s` never exists as a single value in memory,
//! the Hamming-weight / Hamming-distance hypothesis at the core of
//! CPA becomes non-identifiable: correlating the power trace with a
//! guess of any byte of `s` produces no peak since each share alone
//! is uniform.
//!
//! ## 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) |
//! | `masked_ntt` / `masked_ntt_inv` | NTT on each share independently |
//! | `masked_multiply_public` | Multiply masked poly by a public poly |
//! | `masked_add` / `masked_sub` | Add/subtract two masked polys |
//! | `masked_multiply_accumulate` | Fused multiply-add with a public poly |
//!
//! ## References
//!
//! * *Side-channel analysis of the ML-KEM pointwise multiplication*
//!   (IACR ePrint 2025, `doc/papers/eprint2025_sca_mlkem_pointwise.pdf`)
//!   — identifies the pointwise multiplication as the key DPA
//!   target; masking the shares defeats the published attack.
//! * *ML-KEM and ML-DSA on OpenTitan: side-channel evaluation*
//!   (cryptojedi, 2024) — independent evaluation of the trace-count
//!   blow-up induced by masking.
//!
//! ## Where to look next
//!
//! * Countermeasure description and threat analysis:
//!   `doc/sca/countermeasures/ml_kem.rst`, section
//!   *DPA / DEMA / CPA — first-order masking of the KPKE secret*.
//! * Call sites: `keygen_internal_sca`,
//!   `decaps_internal_sca`,
//!   [`crate::ml_kem::kpke::decrypt_sca`].
//!
//! ## Scope and residual risk
//!
//! Masking is **first-order**. A higher-order DPA that combines two
//! independent time samples can still recover the secret with
//! substantially more traces. Tier-4 item `K-SCA3` (not yet
//! scheduled) would extend this to a 3-share scheme (CC EAL4+-grade).

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

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

impl MaskedPoly {
    /// Split a plaintext polynomial into two random additive shares.
    ///
    /// Generates a uniformly random `share1` from the RNG, then computes
    /// `share0 = poly - share1 mod q`. The random bytes are zeroized after use.
    ///
    /// # Arguments
    ///
    /// * `poly` - The secret polynomial to mask (coefficients in `[0, q-1]`).
    /// * `rng` - A cryptographic RNG for generating the random share.
    ///
    /// # Errors
    ///
    /// Returns [`MlKemError::RngFailure`] if the RNG fails.
    pub fn mask(poly: &[i16; N], rng: &mut impl CryptoRng) -> Result<Self, MlKemError> {
        let mut share1 = [0i16; N];
        let mut rand_bytes = [0u8; N * 2];
        rng.fill_bytes(&mut rand_bytes)?;

        for i in 0..N {
            // Random value mod q from 2 random bytes
            let r = u16::from_le_bytes([rand_bytes[2 * i], rand_bytes[2 * i + 1]]);
            share1[i] = (r % (Q as u16)) as i16;
        }
        ntt::zeroize_bytes(&mut rand_bytes);

        let mut share0 = [0i16; N];
        for i in 0..N {
            let mut r = (poly[i] as i32 - share1[i] as i32) % (Q as i32);
            r += (Q as i32) & (r >> 31);
            share0[i] = r as i16;
        }

        Ok(MaskedPoly { share0, share1 })
    }

    /// Reconstruct the plaintext polynomial from the two shares.
    ///
    /// Computes `(share0[i] + share1[i]) mod q` for each coefficient.
    /// The resulting polynomial has coefficients in `[0, q-1]`.
    ///
    /// # Security note
    ///
    /// The returned polynomial is unmasked and should be handled as secret
    /// data (zeroized after use).
    pub fn unmask(&self) -> [i16; N] {
        let mut result = [0i16; N];
        for i in 0..N {
            let mut r = (self.share0[i] as i32 + self.share1[i] as i32) % (Q as i32);
            r += (Q as i32) & (r >> 31);
            result[i] = r as i16;
        }
        result
    }

    /// Securely erase both shares via volatile writes.
    ///
    /// Delegates to [`super::ntt::zeroize_poly`] for each share to ensure
    /// the optimizer cannot elide the writes.
    pub fn zeroize(&mut self) {
        ntt::zeroize_poly(&mut self.share0);
        ntt::zeroize_poly(&mut self.share1);
    }

    /// Re-randomize the shares without changing the unmasked value.
    ///
    /// Draws a fresh random polynomial `r` and updates the shares:
    /// `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 used in multiple operations.
    ///
    /// # Errors
    ///
    /// Returns [`MlKemError::RngFailure`] if the RNG fails.
    pub fn refresh(&mut self, rng: &mut impl CryptoRng) -> Result<(), MlKemError> {
        let mut rand_bytes = [0u8; N * 2];
        rng.fill_bytes(&mut rand_bytes)?;

        for i in 0..N {
            let r = (u16::from_le_bytes([rand_bytes[2 * i], rand_bytes[2 * i + 1]]) % (Q as u16)) as i16;
            let mut s0 = (self.share0[i] as i32 - r as i32) % (Q as i32);
            s0 += (Q as i32) & (s0 >> 31);
            self.share0[i] = s0 as i16;

            let mut s1 = (self.share1[i] as i32 + r as i32) % (Q as i32);
            s1 += (Q as i32) & (s1 >> 31);
            self.share1[i] = s1 as i16;
        }
        ntt::zeroize_bytes(&mut rand_bytes);
        Ok(())
    }
}

// ---- Masked NTT operations ----
// NTT is linear: NTT(s₀ + s₁) = NTT(s₀) + NTT(s₁)
// So we NTT each share independently.

/// Apply the forward NTT to each share of a [`MaskedPoly`] independently.
///
/// Exploits the linearity of the NTT: `NTT(s0 + s1) = NTT(s0) + NTT(s1)`.
/// Each share is transformed independently so the masking invariant is preserved.
pub fn masked_ntt(m: &mut MaskedPoly) {
    ntt::ntt(&mut m.share0);
    ntt::ntt(&mut m.share1);
}

/// Apply the inverse NTT to each share of a [`MaskedPoly`] independently.
///
/// The inverse of `masked_ntt`. Transforms both shares from the NTT domain
/// back to the standard domain while preserving the masking invariant.
pub fn masked_ntt_inv(m: &mut MaskedPoly) {
    ntt::ntt_inv(&mut m.share0);
    ntt::ntt_inv(&mut m.share1);
}

/// Multiply a masked polynomial by a **public** polynomial in NTT domain.
///
/// Computes `(s0 + s1) * p = s0*p + s1*p` by multiplying each share
/// independently. This is secure because `p` is public data -- there is
/// no secret-times-secret interaction that would require second-order masking.
///
/// # Arguments
///
/// * `masked` - The masked secret polynomial (NTT domain).
/// * `public` - A public polynomial (NTT domain).
/// * `out` - Output masked polynomial for the product.
pub fn masked_multiply_public(masked: &MaskedPoly, public: &[i16; N], out: &mut MaskedPoly) {
    ntt::multiply_ntts(&masked.share0, public, &mut out.share0);
    ntt::multiply_ntts(&masked.share1, public, &mut out.share1);
}

/// Add a **public** polynomial to a masked polynomial.
///
/// Adds the public value only to `share0`. Since the public polynomial
/// is not secret, it does not need to be split across shares.
///
/// # Arguments
///
/// * `m` - The masked polynomial to modify in place.
/// * `public` - A public polynomial to add.
pub fn masked_add_public(m: &mut MaskedPoly, public: &[i16; N]) {
    ntt::poly_add(&m.share0, public, &mut m.share0.clone());
    let old = m.share0;
    ntt::poly_add(&old, public, &mut m.share0);
}

/// Add two masked polynomials share-wise.
///
/// Computes `c.share0 = a.share0 + b.share0` and
/// `c.share1 = a.share1 + b.share1`, so
/// `unmask(c) = unmask(a) + unmask(b) mod q`.
pub fn masked_add(a: &MaskedPoly, b: &MaskedPoly, c: &mut MaskedPoly) {
    ntt::poly_add(&a.share0, &b.share0, &mut c.share0);
    ntt::poly_add(&a.share1, &b.share1, &mut c.share1);
}

/// Subtract two masked polynomials share-wise.
///
/// Computes `c.share0 = a.share0 - b.share0` and
/// `c.share1 = a.share1 - b.share1`, so
/// `unmask(c) = unmask(a) - unmask(b) mod q`.
pub fn masked_sub(a: &MaskedPoly, b: &MaskedPoly, c: &mut MaskedPoly) {
    ntt::poly_sub(&a.share0, &b.share0, &mut c.share0);
    ntt::poly_sub(&a.share1, &b.share1, &mut c.share1);
}

/// Fused multiply-accumulate: `out += masked * public` (NTT domain).
///
/// Multiplies a masked polynomial by a public polynomial and adds the
/// result to the existing shares in `out`. Useful for computing matrix-vector
/// products share-wise without allocating temporaries for each column.
///
/// # Arguments
///
/// * `out` - Accumulator masked polynomial (modified in place).
/// * `masked` - The masked secret polynomial (NTT domain).
/// * `public` - A public polynomial (NTT domain).
pub fn masked_multiply_accumulate(out: &mut MaskedPoly, masked: &MaskedPoly, public: &[i16; N]) {
    let mut tmp0 = [0i16; N];
    let mut tmp1 = [0i16; N];
    ntt::multiply_ntts(&masked.share0, public, &mut tmp0);
    ntt::multiply_ntts(&masked.share1, public, &mut tmp1);
    for i in 0..N {
        out.share0[i] = ntt::barrett_reduce((out.share0[i] as i32 + tmp0[i] as i32) as i16);
        out.share1[i] = ntt::barrett_reduce((out.share1[i] as i32 + tmp1[i] as i32) as i16);
    }
}