synta-certificate 0.2.6

//! Composite ML-DSA algorithm spec table and backend-agnostic helpers.
//!
//! Implements draft-ietf-lamps-pq-composite-sigs-19: 18 composite algorithms
//! combining ML-DSA with a traditional algorithm (RSA, ECDSA, EdDSA) under
//! OID arc 1.3.6.1.5.5.7.6, sub-arcs 37–54.
//!
//! This module is feature-gated: `#[cfg(any(feature = "openssl", feature = "nss"))]`
//! because composite key operations always need a crypto backend.

// ── Traditional component algorithm variants ──────────────────────────────────

/// The traditional (non-ML-DSA) component of a composite ML-DSA algorithm.
#[derive(Debug, Clone, Copy)]
pub enum TradAlg {
    /// RSA with RSASSA-PSS padding.  `bits` is the modulus size (2048, 3072, 4096).
    /// `hash` is `"sha256"` or `"sha512"`.
    RsaPss { bits: u32, hash: &'static str },
    /// RSA with PKCS#1 v1.5 padding.  `bits` is the modulus size.
    /// `hash` is `"sha256"` or `"sha512"`.
    RsaPkcs15 { bits: u32, hash: &'static str },
    /// ECDSA on a named curve.  `curve` is the curve name (e.g. `"P-256"`).
    /// `hash` is `"sha256"` or `"sha512"` for the ECDSA digest.
    Ec {
        curve: &'static str,
        hash: &'static str,
    },
    /// Ed25519 (RFC 8032).
    Ed25519,
    /// Ed448 (RFC 8032).
    Ed448,
}

/// The pre-hash function applied to TBS before constructing M'.
#[derive(Debug, Clone, Copy)]
pub enum CompHash {
    /// SHA-256, used by ML-DSA-44 + RSA/ECDSA-P256-SHA256 variants.
    Sha256,
    /// SHA-512, used by most ML-DSA-65 and ML-DSA-87 variants.
    Sha512,
    /// SHAKE256 with 64-byte (512-bit) output, used by MLDSA87-Ed448-SHAKE256.
    Shake256_64,
}

// ── Spec struct ───────────────────────────────────────────────────────────────

/// Per-variant configuration for a composite ML-DSA algorithm.
#[derive(Debug)]
pub struct CompositeMlDsaSpec {
    /// Sub-arc component (37–54) of the composite OID 1.3.6.1.5.5.7.6.<sub_arc>.
    pub sub_arc: u32,
    /// ML-DSA variant string for key generation: `"ML-DSA-44"`, `"ML-DSA-65"`, or `"ML-DSA-87"`.
    pub mldsa_variant: &'static str,
    /// Raw ML-DSA public key size in bytes (for SPKI split).
    pub mldsa_pk_size: usize,
    /// ML-DSA signature size in bytes (for composite signature split).
    pub mldsa_sig_size: usize,
    /// The traditional component algorithm.
    pub trad_alg: TradAlg,
    /// Pre-hash function for M' construction.
    pub hash: CompHash,
    /// Domain-separation label for M' and ML-DSA context string.
    pub label: &'static str,
}

// ── Static spec table ─────────────────────────────────────────────────────────

static COMPOSITE_SPECS: &[CompositeMlDsaSpec] = &[
    // ── ML-DSA-44 variants (sub-arcs 37–40) ──────────────────────────────────
    CompositeMlDsaSpec {
        sub_arc: 37,
        mldsa_variant: "ML-DSA-44",
        mldsa_pk_size: 1312,
        mldsa_sig_size: 2420,
        trad_alg: TradAlg::RsaPss {
            bits: 2048,
            hash: "sha256",
        },
        hash: CompHash::Sha256,
        label: "COMPSIG-MLDSA44-RSA2048-PSS-SHA256",
    },
    CompositeMlDsaSpec {
        sub_arc: 38,
        mldsa_variant: "ML-DSA-44",
        mldsa_pk_size: 1312,
        mldsa_sig_size: 2420,
        trad_alg: TradAlg::RsaPkcs15 {
            bits: 2048,
            hash: "sha256",
        },
        hash: CompHash::Sha256,
        label: "COMPSIG-MLDSA44-RSA2048-PKCS15-SHA256",
    },
    CompositeMlDsaSpec {
        sub_arc: 39,
        mldsa_variant: "ML-DSA-44",
        mldsa_pk_size: 1312,
        mldsa_sig_size: 2420,
        trad_alg: TradAlg::Ed25519,
        hash: CompHash::Sha512,
        label: "COMPSIG-MLDSA44-Ed25519-SHA512",
    },
    CompositeMlDsaSpec {
        sub_arc: 40,
        mldsa_variant: "ML-DSA-44",
        mldsa_pk_size: 1312,
        mldsa_sig_size: 2420,
        trad_alg: TradAlg::Ec {
            curve: "P-256",
            hash: "sha256",
        },
        hash: CompHash::Sha256,
        label: "COMPSIG-MLDSA44-ECDSA-P256-SHA256",
    },
    // ── ML-DSA-65 variants (sub-arcs 41–48) ──────────────────────────────────
    CompositeMlDsaSpec {
        sub_arc: 41,
        mldsa_variant: "ML-DSA-65",
        mldsa_pk_size: 1952,
        mldsa_sig_size: 3309,
        trad_alg: TradAlg::RsaPss {
            bits: 3072,
            hash: "sha512",
        },
        hash: CompHash::Sha512,
        label: "COMPSIG-MLDSA65-RSA3072-PSS-SHA512",
    },
    CompositeMlDsaSpec {
        sub_arc: 42,
        mldsa_variant: "ML-DSA-65",
        mldsa_pk_size: 1952,
        mldsa_sig_size: 3309,
        trad_alg: TradAlg::RsaPkcs15 {
            bits: 3072,
            hash: "sha512",
        },
        hash: CompHash::Sha512,
        label: "COMPSIG-MLDSA65-RSA3072-PKCS15-SHA512",
    },
    CompositeMlDsaSpec {
        sub_arc: 43,
        mldsa_variant: "ML-DSA-65",
        mldsa_pk_size: 1952,
        mldsa_sig_size: 3309,
        trad_alg: TradAlg::RsaPss {
            bits: 4096,
            hash: "sha512",
        },
        hash: CompHash::Sha512,
        label: "COMPSIG-MLDSA65-RSA4096-PSS-SHA512",
    },
    CompositeMlDsaSpec {
        sub_arc: 44,
        mldsa_variant: "ML-DSA-65",
        mldsa_pk_size: 1952,
        mldsa_sig_size: 3309,
        trad_alg: TradAlg::RsaPkcs15 {
            bits: 4096,
            hash: "sha512",
        },
        hash: CompHash::Sha512,
        label: "COMPSIG-MLDSA65-RSA4096-PKCS15-SHA512",
    },
    CompositeMlDsaSpec {
        sub_arc: 45,
        mldsa_variant: "ML-DSA-65",
        mldsa_pk_size: 1952,
        mldsa_sig_size: 3309,
        trad_alg: TradAlg::Ec {
            curve: "P-256",
            hash: "sha512",
        },
        hash: CompHash::Sha512,
        label: "COMPSIG-MLDSA65-ECDSA-P256-SHA512",
    },
    CompositeMlDsaSpec {
        sub_arc: 46,
        mldsa_variant: "ML-DSA-65",
        mldsa_pk_size: 1952,
        mldsa_sig_size: 3309,
        trad_alg: TradAlg::Ec {
            curve: "P-384",
            hash: "sha512",
        },
        hash: CompHash::Sha512,
        label: "COMPSIG-MLDSA65-ECDSA-P384-SHA512",
    },
    CompositeMlDsaSpec {
        sub_arc: 47,
        mldsa_variant: "ML-DSA-65",
        mldsa_pk_size: 1952,
        mldsa_sig_size: 3309,
        trad_alg: TradAlg::Ec {
            curve: "brainpoolP256r1",
            hash: "sha512",
        },
        hash: CompHash::Sha512,
        label: "COMPSIG-MLDSA65-ECDSA-BP256-SHA512",
    },
    CompositeMlDsaSpec {
        sub_arc: 48,
        mldsa_variant: "ML-DSA-65",
        mldsa_pk_size: 1952,
        mldsa_sig_size: 3309,
        trad_alg: TradAlg::Ed25519,
        hash: CompHash::Sha512,
        label: "COMPSIG-MLDSA65-Ed25519-SHA512",
    },
    // ── ML-DSA-87 variants (sub-arcs 49–54) ──────────────────────────────────
    CompositeMlDsaSpec {
        sub_arc: 49,
        mldsa_variant: "ML-DSA-87",
        mldsa_pk_size: 2592,
        mldsa_sig_size: 4627,
        trad_alg: TradAlg::Ec {
            curve: "P-384",
            hash: "sha512",
        },
        hash: CompHash::Sha512,
        label: "COMPSIG-MLDSA87-ECDSA-P384-SHA512",
    },
    CompositeMlDsaSpec {
        sub_arc: 50,
        mldsa_variant: "ML-DSA-87",
        mldsa_pk_size: 2592,
        mldsa_sig_size: 4627,
        trad_alg: TradAlg::Ec {
            curve: "brainpoolP384r1",
            hash: "sha512",
        },
        hash: CompHash::Sha512,
        label: "COMPSIG-MLDSA87-ECDSA-BP384-SHA512",
    },
    CompositeMlDsaSpec {
        sub_arc: 51,
        mldsa_variant: "ML-DSA-87",
        mldsa_pk_size: 2592,
        mldsa_sig_size: 4627,
        trad_alg: TradAlg::Ed448,
        hash: CompHash::Shake256_64,
        label: "COMPSIG-MLDSA87-Ed448-SHAKE256",
    },
    CompositeMlDsaSpec {
        sub_arc: 52,
        mldsa_variant: "ML-DSA-87",
        mldsa_pk_size: 2592,
        mldsa_sig_size: 4627,
        trad_alg: TradAlg::RsaPss {
            bits: 3072,
            hash: "sha512",
        },
        hash: CompHash::Sha512,
        label: "COMPSIG-MLDSA87-RSA3072-PSS-SHA512",
    },
    CompositeMlDsaSpec {
        sub_arc: 53,
        mldsa_variant: "ML-DSA-87",
        mldsa_pk_size: 2592,
        mldsa_sig_size: 4627,
        trad_alg: TradAlg::RsaPss {
            bits: 4096,
            hash: "sha512",
        },
        hash: CompHash::Sha512,
        label: "COMPSIG-MLDSA87-RSA4096-PSS-SHA512",
    },
    CompositeMlDsaSpec {
        sub_arc: 54,
        mldsa_variant: "ML-DSA-87",
        mldsa_pk_size: 2592,
        mldsa_sig_size: 4627,
        trad_alg: TradAlg::Ec {
            curve: "P-521",
            hash: "sha512",
        },
        hash: CompHash::Sha512,
        label: "COMPSIG-MLDSA87-ECDSA-P521-SHA512",
    },
];

/// Look up a composite ML-DSA spec by the sub-arc component (37–54).
///
/// Returns `None` if `sub_arc` is not a recognised composite ML-DSA variant.
pub fn composite_spec(sub_arc: u32) -> Option<&'static CompositeMlDsaSpec> {
    const FIRST: u32 = 37;
    debug_assert!(
        COMPOSITE_SPECS
            .iter()
            .enumerate()
            .all(|(i, s)| s.sub_arc == FIRST + i as u32),
        "COMPOSITE_SPECS table is not sorted by sub_arc starting at 37"
    );
    let idx = sub_arc.checked_sub(FIRST)? as usize;
    COMPOSITE_SPECS.get(idx)
}

/// Look up a [`CompositeMlDsaSpec`] from a full OID component slice.
///
/// Accepts OIDs in the composite ML-DSA arc (1.3.6.1.5.5.7.6.37–54).
/// Returns `None` for unrecognised or shorter OIDs.
pub fn composite_spec_from_oid(comps: &[u32]) -> Option<&'static CompositeMlDsaSpec> {
    let arc = crate::oids::COMPOSITE_MLDSA_ARC;
    if comps.len() != arc.len() + 1 {
        return None;
    }
    if !comps[..arc.len()]
        .iter()
        .zip(arc.iter())
        .all(|(a, b)| a == b)
    {
        return None;
    }
    composite_spec(comps[arc.len()])
}

