krypteia-tessera 0.2.0

// SPDX-License-Identifier: Apache-2.0
// Copyright 2026 Cédric Mesnil <cslashm@pm.me>

//! SHA-3 and SHAKE (FIPS 202) over the `Keccak-f[1600]` sponge.
//!
//! Provides SHA3-224/256/384/512 ([`Digest`]), SHAKE128/256 ([`Xof`]),
//! and cSHAKE128/256 (NIST SP 800-185). This is the single workspace copy
//! of Keccak, shared by `quantica` (PQC) and `arcana` (classical).
//!
//! Endianness: the sponge state is XOR-ed as raw little-endian lane bytes
//! (FIPS 202 lane order). All krypteia targets are little-endian (x86_64,
//! aarch64, thumbv6m/7em/8m, riscv32/64 LE); a big-endian host would need
//! a byte-swap shim in [`KeccakState`].

use crate::{Digest, Xof};
use alloc::vec::Vec;

// ============================================================
// Keccak-f[1600] permutation
// ============================================================

// Keccak-f[1600] runs 24 rounds — one per round constant in `RC`.
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 &rc in &RC {
        // theta
        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];
        }

        // rho and pi
        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;
        }

        // chi
        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);
        }

        // iota
        state[0] ^= rc;
    }
}

// SAFETY: `[u64; 25]` and `[u8; 200]` have identical size (200 bytes); a
// `&[u64; 25]` is more strictly aligned than `&[u8; 200]`, so reborrowing
// it as bytes is sound. On a little-endian target the byte view matches
// FIPS 202 lane order (see the module note). These helpers avoid a
// per-lane to_le_bytes copy on the Keccak hot path.
#[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]) }
}

// ============================================================
// Keccak sponge state
// ============================================================

/// A `Keccak-f[1600]` sponge in absorbing/squeezing mode.
#[derive(Clone)]
pub struct KeccakState {
    state: [u64; 25],
    offset: usize,
    rate: usize,
    suffix: u8,
    squeezing: bool,
}

impl KeccakState {
    /// New sponge with the given `rate` (bytes) and domain `suffix`
    /// (0x06 for SHA-3, 0x1f for SHAKE, 0x04 for cSHAKE).
    pub fn new(rate: usize, suffix: u8) -> Self {
        Self {
            state: [0u64; 25],
            offset: 0,
            rate,
            suffix,
            squeezing: false,
        }
    }

    /// Absorb input. Must not be called after squeezing has begun.
    pub fn absorb(&mut self, data: &[u8]) {
        debug_assert!(!self.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);
            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;
            }
        }
    }

    fn pad_and_squeeze(&mut self) {
        if !self.squeezing {
            let state_bytes = state_as_bytes_mut(&mut self.state);
            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. Pads on the first call; may be called repeatedly.
    pub fn squeeze(&mut self, out: &mut [u8]) {
        self.pad_and_squeeze();
        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);
            let state_bytes = state_as_bytes(&self.state);
            out[pos..pos + to_copy].copy_from_slice(&state_bytes[self.offset..self.offset + to_copy]);
            self.offset += to_copy;
            pos += to_copy;
        }
    }
}

// ============================================================
// SHA3-224/256/384/512 (Digest)
// ============================================================

macro_rules! sha3_digest {
    ($name:ident, $out:expr, $rate:expr, $doc:literal) => {
        #[doc = $doc]
        #[derive(Clone)]
        pub struct $name {
            state: KeccakState,
        }

        impl Digest for $name {
            const OUTPUT_LEN: usize = $out;
            const BLOCK_LEN: usize = $rate;

            fn new() -> Self {
                Self {
                    state: KeccakState::new($rate, 0x06),
                }
            }

            fn update(&mut self, data: &[u8]) {
                self.state.absorb(data);
            }

            fn finalize(mut self, out: &mut [u8]) {
                let mut buf = [0u8; $out];
                self.state.squeeze(&mut buf);
                out[..$out].copy_from_slice(&buf);
            }
        }
    };
}

sha3_digest!(Sha3_224, 28, 144, "SHA3-224 (rate 144, output 28 octets).");
sha3_digest!(Sha3_256, 32, 136, "SHA3-256 (rate 136, output 32 octets).");
sha3_digest!(Sha3_384, 48, 104, "SHA3-384 (rate 104, output 48 octets).");
sha3_digest!(Sha3_512, 64, 72, "SHA3-512 (rate 72, output 64 octets).");

// ============================================================
// SHAKE128 / SHAKE256 (Xof)
// ============================================================

macro_rules! shake_xof {
    ($name:ident, $rate:expr, $doc:literal) => {
        #[doc = $doc]
        #[derive(Clone)]
        pub struct $name {
            state: KeccakState,
        }

        impl Xof for $name {
            const BLOCK_LEN: usize = $rate;

            fn new() -> Self {
                Self {
                    state: KeccakState::new($rate, 0x1f),
                }
            }

            fn update(&mut self, data: &[u8]) {
                self.state.absorb(data);
            }

            fn squeeze(&mut self, out: &mut [u8]) {
                self.state.squeeze(out);
            }
        }
    };
}

shake_xof!(Shake128, 168, "SHAKE128 XOF (rate 168, 128-bit security).");
shake_xof!(Shake256, 136, "SHAKE256 XOF (rate 136, 256-bit security).");

