krypteia-quantica 0.1.0

//! Top-level SLH-DSA algorithms (FIPS 205, Algorithms 18-22, 24).
//!
//! This module implements the complete SLH-DSA signature lifecycle:
//!
//! - **Key generation** (`slh_keygen`, `slh_keygen_internal`): generates the master
//!   seeds and computes the hypertree root to form the public key.
//! - **Signing** (`slh_sign`, `slh_sign_internal`): hashes the message, signs it
//!   with FORS, then signs the FORS public key with the hypertree.
//! - **Verification** (`slh_verify`, `slh_verify_internal`): recomputes the FORS
//!   public key from the FORS signature, then verifies the hypertree signature against
//!   the public key root.
//!
//! The signature format is: `R || SIG_FORS || SIG_HT`, where `R` is the `n`-byte
//! randomizer, `SIG_FORS` is the FORS signature, and `SIG_HT` is the hypertree signature.

use super::SlhDsaError;
use super::address::Adrs;
use super::fors;
use super::hash;
use super::hypertree;
use super::params::{self, Params};
use super::rng::CryptoRng;
use super::xmss;
use alloc::vec::Vec;

/// Generate an SLH-DSA key pair from explicit seed material (deterministic).
///
/// Implements Algorithm 18 of FIPS 205. Computes the top-level XMSS tree root
/// from the provided seeds, then assembles the secret and public keys.
///
/// - `sk_seed`: the master secret seed (`n` bytes), used to derive all WOTS+ and FORS secrets
/// - `sk_prf`: the PRF key (`n` bytes), used to generate per-signature randomizers
/// - `pk_seed`: the public seed (`n` bytes), included in every hash call for domain separation
///
/// Returns `(sk, pk)` where:
/// - `sk = SK.seed || SK.prf || PK.seed || PK.root` (4n bytes)
/// - `pk = PK.seed || PK.root` (2n bytes)
pub fn slh_keygen_internal<P: Params>(sk_seed: &[u8], sk_prf: &[u8], pk_seed: &[u8]) -> (Vec<u8>, Vec<u8>) {
    let mut adrs = Adrs::new();
    adrs.set_layer_address((P::D - 1) as u32);

    // Compute the top-level XMSS tree root
    let pk_root = xmss::xmss_node::<P>(sk_seed, pk_seed, 0, P::H_PRIME as u32, &mut adrs);

    // sk = SK.seed || SK.prf || PK.seed || PK.root
    let mut sk = Vec::with_capacity(P::SK_LEN);
    sk.extend_from_slice(sk_seed);
    sk.extend_from_slice(sk_prf);
    sk.extend_from_slice(pk_seed);
    sk.extend_from_slice(&pk_root);

    // pk = PK.seed || PK.root
    let mut pk = Vec::with_capacity(P::PK_LEN);
    pk.extend_from_slice(pk_seed);
    pk.extend_from_slice(&pk_root);

    (sk, pk)
}