// ── M' construction helper ────────────────────────────────────────────────────

/// Assemble M' from a pre-computed TBS hash.
///
/// ```text
/// M' = "CompositeAlgorithmSignatures2025" || Label || 0x00 || PH(TBS)
/// ```
///
/// The caller is responsible for computing `tbs_hash = PH(TBS)` using the
/// appropriate hash function for the given spec (`spec.hash`).
pub fn build_m_prime_from_hash(tbs_hash: &[u8], spec: &CompositeMlDsaSpec) -> Vec<u8> {
    const PREFIX: &[u8] = b"CompositeAlgorithmSignatures2025";
    let label = spec.label.as_bytes();
    let mut m = Vec::with_capacity(PREFIX.len() + label.len() + 1 + tbs_hash.len());
    m.extend_from_slice(PREFIX);
    m.extend_from_slice(label);
    m.push(0x00);
    m.extend_from_slice(tbs_hash);
    m
}

// ── DER encoding helpers ──────────────────────────────────────────────────────

/// Build a composite SubjectPublicKeyInfo DER from component raw public key bytes.
///
/// The composite OID arc is 1.3.6.1.5.5.7.6.<sub_arc>.
///
/// ```text
/// SubjectPublicKeyInfo ::= SEQUENCE {
///   algorithm  AlgorithmIdentifier,   -- OID only, no params
///   publicKey  BIT STRING             -- 0x00 || mldsa_pk || trad_pk
/// }
/// ```
pub fn encode_composite_spki(
    composite_oid: &[u32],
    mldsa_pk: &[u8],
    trad_pk: &[u8],
) -> Result<Vec<u8>, String> {
    use synta::tag::{Tag, TAG_SEQUENCE};
    use synta::types::string::BitStringRef;
    use synta::{Encoder, Encoding, ObjectIdentifier};

    let oid =
        ObjectIdentifier::new(composite_oid).map_err(|e| format!("invalid composite OID: {e}"))?;

    // Concatenated raw public key bytes.
    let mut raw_pk = Vec::with_capacity(mldsa_pk.len() + trad_pk.len());
    raw_pk.extend_from_slice(mldsa_pk);
    raw_pk.extend_from_slice(trad_pk);

    let pk_bit =
        BitStringRef::new(&raw_pk, 0).map_err(|e| format!("BIT STRING encoding failed: {e}"))?;

    (|| -> synta::Result<Vec<u8>> {
        let mut enc = Encoder::new(Encoding::Der);
        enc.start_constructed_no_guard(Tag::universal_constructed(TAG_SEQUENCE))?;
        // AlgorithmIdentifier SEQUENCE { OID }  — no parameters
        enc.start_constructed_no_guard(Tag::universal_constructed(TAG_SEQUENCE))?;
        enc.encode(&oid)?;
        enc.end_constructed()?;
        // SubjectPublicKey BIT STRING
        enc.encode(&pk_bit)?;
        enc.end_constructed()?;
        enc.finish()
    })()
    .map_err(|e| format!("SPKI DER encoding failed: {e}"))
}

