krypteia-arcana 0.1.0

//! SHA-3 and SHAKE implementations (FIPS 202).
//!
//! Wraps the Keccak-f\[1600\] sponge construction to provide:
//! - SHA3-256, SHA3-384, SHA3-512 (fixed-output, implementing `Hasher`)
//! - SHAKE128, SHAKE256 (extendable-output, implementing `Xof`)

use crate::Hasher;
use crate::Xof;

// ============================================================
// Keccak permutation (copied from ml_kem/src/sha3.rs)
// ============================================================

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 {
        // 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[round];
    }
}

#[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
// ============================================================

#[derive(Clone)]
pub(crate) struct KeccakState {
    state: [u64; 25],
    offset: usize,
    rate: usize,
    suffix: u8,
    squeezing: bool,
}

impl KeccakState {
    pub(crate) fn new(rate: usize, suffix: u8) -> Self {
        Self {
            state: [0u64; 25],
            offset: 0,
            rate,
            suffix,
            squeezing: false,
        }
    }

    pub(crate) 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;
        }
    }

    pub(crate) 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
// ============================================================

/// SHA3-224: rate = 1152 bits (144 bytes), capacity = 448 bits, output = 28 bytes.
#[derive(Clone)]
pub struct Sha3_224 {
    state: KeccakState,
}

impl Hasher for Sha3_224 {
    const OUTPUT_LEN: usize = 28;
    const BLOCK_LEN: usize = 144;

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

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

    fn finalize(self) -> Vec<u8> {
        let mut out = vec![0u8; 28];
        self.finalize_into(&mut out);
        out
    }

    fn finalize_into(mut self, out: &mut [u8]) {
        let len = out.len().min(28);
        let mut buf = [0u8; 28];
        self.state.squeeze(&mut buf);
        out[..len].copy_from_slice(&buf[..len]);
    }
}

// ============================================================
// SHA3-256
// ============================================================

/// SHA3-256: rate = 1088 bits (136 bytes), capacity = 512 bits, output = 32 bytes.
#[derive(Clone)]
pub struct Sha3_256 {
    state: KeccakState,
}

impl Hasher for Sha3_256 {
    const OUTPUT_LEN: usize = 32;
    const BLOCK_LEN: usize = 136;

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

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

    fn finalize(self) -> Vec<u8> {
        let mut out = vec![0u8; 32];
        self.finalize_into(&mut out);
        out
    }

    fn finalize_into(mut self, out: &mut [u8]) {
        let len = out.len().min(32);
        let mut buf = [0u8; 32];
        self.state.squeeze(&mut buf);
        out[..len].copy_from_slice(&buf[..len]);
    }
}

// ============================================================
// SHA3-384
// ============================================================

/// SHA3-384: rate = 832 bits (104 bytes), capacity = 768 bits, output = 48 bytes.
#[derive(Clone)]
pub struct Sha3_384 {
    state: KeccakState,
}

impl Hasher for Sha3_384 {
    const OUTPUT_LEN: usize = 48;
    const BLOCK_LEN: usize = 104;

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

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

    fn finalize(self) -> Vec<u8> {
        let mut out = vec![0u8; 48];
        self.finalize_into(&mut out);
        out
    }

    fn finalize_into(mut self, out: &mut [u8]) {
        let len = out.len().min(48);
        let mut buf = [0u8; 48];
        self.state.squeeze(&mut buf);
        out[..len].copy_from_slice(&buf[..len]);
    }
}

// ============================================================
// SHA3-512
// ============================================================

/// SHA3-512: rate = 576 bits (72 bytes), capacity = 1024 bits, output = 64 bytes.
#[derive(Clone)]
pub struct Sha3_512 {
    state: KeccakState,
}

impl Hasher for Sha3_512 {
    const OUTPUT_LEN: usize = 64;
    const BLOCK_LEN: usize = 72;

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

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

    fn finalize(self) -> Vec<u8> {
        let mut out = vec![0u8; 64];
        self.finalize_into(&mut out);
        out
    }

    fn finalize_into(mut self, out: &mut [u8]) {
        let len = out.len().min(64);
        let mut buf = [0u8; 64];
        self.state.squeeze(&mut buf);
        out[..len].copy_from_slice(&buf[..len]);
    }
}

// ============================================================
// SHAKE128
// ============================================================

