krypteia-quantica 0.1.0

//! WOTS+ one-time signature scheme (FIPS 205, Algorithms 1, 4-8).
//!
//! WOTS+ (Winternitz One-Time Signature Plus) is the foundational one-time signature
//! scheme used at the leaves of every XMSS tree in SLH-DSA. It signs a single `n`-byte
//! message by:
//!
//! 1. Splitting the message into base-`w` digits and computing a checksum.
//! 2. For each digit `d_i`, evaluating a hash chain `F^(d_i)(sk_i)` where `sk_i` is
//!    derived from the secret seed via PRF.
//! 3. Verification completes each chain to step `w - 1` and compresses all endpoints
//!    into a single `n`-byte public key using `T_l`.
//!
//! WOTS+ is purely hash-based: its one-time security relies solely on the second-preimage
//! resistance of the hash function `F`.

use super::address::{Adrs, WOTS_PK, WOTS_PRF};
use super::hash;
use super::params::Params;
use alloc::vec::Vec;

/// Extract base-2^b digits from a byte string.
///
/// Implements Algorithm 4 of FIPS 205. Interprets the byte string `x` as a sequence
/// of `b`-bit unsigned integers in big-endian bit order and returns `out_len` digits,
/// each in the range `[0, 2^b)`.
///
/// For the standard WOTS+ parameters (`lg_w = 4`, `b = 4`), this extracts nibbles.
/// It is also used by FORS to extract `a`-bit indices from the message digest.
pub fn base_2b(x: &[u8], b: usize, out_len: usize) -> Vec<u32> {
    let mut result = Vec::with_capacity(out_len);
    let mut in_idx = 0usize;
    let mut bits = 0u32;
    let mut total = 0usize;
    let mask = (1u32 << b) - 1;

    for _ in 0..out_len {
        while total < b {
            bits = (bits << 8) | (x[in_idx] as u32);
            in_idx += 1;
            total += 8;
        }
        total -= b;
        result.push((bits >> total) & mask);
    }
    result
}

/// Compute the WOTS+ checksum and return `len2` base-`w` digits.
///
/// Implements Algorithm 1 of FIPS 205. The checksum ensures that an attacker cannot
/// forge a signature by only advancing hash chains forward: decreasing any message
/// digit necessarily increases a checksum digit.
fn gen_len2<P: Params>(msg_digits: &[u32]) -> Vec<u32> {
    let w = P::W as u32;
    // csum = sum of (w - 1 - digit) for each digit in msg
    let mut csum: u32 = 0;
    for &d in msg_digits.iter() {
        csum += w - 1 - d;
    }
    // Shift csum left by (8 - ((len2 * lg_w) % 8)) % 8 bits
    let shift = (8 - ((P::LEN2 * P::LG_W) % 8)) % 8;
    csum <<= shift;

    // Convert csum to bytes, then extract len2 base-w digits
    // csum fits in ceil((len2 * lg_w + shift) / 8) bytes
    let num_bytes = ((P::LEN2 * P::LG_W) + shift + 7) / 8;
    let csum_bytes = to_byte(csum as u64, num_bytes);
    base_2b(&csum_bytes, P::LG_W, P::LEN2)
}

/// Apply the WOTS+ chain function `s` times starting from step `i`.
///
/// Implements Algorithm 5 of FIPS 205. Computes `F^s(X)` by iterating the tweakable
/// hash function `F` a total of `s` times, with hash addresses ranging from `i` to
/// `i + s - 1`. This recursive implementation is provided for clarity; the internal
/// signing and verification routines use `chain_iter` (private helper) for efficiency.
///
/// Returns an `n`-byte hash chain value.
pub fn chain<P: Params>(x: &[u8], i: u32, s: u32, pk_seed: &[u8], adrs: &mut Adrs) -> Vec<u8> {
    if s == 0 {
        return x.to_vec();
    }
    let mut tmp = chain::<P>(x, i, s - 1, pk_seed, adrs);
    adrs.set_hash_address(i + s - 1);
    tmp = hash::f_hash::<P>(pk_seed, adrs, &tmp);
    tmp
}

