krypteia-quantica 0.1.0

//! Shared `Keccak-f[1600]` permutation and sponge core (FIPS 202).
//!
//! This module provides the single, deduplicated implementation of the
//! Keccak permutation and the absorb/squeeze sponge state used by all
//! three quantica algorithms (ML-KEM, ML-DSA, SLH-DSA). Each algorithm
//! still ships its own thin `sha3` wrapper module exposing the
//! algorithm-specific high-level functions (`h`, `g`, `j`, `prf`,
//! `Xof` for ML-KEM; `sha3_256`, `sha3_512`, `shake128`, `shake256`,
//! … for ML-DSA; `Shake256`, `shake256_into`, … for SLH-DSA), but they
//! all build on top of the [`KeccakState`] defined here.
//!
//! Pure Rust, no external dependencies.

const KECCAK_ROUNDS: usize = 24;

const RC: [u64; 24] = [
    0x0000000000000001,
    0x0000000000008082,
    0x800000000000808A,
    0x8000000080008000,
    0x000000000000808B,
    0x0000000080000001,
    0x8000000080008081,
    0x8000000000008009,
    0x000000000000008A,
    0x0000000000000088,
    0x0000000080008009,
    0x000000008000000A,
    0x000000008000808B,
    0x800000000000008B,
    0x8000000000008089,
    0x8000000000008003,
    0x8000000000008002,
    0x8000000000000080,
    0x000000000000800A,
    0x800000008000000A,
    0x8000000080008081,
    0x8000000000008080,
    0x0000000080000001,
    0x8000000080008008,
];

const ROTC: [u32; 24] = [
    1, 3, 6, 10, 15, 21, 28, 36, 45, 55, 2, 14, 27, 41, 56, 8, 25, 43, 62, 18, 39, 61, 20, 44,
];

const PI: [usize; 24] = [
    10, 7, 11, 17, 18, 3, 5, 16, 8, 21, 24, 4, 15, 23, 19, 13, 12, 2, 20, 14, 22, 9, 6, 1,
];

#[inline(always)]
fn keccak_f(state: &mut [u64; 25]) {
    for round in 0..KECCAK_ROUNDS {
        // θ step
        let mut c = [0u64; 5];
        for x in 0..5 {
            c[x] = state[x] ^ state[x + 5] ^ state[x + 10] ^ state[x + 15] ^ state[x + 20];
        }
        let mut d = [0u64; 5];
        for x in 0..5 {
            d[x] = c[(x + 4) % 5] ^ c[(x + 1) % 5].rotate_left(1);
        }
        for i in 0..25 {
            state[i] ^= d[i % 5];
        }

        // ρ and π steps
        let mut last = state[1];
        for i in 0..24 {
            let j = PI[i];
            let temp = state[j];
            state[j] = last.rotate_left(ROTC[i]);
            last = temp;
        }

        // χ step
        for y in (0..25).step_by(5) {
            let t0 = state[y];
            let t1 = state[y + 1];
            let t2 = state[y + 2];
            let t3 = state[y + 3];
            let t4 = state[y + 4];
            state[y] = t0 ^ (!t1 & t2);
            state[y + 1] = t1 ^ (!t2 & t3);
            state[y + 2] = t2 ^ (!t3 & t4);
            state[y + 3] = t3 ^ (!t4 & t0);
            state[y + 4] = t4 ^ (!t0 & t1);
        }

        // ι step
        state[0] ^= RC[round];
    }
}

/// Keccak sponge state for SHA-3 and SHAKE constructions.
///
/// Implements the sponge construction over `Keccak-f[1600]` with
/// configurable rate and domain-separation suffix. Supports the
/// absorb-then-squeeze paradigm; once squeezing has begun, further
/// absorption is not permitted.
pub struct KeccakState {
    state: [u64; 25],
    /// Current byte offset in the rate portion.
    offset: usize,
    /// Rate in bytes.
    rate: usize,
    /// Domain separation / padding byte.
    suffix: u8,
    /// Whether we've already switched to squeezing.
    squeezing: bool,
}

impl KeccakState {
    /// Create a new Keccak sponge state.
    ///
    /// * `rate`   — the sponge rate in bytes
    ///   (e.g., 168 for SHAKE128, 136 for SHA3-256/SHAKE256, 72 for SHA3-512).
    /// * `suffix` — the domain-separation/padding byte (0x06 for SHA-3, 0x1f for SHAKE).
    pub fn new(rate: usize, suffix: u8) -> Self {
        Self {
            state: [0u64; 25],
            offset: 0,
            rate,
            suffix,
            squeezing: false,
        }
    }

    /// Absorb input bytes into the sponge.
    ///
    /// May be called multiple times before squeezing to incrementally
    /// feed data. XORs input into the rate portion of the state and
    /// applies the Keccak-f permutation whenever a full rate block is
    /// filled.
    ///
    /// # Panics
    ///
    /// Debug-panics if called after squeezing has begun.
    pub fn absorb(&mut self, data: &[u8]) {
        debug_assert!(!self.squeezing, "Cannot absorb after squeezing");
        let mut pos = 0;
        while pos < data.len() {
            let block_remaining = self.rate - self.offset;
            let to_copy = block_remaining.min(data.len() - pos);

            // XOR input into state bytes.
            let state_bytes = state_as_bytes_mut(&mut self.state);
            for i in 0..to_copy {
                state_bytes[self.offset + i] ^= data[pos + i];
            }
            self.offset += to_copy;
            pos += to_copy;

            if self.offset == self.rate {
                keccak_f(&mut self.state);
                self.offset = 0;
            }
        }
    }

