krypteia-quantica 0.1.0

use super::MlKemError;
use super::encode;
use super::kpke;
use super::ntt;
use super::params::N;
use super::params::Params;
use super::rng::CryptoRng;
use super::sha3;
/// ML-KEM Key-Encapsulation Mechanism (FIPS 203 Sections 6-7).
///
/// Implements Algorithms 16-21 with the following side-channel countermeasures:
///
/// - **Constant-time**: no secret-dependent branches or memory access patterns
/// - **Zeroization**: all secret intermediates erased via volatile writes
/// - **Double decaps**: fault detection by running decapsulation twice (DFA protection)
/// - **dk integrity**: `H(ek)` verification at decapsulation time (DFA protection on key storage)
///
/// The public API ([`keygen`], [`encaps`], [`decaps`]) includes input validation
/// and full countermeasures. The internal variants ([`keygen_internal`],
/// [`encaps_internal`], [`decaps_internal`]) are deterministic and intended
/// for CAVP testing.
use alloc::vec::Vec;

/// Maximum module rank across all ML-KEM parameter sets.
const MAX_K: usize = 4;
/// Maximum encapsulation key size: 384*4 + 32 = 1568 bytes.
const MAX_EK_LEN: usize = 384 * MAX_K + 32;
/// Maximum ciphertext size: 32*(11*4 + 5) = 1568 bytes.
const MAX_CT_LEN: usize = 1568;
/// Maximum j_input size: 32 + max ciphertext = 1600 bytes.
const MAX_J_INPUT_LEN: usize = 32 + MAX_CT_LEN;

// =====================================================================
// Internal (deterministic) functions — used for CAVP testing
// =====================================================================

/// Deterministic ML-KEM key generation (Algorithm 16).
///
/// Generates an encapsulation/decapsulation key pair from explicit seeds.
/// The decapsulation key is structured as `dk_pke || ek || H(ek) || z`.
///
/// # Arguments
///
/// * `d` - 32-byte seed passed to [`kpke::keygen`] for K-PKE key generation.
/// * `z` - 32-byte implicit rejection value embedded in the decapsulation key.
///
/// # Returns
///
/// A tuple `(ek, dk)` of byte vectors.
pub fn keygen_internal<P: Params>(d: &[u8; 32], z: &[u8; 32]) -> (Vec<u8>, Vec<u8>) {
    let k = P::K;
    let ek_len = 384 * k + 32;
    let dk_pke_len = 384 * k;

    let mut ek_buf = [0u8; MAX_EK_LEN];
    let mut dk_pke_buf = [0u8; 384 * MAX_K];
    kpke::keygen::<P>(d, &mut ek_buf, &mut dk_pke_buf);

    let ek = &ek_buf[..ek_len];
    let h_ek = sha3::h(ek);

    let mut dk = Vec::with_capacity(P::DK_LEN);
    dk.extend_from_slice(&dk_pke_buf[..dk_pke_len]);
    dk.extend_from_slice(ek);
    dk.extend_from_slice(&h_ek);
    dk.extend_from_slice(z);

    (ek[..ek_len].to_vec(), dk)
}

/// SCA-protected ML-KEM key generation.
///
/// Uses [`kpke::keygen_sca`] which applies shuffled NTT on secret polynomials,
/// then assembles the full decapsulation key. The deterministic seed derivation
/// is identical to [`keygen_internal`].
#[cfg(feature = "sca-protected")]
fn keygen_internal_sca<P: Params>(
    d: &[u8; 32],
    z: &[u8; 32],
    rng: &mut impl CryptoRng,
) -> Result<(Vec<u8>, Vec<u8>), MlKemError> {
    let k = P::K;
    let ek_len = 384 * k + 32;
    let dk_pke_len = 384 * k;

    let mut ek_buf = [0u8; MAX_EK_LEN];
    let mut dk_pke_buf = [0u8; 384 * MAX_K];
    kpke::keygen_sca::<P>(d, &mut ek_buf, &mut dk_pke_buf, rng)?;

    let ek = &ek_buf[..ek_len];
    let h_ek = sha3::h(ek);

    let mut dk = Vec::with_capacity(P::DK_LEN);
    dk.extend_from_slice(&dk_pke_buf[..dk_pke_len]);
    dk.extend_from_slice(ek);
    dk.extend_from_slice(&h_ek);
    dk.extend_from_slice(z);

    Ok((ek[..ek_len].to_vec(), dk))
}