/// Internal SLH-DSA signing function.
///
/// Implements Algorithm 19 of FIPS 205. Performs the full signing pipeline:
///
/// 1. Generates the per-signature randomizer `R` via `PRF_msg(SK.prf, addrnd, M)`.
/// 2. Hashes the message with `H_msg` to produce a digest, then splits it into
///    FORS indices (`md`), a tree index (`idx_tree`), and a leaf index (`idx_leaf`).
/// 3. Signs `md` with FORS to obtain `SIG_FORS`, then computes the FORS public key.
/// 4. Signs the FORS public key with the hypertree to obtain `SIG_HT`.
///
/// The `addrnd` parameter is `n` bytes of additional randomness for hedged signing.
/// Passing `PK.seed` instead yields deterministic signing.
///
/// Returns the full signature: `R || SIG_FORS || SIG_HT`.
///
/// # Feature: `sca-fors-redundancy`
///
/// When this cargo feature is enabled, the FORS sub-signature is produced
/// by [`fors::fors_sign_into_redundant`] (T1-C — recompute-and-compare
/// against single-fault grafting-tree forgery, per
/// :cite:`genet2023_protecting_sphincs_faults`,
/// :cite:`castelnovi2018_grafting_trees`,
/// :cite:`adiletta2025_slashdsa_rowhammer`). On a detected divergence the
/// function returns [`Err(SlhDsaError::FaultDetected)`][SlhDsaError] —
/// the faulted signature never leaves the device. The validated FORS
/// public key returned by the redundant routine is fed straight into the
/// hypertree, saving a third FORS-root derivation.
///
/// Without the feature, the result is always [`Ok`] — the unified `Result`
/// envelope is paid only to keep one signature across the two
/// configurations. CAVP / KAT output bytes are unchanged.
pub fn slh_sign_internal<P: Params>(m: &[u8], sk: &[u8], addrnd: &[u8]) -> Result<Vec<u8>, SlhDsaError> {
    let n = P::N;
    // Parse SK
    let sk_seed = &sk[..n];
    let sk_prf = &sk[n..2 * n];
    let pk_seed = &sk[2 * n..3 * n];
    let pk_root = &sk[3 * n..4 * n];

    // Generate randomizer R
    let r = hash::prf_msg::<P>(sk_prf, addrnd, m);

    // One allocation for the full signature. All sub-components
    // (R || SIG_FORS || SIG_HT) are written directly into slices of
    // this buffer — no transient per-component `Vec<u8>` is held.
    let sig_len = params::sig_len::<P>();
    let fors_sig_len = P::K * (1 + P::A) * n;
    let mut sig = alloc::vec![0u8; sig_len];
    sig[..n].copy_from_slice(&r);

    // Compute message digest
    let digest = hash::h_msg::<P>(&r, pk_seed, pk_root, m);

    // Split digest into: md (k*a/8 bytes), idx_tree, idx_leaf
    let md_len = (P::K * P::A + 7) / 8;
    let tree_bits = P::H - P::H_PRIME; // bits for idx_tree
    let tree_bytes = (tree_bits + 7) / 8;
    let leaf_bytes = (P::H_PRIME + 7) / 8;

    let md = &digest[..md_len];
    let tree_part = &digest[md_len..md_len + tree_bytes];
    let leaf_part = &digest[md_len + tree_bytes..md_len + tree_bytes + leaf_bytes];

    // Extract idx_tree and idx_leaf
    let idx_tree = bytes_to_u64(tree_part) & ((1u64 << tree_bits) - 1);
    let idx_leaf = bytes_to_u32(leaf_part) & ((1u32 << P::H_PRIME) - 1);

    // FORS signature — streamed into sig[n..n + fors_sig_len].
    // `adrs` is the immutable template passed to all FORS routines; each
    // callee clones internally for its own scratch state.
    let mut adrs = Adrs::new();
    adrs.set_tree_address(idx_tree);
    adrs.set_type_and_clear(super::address::FORS_TREE);
    adrs.set_key_pair_address(idx_leaf);
    let adrs = adrs;

    let pk_fors = {
        let fors_slot = &mut sig[n..n + fors_sig_len];
        #[cfg(feature = "sca-fors-redundancy")]
        {
            // T1-C — sign FORS twice, CT-compare both signatures and both
            // derived pks, abort with `FaultDetected` on any mismatch. The
            // returned pk is reused for the hypertree signing below.
            fors::fors_sign_into_redundant::<P>(md, sk_seed, pk_seed, &adrs, fors_slot)?
        }
        #[cfg(not(feature = "sca-fors-redundancy"))]
        {
            fors::fors_sign_into::<P>(md, sk_seed, pk_seed, &adrs, fors_slot)?;
            // Re-read SIG_FORS from the slice we just filled to compute
            // the FORS public key — no extra buffer needed.
            fors::fors_pk_from_sig::<P>(&sig[n..n + fors_sig_len], md, pk_seed, &adrs)
        }
    };

    // Hypertree signature — streamed into sig[n + fors_sig_len..].
    {
        let ht_slot = &mut sig[n + fors_sig_len..];
        hypertree::ht_sign_into::<P>(&pk_fors, sk_seed, pk_seed, idx_tree, idx_leaf, ht_slot);
    }

    Ok(sig)
}

