krypteia-quantica 0.1.0

#[cfg(feature = "sca-protected")]
use super::MlKemError;
use super::encode;
#[cfg(feature = "sca-protected")]
use super::masked::{self, MaskedPoly};
use super::ntt;
/// K-PKE lattice-based public-key encryption component (FIPS 203 Section 5).
///
/// Provides the three core algorithms that underlie ML-KEM:
///
/// - [`keygen`] -- Algorithm 13: K-PKE.KeyGen
/// - [`encrypt`] -- Algorithm 14: K-PKE.Encrypt
/// - [`decrypt`] -- Algorithm 15: K-PKE.Decrypt
///
/// All operations are constant-time on secret data. Secret intermediates
/// (private key polynomials, error vectors, randomness vectors) are zeroized
/// via volatile writes after use.
///
/// When the `sca-protected` feature is enabled, additional side-channel
/// countermeasures are applied:
/// - **Masked decryption**: the secret key is split into additive shares
///   and the inner product is computed share-wise (DPA protection).
/// - **Shuffled NTT**: butterfly operations on secret data use randomized
///   ordering (SPA protection).
use super::params::{N, Params};
#[cfg(feature = "sca-protected")]
use super::rng::CryptoRng;
use super::sample;
use super::sha3;
#[cfg(feature = "sca-protected")]
use super::shuffle;

/// Maximum module rank across all ML-KEM parameter sets.
const MAX_K: usize = 4;
/// Maximum PRF output size: 64 * max(eta) = 64 * 3 = 192 bytes.
const MAX_PRF_LEN: usize = 192;

/// Generate a K-PKE key pair (Algorithm 13).
///
/// From a 32-byte seed `d`, derives the public matrix A (in NTT domain),
/// secret vector s, and error vector e using SHA3-512 and SHAKE256-based
/// CBD sampling. Writes the encoded encapsulation key `ek_pke = ByteEncode_12(t_hat) || rho`
/// into `ek_out` and the encoded decapsulation key `dk_pke = ByteEncode_12(s_hat)`
/// into `dk_out`.
///
/// Secret polynomials `s` and `e` are zeroized after use.
///
/// # Arguments
///
/// * `d` - 32-byte seed for deterministic key generation.
/// * `ek_out` - Output slice for ek_pke, must be at least `384*k + 32` bytes.
/// * `dk_out` - Output slice for dk_pke, must be at least `384*k` bytes.
///
/// # Returns
///
/// A tuple `(ek_len, dk_len)` with the actual lengths written.
pub fn keygen<P: Params>(d: &[u8; 32], ek_out: &mut [u8], dk_out: &mut [u8]) -> (usize, usize) {
    let k = P::K;
    let ek_len = 384 * k + 32;
    let dk_len = 384 * k;

    // (ρ, σ) ← G(d || k)
    let mut g_input = [0u8; 33];
    g_input[..32].copy_from_slice(d);
    g_input[32] = k as u8;
    let (rho, sigma) = sha3::g(&g_input);
    ntt::zeroize_bytes(&mut g_input);

    // Generate matrix  (in NTT domain) — public data
    let mut a_hat = [[0i16; N]; MAX_K * MAX_K];
    for i in 0..k {
        for j in 0..k {
            let mut seed = [0u8; 34];
            seed[..32].copy_from_slice(&rho);
            seed[32] = j as u8;
            seed[33] = i as u8;
            a_hat[i * k + j] = sample::sample_ntt(&seed);
        }
    }

    // Generate s (secret) and e (secret) from CBD
    let mut n_counter = 0u8;
    let mut s_hat = [[0i16; N]; MAX_K];
    let mut prf_buf = [0u8; MAX_PRF_LEN];
    for i in 0..k {
        sha3::prf(P::ETA1, &sigma, n_counter, &mut prf_buf);
        s_hat[i] = sample::sample_poly_cbd(P::ETA1, &prf_buf[..64 * P::ETA1]);
        ntt::zeroize_bytes(&mut prf_buf[..64 * P::ETA1]);
        ntt::ntt(&mut s_hat[i]);
        n_counter += 1;
    }

    let mut e_hat = [[0i16; N]; MAX_K];
    for i in 0..k {
        sha3::prf(P::ETA1, &sigma, n_counter, &mut prf_buf);
        e_hat[i] = sample::sample_poly_cbd(P::ETA1, &prf_buf[..64 * P::ETA1]);
        ntt::zeroize_bytes(&mut prf_buf[..64 * P::ETA1]);
        ntt::ntt(&mut e_hat[i]);
        n_counter += 1;
    }

    // t̂ = Â ∘ ŝ + ê (all in NTT domain)
    // multiply_ntts output is in /R domain (Montgomery basemul).
    // to_mont_poly converts /R → normal before adding ê (normal domain from NTT).
    let mut t_hat = [[0i16; N]; MAX_K];
    for i in 0..k {
        for j in 0..k {
            let mut tmp = [0i16; N];
            ntt::multiply_ntts(&a_hat[i * k + j], &s_hat[j], &mut tmp);
            for l in 0..N {
                t_hat[i][l] = t_hat[i][l] + tmp[l];
            }
        }
        ntt::to_mont_poly(&mut t_hat[i]); // /R → normal
        for l in 0..N {
            t_hat[i][l] = t_hat[i][l] + e_hat[i][l];
        }
    }

    // ek_pke = ByteEncode_12(t̂) || ρ
    for i in 0..k {
        let mut t_u16 = [0u16; N];
        for l in 0..N {
            t_u16[l] = ntt::barrett_reduce(t_hat[i][l]) as u16;
        }
        encode::byte_encode(12, &t_u16, &mut ek_out[384 * i..384 * (i + 1)]);
    }
    ek_out[384 * k..384 * k + 32].copy_from_slice(&rho);

    // dk_pke = ByteEncode_12(ŝ)
    for i in 0..k {
        let mut s_u16 = [0u16; N];
        for l in 0..N {
            s_u16[l] = ntt::barrett_reduce(s_hat[i][l]) as u16;
        }
        encode::byte_encode(12, &s_u16, &mut dk_out[384 * i..384 * (i + 1)]);
    }

    // Zeroize secrets
    for poly in s_hat[..k].iter_mut() {
        ntt::zeroize_poly(poly);
    }
    for poly in e_hat[..k].iter_mut() {
        ntt::zeroize_poly(poly);
    }

    (ek_len, dk_len)
}