/// Deterministic ML-KEM encapsulation (Algorithm 17).
///
/// Computes the shared secret and ciphertext from an encapsulation key
/// and an explicit 32-byte message `m`. The shared secret and encryption
/// randomness are derived as `(K, r) = G(m || H(ek))`.
///
/// No input validation is performed on `ek`; callers should use [`encaps`]
/// for production use.
///
/// # Arguments
///
/// * `ek` - The encapsulation (public) key.
/// * `m` - 32-byte random message seed.
///
/// # Returns
///
/// A tuple `(shared_secret, ciphertext)`.
pub fn encaps_internal<P: Params>(ek: &[u8], m: &[u8; 32]) -> ([u8; 32], Vec<u8>) {
    let h_ek = sha3::h(ek);
    let mut g_input = [0u8; 64];
    g_input[..32].copy_from_slice(m);
    g_input[32..64].copy_from_slice(&h_ek);
    let (shared_key, r) = sha3::g(&g_input);
    ntt::zeroize_bytes(&mut g_input);

    let mut ct_buf = [0u8; MAX_CT_LEN];
    let ct_len = kpke::encrypt::<P>(ek, m, &r, &mut ct_buf);
    (shared_key, ct_buf[..ct_len].to_vec())
}

/// Deterministic ML-KEM decapsulation (Algorithm 18).
///
/// Implements the Fujisaki-Okamoto transform: decrypts the ciphertext to
/// recover `m'`, re-encrypts to get `c'`, then uses a constant-time
/// comparison to select either the real shared key `K'` (if `c == c'`)
/// or an implicit rejection key `J(z || c)` (otherwise).
///
/// All operations are constant-time with no secret-dependent branches.
/// Intermediates (`m'`, `g_input`, `j_input`) are zeroized after use.
///
/// # Arguments
///
/// * `dk` - The full decapsulation key (layout: `dk_pke || ek || H(ek) || z`).
/// * `c`  - The ciphertext.
///
/// # Returns
///
/// The 32-byte shared secret.
pub fn decaps_internal<P: Params>(dk: &[u8], c: &[u8]) -> [u8; 32] {
    let k = P::K;
    let dk_pke = &dk[0..384 * k];
    let ek_pke = &dk[384 * k..768 * k + 32];
    let h = &dk[768 * k + 32..768 * k + 64];
    let z = &dk[768 * k + 64..768 * k + 96];

    let mut m_prime = kpke::decrypt::<P>(dk_pke, c);

    let mut g_input = [0u8; 64];
    g_input[..32].copy_from_slice(&m_prime);
    g_input[32..64].copy_from_slice(h);
    let (k_prime, r_prime) = sha3::g(&g_input);
    ntt::zeroize_bytes(&mut g_input);

    let j_input_len = 32 + c.len();
    let mut j_input = [0u8; MAX_J_INPUT_LEN];
    j_input[..32].copy_from_slice(z);
    j_input[32..j_input_len].copy_from_slice(c);
    let k_bar = sha3::j(&j_input[..j_input_len]);
    ntt::zeroize_bytes(&mut j_input[..j_input_len]);

    let mut c_prime = [0u8; MAX_CT_LEN];
    let ct_len = kpke::encrypt::<P>(ek_pke, &m_prime, &r_prime, &mut c_prime);

    let eq = ct_eq(c, &c_prime[..ct_len]);
    let mut result = [0u8; 32];
    ct_select(&mut result, &k_prime, &k_bar, eq);

    ntt::zeroize_bytes(&mut m_prime);
    result
}

