krypteia-quantica 0.2.0

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

//! Tweakable hash functions for SLH-DSA (FIPS 205, Section 11).
//!
//! SLH-DSA defines six tweakable hash functions (`H_msg`, `PRF`,
//! `PRF_msg`, `T_l`, `H`, `F`) and two instantiations: SHAKE (§11.1) and
//! SHA-2 (§11.2). The family is selected per parameter set via the
//! [`Params`](crate::slh_dsa::params::Params) `Hash` associated type,
//! which names a [`TweakableHash`](crate::slh_dsa::hash::TweakableHash)
//! implementation ([`ShakeFamily`](crate::slh_dsa::hash::ShakeFamily)
//! here, [`Sha2Family`](crate::slh_dsa::hash_sha2::Sha2Family) for SHA-2).
//!
//! The free functions below are the stable call surface used throughout
//! `slh`/`fors`/`wots`/`xmss`; they simply dispatch to `P::Hash`. The
//! address (`Adrs`) is mixed into every keyed call for domain separation.
//! The `_into` variants write to a caller-provided buffer (zero heap);
//! the plain variants return `Vec<u8>` for convenience.
//!
//! # Stability (hazmat tier)
//!
//! This module is part of the `hazmat` expert building-block layer: it is
//! public only with the `hazmat` Cargo feature, for KAT/test-vector tooling
//! and research. Read each function's invariants before use. **No stability,
//! misuse-resistance, or side-channel promise** applies to this tier — its
//! API may change in any release. The typed facade is the guaranteed API.

use super::address::Adrs;
use super::params::Params;
use super::sha3::{KeccakState, SHAKE256_RATE};
use alloc::vec::Vec;

/// One instantiation of the six SLH-DSA tweakable hash functions
/// (FIPS 205 §11). Implemented by [`ShakeFamily`] (§11.1) and
/// [`Sha2Family`](super::hash_sha2::Sha2Family) (§11.2). Every method is
/// generic over the parameter set `P`, which supplies the output lengths
/// (`P::N`, `P::M`) and — for SHA-2 — the security-category split.
///
/// All outputs are written into the caller's buffer. `H_msg` writes
/// `P::M` bytes; the others write `P::N` bytes.
pub trait TweakableHash {
    /// `H_msg(R, PK.seed, PK.root, M)` — message hash, `P::M` bytes.
    fn h_msg_into<P: Params>(r: &[u8], pk_seed: &[u8], pk_root: &[u8], m: &[u8], out: &mut [u8]);
    /// `PRF(PK.seed, SK.seed, ADRS)` — secret value derivation, `P::N` bytes.
    fn prf_into<P: Params>(pk_seed: &[u8], sk_seed: &[u8], adrs: &Adrs, out: &mut [u8]);
    /// `PRF_msg(SK.prf, opt_rand, M)` — message randomizer, `P::N` bytes.
    fn prf_msg_into<P: Params>(sk_prf: &[u8], opt_rand: &[u8], m: &[u8], out: &mut [u8]);
    /// `T_l(PK.seed, ADRS, M)` — multi-input compression, `P::N` bytes.
    fn t_l_into<P: Params>(pk_seed: &[u8], adrs: &Adrs, m: &[u8], out: &mut [u8]);
    /// `H(PK.seed, ADRS, M1 ‖ M2)` — two-input Merkle node hash, `P::N` bytes.
    fn hash_h_into<P: Params>(pk_seed: &[u8], adrs: &Adrs, m1: &[u8], m2: &[u8], out: &mut [u8]);
    /// `F(PK.seed, ADRS, M1)` — single-input chain hash, `P::N` bytes.
    fn f_hash_into<P: Params>(pk_seed: &[u8], adrs: &Adrs, m1: &[u8], out: &mut [u8]);
}

// ============================================================
// SHAKE instantiation (FIPS 205 §11.1) — all six on SHAKE256.
// ============================================================

/// SHAKE256 instantiation of the tweakable hash family (FIPS 205 §11.1).
pub struct ShakeFamily;

/// Internal: SHAKE256 context with `PK.seed ‖ ADRS` absorbed. Common
/// prefix for PRF, T_l, H, and F.
#[inline(always)]
fn shake256_pk_adrs(pk_seed: &[u8], adrs: &Adrs) -> KeccakState {
    let mut ctx = KeccakState::new(SHAKE256_RATE, 0x1f);
    ctx.absorb(pk_seed);
    ctx.absorb(adrs.as_bytes());
    ctx
}

impl TweakableHash for ShakeFamily {
    fn h_msg_into<P: Params>(r: &[u8], pk_seed: &[u8], pk_root: &[u8], m: &[u8], out: &mut [u8]) {
        let mut ctx = KeccakState::new(SHAKE256_RATE, 0x1f);
        ctx.absorb(r);
        ctx.absorb(pk_seed);
        ctx.absorb(pk_root);
        ctx.absorb(m);
        ctx.squeeze(out);
    }

    fn prf_into<P: Params>(pk_seed: &[u8], sk_seed: &[u8], adrs: &Adrs, out: &mut [u8]) {
        let mut ctx = shake256_pk_adrs(pk_seed, adrs);
        ctx.absorb(sk_seed);
        ctx.squeeze(out);
    }

    fn prf_msg_into<P: Params>(sk_prf: &[u8], opt_rand: &[u8], m: &[u8], out: &mut [u8]) {
        let mut ctx = KeccakState::new(SHAKE256_RATE, 0x1f);
        ctx.absorb(sk_prf);
        ctx.absorb(opt_rand);
        ctx.absorb(m);
        ctx.squeeze(out);
    }