// ============================================================
// cSHAKE128 / cSHAKE256 (NIST SP 800-185)
// ============================================================

macro_rules! cshake {
    ($name:ident, $rate:expr, $doc:literal) => {
        #[doc = $doc]
        #[derive(Clone)]
        pub struct $name {
            state: KeccakState,
        }

        impl $name {
            /// Rate in octets.
            pub const RATE: usize = $rate;

            /// New cSHAKE with a function name `N` and customization `S`.
            /// With both empty it reduces to the plain SHAKE of the same
            /// security level.
            pub fn new(function_name: &[u8], customization: &[u8]) -> Self {
                let suffix = if function_name.is_empty() && customization.is_empty() {
                    0x1f
                } else {
                    0x04
                };
                let mut state = KeccakState::new($rate, suffix);
                if !function_name.is_empty() || !customization.is_empty() {
                    let prefix = cshake_prefix($rate, function_name, customization);
                    state.absorb(&prefix);
                }
                Self { state }
            }

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

            /// Squeeze output bytes; may be called repeatedly.
            pub fn squeeze(&mut self, out: &mut [u8]) {
                self.state.squeeze(out);
            }
        }
    };
}

cshake!(CShake128, 168, "cSHAKE128 (NIST SP 800-185).");
cshake!(CShake256, 136, "cSHAKE256 (NIST SP 800-185).");

/// `bytepad(encode_string(N) || encode_string(S), rate)`.
///
/// NIST SP 800-185 §3.2 (cSHAKE definition).
fn cshake_prefix(rate: usize, function_name: &[u8], customization: &[u8]) -> Vec<u8> {
    let mut buf = Vec::new();
    append_encode_string(&mut buf, function_name);
    append_encode_string(&mut buf, customization);
    bytepad(&buf, rate)
}

/// `left_encode(x)` — NIST SP 800-185 §2.3.1.
///
/// Prefixes the minimal big-endian byte representation of `x` with a
/// single byte giving that representation's length; `left_encode(0)`
/// is `[0x01, 0x00]`.
pub(crate) fn left_encode(x: u64) -> Vec<u8> {
    let mut bytes = Vec::new();
    let be = x.to_be_bytes();
    let first = be.iter().position(|&b| b != 0).unwrap_or(7);
    let n = (8 - first) as u8;
    bytes.push(n);
    bytes.extend_from_slice(&be[first..]);
    bytes
}

/// `right_encode(x)` — NIST SP 800-185 §2.3.1.
///
/// Same as [`left_encode`] but with the length byte appended after the
/// big-endian representation; `right_encode(0)` is `[0x00, 0x01]`.
pub(crate) fn right_encode(x: u64) -> Vec<u8> {
    let mut bytes = Vec::new();
    let be = x.to_be_bytes();
    let first = be.iter().position(|&b| b != 0).unwrap_or(7);
    let n = (8 - first) as u8;
    bytes.extend_from_slice(&be[first..]);
    bytes.push(n);
    bytes
}

/// `encode_string(S)` = `left_encode(len_in_bits(S)) || S`.
///
/// NIST SP 800-185 §2.3.2. Note the length is encoded in **bits**.
pub(crate) fn append_encode_string(out: &mut Vec<u8>, s: &[u8]) {
    out.extend_from_slice(&left_encode((s.len() as u64) * 8));
    out.extend_from_slice(s);
}

/// `bytepad(data, w)` — NIST SP 800-185 §2.3.3.
///
/// Returns `left_encode(w) || data`, zero-padded up to a multiple of
/// `w` octets. `w` (the rate) must be non-zero.
pub(crate) fn bytepad(data: &[u8], w: usize) -> Vec<u8> {
    let prefix = left_encode(w as u64);
    let mut out = Vec::with_capacity(prefix.len() + data.len());
    out.extend_from_slice(&prefix);
    out.extend_from_slice(data);
    let pad = (w - (out.len() % w)) % w;
    out.resize(out.len() + pad, 0u8);
    out
}

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

    // Smoke KATs. Exhaustive SHA3-224/256/384/512, SHAKE128/256 and
    // cSHAKE coverage (NIST ACVP + SP 800-185) lives in the integration
    // suites `tests/acvp_hashes.rs` and `tests/sp800_185.rs`.

    #[test]
    fn sha3_256_smoke() {
        // FIPS 202 "abc".
        let mut out = [0u8; 32];
        Sha3_256::digest(b"abc", &mut out);
        assert_eq!(
            out,
            [
                0x3a, 0x98, 0x5d, 0xa7, 0x4f, 0xe2, 0x25, 0xb2, 0x04, 0x5c, 0x17, 0x2d, 0x6b, 0xd3, 0x90, 0xbd, 0x85,
                0x5f, 0x08, 0x6e, 0x3e, 0x9d, 0x52, 0x5b, 0x46, 0xbf, 0xe2, 0x45, 0x11, 0x43, 0x15, 0x32,
            ]
        );
    }

    #[test]
    fn shake256_smoke() {
        // SHAKE256("", 32) known prefix (FIPS 202).
        let mut x = Shake256::new();
        x.update(b"");
        let mut out = [0u8; 32];
        x.squeeze(&mut out);
        assert_eq!(out[..4], [0x46, 0xb9, 0xdd, 0x2b]);
    }
}