/// SCA-protected K-PKE key generation with shuffled NTT on secret polynomials.
///
/// Functionally identical to [`keygen`] but uses [`shuffle::ntt_shuffled`] for
/// the forward NTT on secret polynomials `s` and `e`, randomizing the butterfly
/// execution order to defeat Simple Power Analysis.
///
/// The NTT on public data (matrix A via `SampleNTT`) is unaffected.
///
/// # Arguments
///
/// * `d` - 32-byte seed for deterministic key generation.
/// * `ek_out` - Output slice for ek_pke, must be at least `384*k + 32` bytes.
/// * `dk_out` - Output slice for dk_pke, must be at least `384*k` bytes.
/// * `rng` - A cryptographic RNG for shuffle permutation randomness.
///
/// # Errors
///
/// Returns [`MlKemError::RngFailure`] if the RNG fails.
#[cfg(feature = "sca-protected")]
pub fn keygen_sca<P: Params>(
    d: &[u8; 32],
    ek_out: &mut [u8],
    dk_out: &mut [u8],
    rng: &mut impl CryptoRng,
) -> Result<(usize, usize), MlKemError> {
    let k = P::K;
    let ek_len = 384 * k + 32;
    let dk_len = 384 * k;

    // (ρ, σ) ← G(d || k)
    let mut g_input = [0u8; 33];
    g_input[..32].copy_from_slice(d);
    g_input[32] = k as u8;
    let (rho, sigma) = sha3::g(&g_input);
    ntt::zeroize_bytes(&mut g_input);

    // Generate matrix  (in NTT domain) — public data, no shuffling needed
    let mut a_hat = [[0i16; N]; MAX_K * MAX_K];
    for i in 0..k {
        for j in 0..k {
            let mut seed = [0u8; 34];
            seed[..32].copy_from_slice(&rho);
            seed[32] = j as u8;
            seed[33] = i as u8;
            a_hat[i * k + j] = sample::sample_ntt(&seed);
        }
    }

    // Generate s (secret) and e (secret) from CBD — shuffled NTT for SPA protection
    let mut n_counter = 0u8;
    let mut s_hat = [[0i16; N]; MAX_K];
    let mut prf_buf = [0u8; MAX_PRF_LEN];
    for i in 0..k {
        sha3::prf(P::ETA1, &sigma, n_counter, &mut prf_buf);
        s_hat[i] = sample::sample_poly_cbd(P::ETA1, &prf_buf[..64 * P::ETA1]);
        ntt::zeroize_bytes(&mut prf_buf[..64 * P::ETA1]);
        shuffle::ntt_shuffled(&mut s_hat[i], rng)?;
        n_counter += 1;
    }

    let mut e_hat = [[0i16; N]; MAX_K];
    for i in 0..k {
        sha3::prf(P::ETA1, &sigma, n_counter, &mut prf_buf);
        e_hat[i] = sample::sample_poly_cbd(P::ETA1, &prf_buf[..64 * P::ETA1]);
        ntt::zeroize_bytes(&mut prf_buf[..64 * P::ETA1]);
        shuffle::ntt_shuffled(&mut e_hat[i], rng)?;
        n_counter += 1;
    }

    // t̂ = Â ∘ ŝ + ê (all in NTT domain)
    let mut t_hat = [[0i16; N]; MAX_K];
    for i in 0..k {
        for j in 0..k {
            let mut tmp = [0i16; N];
            ntt::multiply_ntts(&a_hat[i * k + j], &s_hat[j], &mut tmp);
            for l in 0..N {
                t_hat[i][l] = t_hat[i][l] + tmp[l];
            }
        }
        ntt::to_mont_poly(&mut t_hat[i]);
        for l in 0..N {
            t_hat[i][l] = t_hat[i][l] + e_hat[i][l];
        }
    }

    // ek_pke = ByteEncode_12(t̂) || ρ
    for i in 0..k {
        let mut t_u16 = [0u16; N];
        for l in 0..N {
            t_u16[l] = ntt::barrett_reduce(t_hat[i][l]) as u16;
        }
        encode::byte_encode(12, &t_u16, &mut ek_out[384 * i..384 * (i + 1)]);
    }
    ek_out[384 * k..384 * k + 32].copy_from_slice(&rho);

    // dk_pke = ByteEncode_12(ŝ)
    for i in 0..k {
        let mut s_u16 = [0u16; N];
        for l in 0..N {
            s_u16[l] = ntt::barrett_reduce(s_hat[i][l]) as u16;
        }
        encode::byte_encode(12, &s_u16, &mut dk_out[384 * i..384 * (i + 1)]);
    }

    // Zeroize secrets
    for poly in s_hat[..k].iter_mut() {
        ntt::zeroize_poly(poly);
    }
    for poly in e_hat[..k].iter_mut() {
        ntt::zeroize_poly(poly);
    }

    Ok((ek_len, dk_len))
}