    fn t_l_into<P: Params>(pk_seed: &[u8], adrs: &Adrs, m: &[u8], out: &mut [u8]) {
        let mut ctx = shake256_pk_adrs(pk_seed, adrs);
        ctx.absorb(m);
        ctx.squeeze(out);
    }

    fn hash_h_into<P: Params>(pk_seed: &[u8], adrs: &Adrs, m1: &[u8], m2: &[u8], out: &mut [u8]) {
        let mut ctx = shake256_pk_adrs(pk_seed, adrs);
        ctx.absorb(m1);
        ctx.absorb(m2);
        ctx.squeeze(out);
    }

    fn f_hash_into<P: Params>(pk_seed: &[u8], adrs: &Adrs, m1: &[u8], out: &mut [u8]) {
        let mut ctx = shake256_pk_adrs(pk_seed, adrs);
        ctx.absorb(m1);
        ctx.squeeze(out);
    }
}

// ============================================================
// Free dispatch functions — the stable call surface.
// Each delegates to the parameter set's hash family `P::Hash`.
// ============================================================

/// `H_msg(R, PK.seed, PK.root, M)` — message hash. Output (`P::M` bytes) to `out`.
pub fn h_msg_into<P: Params>(r: &[u8], pk_seed: &[u8], pk_root: &[u8], m: &[u8], out: &mut [u8]) {
    P::Hash::h_msg_into::<P>(r, pk_seed, pk_root, m, out);
}

/// `H_msg` — convenience wrapper returning `Vec`.
pub fn h_msg<P: Params>(r: &[u8], pk_seed: &[u8], pk_root: &[u8], m: &[u8]) -> Vec<u8> {
    let mut out = vec![0u8; P::M];
    h_msg_into::<P>(r, pk_seed, pk_root, m, &mut out);
    out
}

/// `PRF(PK.seed, SK.seed, ADRS)` — secret value derivation. Output (`P::N` bytes) to `out`.
pub fn prf_into<P: Params>(pk_seed: &[u8], sk_seed: &[u8], adrs: &Adrs, out: &mut [u8]) {
    P::Hash::prf_into::<P>(pk_seed, sk_seed, adrs, out);
}

/// `PRF` — convenience wrapper returning `Vec`.
pub fn prf<P: Params>(pk_seed: &[u8], sk_seed: &[u8], adrs: &Adrs) -> Vec<u8> {
    let mut out = vec![0u8; P::N];
    prf_into::<P>(pk_seed, sk_seed, adrs, &mut out);
    out
}

/// `PRF_msg(SK.prf, opt_rand, M)` — message randomizer. Output (`P::N` bytes) to `out`.
pub fn prf_msg_into<P: Params>(sk_prf: &[u8], opt_rand: &[u8], m: &[u8], out: &mut [u8]) {
    P::Hash::prf_msg_into::<P>(sk_prf, opt_rand, m, out);
}

/// `PRF_msg` — convenience wrapper returning `Vec`.
pub fn prf_msg<P: Params>(sk_prf: &[u8], opt_rand: &[u8], m: &[u8]) -> Vec<u8> {
    let mut out = vec![0u8; P::N];
    prf_msg_into::<P>(sk_prf, opt_rand, m, &mut out);
    out
}

/// `T_l(PK.seed, ADRS, M)` — multi-input compression. Output (`P::N` bytes) to `out`.
pub fn t_l_into<P: Params>(pk_seed: &[u8], adrs: &Adrs, m: &[u8], out: &mut [u8]) {
    P::Hash::t_l_into::<P>(pk_seed, adrs, m, out);
}

/// `T_l` — convenience wrapper returning `Vec`.
pub fn t_l<P: Params>(pk_seed: &[u8], adrs: &Adrs, m: &[u8]) -> Vec<u8> {
    let mut out = vec![0u8; P::N];
    t_l_into::<P>(pk_seed, adrs, m, &mut out);
    out
}

/// `H(PK.seed, ADRS, M1 ‖ M2)` — two-input Merkle node hash. Output (`P::N` bytes) to `out`.
pub fn hash_h_into<P: Params>(pk_seed: &[u8], adrs: &Adrs, m1: &[u8], m2: &[u8], out: &mut [u8]) {
    P::Hash::hash_h_into::<P>(pk_seed, adrs, m1, m2, out);
}

/// `H` — convenience wrapper returning `Vec`.
pub fn hash_h<P: Params>(pk_seed: &[u8], adrs: &Adrs, m1: &[u8], m2: &[u8]) -> Vec<u8> {
    let mut out = vec![0u8; P::N];
    hash_h_into::<P>(pk_seed, adrs, m1, m2, &mut out);
    out
}

/// `F(PK.seed, ADRS, M1)` — single-input chain hash. Output (`P::N` bytes) to `out`.
///
/// The most frequently called hash function in SLH-DSA (WOTS+ chains).
pub fn f_hash_into<P: Params>(pk_seed: &[u8], adrs: &Adrs, m1: &[u8], out: &mut [u8]) {
    P::Hash::f_hash_into::<P>(pk_seed, adrs, m1, out);
}

/// `F` — convenience wrapper returning `Vec`.
pub fn f_hash<P: Params>(pk_seed: &[u8], adrs: &Adrs, m1: &[u8]) -> Vec<u8> {
    let mut out = vec![0u8; P::N];
    f_hash_into::<P>(pk_seed, adrs, m1, &mut out);
    out
}