krypteia-quantica 0.2.0

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

//! Pre-hash selectors and `M′` message encoding for the external
//! signature APIs (FIPS 204 §5.4 / FIPS 205 §10.2).
//!
//! Both ML-DSA (HashML-DSA) and SLH-DSA (HashSLH-DSA) wrap the message
//! into `M′` for domain separation before calling the internal signer:
//!
//! - pure:    `M′ = 0x00 ‖ len(ctx) ‖ ctx ‖ M`
//! - prehash: `M′ = 0x01 ‖ len(ctx) ‖ ctx ‖ OID(PH) ‖ PH(M)`
//!
//! `OID(PH)` is the full DER object identifier (tag + length + content,
//! 11 bytes for every NIST CSOR hash/XOF) per FIPS 204 Algorithm 4. The
//! pre-hash digest is computed from the shared `krypteia-tessera`
//! primitives; SHAKE128 is squeezed to 256 bits and SHAKE256 to 512 bits.
//!
//! `M′` is built over public data (domain tag, context, OID, message or
//! its digest) — no secret-dependent branch, index, or length. The only
//! decision is the public `|ctx| ≤ 255` check, done by the callers.

use alloc::vec::Vec;
use tessera::{
    Digest, Sha3_224, Sha3_256, Sha3_384, Sha3_512, Sha224, Sha256, Sha384, Sha512, Sha512_224, Sha512_256, Shake128,
    Shake256, Xof,
};

/// Approved pre-hash function for HashML-DSA / HashSLH-DSA
/// (FIPS 204 §5.4, FIPS 205 §10.2). Selects both the digest computed over
/// the message and the DER OID embedded in `M′`.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum PreHash {
    /// SHA-256 (OID 2.16.840.1.101.3.4.2.1).
    Sha256,
    /// SHA-384 (OID 2.16.840.1.101.3.4.2.2).
    Sha384,
    /// SHA-512 (OID 2.16.840.1.101.3.4.2.3).
    Sha512,
    /// SHA-224 (OID 2.16.840.1.101.3.4.2.4).
    Sha224,
    /// SHA-512/224 (OID 2.16.840.1.101.3.4.2.5).
    Sha512_224,
    /// SHA-512/256 (OID 2.16.840.1.101.3.4.2.6).
    Sha512_256,
    /// SHA3-224 (OID 2.16.840.1.101.3.4.2.7).
    Sha3_224,
    /// SHA3-256 (OID 2.16.840.1.101.3.4.2.8).
    Sha3_256,
    /// SHA3-384 (OID 2.16.840.1.101.3.4.2.9).
    Sha3_384,
    /// SHA3-512 (OID 2.16.840.1.101.3.4.2.10).
    Sha3_512,
    /// SHAKE128 squeezed to 256 bits (OID 2.16.840.1.101.3.4.2.11).
    Shake128,
    /// SHAKE256 squeezed to 512 bits (OID 2.16.840.1.101.3.4.2.12).
    Shake256,
}