/// Encrypt a 32-byte message under a K-PKE public key (Algorithm 14).
///
/// Computes `u = NTT_inv(A^T * y_hat) + e1` and
/// `v = NTT_inv(t_hat^T * y_hat) + e2 + Decompress_1(m)`, then compresses
/// and encodes both into a ciphertext of [`Params::CT_LEN`] bytes.
///
/// The randomness vectors `y`, `e1`, and `e2` are derived deterministically
/// from the 32-byte seed `r` via SHAKE256-based CBD sampling, and are
/// zeroized after use.
///
/// # Arguments
///
/// * `ek_pke` - The K-PKE encapsulation key (`ByteEncode_12(t_hat) || rho`).
/// * `m` - The 32-byte message to encrypt (one bit per coefficient).
/// * `r` - 32-byte encryption randomness seed.
/// * `ct_out` - Output slice for the ciphertext, must be at least `P::CT_LEN` bytes.
///
/// # Returns
///
/// The actual ciphertext length written.
pub fn encrypt<P: Params>(ek_pke: &[u8], m: &[u8; 32], r: &[u8; 32], ct_out: &mut [u8]) -> usize {
    let k = P::K;
    let du = P::DU;
    let dv = P::DV;

    // Decode t̂
    let mut t_hat = [[0i16; N]; MAX_K];
    for i in 0..k {
        let mut t_decoded = [0u16; N];
        encode::byte_decode(12, &ek_pke[384 * i..384 * (i + 1)], &mut t_decoded);
        for l in 0..N {
            t_hat[i][l] = t_decoded[l] as i16;
        }
    }

    // Extract ρ and re-generate Â
    let rho = &ek_pke[384 * k..384 * k + 32];
    let mut a_hat = [[0i16; N]; MAX_K * MAX_K];
    for i in 0..k {
        for j in 0..k {
            let mut seed = [0u8; 34];
            seed[..32].copy_from_slice(rho);
            seed[32] = j as u8;
            seed[33] = i as u8;
            a_hat[i * k + j] = sample::sample_ntt(&seed);
        }
    }

    // Generate y, e1, e2 from CBD
    let mut n_counter = 0u8;
    let mut y_hat = [[0i16; N]; MAX_K];
    let mut prf_buf = [0u8; MAX_PRF_LEN];
    for i in 0..k {
        sha3::prf(P::ETA1, r, n_counter, &mut prf_buf);
        y_hat[i] = sample::sample_poly_cbd(P::ETA1, &prf_buf[..64 * P::ETA1]);
        ntt::zeroize_bytes(&mut prf_buf[..64 * P::ETA1]);
        ntt::ntt(&mut y_hat[i]);
        n_counter += 1;
    }

    let mut e1 = [[0i16; N]; MAX_K];
    for i in 0..k {
        sha3::prf(P::ETA2, r, n_counter, &mut prf_buf);
        e1[i] = sample::sample_poly_cbd(P::ETA2, &prf_buf[..64 * P::ETA2]);
        ntt::zeroize_bytes(&mut prf_buf[..64 * P::ETA2]);
        n_counter += 1;
    }

    sha3::prf(P::ETA2, r, n_counter, &mut prf_buf);
    let mut e2 = sample::sample_poly_cbd(P::ETA2, &prf_buf[..64 * P::ETA2]);
    ntt::zeroize_bytes(&mut prf_buf[..64 * P::ETA2]);

    // u = NTT⁻¹(Â^T ∘ ŷ) + e1
    let mut u = [[0i16; N]; MAX_K];
    for i in 0..k {
        let mut acc = [0i16; N];
        for j in 0..k {
            let mut tmp = [0i16; N];
            ntt::multiply_ntts(&a_hat[j * k + i], &y_hat[j], &mut tmp);
            for l in 0..N {
                acc[l] = acc[l].wrapping_add(tmp[l]);
            }
        }
        ntt::ntt_inv(&mut acc);
        ntt::poly_add(&acc, &e1[i], &mut u[i]);
    }

    // Decode message
    let mut mu = [0u16; N];
    encode::byte_decode(1, m, &mut mu);

    // v = NTT⁻¹(t̂^T ∘ ŷ) + e2 + Decompress_1(mu)
    let mut v = [0i16; N];
    {
        let mut acc = [0i16; N];
        for j in 0..k {
            let mut tmp = [0i16; N];
            ntt::multiply_ntts(&t_hat[j], &y_hat[j], &mut tmp);
            for l in 0..N {
                acc[l] = acc[l].wrapping_add(tmp[l]);
            }
        }
        ntt::ntt_inv(&mut acc);
        for l in 0..N {
            let mu_dec = encode::decompress(1, mu[l]) as i16;
            v[l] = acc[l].wrapping_add(e2[l]).wrapping_add(mu_dec);
        }
    }

    // Compress and encode
    let ct_len = 32 * (du * k + dv);

    for i in 0..k {
        let mut u_comp = [0u16; N];
        for l in 0..N {
            u_comp[l] = encode::compress(du as u32, ntt::barrett_reduce(u[i][l]) as u16);
        }
        encode::byte_encode(du, &u_comp, &mut ct_out[32 * du * i..32 * du * (i + 1)]);
    }

    let c2_off = 32 * du * k;
    let mut v_comp = [0u16; N];
    for l in 0..N {
        v_comp[l] = encode::compress(dv as u32, ntt::barrett_reduce(v[l]) as u16);
    }
    encode::byte_encode(dv, &v_comp, &mut ct_out[c2_off..c2_off + 32 * dv]);

    // Zeroize secrets
    for poly in y_hat[..k].iter_mut() {
        ntt::zeroize_poly(poly);
    }
    for poly in e1[..k].iter_mut() {
        ntt::zeroize_poly(poly);
    }
    ntt::zeroize_poly(&mut e2);

    ct_len
}