/// Internal SLH-DSA verification function.
///
/// Implements Algorithm 20 of FIPS 205. Reverses the signing pipeline:
///
/// 1. Parses `R`, `SIG_FORS`, and `SIG_HT` from the signature.
/// 2. Recomputes the message digest using `H_msg(R, PK.seed, PK.root, M)`.
/// 3. Recovers the FORS public key from `SIG_FORS` and the digest indices.
/// 4. Verifies the hypertree signature `SIG_HT` over the recovered FORS public key.
///
/// Returns `true` if the hypertree root matches `PK.root`, `false` otherwise.
pub fn slh_verify_internal<P: Params>(m: &[u8], sig: &[u8], pk: &[u8]) -> bool {
    let n = P::N;
    let expected_sig_len = params::sig_len::<P>();
    if sig.len() != expected_sig_len {
        return false;
    }
    if pk.len() != P::PK_LEN {
        return false;
    }

    // Parse PK
    let pk_seed = &pk[..n];
    let pk_root = &pk[n..2 * n];

    // Parse signature: R || SIG_FORS || SIG_HT
    let r = &sig[..n];
    let fors_sig_len = P::K * (1 + P::A) * n;
    let sig_fors = &sig[n..n + fors_sig_len];
    let sig_ht = &sig[n + fors_sig_len..];

    // Compute message digest
    let digest = hash::h_msg::<P>(r, pk_seed, pk_root, m);

    let md_len = (P::K * P::A + 7) / 8;
    let tree_bits = P::H - P::H_PRIME;
    let tree_bytes = (tree_bits + 7) / 8;
    let leaf_bytes = (P::H_PRIME + 7) / 8;

    let md = &digest[..md_len];
    let tree_part = &digest[md_len..md_len + tree_bytes];
    let leaf_part = &digest[md_len + tree_bytes..md_len + tree_bytes + leaf_bytes];

    let idx_tree = bytes_to_u64(tree_part) & ((1u64 << tree_bits) - 1);
    let idx_leaf = bytes_to_u32(leaf_part) & ((1u32 << P::H_PRIME) - 1);

    // Compute FORS public key from signature
    let mut adrs = Adrs::new();
    adrs.set_tree_address(idx_tree);
    adrs.set_type_and_clear(super::address::FORS_TREE);
    adrs.set_key_pair_address(idx_leaf);
    let adrs = adrs;

    let pk_fors = fors::fors_pk_from_sig::<P>(sig_fors, md, pk_seed, &adrs);

    // Verify hypertree signature
    hypertree::ht_verify::<P>(&pk_fors, sig_ht, pk_seed, idx_tree, idx_leaf, pk_root)
}

/// Generate an SLH-DSA key pair using random seeds from the provided RNG.
///
/// Implements Algorithm 21 of FIPS 205. Draws three independent `n`-byte random
/// values (`SK.seed`, `SK.prf`, `PK.seed`) from `rng`, then delegates to
/// `slh_keygen_internal` to compute the hypertree root and assemble the keys.
///
/// Returns `(secret_key, public_key)` or an error if the RNG fails.
pub fn slh_keygen<P: Params>(rng: &mut dyn CryptoRng) -> Result<(Vec<u8>, Vec<u8>), SlhDsaError> {
    let n = P::N;
    let mut sk_seed = vec![0u8; n];
    let mut sk_prf = vec![0u8; n];
    let mut pk_seed = vec![0u8; n];

    rng.fill_bytes(&mut sk_seed)?;
    rng.fill_bytes(&mut sk_prf)?;
    rng.fill_bytes(&mut pk_seed)?;

    Ok(slh_keygen_internal::<P>(&sk_seed, &sk_prf, &pk_seed))
}

/// Sign a message with randomized (hedged) signing.
///
/// Implements Algorithm 22 of FIPS 205. Draws `n` bytes of randomness from `rng` to
/// use as the `addrnd` parameter, providing hedged signing that remains secure even
/// if the RNG is somewhat predictable.
///
/// Returns the full SLH-DSA signature or an error if the key is invalid, the RNG
/// fails, or — when `sca-fors-redundancy` is enabled — the T1-C FORS
/// recompute-and-compare check detects a fault.
pub fn slh_sign<P: Params>(m: &[u8], sk: &[u8], rng: &mut dyn CryptoRng) -> Result<Vec<u8>, SlhDsaError> {
    let n = P::N;
    if sk.len() != P::SK_LEN {
        return Err(SlhDsaError::InvalidKey);
    }

    let mut addrnd = vec![0u8; n];
    rng.fill_bytes(&mut addrnd)?;

    slh_sign_internal::<P>(m, sk, &addrnd)
}

/// Verify an SLH-DSA signature against a message and public key.
///
/// Implements Algorithm 24 of FIPS 205. Validates the key and signature lengths,
/// then delegates to `slh_verify_internal`.
///
/// Returns `Ok(true)` if valid, `Ok(false)` if the signature does not verify,
/// or `Err` if the key or signature has an invalid length.
pub fn slh_verify<P: Params>(m: &[u8], sig: &[u8], pk: &[u8]) -> Result<bool, SlhDsaError> {
    if pk.len() != P::PK_LEN {
        return Err(SlhDsaError::InvalidKey);
    }
    let expected_sig_len = params::sig_len::<P>();
    if sig.len() != expected_sig_len {
        return Err(SlhDsaError::InvalidSignature);
    }
    Ok(slh_verify_internal::<P>(m, sig, pk))
}

// ---------- Helpers ----------

/// Convert a big-endian byte slice to u64.
fn bytes_to_u64(b: &[u8]) -> u64 {
    let mut val = 0u64;
    for &byte in b {
        val = (val << 8) | byte as u64;
    }
    val
}

/// Convert a big-endian byte slice to u32.
fn bytes_to_u32(b: &[u8]) -> u32 {
    let mut val = 0u32;
    for &byte in b {
        val = (val << 8) | byte as u32;
    }
    val
}