/// Build a composite OneAsymmetricKey (PKCS#8) DER.
///
/// ```text
/// OneAsymmetricKey ::= SEQUENCE {
///   version            INTEGER 0,
///   privateKeyAlg      AlgorithmIdentifier,   -- OID only, no params
///   privateKey         OCTET STRING           -- mldsa_seed (32) || trad_sk
/// }
/// ```
pub fn encode_composite_pkcs8(
    composite_oid: &[u32],
    mldsa_seed: &[u8],
    trad_sk: &[u8],
) -> Result<Vec<u8>, String> {
    use synta::types::string::OctetStringRef;
    use synta::ObjectIdentifier;

    let oid =
        ObjectIdentifier::new(composite_oid).map_err(|e| format!("invalid composite OID: {e}"))?;

    let alg = crate::AlgorithmIdentifier {
        algorithm: oid,
        parameters: None,
    };

    let mut raw_sk = Vec::with_capacity(mldsa_seed.len() + trad_sk.len());
    raw_sk.extend_from_slice(mldsa_seed);
    raw_sk.extend_from_slice(trad_sk);

    let pki = crate::pkcs8_types::OneAsymmetricKey {
        version: synta::Integer::from_i64(0),
        private_key_algorithm: alg,
        private_key: OctetStringRef::new(&raw_sk),
        attributes: None,
        public_key: None,
    };

    pki.to_der()
        .map_err(|e| format!("composite PKCS#8 DER encoding failed: {e}"))
}

