krypteia-tessera 0.2.0

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

//! BLAKE3 cryptographic hash / XOF (BLAKE3 spec, 2020-01-09; reference
//! `blake3.py` / `reference_impl.rs` from the BLAKE3 team).
//!
//! BLAKE3 is a Merkle-tree hash built on a keyed compression function
//! derived from BLAKE2s (ARX, 7 rounds). It provides three modes:
//!
//! * **hash** — unkeyed ([`hash`], [`Blake3::new`]).
//! * **keyed_hash** — 32-byte key ([`keyed_hash`], [`Blake3::new_keyed`]),
//!   a PRF / MAC.
//! * **derive_key** — context-string KDF ([`derive_key`],
//!   [`Blake3::new_derive_key`]).
//!
//! The output is an XOF: the default is 32 bytes ([`Digest`]), but
//! [`Blake3::finalize_xof`] streams an arbitrary-length, seekable output.
//!
//! # Side-channel posture
//!
//! BLAKE3 is a pure ARX design (add / rotate / xor over 32-bit words). The
//! whole state transition, the message schedule (a fixed permutation), the
//! chunk/tree bookkeeping, and the output expansion are **data-oblivious**:
//! no secret-dependent branch, no secret-dependent memory index, no table
//! lookup. In `keyed_hash` / `derive_key` the key is only folded into the
//! IV / chaining state before the same oblivious permutation runs, so the
//! secret never steers control flow or addressing. Timing is a function of
//! input *length* only, which is public. No `silentops` CT primitive is
//! required here; secret-dependent comparison / zeroization of a BLAKE3
//! output (e.g. MAC verification) is the caller's responsibility and must
//! use `silentops::ct_eq`.

use crate::Digest;

// ====================================================================
// Constants (BLAKE3 spec §2.1, §5)
// ====================================================================

/// Initialization vector — the SHA-256 IV (first 32 bits of the fractional
/// parts of the square roots of the first 8 primes). BLAKE3 spec §2.1.
const IV: [u32; 8] = [
    0x6A09E667, 0xBB67AE85, 0x3C6EF372, 0xA54FF53A, 0x510E527F, 0x9B05688C, 0x1F83D9AB, 0x5BE0CD19,
];

/// Message word permutation applied between rounds. BLAKE3 spec §2.2.
const MSG_PERMUTATION: [usize; 16] = [2, 6, 3, 10, 7, 0, 4, 13, 1, 11, 12, 5, 9, 14, 15, 8];

/// Bytes per chunk (leaf of the tree).
const CHUNK_LEN: usize = 1024;
/// Bytes per compression block.
const BLOCK_LEN: usize = 64;
/// Output / chaining-value length in bytes.
const OUT_LEN: usize = 32;

// Domain-separation flags (BLAKE3 spec §2.1).
const CHUNK_START: u32 = 1 << 0;
const CHUNK_END: u32 = 1 << 1;
const PARENT: u32 = 1 << 2;
const ROOT: u32 = 1 << 3;
const KEYED_HASH: u32 = 1 << 4;
const DERIVE_KEY_CONTEXT: u32 = 1 << 5;
const DERIVE_KEY_MATERIAL: u32 = 1 << 6;

// ====================================================================
// Compression function (BLAKE3 spec §2.2)
// ====================================================================

/// The G mixing function of BLAKE3 (BLAKE2s rotations 16/12/8/7).
#[inline(always)]
fn g(state: &mut [u32; 16], a: usize, b: usize, c: usize, d: usize, mx: u32, my: u32) {
    state[a] = state[a].wrapping_add(state[b]).wrapping_add(mx);
    state[d] = (state[d] ^ state[a]).rotate_right(16);
    state[c] = state[c].wrapping_add(state[d]);
    state[b] = (state[b] ^ state[c]).rotate_right(12);
    state[a] = state[a].wrapping_add(state[b]).wrapping_add(my);
    state[d] = (state[d] ^ state[a]).rotate_right(8);
    state[c] = state[c].wrapping_add(state[d]);
    state[b] = (state[b] ^ state[c]).rotate_right(7);
}