/// Iterative (non-recursive) variant of the WOTS+ chain function.
///
/// Functionally equivalent to [`chain`] but avoids deep recursion by using a simple loop.
/// This is the version used by all internal signing and verification routines.
fn chain_iter<P: Params>(x: &[u8], i: u32, s: u32, pk_seed: &[u8], adrs: &mut Adrs) -> Vec<u8> {
    // Use two stack buffers to avoid heap allocation in the inner loop.
    // MAX_N = 32 covers all SLH-DSA parameter sets.
    let n = P::N;
    let mut buf_a = [0u8; 32];
    let mut buf_b = [0u8; 32];
    buf_a[..n].copy_from_slice(&x[..n]);

    for j in 0..s {
        adrs.set_hash_address(i + j);
        if j % 2 == 0 {
            hash::f_hash_into::<P>(pk_seed, adrs, &buf_a[..n], &mut buf_b[..n]);
        } else {
            hash::f_hash_into::<P>(pk_seed, adrs, &buf_b[..n], &mut buf_a[..n]);
        }
    }
    if s == 0 {
        buf_a[..n].to_vec()
    } else if s % 2 == 1 {
        buf_b[..n].to_vec()
    } else {
        buf_a[..n].to_vec()
    }
}

/// Generate a WOTS+ public key.
///
/// Implements Algorithm 6 of FIPS 205. Derives `len` secret chain values from `SK.seed`
/// via PRF, evaluates each chain to its full length (`w - 1` steps), and compresses
/// all `len` endpoints into a single `n`-byte public key using `T_len`.
///
/// The address `adrs` must have its key pair address set before calling.
pub fn wots_pk_gen<P: Params>(sk_seed: &[u8], pk_seed: &[u8], adrs: &mut Adrs) -> Vec<u8> {
    let mut sk_adrs = adrs.clone();
    sk_adrs.set_type_and_clear(WOTS_PRF);
    sk_adrs.set_key_pair_address(adrs.get_key_pair_address());

    let mut tmp = Vec::with_capacity(P::LEN * P::N);

    for i in 0..P::LEN {
        sk_adrs.set_chain_address(i as u32);
        let sk = hash::prf::<P>(pk_seed, sk_seed, &sk_adrs);
        adrs.set_chain_address(i as u32);
        let pk_i = chain_iter::<P>(&sk, 0, (P::W - 1) as u32, pk_seed, adrs);
        tmp.extend_from_slice(&pk_i);
    }

    // Compress: T_len(PK.seed, ADRS', tmp)
    let mut wots_pk_adrs = adrs.clone();
    wots_pk_adrs.set_type_and_clear(WOTS_PK);
    wots_pk_adrs.set_key_pair_address(adrs.get_key_pair_address());

    hash::t_l::<P>(pk_seed, &wots_pk_adrs, &tmp)
}

/// Sign an `n`-byte message using WOTS+.
///
/// Implements Algorithm 7 of FIPS 205. Converts the message `m` to base-`w` digits,
/// appends the checksum digits, and for each digit `d_i` outputs the chain value
/// `F^(d_i)(sk_i)`. The resulting signature is `len * n` bytes.
///
/// This is a one-time signature: using the same WOTS+ key pair to sign two different
/// messages leaks enough information to allow forgery.
pub fn wots_sign<P: Params>(m: &[u8], sk_seed: &[u8], pk_seed: &[u8], adrs: &mut Adrs) -> Vec<u8> {
    let mut sig = alloc::vec![0u8; P::LEN * P::N];
    wots_sign_into::<P>(m, sk_seed, pk_seed, adrs, &mut sig);
    sig
}