// ── DER split helpers ─────────────────────────────────────────────────────────

/// Extract the BIT STRING content (raw public key bytes) from a SPKI DER.
///
/// Skips the outer SEQUENCE and AlgorithmIdentifier, reads the BIT STRING,
/// and returns the content without the leading unused-bits byte.
pub fn extract_spki_bitstring_payload(spki_der: &[u8]) -> Result<Vec<u8>, String> {
    use synta::types::string::BitStringRef;
    use synta::{Decoder, Encoding};

    let mut dec = Decoder::new(spki_der, Encoding::Der);
    // Outer SEQUENCE
    dec.read_tag().map_err(|e| e.to_string())?;
    dec.read_length()
        .and_then(|l| l.definite())
        .map_err(|e| e.to_string())?;
    // Skip AlgorithmIdentifier
    dec.decode::<synta::Element>().map_err(|e| e.to_string())?;
    // BIT STRING
    let bs: BitStringRef = dec.decode().map_err(|e| e.to_string())?;
    Ok(bs.as_bytes().to_vec())
}

/// Split a composite SPKI BIT STRING payload into (mldsa_pk, trad_pk).
pub fn split_composite_spki_content<'a>(
    payload: &'a [u8],
    spec: &CompositeMlDsaSpec,
) -> Result<(&'a [u8], &'a [u8]), String> {
    if payload.len() < spec.mldsa_pk_size {
        return Err(format!(
            "composite SPKI payload too short for {} (mldsa_pk_size={}): got {} bytes",
            spec.label,
            spec.mldsa_pk_size,
            payload.len()
        ));
    }
    Ok(payload.split_at(spec.mldsa_pk_size))
}