impl PreHash {
    /// DER encoding of the function's OID (tag `0x06`, length `0x09`, then
    /// the nine content bytes), per FIPS 204 Algorithm 4 — verified against
    /// the NIST.FIPS.204 primary source. The arcs differ only in the final
    /// byte (the CSOR `hashAlgs` index).
    pub(crate) fn oid(self) -> &'static [u8] {
        // DER TLV `06 09 60 86 48 01 65 03 04 02 XX`; only the final CSOR
        // `hashAlgs` index XX differs. Array-literal refs promote to 'static.
        match self {
            PreHash::Sha256 => &[0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x01],
            PreHash::Sha384 => &[0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x02],
            PreHash::Sha512 => &[0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x03],
            PreHash::Sha224 => &[0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x04],
            PreHash::Sha512_224 => &[0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x05],
            PreHash::Sha512_256 => &[0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x06],
            PreHash::Sha3_224 => &[0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x07],
            PreHash::Sha3_256 => &[0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x08],
            PreHash::Sha3_384 => &[0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x09],
            PreHash::Sha3_512 => &[0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x0A],
            PreHash::Shake128 => &[0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x0B],
            PreHash::Shake256 => &[0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x0C],
        }
    }

    /// Compute `PH(M)`. Fixed-output hashes return their digest; SHAKE128
    /// is squeezed to 256 bits and SHAKE256 to 512 bits (FIPS 204 §5.4).
    pub(crate) fn hash(self, m: &[u8]) -> Vec<u8> {
        fn fixed<H: Digest>(m: &[u8]) -> Vec<u8> {
            let mut out = alloc::vec![0u8; H::OUTPUT_LEN];
            H::digest(m, &mut out);
            out
        }
        fn xof<X: Xof>(m: &[u8], out_len: usize) -> Vec<u8> {
            let mut x = X::new();
            x.update(m);
            let mut out = alloc::vec![0u8; out_len];
            x.squeeze(&mut out);
            out
        }
        match self {
            PreHash::Sha256 => fixed::<Sha256>(m),
            PreHash::Sha384 => fixed::<Sha384>(m),
            PreHash::Sha512 => fixed::<Sha512>(m),
            PreHash::Sha224 => fixed::<Sha224>(m),
            PreHash::Sha512_224 => fixed::<Sha512_224>(m),
            PreHash::Sha512_256 => fixed::<Sha512_256>(m),
            PreHash::Sha3_224 => fixed::<Sha3_224>(m),
            PreHash::Sha3_256 => fixed::<Sha3_256>(m),
            PreHash::Sha3_384 => fixed::<Sha3_384>(m),
            PreHash::Sha3_512 => fixed::<Sha3_512>(m),
            PreHash::Shake128 => xof::<Shake128>(m, 32),
            PreHash::Shake256 => xof::<Shake256>(m, 64),
        }
    }
}

/// Pure external `M′ = 0x00 ‖ len(ctx) ‖ ctx ‖ M`
/// (FIPS 204 Algorithm 2 / FIPS 205 Algorithm 22).
///
/// The caller must ensure `ctx.len() <= 255` (mapped to the caller's own
/// `ContextTooLong` error) before calling.
pub(crate) fn m_prime_pure(ctx: &[u8], msg: &[u8]) -> Vec<u8> {
    let mut out = Vec::with_capacity(2 + ctx.len() + msg.len());
    out.push(0x00);
    out.push(ctx.len() as u8);
    out.extend_from_slice(ctx);
    out.extend_from_slice(msg);
    out
}