/// SHAKE128: rate = 1344 bits (168 bytes), capacity = 256 bits, extendable output.
#[derive(Clone)]
pub struct Shake128 {
    state: KeccakState,
}

impl Xof for Shake128 {
    const BLOCK_LEN: usize = 168;

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

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

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

// ============================================================
// SHAKE256
// ============================================================

/// SHAKE256: rate = 1088 bits (136 bytes), capacity = 512 bits, extendable output.
#[derive(Clone)]
pub struct Shake256 {
    state: KeccakState,
}

impl Xof for Shake256 {
    const BLOCK_LEN: usize = 136;

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

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

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

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

/// cSHAKE128: customizable SHAKE with 128-bit security (NIST SP 800-185).
///
/// If both `function_name` and `customization` are empty, cSHAKE128
/// reduces to SHAKE128 (the domain-separation prefix is omitted).
#[derive(Clone)]
pub struct CShake128 {
    state: KeccakState,
}

impl CShake128 {
    /// Rate in bytes (same as SHAKE128).
    pub const RATE: usize = 168;

    /// Create a new cSHAKE128 instance with a function name and
    /// customization string.
    pub fn new(function_name: &[u8], customization: &[u8]) -> Self {
        let mut state = KeccakState::new(
            168,
            if function_name.is_empty() && customization.is_empty() {
                0x1F
            } else {
                0x04
            },
        );
        if !function_name.is_empty() || !customization.is_empty() {
            let prefix = cshake_prefix(168, 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. Can be called multiple times.
    pub fn squeeze(&mut self, out: &mut [u8]) {
        self.state.squeeze(out);
    }
}

/// cSHAKE256: customizable SHAKE with 256-bit security (NIST SP 800-185).
///
/// If both `function_name` and `customization` are empty, cSHAKE256
/// reduces to SHAKE256 (the domain-separation prefix is omitted).
#[derive(Clone)]
pub struct CShake256 {
    state: KeccakState,
}

impl CShake256 {
    /// Rate in bytes (same as SHAKE256).
    pub const RATE: usize = 136;

    /// Create a new cSHAKE256 instance with a function name and
    /// customization string.
    pub fn new(function_name: &[u8], customization: &[u8]) -> Self {
        let mut state = KeccakState::new(
            136,
            if function_name.is_empty() && customization.is_empty() {
                0x1F
            } else {
                0x04
            },
        );
        if !function_name.is_empty() || !customization.is_empty() {
            let prefix = cshake_prefix(136, 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. Can be called multiple times.
    pub fn squeeze(&mut self, out: &mut [u8]) {
        self.state.squeeze(out);
    }
}

/// Build the cSHAKE domain-separation prefix:
/// `bytepad(encode_string(N) || encode_string(S), rate)`.
fn cshake_prefix(rate: usize, function_name: &[u8], customization: &[u8]) -> Vec<u8> {
    let mut buf = Vec::new();
    // left_encode(rate)
    let le = left_encode(rate as u64);
    buf.extend_from_slice(&le);
    // encode_string(N) = left_encode(len_bits(N)) || N
    append_encode_string(&mut buf, function_name);
    // encode_string(S) = left_encode(len_bits(S)) || S
    append_encode_string(&mut buf, customization);
    // Pad to a multiple of rate.
    let pad = (rate - (buf.len() % rate)) % rate;
    buf.extend(core::iter::repeat(0u8).take(pad));
    buf
}

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
}

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

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

    #[test]
    fn test_sha3_224_empty() {
        let digest = Sha3_224::hash(b"");
        let expected: [u8; 28] = [
            0x6b, 0x4e, 0x03, 0x42, 0x36, 0x67, 0xdb, 0xb7, 0x3b, 0x6e, 0x15, 0x45, 0x4f, 0x0e, 0xb1, 0xab, 0xd4, 0x59,
            0x7f, 0x9a, 0x1b, 0x07, 0x8e, 0x3f, 0x5b, 0x5a, 0x6b, 0xc7,
        ];
        assert_eq!(&digest[..], &expected[..]);
    }

    #[test]
    fn test_sha3_224_abc() {
        let digest = Sha3_224::hash(b"abc");
        let expected: [u8; 28] = [
            0xe6, 0x42, 0x82, 0x4c, 0x3f, 0x8c, 0xf2, 0x4a, 0xd0, 0x92, 0x34, 0xee, 0x7d, 0x3c, 0x76, 0x6f, 0xc9, 0xa3,
            0xa5, 0x16, 0x8d, 0x0c, 0x94, 0xad, 0x73, 0xb4, 0x6f, 0xdf,
        ];
        assert_eq!(&digest[..], &expected[..]);
    }

    #[test]
    fn test_sha3_256_empty() {
        let digest = Sha3_256::hash(b"");
        let expected: [u8; 32] = [
            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!(&digest[..], &expected[..]);
    }

    #[test]
    fn test_sha3_256_abc() {
        let digest = Sha3_256::hash(b"abc");
        let expected: [u8; 32] = [
            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,
        ];
        assert_eq!(&digest[..], &expected[..]);
    }

    #[test]
    fn test_sha3_512_empty() {
        let digest = Sha3_512::hash(b"");
        let expected_first_8: [u8; 8] = [0xa6, 0x9f, 0x73, 0xcc, 0xa2, 0x3a, 0x9a, 0xc5];
        assert_eq!(&digest[..8], &expected_first_8[..]);
    }

    #[test]
    fn test_sha3_384_abc() {
        let digest = Sha3_384::hash(b"abc");
        let expected: [u8; 48] = [
            0xec, 0x01, 0x49, 0x82, 0x88, 0x51, 0x6f, 0xc9, 0x26, 0x45, 0x9f, 0x58, 0xe2, 0xc6, 0xad, 0x8d, 0xf9, 0xb4,
            0x73, 0xcb, 0x0f, 0xc0, 0x8c, 0x25, 0x96, 0xda, 0x7c, 0xf0, 0xe4, 0x9b, 0xe4, 0xb2, 0x98, 0xd8, 0x8c, 0xea,
            0x92, 0x7a, 0xc7, 0xf5, 0x39, 0xf1, 0xed, 0xf2, 0x28, 0x37, 0x6d, 0x25,
        ];
        assert_eq!(&digest[..], &expected[..]);
    }

    // cSHAKE test vectors from NIST SP 800-185 §4.

    #[test]
    fn test_cshake128_empty_custom_is_shake128() {
        // With empty N and S, cSHAKE128 == SHAKE128.
        let data = [0x00u8; 4];
        let mut cshake = CShake128::new(b"", b"");
        cshake.update(&data);
        let mut out1 = [0u8; 32];
        cshake.squeeze(&mut out1);

        let mut shake = Shake128::new();
        shake.update(&data);
        let mut out2 = [0u8; 32];
        shake.squeeze(&mut out2);

        assert_eq!(out1, out2);
    }

    #[test]
    fn test_cshake128_sample1() {
        // NIST SP 800-185 §4, Sample #3: N="", S="Email Signature",
        // X=0x00010203, output 256 bits.
        let data = [0x00, 0x01, 0x02, 0x03];
        let mut c = CShake128::new(b"", b"Email Signature");
        c.update(&data);
        let mut out = [0u8; 32];
        c.squeeze(&mut out);
        let expected: [u8; 32] = [
            0xC1, 0xC3, 0x69, 0x25, 0xB6, 0x40, 0x9A, 0x04, 0xF1, 0xB5, 0x04, 0xFC, 0xBC, 0xA9, 0xD8, 0x2B, 0x40, 0x17,
            0x27, 0x7C, 0xB5, 0xED, 0x2B, 0x20, 0x65, 0xFC, 0x1D, 0x38, 0x14, 0xD5, 0xAA, 0xF5,
        ];
        assert_eq!(out, expected);
    }

    #[test]
    fn test_cshake256_sample1() {
        // NIST SP 800-185 §4, Sample #5: N="", S="Email Signature",
        // X=0x00010203, output 256 bits.
        let data = [0x00, 0x01, 0x02, 0x03];
        let mut c = CShake256::new(b"", b"Email Signature");
        c.update(&data);
        let mut out = [0u8; 32];
        c.squeeze(&mut out);
        let expected: [u8; 32] = [
            0xD0, 0x08, 0x82, 0x8E, 0x2B, 0x80, 0xAC, 0x9D, 0x22, 0x18, 0xFF, 0xEE, 0x1D, 0x07, 0x0C, 0x48, 0xB8, 0xE4,
            0xC8, 0x7B, 0xFF, 0x32, 0xC9, 0x69, 0x9D, 0x5B, 0x68, 0x96, 0xEE, 0xE0, 0xED, 0xD1,
        ];
        assert_eq!(out, expected);
    }
}