/// SCA-protected ML-KEM decapsulation (internal).
///
/// Identical to [`decaps_internal`] but uses [`kpke::decrypt_sca`] which applies
/// masked multiplication and shuffled NTT during decryption.
#[cfg(feature = "sca-protected")]
fn decaps_internal_sca<P: Params>(dk: &[u8], c: &[u8], rng: &mut impl CryptoRng) -> Result<[u8; 32], MlKemError> {
    let k = P::K;
    let dk_pke = &dk[0..384 * k];
    let ek_pke = &dk[384 * k..768 * k + 32];
    let h = &dk[768 * k + 32..768 * k + 64];
    let z = &dk[768 * k + 64..768 * k + 96];

    let mut m_prime = kpke::decrypt_sca::<P>(dk_pke, c, rng)?;

    let mut g_input = [0u8; 64];
    g_input[..32].copy_from_slice(&m_prime);
    g_input[32..64].copy_from_slice(h);
    let (k_prime, r_prime) = sha3::g(&g_input);
    ntt::zeroize_bytes(&mut g_input);

    let j_input_len = 32 + c.len();
    let mut j_input = [0u8; MAX_J_INPUT_LEN];
    j_input[..32].copy_from_slice(z);
    j_input[32..j_input_len].copy_from_slice(c);
    let k_bar = sha3::j(&j_input[..j_input_len]);
    ntt::zeroize_bytes(&mut j_input[..j_input_len]);

    let mut c_prime = [0u8; MAX_CT_LEN];
    let ct_len = kpke::encrypt::<P>(ek_pke, &m_prime, &r_prime, &mut c_prime);

    let eq = ct_eq(c, &c_prime[..ct_len]);
    let mut result = [0u8; 32];
    ct_select(&mut result, &k_prime, &k_bar, eq);

    ntt::zeroize_bytes(&mut m_prime);
    Ok(result)
}

// =====================================================================
// Public API with full side-channel countermeasures
// =====================================================================

/// Generate an ML-KEM key pair (Algorithm 19).
///
/// Draws 32-byte seeds `d` and `z` from the RNG, delegates to
/// [`keygen_internal`], then zeroizes both seeds. This is the
/// recommended entry point for key generation.
///
/// # Arguments
///
/// * `rng` - A cryptographic random number generator implementing [`CryptoRng`].
///
/// # Returns
///
/// A tuple `(encapsulation_key, decapsulation_key)` as byte vectors.
///
/// # Errors
///
/// Returns [`MlKemError::RngFailure`] if the RNG fails to produce bytes.
pub fn keygen<P: Params>(rng: &mut impl CryptoRng) -> Result<(Vec<u8>, Vec<u8>), MlKemError> {
    let mut d = [0u8; 32];
    let mut z = [0u8; 32];
    rng.fill_bytes(&mut d)?;
    rng.fill_bytes(&mut z)?;

    #[cfg(feature = "sca-protected")]
    let result = keygen_internal_sca::<P>(&d, &z, rng)?;

    #[cfg(not(feature = "sca-protected"))]
    let result = keygen_internal::<P>(&d, &z);

    ntt::zeroize_bytes(&mut d);
    ntt::zeroize_bytes(&mut z);
    Ok(result)
}

/// Encapsulate a shared secret (Algorithm 20) with input validation.
///
/// Validates the encapsulation key (length check and constant-time modulus
/// check per FIPS 203 Section 7.2), then draws a random 32-byte message
/// and delegates to [`encaps_internal`].
///
/// The modulus check ensures each 12-bit coefficient in `ek` is less than
/// q = 3329 by round-tripping through encode/decode.
///
/// # Arguments
///
/// * `ek` - The encapsulation (public) key, exactly [`Params::EK_LEN`] bytes.
/// * `rng` - A cryptographic random number generator implementing [`CryptoRng`].
///
/// # Returns
///
/// A tuple `(shared_secret, ciphertext)`.
///
/// # Errors
///
/// * [`MlKemError::InvalidEncapsulationKey`] if the key has wrong length or
///   fails the modulus check.
/// * [`MlKemError::RngFailure`] if the RNG fails.
pub fn encaps<P: Params>(ek: &[u8], rng: &mut impl CryptoRng) -> Result<([u8; 32], Vec<u8>), MlKemError> {
    if ek.len() != P::EK_LEN {
        return Err(MlKemError::InvalidEncapsulationKey);
    }
    // Modulus check (constant-time)
    let k = P::K;
    for i in 0..k {
        let slice = &ek[384 * i..384 * (i + 1)];
        let mut decoded = [0u16; N];
        encode::byte_decode(12, slice, &mut decoded);
        let mut reencoded = [0u8; 384];
        encode::byte_encode(12, &decoded, &mut reencoded);
        if !ct_eq(slice, &reencoded) {
            return Err(MlKemError::InvalidEncapsulationKey);
        }
    }

    let mut m = [0u8; 32];
    rng.fill_bytes(&mut m)?;
    let (shared_key, ct) = encaps_internal::<P>(ek, &m);
    ntt::zeroize_bytes(&mut m);
    Ok((shared_key, ct))
}