    /// Finalize absorption and switch to squeezing.
    fn finalize(&mut self) {
        if !self.squeezing {
            let state_bytes = state_as_bytes_mut(&mut self.state);
            // Padding: suffix byte then final bit.
            state_bytes[self.offset] ^= self.suffix;
            state_bytes[self.rate - 1] ^= 0x80;
            keccak_f(&mut self.state);
            self.offset = 0;
            self.squeezing = true;
        }
    }

    /// Squeeze output bytes from the sponge.
    ///
    /// On the first call, finalizes absorption by applying padding and
    /// running Keccak-f. Subsequent calls continue squeezing, applying
    /// Keccak-f each time the rate buffer is exhausted. May be called
    /// multiple times to produce arbitrary-length output (XOF mode).
    pub fn squeeze(&mut self, out: &mut [u8]) {
        self.finalize();
        let mut pos = 0;
        while pos < out.len() {
            if self.offset == self.rate {
                keccak_f(&mut self.state);
                self.offset = 0;
            }
            let available = self.rate - self.offset;
            let to_copy = available.min(out.len() - pos);

            // Fast path: read by u64 lanes when aligned.
            let mut i = 0;
            while i + 8 <= to_copy && (self.offset + i) % 8 == 0 {
                let lane = (self.offset + i) / 8;
                let bytes = self.state[lane].to_le_bytes();
                out[pos + i..pos + i + 8].copy_from_slice(&bytes);
                i += 8;
            }
            // Slow path: remaining bytes.
            let state_bytes = state_as_bytes(&self.state);
            while i < to_copy {
                out[pos + i] = state_bytes[self.offset + i];
                i += 1;
            }

            self.offset += to_copy;
            pos += to_copy;
        }
    }
}

#[inline]
fn state_as_bytes(state: &[u64; 25]) -> &[u8; 200] {
    unsafe { &*(state.as_ptr() as *const [u8; 200]) }
}

#[inline]
fn state_as_bytes_mut(state: &mut [u64; 25]) -> &mut [u8; 200] {
    unsafe { &mut *(state.as_mut_ptr() as *mut [u8; 200]) }
}

// ============================================================
// Standard FIPS 202 rates (in bytes).
// ============================================================

/// SHA3-256: rate = (1600 − 2·256) / 8 = 136 bytes, suffix = 0x06.
pub const SHA3_256_RATE: usize = 136;
/// SHA3-512: rate = (1600 − 2·512) / 8 = 72 bytes, suffix = 0x06.
pub const SHA3_512_RATE: usize = 72;
/// SHAKE128: rate = (1600 − 2·128) / 8 = 168 bytes, suffix = 0x1f.
pub const SHAKE128_RATE: usize = 168;
/// SHAKE256: rate = (1600 − 2·256) / 8 = 136 bytes, suffix = 0x1f.
pub const SHAKE256_RATE: usize = 136;

#[cfg(test)]
mod tests {
    use super::*;

    fn sha3_256(input: &[u8]) -> [u8; 32] {
        let mut s = KeccakState::new(SHA3_256_RATE, 0x06);
        s.absorb(input);
        let mut out = [0u8; 32];
        s.squeeze(&mut out);
        out
    }

    fn sha3_512(input: &[u8]) -> [u8; 64] {
        let mut s = KeccakState::new(SHA3_512_RATE, 0x06);
        s.absorb(input);
        let mut out = [0u8; 64];
        s.squeeze(&mut out);
        out
    }

    #[test]
    fn sha3_256_empty_kat() {
        // SHA3-256("") = a7ffc6f8bf1ed76651c14756a061d662f580ff4de43b49fa82d80a4b80f8434a
        let out = sha3_256(b"");
        let expected = [
            0xa7, 0xff, 0xc6, 0xf8, 0xbf, 0x1e, 0xd7, 0x66, 0x51, 0xc1, 0x47, 0x56, 0xa0, 0x61, 0xd6, 0x62, 0xf5, 0x80,
            0xff, 0x4d, 0xe4, 0x3b, 0x49, 0xfa, 0x82, 0xd8, 0x0a, 0x4b, 0x80, 0xf8, 0x43, 0x4a,
        ];
        assert_eq!(out, expected);
    }

    #[test]
    fn sha3_512_empty_kat_first_32() {
        let out = sha3_512(b"");
        let expected_first_32 = [
            0xa6, 0x9f, 0x73, 0xcc, 0xa2, 0x3a, 0x9a, 0xc5, 0xc8, 0xb5, 0x67, 0xdc, 0x18, 0x5a, 0x75, 0x6e, 0x97, 0xc9,
            0x82, 0x16, 0x4f, 0xe2, 0x58, 0x59, 0xe0, 0xd1, 0xdc, 0xc1, 0x47, 0x5c, 0x80, 0xa6,
        ];
        assert_eq!(&out[..32], &expected_first_32);
    }
}