/// One of the 7 rounds: 4 column mixes then 4 diagonal mixes.
#[inline(always)]
fn round(state: &mut [u32; 16], m: &[u32; 16]) {
    // Columns.
    g(state, 0, 4, 8, 12, m[0], m[1]);
    g(state, 1, 5, 9, 13, m[2], m[3]);
    g(state, 2, 6, 10, 14, m[4], m[5]);
    g(state, 3, 7, 11, 15, m[6], m[7]);
    // Diagonals.
    g(state, 0, 5, 10, 15, m[8], m[9]);
    g(state, 1, 6, 11, 12, m[10], m[11]);
    g(state, 2, 7, 8, 13, m[12], m[13]);
    g(state, 3, 4, 9, 14, m[14], m[15]);
}

#[inline(always)]
fn permute(m: &mut [u32; 16]) {
    let old = *m;
    for i in 0..16 {
        m[i] = old[MSG_PERMUTATION[i]];
    }
}

/// The keyed compression function. Returns the full 16-word state; the
/// low 8 words are the chaining value / output, the high 8 words feed the
/// extended (XOF) output. BLAKE3 spec §2.2.
fn compress(chaining_value: &[u32; 8], block_words: &[u32; 16], counter: u64, block_len: u32, flags: u32) -> [u32; 16] {
    let mut state: [u32; 16] = [
        chaining_value[0],
        chaining_value[1],
        chaining_value[2],
        chaining_value[3],
        chaining_value[4],
        chaining_value[5],
        chaining_value[6],
        chaining_value[7],
        IV[0],
        IV[1],
        IV[2],
        IV[3],
        counter as u32,
        (counter >> 32) as u32,
        block_len,
        flags,
    ];

    let mut m = *block_words;
    round(&mut state, &m); // round 1
    permute(&mut m);
    round(&mut state, &m); // round 2
    permute(&mut m);
    round(&mut state, &m); // round 3
    permute(&mut m);
    round(&mut state, &m); // round 4
    permute(&mut m);
    round(&mut state, &m); // round 5
    permute(&mut m);
    round(&mut state, &m); // round 6
    permute(&mut m);
    round(&mut state, &m); // round 7

    // Feed-forward: XOR the two halves of the state.
    for i in 0..8 {
        state[i] ^= state[i + 8];
        state[i + 8] ^= chaining_value[i];
    }
    state
}

#[inline(always)]
fn first_8(words: &[u32; 16]) -> [u32; 8] {
    [
        words[0], words[1], words[2], words[3], words[4], words[5], words[6], words[7],
    ]
}

/// Read a 64-byte block into 16 little-endian words. The block is
/// zero-padded to 64 bytes by the caller.
#[inline(always)]
fn words_from_block(block: &[u8; BLOCK_LEN]) -> [u32; 16] {
    let mut out = [0u32; 16];
    for (i, w) in out.iter_mut().enumerate() {
        *w = u32::from_le_bytes([block[4 * i], block[4 * i + 1], block[4 * i + 2], block[4 * i + 3]]);
    }
    out
}

// ====================================================================
// Output node — the XOF root / any chaining value (BLAKE3 spec §2.3)
// ====================================================================

/// A node whose output has not yet been produced. Re-compressing it with
/// an incrementing counter and the `ROOT` flag yields the XOF stream.
#[derive(Clone)]
struct Output {
    input_chaining_value: [u32; 8],
    block_words: [u32; 16],
    counter: u64,
    block_len: u32,
    flags: u32,
}

impl Output {
    /// Chaining value of a non-root node (counter 0, no ROOT flag).
    fn chaining_value(&self) -> [u32; 8] {
        first_8(&compress(
            &self.input_chaining_value,
            &self.block_words,
            self.counter,
            self.block_len,
            self.flags,
        ))
    }
}

// ====================================================================
// Chunk state (BLAKE3 spec §2.3) — one leaf of the tree
// ====================================================================

#[derive(Clone)]
struct ChunkState {
    chaining_value: [u32; 8],
    chunk_counter: u64,
    block: [u8; BLOCK_LEN],
    block_len: usize,
    blocks_compressed: u64,
    flags: u32,
}