/// Algorithm 21: ML-KEM.Decaps(dk, c) with input validation and DFA protection.
///
/// # Side-channel / fault countermeasures
///
/// Three independent countermeasures are layered inside this function.
/// Full threat-model context in `doc/sca/countermeasures/ml_kem.rst`,
/// sections *DFA — double computation + CT fault fallback* and
/// *DFA on `dk` — `H(ek)` integrity check*.
///
/// ## 1. `dk` integrity check — fault on `dk` in storage
///
/// The decapsulation key's FIPS 203 layout is
/// `dk_pke ‖ ek ‖ H(ek) ‖ z`. A fault that alters `dk` in memory
/// (for example a hot-carrier-induced bit flip in flash) would
/// undermine the FO security argument because the attacker could
/// coerce decapsulation into using a crafted `dk_pke`. We recompute
/// `H(ek_in_dk)` and compare it with the stored `H(ek)` using
/// `silentops::ct_eq`; a mismatch aborts with
/// [`MlKemError::InvalidDecapsulationKey`] before `decaps_internal`
/// runs.
///
/// ## 2. Double computation — fault on FO re-encryption
///
/// ML-KEM decapsulation is vulnerable to a classical DFA on the FO
/// re-encryption step: if an attacker can make the re-encryption
/// return a value close to the real ciphertext on one specific
/// input, the implicit-rejection path is bypassed and the KEM acts
/// as a decryption oracle (Boneh–DeMillo–Lipton, EUROCRYPT 1997).
///
/// We run `decaps_internal_sca` twice. A single-fault attacker can
/// only affect one execution; if the two shared secrets differ,
/// we conclude a fault happened and switch to the constant-time
/// fallback described below. In the no-fault path both runs agree
/// and we return either result.
///
/// ## 3. Constant-time fault fallback — leakage on the fault branch
///
/// Naive code would write
/// `if !results_match { return k_fault }`, which introduces a
/// conditional jump on a secret-derived bit (`results_match` is
/// derived from `k1`, `k2` which both depend on the sk). An
/// attacker able to inject a fault AND measure timing learns
/// "fault was detected" from the branch timing alone.
///
/// To close this, we **always** compute `k_fault = SHA3(z ‖ 0xFF)`
/// and select between `k1` and `k_fault` with
/// [`silentops::ct_select_u8`] (via the local 32-byte `ct_select`
/// wrapper). The branch is gone; timing is identical in both
/// cases. `k_fault` is:
/// * deterministic for a given `(dk, c)` so a repeated faulted
///   call returns the same value (prevents oracle-by-repetition);
/// * distinct from both the legitimate FO output and the implicit-
///   rejection output (`J(z ‖ c)`), so the attacker cannot
///   distinguish "fault detected" from either correct branch.
///
/// Recommended for embedded and high-security contexts where
/// physical fault attacks are in the threat model.
pub fn decaps<P: Params>(dk: &[u8], c: &[u8], _rng: &mut impl CryptoRng) -> Result<[u8; 32], MlKemError> {
    let k = P::K;

    // Length checks
    if c.len() != P::CT_LEN {
        return Err(MlKemError::InvalidCiphertext);
    }
    if dk.len() != P::DK_LEN {
        return Err(MlKemError::InvalidDecapsulationKey);
    }

    // ---- DFA countermeasure 1: dk integrity check ----
    // Verify H(ek) stored in dk matches recomputed value.
    // Detects fault injection on dk in memory.
    let ek_in_dk = &dk[384 * k..768 * k + 32];
    let h_stored = &dk[768 * k + 32..768 * k + 64];
    let h_computed = sha3::h(ek_in_dk);
    if !ct_eq(&h_computed, h_stored) {
        return Err(MlKemError::InvalidDecapsulationKey);
    }

    // ---- DFA countermeasure 2: double computation ----
    // Run decaps_internal twice. A single-fault attack can only affect
    // one execution. If results differ, a fault was detected.
    #[cfg(feature = "sca-protected")]
    let k1 = decaps_internal_sca::<P>(dk, c, _rng)?;
    #[cfg(feature = "sca-protected")]
    let k2 = decaps_internal_sca::<P>(dk, c, _rng)?;

    #[cfg(not(feature = "sca-protected"))]
    let k1 = decaps_internal::<P>(dk, c);
    #[cfg(not(feature = "sca-protected"))]
    let k2 = decaps_internal::<P>(dk, c);

    // Always compute the fault-fallback key, then select between `k1`
    // and `k_fault` in constant time. A `if !match { return k_fault }`
    // branch would leak, through timing, whether a fault was injected
    // — measurable under ctgrind, and exploitable when the attacker
    // can both trigger faults and observe timing. Constant-time
    // selection costs one extra SHA3 (negligible compared to the two
    // full decaps calls already performed).
    let results_match = ct_eq(&k1, &k2);
    let z = &dk[768 * k + 64..768 * k + 96];
    let mut fault_input = [0u8; 33];
    fault_input[..32].copy_from_slice(z);
    fault_input[32] = 0xFF; // domain separator: distinct from J(z||c)
    let k_fault = sha3::h(&fault_input);
    ntt::zeroize_bytes(&mut fault_input);

    let mut out = [0u8; 32];
    // out = if results_match { k1 } else { k_fault }
    ct_select(&mut out, &k1, &k_fault, results_match);
    Ok(out)
}