/// Streaming variant of [`wots_sign`] — writes the `LEN * N`-byte
/// signature into the start of `out` (which must be at least that
/// size) instead of returning a freshly-allocated `Vec<u8>`.
///
/// Used by the streaming sign path (`xmss_sign_into` → `ht_sign_into`
/// → `slh_sign_into`) to avoid the transient per-layer heap buffers.
pub fn wots_sign_into<P: Params>(m: &[u8], sk_seed: &[u8], pk_seed: &[u8], adrs: &mut Adrs, out: &mut [u8]) {
    debug_assert!(out.len() >= P::LEN * P::N);

    // Convert M to base-w, then append the checksum digits.
    let msg_digits = base_2b(m, P::LG_W, P::LEN1);
    let csum_digits = gen_len2::<P>(&msg_digits);
    let mut digits = msg_digits;
    digits.extend_from_slice(&csum_digits);

    let mut sk_adrs = adrs.clone();
    sk_adrs.set_type_and_clear(WOTS_PRF);
    sk_adrs.set_key_pair_address(adrs.get_key_pair_address());

    for i in 0..P::LEN {
        sk_adrs.set_chain_address(i as u32);
        let sk = hash::prf::<P>(pk_seed, sk_seed, &sk_adrs);
        adrs.set_chain_address(i as u32);
        let sig_i = chain_iter::<P>(&sk, 0, digits[i], pk_seed, adrs);
        out[i * P::N..(i + 1) * P::N].copy_from_slice(&sig_i);
    }
}

/// Compute a WOTS+ public key candidate from a signature.
///
/// Implements Algorithm 8 of FIPS 205. For each digit `d_i` of the message (including
/// checksum), completes the hash chain from the signature value (at step `d_i`) to the
/// full chain endpoint (step `w - 1`), then compresses all endpoints with `T_len`.
///
/// If the signature is valid, the returned value equals the original public key.
pub fn wots_pk_from_sig<P: Params>(sig: &[u8], m: &[u8], pk_seed: &[u8], adrs: &mut Adrs) -> Vec<u8> {
    let msg_digits = base_2b(m, P::LG_W, P::LEN1);
    let csum_digits = gen_len2::<P>(&msg_digits);
    let mut digits = msg_digits;
    digits.extend_from_slice(&csum_digits);

    let mut tmp = Vec::with_capacity(P::LEN * P::N);

    for i in 0..P::LEN {
        adrs.set_chain_address(i as u32);
        let sig_i = &sig[i * P::N..(i + 1) * P::N];
        let w_minus_1 = (P::W - 1) as u32;
        let pk_i = chain_iter::<P>(sig_i, digits[i], w_minus_1 - digits[i], pk_seed, adrs);
        tmp.extend_from_slice(&pk_i);
    }

    let mut wots_pk_adrs = adrs.clone();
    wots_pk_adrs.set_type_and_clear(WOTS_PK);
    wots_pk_adrs.set_key_pair_address(adrs.get_key_pair_address());

    hash::t_l::<P>(pk_seed, &wots_pk_adrs, &tmp)
}

/// Convert an integer to a big-endian byte vector of length `n`.
///
/// The integer `val` is encoded in exactly `n` bytes, zero-padded on the left.
/// Used internally to serialize the checksum value before extracting its base-`w` digits.
/// Convert an integer to a big-endian byte array.
/// Returns a stack-allocated [u8; 8] (max needed is 4 bytes for checksum).
pub fn to_byte(val: u64, n: usize) -> Vec<u8> {
    let mut result = vec![0u8; n];
    let mut v = val;
    for i in (0..n).rev() {
        result[i] = (v & 0xff) as u8;
        v >>= 8;
    }
    result
}

/// Stack-allocated to_byte for small sizes (up to 8 bytes).
pub fn to_byte_stack(val: u64, n: usize) -> [u8; 8] {
    let mut result = [0u8; 8];
    let mut v = val;
    let start = 8 - n;
    for i in (start..8).rev() {
        result[i] = (v & 0xff) as u8;
        v >>= 8;
    }
    result
}