impl ChunkState {
    fn new(key: &[u32; 8], chunk_counter: u64, flags: u32) -> Self {
        Self {
            chaining_value: *key,
            chunk_counter,
            block: [0u8; BLOCK_LEN],
            block_len: 0,
            blocks_compressed: 0,
            flags,
        }
    }

    fn len(&self) -> usize {
        BLOCK_LEN * self.blocks_compressed as usize + self.block_len
    }

    #[inline(always)]
    fn start_flag(&self) -> u32 {
        if self.blocks_compressed == 0 { CHUNK_START } else { 0 }
    }

    fn update(&mut self, mut input: &[u8]) {
        while !input.is_empty() {
            if self.block_len == BLOCK_LEN {
                // Compress the full block; it is not the chunk's last.
                let block_words = words_from_block(&self.block);
                self.chaining_value = first_8(&compress(
                    &self.chaining_value,
                    &block_words,
                    self.chunk_counter,
                    BLOCK_LEN as u32,
                    self.flags | self.start_flag(),
                ));
                self.blocks_compressed += 1;
                self.block = [0u8; BLOCK_LEN];
                self.block_len = 0;
            }
            let want = BLOCK_LEN - self.block_len;
            let take = want.min(input.len());
            self.block[self.block_len..self.block_len + take].copy_from_slice(&input[..take]);
            self.block_len += take;
            input = &input[take..];
        }
    }

    /// Finalize the chunk into an [`Output`] (its chaining value, or the
    /// XOF root if this chunk is the whole message).
    fn output(&self) -> Output {
        let block_words = words_from_block(&self.block);
        Output {
            input_chaining_value: self.chaining_value,
            block_words,
            counter: self.chunk_counter,
            block_len: self.block_len as u32,
            flags: self.flags | self.start_flag() | CHUNK_END,
        }
    }
}

/// Parent node output: two 8-word children -> one [`Output`]. BLAKE3
/// spec §2.3.
fn parent_output(left: &[u32; 8], right: &[u32; 8], key: &[u32; 8], flags: u32) -> Output {
    let mut block_words = [0u32; 16];
    block_words[..8].copy_from_slice(left);
    block_words[8..].copy_from_slice(right);
    Output {
        input_chaining_value: *key,
        block_words,
        counter: 0,
        block_len: BLOCK_LEN as u32,
        flags: flags | PARENT,
    }
}

fn parent_cv(left: &[u32; 8], right: &[u32; 8], key: &[u32; 8], flags: u32) -> [u32; 8] {
    parent_output(left, right, key, flags).chaining_value()
}

// ====================================================================
// Streaming hasher (BLAKE3 spec §2.3 — the tree / CV-stack)
// ====================================================================

/// Streaming BLAKE3 hasher / XOF.
///
/// Feed input with [`update`](Blake3::update); finalize with
/// [`finalize`](Digest::finalize) (32 bytes) or
/// [`finalize_xof`](Blake3::finalize_xof) (arbitrary length). BLAKE3 spec
/// §2.3.
#[derive(Clone)]
pub struct Blake3 {
    chunk_state: ChunkState,
    key: [u32; 8],
    // Subtree chaining-value stack: at most 54 entries (a 2^64-byte input
    // has depth < 54). Fixed array keeps the hasher `no_std`/stack-only.
    cv_stack: [[u32; 8]; 54],
    cv_stack_len: u8,
    flags: u32,
}

impl Blake3 {
    fn new_internal(key: [u32; 8], flags: u32) -> Self {
        Self {
            chunk_state: ChunkState::new(&key, 0, flags),
            key,
            cv_stack: [[0u32; 8]; 54],
            cv_stack_len: 0,
            flags,
        }
    }

    /// Unkeyed BLAKE3 hasher (default mode). BLAKE3 spec §2.1.
    // `new` intentionally mirrors `new_keyed`/`new_derive_key` (the three BLAKE3
    // modes); a `Default` impl would obscure that symmetry. Scoped lint waiver,
    // not a blanket suppression (CLAUDE.md §6).
    #[allow(clippy::should_implement_trait)]
    pub fn new() -> Self {
        Self::new_internal(IV, 0)
    }