/// Pre-hash external `M′ = 0x01 ‖ len(ctx) ‖ ctx ‖ OID(PH) ‖ PH(M)`
/// (FIPS 204 Algorithm 4 / FIPS 205 Algorithm 23).
///
/// The caller must ensure `ctx.len() <= 255` before calling.
pub(crate) fn m_prime_prehash(ctx: &[u8], ph: PreHash, msg: &[u8]) -> Vec<u8> {
    let oid = ph.oid();
    let digest = ph.hash(msg);
    let mut out = Vec::with_capacity(2 + ctx.len() + oid.len() + digest.len());
    out.push(0x01);
    out.push(ctx.len() as u8);
    out.extend_from_slice(ctx);
    out.extend_from_slice(oid);
    out.extend_from_slice(&digest);
    out
}

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

    #[test]
    fn oids_are_11_byte_der_tlv() {
        for ph in [
            PreHash::Sha256,
            PreHash::Sha384,
            PreHash::Sha512,
            PreHash::Sha224,
            PreHash::Sha512_224,
            PreHash::Sha512_256,
            PreHash::Sha3_224,
            PreHash::Sha3_256,
            PreHash::Sha3_384,
            PreHash::Sha3_512,
            PreHash::Shake128,
            PreHash::Shake256,
        ] {
            let oid = ph.oid();
            assert_eq!(oid.len(), 11);
            assert_eq!(oid[0], 0x06); // DER OBJECT IDENTIFIER tag
            assert_eq!(oid[1], 0x09); // length
            assert_eq!(&oid[2..10], &[0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02]);
        }
        // Spot-check the FIPS 204 Algorithm 4 examples (verbatim).
        assert_eq!(PreHash::Sha256.oid()[10], 0x01);
        assert_eq!(PreHash::Sha512.oid()[10], 0x03);
        assert_eq!(PreHash::Shake128.oid()[10], 0x0B);
    }

    #[test]
    fn shake_prehash_output_lengths() {
        assert_eq!(PreHash::Shake128.hash(b"x").len(), 32); // 256 bits
        assert_eq!(PreHash::Shake256.hash(b"x").len(), 64); // 512 bits
        assert_eq!(PreHash::Sha256.hash(b"x").len(), 32);
        assert_eq!(PreHash::Sha512.hash(b"x").len(), 64);
    }

    #[test]
    fn pure_m_prime_layout() {
        let mp = m_prime_pure(b"ctx", b"msg");
        assert_eq!(mp[0], 0x00);
        assert_eq!(mp[1], 3);
        assert_eq!(&mp[2..5], b"ctx");
        assert_eq!(&mp[5..], b"msg");
    }

    #[test]
    fn prehash_m_prime_layout() {
        let mp = m_prime_prehash(b"", PreHash::Sha256, b"msg");
        assert_eq!(mp[0], 0x01);
        assert_eq!(mp[1], 0); // empty ctx
        assert_eq!(&mp[2..13], PreHash::Sha256.oid());
        assert_eq!(mp.len(), 2 + 11 + 32);
    }

    fn hexd(s: &str) -> Vec<u8> {
        (0..s.len())
            .step_by(2)
            .map(|i| u8::from_str_radix(&s[i..i + 2], 16).unwrap())
            .collect()
    }

    #[test]
    fn prehash_m_prime_matches_fips_reference() {
        // Independently-computed references for every PreHash (ctx empty,
        // msg = b"abc"): `0x01 ‖ 0x00 ‖ OID(PH) ‖ PH("abc")`. Cross-checks the
        // OID byte strings *and* the underlying digests in one shot.
        let refs: &[(PreHash, &str)] = &[
            (
                PreHash::Sha384,
                "01000609608648016503040202cb00753f45a35e8bb5a03d699ac65007272c32ab0eded1631a8b605a43ff5bed8086072ba1e7cc2358baeca134c825a7",
            ),
            (
                PreHash::Shake256,
                "0100060960864801650304020c483366601360a8771c6863080cc4114d8db44530f8f1e1ee4f94ea37e78b5739d5a15bef186a5386c75744c0527e1faa9f8726e462a12a4feb06bd8801e751e4",
            ),
            (
                PreHash::Sha3_224,
                "01000609608648016503040207e642824c3f8cf24ad09234ee7d3c766fc9a3a5168d0c94ad73b46fdf",
            ),
            (
                PreHash::Sha3_256,
                "010006096086480165030402083a985da74fe225b2045c172d6bd390bd855f086e3e9d525b46bfe24511431532",
            ),
            (
                PreHash::Sha3_384,
                "01000609608648016503040209ec01498288516fc926459f58e2c6ad8df9b473cb0fc08c2596da7cf0e49be4b298d88cea927ac7f539f1edf228376d25",
            ),
            (
                PreHash::Sha3_512,
                "0100060960864801650304020ab751850b1a57168a5693cd924b6b096e08f621827444f70d884f5d0240d2712e10e116e9192af3c91a7ec57647e3934057340b4cf408d5a56592f8274eec53f0",
            ),
            (
                PreHash::Sha224,
                "0100060960864801650304020423097d223405d8228642a477bda255b32aadbce4bda0b3f7e36c9da7",
            ),
            (
                PreHash::Sha256,
                "01000609608648016503040201ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad",
            ),
            (
                PreHash::Sha512,
                "01000609608648016503040203ddaf35a193617abacc417349ae20413112e6fa4e89a97ea20a9eeee64b55d39a2192992a274fc1a836ba3c23a3feebbd454d4423643ce80e2a9ac94fa54ca49f",
            ),
            (
                PreHash::Sha512_224,
                "010006096086480165030402054634270f707b6a54daae7530460842e20e37ed265ceee9a43e8924aa",
            ),
            (
                PreHash::Sha512_256,
                "0100060960864801650304020653048e2681941ef99b2e29b76b4c7dabe4c2d0c634fc6d46e0e2f13107e7af23",
            ),
            (
                PreHash::Shake128,
                "0100060960864801650304020b5881092dd818bf5cf8a3ddb793fbcba74097d5c526a6d35f97b83351940f2cc8",
            ),
        ];
        for (ph, want) in refs {
            assert_eq!(m_prime_prehash(b"", *ph, b"abc"), hexd(want), "{ph:?} M' mismatch");
        }
    }
}