/// Decrypt a K-PKE ciphertext to recover the 32-byte message (Algorithm 15).
///
/// Decompresses `u` and `v` from the ciphertext, decodes the secret key `s`,
/// then computes `w = v - NTT_inv(s_hat^T * NTT(u))` and compresses `w`
/// to recover the original one-bit-per-coefficient message.
///
/// The secret key polynomial `s_hat` and the accumulator are zeroized
/// after use.
///
/// # Arguments
///
/// * `dk_pke` - The K-PKE decapsulation key (`ByteEncode_12(s_hat)`).
/// * `c` - The ciphertext.
///
/// # Returns
///
/// The recovered 32-byte message.
pub fn decrypt<P: Params>(dk_pke: &[u8], c: &[u8]) -> [u8; 32] {
    let k = P::K;
    let du = P::DU;
    let dv = P::DV;

    // Decompress u
    let mut u = [[0i16; N]; MAX_K];
    for i in 0..k {
        let mut u_comp = [0u16; N];
        encode::byte_decode(du, &c[32 * du * i..32 * du * (i + 1)], &mut u_comp);
        for l in 0..N {
            u[i][l] = encode::decompress(du as u32, u_comp[l]) as i16;
        }
    }

    // Decompress v
    let c2_off = 32 * du * k;
    let mut v_comp = [0u16; N];
    encode::byte_decode(dv, &c[c2_off..c2_off + 32 * dv], &mut v_comp);
    let mut v = [0i16; N];
    for l in 0..N {
        v[l] = encode::decompress(dv as u32, v_comp[l]) as i16;
    }

    // Decode ŝ
    let mut s_hat = [[0i16; N]; MAX_K];
    for i in 0..k {
        let mut s_dec = [0u16; N];
        encode::byte_decode(12, &dk_pke[384 * i..384 * (i + 1)], &mut s_dec);
        for l in 0..N {
            s_hat[i][l] = s_dec[l] as i16;
        }
    }

    // NTT(u)
    for poly in u[..k].iter_mut() {
        ntt::ntt(poly);
    }

    // ŝ^T ∘ NTT(u)
    let mut acc = [0i16; N];
    for j in 0..k {
        let mut tmp = [0i16; N];
        ntt::multiply_ntts(&s_hat[j], &u[j], &mut tmp);
        for l in 0..N {
            acc[l] = acc[l].wrapping_add(tmp[l]);
        }
    }
    ntt::reduce(&mut acc);
    ntt::ntt_inv(&mut acc);

    // w = v - acc
    let mut w = [0i16; N];
    ntt::poly_sub(&v, &acc, &mut w);

    // m = ByteEncode_1(Compress_1(w))
    let mut w_comp = [0u16; N];
    for l in 0..N {
        w_comp[l] = encode::compress(1, ntt::barrett_reduce(w[l]) as u16);
    }
    let mut m = [0u8; 32];
    encode::byte_encode(1, &w_comp, &mut m);

    // Zeroize secrets
    for poly in s_hat[..k].iter_mut() {
        ntt::zeroize_poly(poly);
    }
    ntt::zeroize_poly(&mut acc);

    m
}