    /// Keyed BLAKE3 (`keyed_hash`) — a PRF / MAC over a 32-byte key.
    /// The key replaces the IV as the chaining state. BLAKE3 spec §2.1.
    pub fn new_keyed(key: &[u8; 32]) -> Self {
        let mut kw = [0u32; 8];
        for (i, w) in kw.iter_mut().enumerate() {
            *w = u32::from_le_bytes([key[4 * i], key[4 * i + 1], key[4 * i + 2], key[4 * i + 3]]);
        }
        Self::new_internal(kw, KEYED_HASH)
    }

    /// `derive_key` KDF context phase. The context string is hashed under
    /// the `DERIVE_KEY_CONTEXT` flag; the resulting 32-byte value becomes
    /// the key for the `DERIVE_KEY_MATERIAL` phase (fed via `update`).
    /// BLAKE3 spec §2.1 / §5.
    pub fn new_derive_key(context: &str) -> Self {
        // Phase 1: hash the context string with the IV, DERIVE_KEY_CONTEXT.
        let mut ctx_hasher = Self::new_internal(IV, DERIVE_KEY_CONTEXT);
        ctx_hasher.update(context.as_bytes());
        let mut context_key = [0u8; 32];
        ctx_hasher.finalize_xof(0, &mut context_key);
        let mut kw = [0u32; 8];
        for (i, w) in kw.iter_mut().enumerate() {
            *w = u32::from_le_bytes([
                context_key[4 * i],
                context_key[4 * i + 1],
                context_key[4 * i + 2],
                context_key[4 * i + 3],
            ]);
        }
        // Phase 2: hasher keyed with the context key, DERIVE_KEY_MATERIAL.
        Self::new_internal(kw, DERIVE_KEY_MATERIAL)
    }

    #[inline(always)]
    fn push_cv(&mut self, cv: [u32; 8]) {
        self.cv_stack[self.cv_stack_len as usize] = cv;
        self.cv_stack_len += 1;
    }

    #[inline(always)]
    fn pop_cv(&mut self) -> [u32; 8] {
        self.cv_stack_len -= 1;
        self.cv_stack[self.cv_stack_len as usize]
    }

    /// Add a completed chunk's chaining value to the tree, merging any
    /// complete subtrees. `total_chunks` is the count of chunks *after*
    /// this one is added; a subtree merges when its size divides the count
    /// (i.e. when the low bit currently being carried is set). BLAKE3
    /// spec §2.3 (the "add_chunk_chaining_value" carry logic).
    fn add_chunk_cv(&mut self, mut new_cv: [u32; 8], mut total_chunks: u64) {
        // While the current total has an even number of subtrees at this
        // level, merge the top of the stack with the new CV.
        while total_chunks & 1 == 0 {
            let left = self.pop_cv();
            new_cv = parent_cv(&left, &new_cv, &self.key, self.flags);
            total_chunks >>= 1;
        }
        self.push_cv(new_cv);
    }

    /// Absorb input. May be called repeatedly. BLAKE3 spec §2.3.
    pub fn update(&mut self, mut input: &[u8]) {
        while !input.is_empty() {
            // If the current chunk is full, finalize it into the tree and
            // start a new one.
            if self.chunk_state.len() == CHUNK_LEN {
                let chunk_cv = self.chunk_state.output().chaining_value();
                let total_chunks = self.chunk_state.chunk_counter + 1;
                self.add_chunk_cv(chunk_cv, total_chunks);
                self.chunk_state = ChunkState::new(&self.key, total_chunks, self.flags);
            }
            let want = CHUNK_LEN - self.chunk_state.len();
            let take = want.min(input.len());
            self.chunk_state.update(&input[..take]);
            input = &input[take..];
        }
    }

    /// Build the root [`Output`]: fold the CV stack down over the final
    /// chunk. BLAKE3 spec §2.3.
    fn root_output(&self) -> Output {
        let mut output = self.chunk_state.output();
        let mut parent_nodes_remaining = self.cv_stack_len as usize;
        while parent_nodes_remaining > 0 {
            parent_nodes_remaining -= 1;
            let left = self.cv_stack[parent_nodes_remaining];
            let right = output.chaining_value();
            output = parent_output(&left, &right, &self.key, self.flags);
        }
        output
    }