/// Split a composite signature into (mldsa_sig, trad_sig).
pub fn split_composite_sig<'a>(
    sig: &'a [u8],
    spec: &CompositeMlDsaSpec,
) -> Result<(&'a [u8], &'a [u8]), String> {
    if sig.len() < spec.mldsa_sig_size {
        return Err(format!(
            "composite signature too short for {} (mldsa_sig_size={}): got {} bytes",
            spec.label,
            spec.mldsa_sig_size,
            sig.len()
        ));
    }
    Ok(sig.split_at(spec.mldsa_sig_size))
}

/// Split a composite private key content into (mldsa_seed, trad_sk).
///
/// The ML-DSA seed is always 32 bytes; `trad_sk` is the remainder.
pub fn split_composite_privkey(privkey_content: &[u8]) -> Result<(&[u8], &[u8]), String> {
    const SEED_LEN: usize = 32;
    if privkey_content.len() <= SEED_LEN {
        return Err(format!(
            "composite private key content too short: {} <= {} (no traditional key material)",
            privkey_content.len(),
            SEED_LEN
        ));
    }
    Ok(privkey_content.split_at(SEED_LEN))
}

/// Extract the `privateKey` OCTET STRING content from a PKCS#8 DER buffer.
pub fn pkcs8_private_key_content(pkcs8_der: &[u8]) -> Result<Vec<u8>, String> {
    crate::pkcs8_types::PrivateKeyInfo::from_der(pkcs8_der)
        .map(|pki| pki.private_key.as_bytes().to_vec())
        .map_err(|e| format!("PKCS#8 parse error: {e}"))
}

/// Build the composite OID component array for a given sub-arc.
///
/// Returns the full OID `1.3.6.1.5.5.7.6.<sub_arc>` as a fixed-size stack array.
pub fn composite_oid_components(sub_arc: u32) -> [u32; 9] {
    let mut comps = [0u32; 9];
    comps[..8].copy_from_slice(crate::oids::COMPOSITE_MLDSA_ARC);
    comps[8] = sub_arc;
    comps
}