/// SCA-protected K-PKE decryption with masked secret key and shuffled NTT.
///
/// Functionally identical to [`decrypt`] but applies two side-channel
/// countermeasures on the critical `ŝ^T · NTT(u)` inner product:
///
/// 1. **Shuffled NTT**: the forward NTT on each ciphertext polynomial `u[i]`
///    uses randomized butterfly ordering to defeat Simple Power Analysis.
/// 2. **Masked multiplication**: the secret key `ŝ` is split into two additive
///    shares `(s0, s1)` and the inner product is computed as
///    `s0·NTT(u) + s1·NTT(u)`, so no single intermediate reveals the secret
///    (Differential Power Analysis protection).
///
/// The RNG is needed for both generating mask shares and shuffle permutations.
///
/// # Arguments
///
/// * `dk_pke` - The K-PKE decapsulation key (`ByteEncode_12(s_hat)`).
/// * `c` - The ciphertext.
/// * `rng` - A cryptographic RNG for masking and shuffle randomness.
///
/// # Errors
///
/// Returns [`MlKemError::RngFailure`] if the RNG fails.
#[cfg(feature = "sca-protected")]
pub fn decrypt_sca<P: Params>(dk_pke: &[u8], c: &[u8], rng: &mut impl CryptoRng) -> Result<[u8; 32], MlKemError> {
    let k = P::K;
    let du = P::DU;
    let dv = P::DV;

    // Decompress u
    let mut u = [[0i16; N]; MAX_K];
    for i in 0..k {
        let mut u_comp = [0u16; N];
        encode::byte_decode(du, &c[32 * du * i..32 * du * (i + 1)], &mut u_comp);
        for l in 0..N {
            u[i][l] = encode::decompress(du as u32, u_comp[l]) as i16;
        }
    }

    // Decompress v
    let c2_off = 32 * du * k;
    let mut v_comp = [0u16; N];
    encode::byte_decode(dv, &c[c2_off..c2_off + 32 * dv], &mut v_comp);
    let mut v = [0i16; N];
    for l in 0..N {
        v[l] = encode::decompress(dv as u32, v_comp[l]) as i16;
    }

    // Decode ŝ
    let mut s_hat = [[0i16; N]; MAX_K];
    for i in 0..k {
        let mut s_dec = [0u16; N];
        encode::byte_decode(12, &dk_pke[384 * i..384 * (i + 1)], &mut s_dec);
        for l in 0..N {
            s_hat[i][l] = s_dec[l] as i16;
        }
    }

    // NTT(u) — shuffled NTT for SPA protection (u is derived from ciphertext
    // but the access pattern during NTT can leak info about the secret when
    // combined with the subsequent multiplication)
    for poly in u[..k].iter_mut() {
        shuffle::ntt_shuffled(poly, rng)?;
    }

    // Masked inner product: split ŝ into shares, compute each share's
    // contribution independently.
    // s_hat is in NTT domain (Barrett-reduced coefficients in [0, q-1]),
    // so MaskedPoly::mask works directly.
    let mut acc_masked = MaskedPoly {
        share0: [0i16; N],
        share1: [0i16; N],
    };
    for j in 0..k {
        let s_masked = MaskedPoly::mask(&s_hat[j], rng)?;
        masked::masked_multiply_accumulate(&mut acc_masked, &s_masked, &u[j]);
    }

    // Unmask and reduce
    let mut acc = acc_masked.unmask();
    acc_masked.zeroize();
    ntt::reduce(&mut acc);
    ntt::ntt_inv(&mut acc);

    // w = v - acc
    let mut w = [0i16; N];
    ntt::poly_sub(&v, &acc, &mut w);

    // m = ByteEncode_1(Compress_1(w))
    let mut w_comp = [0u16; N];
    for l in 0..N {
        w_comp[l] = encode::compress(1, ntt::barrett_reduce(w[l]) as u16);
    }
    let mut m = [0u8; 32];
    encode::byte_encode(1, &w_comp, &mut m);

    // Zeroize secrets
    for poly in s_hat[..k].iter_mut() {
        ntt::zeroize_poly(poly);
    }
    ntt::zeroize_poly(&mut acc);

    Ok(m)
}