/// Decapsulate without double computation (single-pass variant).
///
/// Performs length validation and the `H(ek)` integrity check on the
/// decapsulation key, then runs [`decaps_internal`] once. This is faster
/// than [`decaps`] but does not detect single-fault injection attacks.
///
/// Suitable for software-only environments where physical fault attacks
/// are not in the threat model.
///
/// # Arguments
///
/// * `dk` - The decapsulation (private) key, exactly [`Params::DK_LEN`] bytes.
/// * `c`  - The ciphertext, exactly [`Params::CT_LEN`] bytes.
///
/// # Returns
///
/// The 32-byte shared secret.
///
/// # Errors
///
/// * [`MlKemError::InvalidDecapsulationKey`] if `dk` has wrong length or
///   fails the integrity check.
/// * [`MlKemError::InvalidCiphertext`] if `c` has wrong length.
pub fn decaps_single<P: Params>(dk: &[u8], c: &[u8]) -> Result<[u8; 32], MlKemError> {
    let k = P::K;
    if c.len() != P::CT_LEN {
        return Err(MlKemError::InvalidCiphertext);
    }
    if dk.len() != P::DK_LEN {
        return Err(MlKemError::InvalidDecapsulationKey);
    }
    let ek_in_dk = &dk[384 * k..768 * k + 32];
    let h_stored = &dk[768 * k + 32..768 * k + 64];
    let h_computed = sha3::h(ek_in_dk);
    if !ct_eq(&h_computed, h_stored) {
        return Err(MlKemError::InvalidDecapsulationKey);
    }
    Ok(decaps_internal::<P>(dk, c))
}

// =====================================================================
// Constant-time primitives (delegated to silentops crate)
// =====================================================================

/// Constant-time equality. No early exit, no secret-dependent branch.
/// Delegates to [`silentops::ct_eq`] and converts the u8 result to bool.
#[inline(never)]
fn ct_eq(a: &[u8], b: &[u8]) -> bool {
    silentops::ct_eq(a, b) == 1
}

/// Constant-time select: out = condition ? a : b. Branchless via arithmetic mask.
/// Delegates to [`silentops::ct_select_u8`] per byte.
#[inline(never)]
fn ct_select(out: &mut [u8; 32], a: &[u8; 32], b: &[u8; 32], condition: bool) {
    let cond = condition as u8;
    for i in 0..32 {
        out[i] = silentops::ct_select_u8(a[i], b[i], cond);
    }
}