    /// Write `out.len()` bytes of XOF output starting at byte offset
    /// `seek`. Each root re-compression yields one 64-byte output block; a
    /// `seek` that does not fall on a block boundary is handled by skipping
    /// the leading bytes of the first block. BLAKE3 spec §2.3.
    pub fn finalize_xof(&self, seek: u64, out: &mut [u8]) {
        if out.is_empty() {
            return;
        }
        let output = self.root_output();
        const BLK: u64 = 2 * OUT_LEN as u64; // 64-byte output block
        let mut block_counter = seek / BLK;
        let mut skip = (seek % BLK) as usize;
        let mut produced = 0usize;
        while produced < out.len() {
            let words = compress(
                &output.input_chaining_value,
                &output.block_words,
                block_counter,
                output.block_len,
                output.flags | ROOT,
            );
            let mut blk = [0u8; BLK as usize];
            for (i, &word) in words.iter().enumerate() {
                blk[i * 4..i * 4 + 4].copy_from_slice(&word.to_le_bytes());
            }
            let avail = &blk[skip..];
            let take = avail.len().min(out.len() - produced);
            out[produced..produced + take].copy_from_slice(&avail[..take]);
            produced += take;
            skip = 0;
            block_counter += 1;
        }
    }
}

impl Default for Blake3 {
    fn default() -> Self {
        Self::new()
    }
}

impl Digest for Blake3 {
    const OUTPUT_LEN: usize = OUT_LEN;
    const BLOCK_LEN: usize = BLOCK_LEN;

    fn new() -> Self {
        Blake3::new()
    }

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

    fn finalize(self, out: &mut [u8]) {
        let len = out.len().min(OUT_LEN);
        let mut full = [0u8; OUT_LEN];
        self.finalize_xof(0, &mut full);
        out[..len].copy_from_slice(&full[..len]);
    }
}

// ====================================================================
// One-shot convenience API (BLAKE3 spec §5)
// ====================================================================

/// One-shot unkeyed BLAKE3, arbitrary-length XOF output. BLAKE3 spec §2.1.
pub fn hash(input: &[u8], out: &mut [u8]) {
    let mut h = Blake3::new();
    h.update(input);
    h.finalize_xof(0, out);
}

/// One-shot keyed BLAKE3 (`keyed_hash`) over a 32-byte key, arbitrary
/// output. A PRF / MAC. BLAKE3 spec §2.1.
pub fn keyed_hash(key: &[u8; 32], input: &[u8], out: &mut [u8]) {
    let mut h = Blake3::new_keyed(key);
    h.update(input);
    h.finalize_xof(0, out);
}

/// One-shot `derive_key` KDF: bind `key_material` to `context`, arbitrary
/// output. The context string should be a hardcoded, application-unique,
/// globally-unique constant. BLAKE3 spec §2.1 / §5.
pub fn derive_key(context: &str, key_material: &[u8], out: &mut [u8]) {
    let mut h = Blake3::new_derive_key(context);
    h.update(key_material);
    h.finalize_xof(0, out);
}

// ====================================================================
// Tests — smoke KAT.
//
// The full pinned upstream corpus
// (github.com/BLAKE3-team/BLAKE3 test_vectors/test_vectors.json) — all
// three modes, every input length, plus mid-block / large-offset XOF
// seek checks — lives in the integration suite `tests/blake3.rs`.
// ====================================================================

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

    fn hex(bytes: &[u8]) -> alloc::string::String {
        use core::fmt::Write;
        let mut s = alloc::string::String::new();
        for b in bytes {
            write!(s, "{:02x}", b).unwrap();
        }
        s
    }

    // BLAKE3 hash of the empty input (default 32-byte output), from the
    // official test_vectors.json input_len = 0 case.
    #[test]
    fn blake3_empty_smoke() {
        let mut out = [0u8; 32];
        hash(b"", &mut out);
        assert_eq!(
            hex(&out),
            "af1349b9f5f9a1a6a0404dea36dcc9499bcb25c9adc112b7cc9a93cae41f3262"
        );
    }
}