krypteia-quantica 0.1.0

//! XMSS: eXtended Merkle Signature Scheme (FIPS 205, Algorithms 9-11).
//!
//! XMSS combines multiple WOTS+ one-time key pairs into a few-time signature scheme
//! by organizing them as leaves of a binary Merkle tree of height `h'`. Each XMSS tree
//! can sign `2^h'` messages (one per WOTS+ leaf).
//!
//! In the SLH-DSA hierarchy, XMSS trees form the building blocks of the hypertree:
//! each layer of the hypertree consists of XMSS trees, where the leaves of upper-layer
//! trees certify the roots of lower-layer trees.
//!
//! The Merkle tree uses the tweakable hash `H` ([`hash::hash_h`]) for internal nodes
//! and WOTS+ public keys ([`wots::wots_pk_gen`]) as leaf values.

use super::address::{Adrs, TREE, WOTS_HASH};
use super::hash;
use super::params::Params;
use super::wots;
use alloc::vec::Vec;

/// Compute the root of an XMSS Merkle subtree.
///
/// Implements Algorithm 9 of FIPS 205. Computes the root node of a subtree of height
/// `z` whose leftmost leaf is at index `i` within the current XMSS tree. When `z`
/// equals `h'` and `i` equals 0, this returns the full XMSS tree root.
///
/// For `z = 0`, returns the WOTS+ public key at leaf `i`. For `z > 0`, builds the
/// tree iteratively: first generating all `2^z` leaf public keys, then hashing pairs
/// of sibling nodes up to the root using `H`.
pub fn xmss_node<P: Params>(sk_seed: &[u8], pk_seed: &[u8], i: u32, z: u32, adrs: &mut Adrs) -> Vec<u8> {
    if z == 0 {
        // Leaf: compute WOTS+ public key for leaf i
        adrs.set_type_and_clear(WOTS_HASH);
        adrs.set_key_pair_address(i);
        return wots::wots_pk_gen::<P>(sk_seed, pk_seed, adrs);
    }

    // Iterative Merkle tree computation.
    // Compute all 2^z leaves, then hash up.
    let num_leaves = 1u32 << z;
    let base = i; // i is the leftmost leaf index at this level

    // Compute leaf nodes
    let mut nodes: Vec<Vec<u8>> = Vec::with_capacity(num_leaves as usize);
    for j in 0..num_leaves {
        adrs.set_type_and_clear(WOTS_HASH);
        adrs.set_key_pair_address(base + j);
        let leaf = wots::wots_pk_gen::<P>(sk_seed, pk_seed, adrs);
        nodes.push(leaf);
    }

    // Hash up the tree
    for height in 1..=z {
        let mut new_nodes = Vec::with_capacity(nodes.len() / 2);
        for j in 0..(nodes.len() / 2) {
            adrs.set_type_and_clear(TREE);
            adrs.set_tree_height(height);
            adrs.set_tree_index(base / (1 << height) + j as u32);
            let parent = hash::hash_h::<P>(pk_seed, adrs, &nodes[2 * j], &nodes[2 * j + 1]);
            new_nodes.push(parent);
        }
        nodes = new_nodes;
    }

    nodes.into_iter().next().unwrap()
}

/// Create an XMSS signature for an `n`-byte message.
///
/// Implements Algorithm 10 of FIPS 205. Signs message `m` using the WOTS+ key pair
/// at leaf index `idx` and produces an authentication path of `h'` sibling nodes
/// that allows the verifier to recompute the tree root.
///
/// The returned signature is `(len + h') * n` bytes: a WOTS+ signature (`len * n`)
/// followed by the Merkle authentication path (`h' * n`).
pub fn xmss_sign<P: Params>(m: &[u8], sk_seed: &[u8], idx: u32, pk_seed: &[u8], adrs: &mut Adrs) -> Vec<u8> {
    let mut sig = alloc::vec![0u8; (P::LEN + P::H_PRIME) * P::N];
    xmss_sign_into::<P>(m, sk_seed, idx, pk_seed, adrs, &mut sig);
    sig
}

/// Streaming variant of [`xmss_sign`] — writes the `(LEN + H') * N`
/// byte signature into the start of `out` (which must be at least
/// that size) instead of returning a freshly-allocated `Vec<u8>`.
///
/// Layout of `out` (after the call): `WOTS sig (LEN * N) || auth path (H' * N)`.
pub fn xmss_sign_into<P: Params>(m: &[u8], sk_seed: &[u8], idx: u32, pk_seed: &[u8], adrs: &mut Adrs, out: &mut [u8]) {
    let wots_sig_len = P::LEN * P::N;
    let total_len = wots_sig_len + P::H_PRIME * P::N;
    debug_assert!(out.len() >= total_len);
    let (sig_wots_slot, auth_slot) = out[..total_len].split_at_mut(wots_sig_len);

    // Authentication path: one node per height.
    for j in 0..P::H_PRIME {
        let k = (idx >> j) ^ 1;
        let node = xmss_node::<P>(sk_seed, pk_seed, k * (1 << j), j as u32, adrs);
        auth_slot[j * P::N..(j + 1) * P::N].copy_from_slice(&node);
    }

    // WOTS+ signature written straight into the first block.
    adrs.set_type_and_clear(WOTS_HASH);
    adrs.set_key_pair_address(idx);
    wots::wots_sign_into::<P>(m, sk_seed, pk_seed, adrs, sig_wots_slot);
}

/// Compute an XMSS public key (tree root) from an XMSS signature.
///
/// Implements Algorithm 11 of FIPS 205. Recovers the WOTS+ public key from the
/// WOTS+ signature component, then walks up the authentication path to recompute
/// the Merkle tree root. If the signature is valid, the returned root matches the
/// original XMSS public key.
///
/// The `sig_xmss` input must be `(len + h') * n` bytes, and `idx` is the leaf index
/// that was used during signing.
pub fn xmss_pk_from_sig<P: Params>(idx: u32, sig_xmss: &[u8], m: &[u8], pk_seed: &[u8], adrs: &mut Adrs) -> Vec<u8> {
    // sig_xmss = sig_wots (len*n bytes) || auth (h'*n bytes)
    let wots_sig_len = P::LEN * P::N;
    let sig_wots = &sig_xmss[..wots_sig_len];
    let auth = &sig_xmss[wots_sig_len..];

    // Compute WOTS+ public key candidate
    adrs.set_type_and_clear(WOTS_HASH);
    adrs.set_key_pair_address(idx);
    let mut node = wots::wots_pk_from_sig::<P>(sig_wots, m, pk_seed, adrs);

    // Climb the authentication path
    adrs.set_type_and_clear(TREE);
    adrs.set_tree_index(idx);

    for k in 0..P::H_PRIME {
        let auth_k = &auth[k * P::N..(k + 1) * P::N];
        adrs.set_tree_height((k + 1) as u32);

        if ((idx >> k) & 1) == 0 {
            // idx/2^k is even: current node is left child
            adrs.set_tree_index(idx >> (k + 1));
            node = hash::hash_h::<P>(pk_seed, adrs, &node, auth_k);
        } else {
            // idx/2^k is odd: current node is right child
            adrs.set_tree_index(idx >> (k + 1));
            node = hash::hash_h::<P>(pk_seed, adrs, auth_k, &node);
        }
    }

    node
}