krypteia-quantica 0.1.0

//! SHA-3 / SHAKE high-level wrappers used by ML-KEM (FIPS 203).
//!
//! This module no longer carries its own copy of the Keccak permutation:
//! it builds on top of the shared sponge core in `crate::sha3`. The
//! ML-KEM-specific surface is unchanged:
//!
//! | Function  | Primitive   | Usage in FIPS 203                                       |
//! |-----------|-------------|---------------------------------------------------------|
//! | `h`       | SHA3-256    | H — hash encapsulation key, dk integrity                |
//! | `g`       | SHA3-512    | G — derive shared key and encryption randomness         |
//! | `j`       | SHAKE-256   | J — implicit rejection key derivation                   |
//! | `prf`     | SHAKE-256   | PRF — CBD sampling randomness                           |
//! | `Xof`     | SHAKE-128   | XOF — matrix sampling via `super::sample::sample_ntt`   |

use crate::sha3::{KeccakState, SHA3_256_RATE, SHA3_512_RATE, SHAKE128_RATE, SHAKE256_RATE};

/// H(s) — SHA3-256 hash returning 32 bytes.
///
/// Used in FIPS 203 for hashing the encapsulation key (`H(ek)`) and
/// for the dk integrity check.
pub fn h(input: &[u8]) -> [u8; 32] {
    let mut state = KeccakState::new(SHA3_256_RATE, 0x06);
    state.absorb(input);
    let mut out = [0u8; 32];
    state.squeeze(&mut out);
    out
}

/// G(c) — SHA3-512 hash returning two 32-byte values.
///
/// Computes `SHA3-512(c)` and splits the 64-byte output into two halves.
/// Used in FIPS 203 to derive both the shared secret and the encryption
/// randomness: `(K, r) = G(m || H(ek))`.
pub fn g(input: &[u8]) -> ([u8; 32], [u8; 32]) {
    let mut state = KeccakState::new(SHA3_512_RATE, 0x06);
    state.absorb(input);
    let mut out = [0u8; 64];
    state.squeeze(&mut out);
    let mut a = [0u8; 32];
    let mut b = [0u8; 32];
    a.copy_from_slice(&out[..32]);
    b.copy_from_slice(&out[32..]);
    (a, b)
}

/// J(s) — SHAKE-256 producing 32 bytes of output.
///
/// Used in FIPS 203 for implicit rejection: `K_bar = J(z || c)` provides
/// the fallback shared secret when ciphertext re-encryption does not match.
pub fn j(input: &[u8]) -> [u8; 32] {
    let mut state = KeccakState::new(SHAKE256_RATE, 0x1f);
    state.absorb(input);
    let mut out = [0u8; 32];
    state.squeeze(&mut out);
    out
}

/// PRF_eta(s, b) — SHAKE-256 pseudo-random function producing `64 * eta` bytes.
///
/// Used in FIPS 203 to generate the random bytes consumed by
/// [`super::sample::sample_poly_cbd`] for secret and error polynomial
/// sampling.
pub fn prf(eta: usize, s: &[u8; 32], b: u8, out: &mut [u8]) {
    debug_assert!(out.len() >= 64 * eta);
    let mut state = KeccakState::new(SHAKE256_RATE, 0x1f);
    state.absorb(s);
    state.absorb(&[b]);
    let len = 64 * eta;
    state.squeeze(&mut out[..len]);
}

/// Extendable Output Function (XOF) context wrapping SHAKE-128.
///
/// Provides an incremental absorb/squeeze interface for generating
/// arbitrary-length pseudo-random output. Used by
/// [`super::sample::sample_ntt`] to sample the public matrix A.
pub struct Xof {
    state: KeccakState,
}

impl Xof {
    /// Create a new SHAKE-128 XOF context with an empty state.
    pub fn new() -> Self {
        Self {
            state: KeccakState::new(SHAKE128_RATE, 0x1f),
        }
    }

    /// Absorb input data into the XOF.
    pub fn absorb(&mut self, data: &[u8]) {
        self.state.absorb(data);
    }

    /// Squeeze output bytes from the XOF.
    pub fn squeeze(&mut self, out: &mut [u8]) {
        self.state.squeeze(out);
    }
}