qs-crypto 0.1.0

#![cfg_attr(not(feature = "std"), no_std)]
//! Quantum-Sign cryptographic module boundary.
//! Provides deterministic randomness interfaces and (eventually) signature/KEM glue.
#![forbid(unsafe_code)]
#![deny(missing_docs)]

#[cfg(not(feature = "std"))]
extern crate alloc;

/// Utilities for canonical SPKI handling and key identifiers.
pub mod public;
#[cfg(feature = "std")]
use std::{format, string::String, vec, vec::Vec};

#[cfg(not(feature = "std"))]
use alloc::{format, string::String, vec, vec::Vec};
use core::{convert::TryFrom, fmt, mem, str::FromStr};
use ml_dsa::{
    EncodedSignature, EncodedSigningKey, EncodedVerifyingKey, KeyGen, MlDsa87,
    Signature as MlSignature, SigningKey, VerifyingKey,
};
use pkcs8::spki::EncodePublicKey;
pub use public::{
    kid_from_spki_der, spki_der_canonical, spki_mldsa_paramset, spki_subject_key_bytes,
};
use qs_drbg::rand_adapter::DrbgRng;
use qs_drbg::{Error as InnerError, HmacDrbg};
use sha2::{Digest, Sha256};
use zeroize::{Zeroize, ZeroizeOnDrop};

/// FIPS 204 ML-DSA-87 canonical lengths (bytes).
pub mod mldsa87 {
    /// Public key length as mandated by FIPS 204 Table 2.
    pub const PUBLIC_KEY_LEN: usize = 2592;
    /// Secret key length as mandated by FIPS 204 Table 2.
    pub const SECRET_KEY_LEN: usize = 4896;
    /// Signature length as mandated by FIPS 204 Table 2.
    pub const SIGNATURE_LEN: usize = 4627;
}

/// Length in bytes of an ML-DSA-87 signing key.
pub const MLDSA87_SECRET_KEY_LEN: usize = mem::size_of::<EncodedSigningKey<MlDsa87>>();
/// Length in bytes of an ML-DSA-87 verifying key.
pub const MLDSA87_PUBLIC_KEY_LEN: usize = mem::size_of::<EncodedVerifyingKey<MlDsa87>>();
/// Length in bytes of an ML-DSA-87 signature.
pub const MLDSA87_SIGNATURE_LEN: usize = mem::size_of::<EncodedSignature<MlDsa87>>();

const SIGNING_CONTEXT: &[u8] = b"quantum-sign.v1";
/// Domain separator used when binding algorithm + policy into the signed transcript.
const TRANSCRIPT_DOMAIN: &[u8] = b"quantum-sign:v1";
/// Length of the transcript digest (SHA-256 output).
pub const TRANSCRIPT_DIGEST_LEN: usize = 32;

/// Supported artifact digest algorithms.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum DigestAlg {
    /// SHA-256 with 32-byte output.
    Sha256,
    /// SHA-512 with 64-byte output.
    Sha512,
    /// SHAKE256 XOF truncated to 64 bytes.
    Shake256_64,
}

impl DigestAlg {
    /// String representation used in policy and intent metadata.
    pub fn as_str(self) -> &'static str {
        match self {
            DigestAlg::Sha256 => "sha256",
            DigestAlg::Sha512 => "sha512",
            DigestAlg::Shake256_64 => "shake256-64",
        }
    }

    /// Output length of the digest in bytes.
    pub fn output_len(self) -> usize {
        match self {
            DigestAlg::Sha256 => 32,
            DigestAlg::Sha512 | DigestAlg::Shake256_64 => 64,
        }
    }
}

impl FromStr for DigestAlg {
    type Err = &'static str;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_ascii_lowercase().as_str() {
            "sha256" => Ok(DigestAlg::Sha256),
            "sha512" => Ok(DigestAlg::Sha512),
            "shake256-64" | "shake256" => Ok(DigestAlg::Shake256_64),
            _ => Err("unsupported digest algorithm"),
        }
    }
}

/// Errors surfaced by deterministic random bit generators inside the crypto module.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DrbgError {
    /// Generated output exceeded the maximum request length.
    RequestTooLarge,
    /// Generator must be reseeded with fresh entropy before more output is produced.
    ReseedRequired,
    /// Underlying entropy source failed.
    EntropyUnavailable,
    /// Entropy health check failed (weak or repeated seed material).
    EntropyHealthFailed,
}

impl fmt::Display for DrbgError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            DrbgError::RequestTooLarge => write!(f, "DRBG request exceeds per-call limit"),
            DrbgError::ReseedRequired => write!(f, "DRBG reseed required"),
            DrbgError::EntropyUnavailable => write!(f, "OS entropy unavailable"),
            DrbgError::EntropyHealthFailed => write!(f, "entropy health check failed"),
        }
    }
}

#[cfg(feature = "std")]
impl std::error::Error for DrbgError {}

impl From<InnerError> for DrbgError {
    fn from(value: InnerError) -> Self {
        match value {
            InnerError::RequestTooLarge => DrbgError::RequestTooLarge,
            InnerError::ReseedRequired => DrbgError::ReseedRequired,
            InnerError::EntropyUnavailable => DrbgError::EntropyUnavailable,
            InnerError::EntropyHealthFailed => DrbgError::EntropyHealthFailed,
        }
    }
}

/// Compute the policy-bound transcript digest consumed by ML-DSA signatures.
pub fn transcript_digest(
    sign_alg: &str,
    digest_alg: &str,
    message_digest: &[u8],
    policy_hash: Option<&[u8]>,
) -> [u8; TRANSCRIPT_DIGEST_LEN] {
    let mut hasher = Sha256::new();
    hasher.update(TRANSCRIPT_DOMAIN);
    hasher.update(b"|alg:");
    hasher.update(sign_alg.as_bytes());
    hasher.update(b"|hash:");
    hasher.update(digest_alg.as_bytes());
    if let Some(ph) = policy_hash {
        hasher.update(b"|policy:");
        hasher.update(ph);
    }
    hasher.update(b"|msg:");
    hasher.update(message_digest);
    let digest = hasher.finalize();
    let mut out = [0u8; TRANSCRIPT_DIGEST_LEN];
    out.copy_from_slice(&digest);
    out
}

/// Trait implemented by deterministic random bit generators used by Quantum-Sign.
pub trait DeterministicRng {
    /// Fill `out` with pseudorandom bytes.
    fn fill_bytes(&mut self, out: &mut [u8]) -> Result<(), DrbgError>;

    /// Reseed the generator with new entropy and optional additional input.
    fn reseed(&mut self, entropy: &[u8], additional_input: Option<&[u8]>) -> Result<(), DrbgError>;

    /// Adjust the reseed interval (number of generate calls permitted before mandatory reseed).
    fn set_reseed_interval(&mut self, interval: u64);

    /// Adjust the byte budget that forces a reseed.
    fn set_max_bytes_between_reseed(&mut self, bytes: u128);
}

/// HMAC-DRBG (SHA-512) wrapper implementing [`DeterministicRng`].
#[derive(Zeroize, ZeroizeOnDrop)]
pub struct HmacSha512Drbg {
    inner: HmacDrbg,
}

impl HmacSha512Drbg {
    /// Instantiate DRBG from caller-provided entropy/nonce/personalization strings.
    pub fn new(
        entropy: &[u8],
        nonce: &[u8],
        personalization: Option<&[u8]>,
    ) -> Result<Self, DrbgError> {
        let inner = HmacDrbg::new(entropy, nonce, personalization).map_err(DrbgError::from)?;
        Ok(Self { inner })
    }

    /// Instantiate DRBG using the operating system CSPRNG for entropy and nonce.
    pub fn from_os(personalization: Option<&[u8]>) -> Result<Self, DrbgError> {
        let inner = HmacDrbg::from_os(personalization).map_err(DrbgError::from)?;
        Ok(Self { inner })
    }

    /// Borrow the inner DRBG mutably (for adapters).
    pub fn inner_mut(&mut self) -> &mut HmacDrbg {
        &mut self.inner
    }

    /// Expose the internal state for testing; callers must not use this outside tests.
    #[cfg(test)]
    pub fn inner(&self) -> &HmacDrbg {
        &self.inner
    }
}

impl fmt::Debug for HmacSha512Drbg {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("HmacSha512Drbg").finish_non_exhaustive()
    }
}

impl DeterministicRng for HmacSha512Drbg {
    fn fill_bytes(&mut self, out: &mut [u8]) -> Result<(), DrbgError> {
        self.inner.generate(out, None).map_err(DrbgError::from)
    }

    fn reseed(&mut self, entropy: &[u8], additional_input: Option<&[u8]>) -> Result<(), DrbgError> {
        self.inner
            .reseed(entropy, additional_input)
            .map_err(DrbgError::from)
    }

    fn set_reseed_interval(&mut self, interval: u64) {
        self.inner.set_reseed_interval(interval);
    }

    fn set_max_bytes_between_reseed(&mut self, bytes: u128) {
        self.inner.set_max_bytes_between_reseed(bytes);
    }
}

/// Utility for generating a fixed number of bytes using a fresh OS-seeded DRBG.
pub fn random_bytes(len: usize) -> Result<Vec<u8>, DrbgError> {
    let mut drbg = HmacSha512Drbg::from_os(None)?;
    let mut buf = vec![0u8; len];
    drbg.fill_bytes(&mut buf)?;
    Ok(buf)
}

/// Simple keypair container.
#[derive(Debug, Clone)]
pub struct Keypair {
    /// Secret signing key bytes.
    pub secret: Vec<u8>,
    /// Public verifying key bytes.
    pub public: Vec<u8>,
}

impl Drop for Keypair {
    fn drop(&mut self) {
        self.zeroize();
    }
}

impl Zeroize for Keypair {
    fn zeroize(&mut self) {
        self.secret.zeroize();
        self.public.zeroize();
    }
}

impl ZeroizeOnDrop for Keypair {}

/// Cryptographic operation errors.
#[derive(Debug)]
pub enum CryptoError {
    /// Deterministic RNG failure.
    Drbg(DrbgError),
    /// Provided key bytes were malformed.
    InvalidKey,
    /// Provided signature bytes were malformed.
    InvalidSignature,
    /// Signature generation failed.
    SigningFailed,
    /// The supplied digest length was incorrect.
    BadDigestLen {
        /// Expected digest length in bytes.
        expected: usize,
        /// Actual digest length supplied by the caller.
        got: usize,
    },
    /// SPKI or public key parsing failed.
    PublicKey(String),
}

impl fmt::Display for CryptoError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            CryptoError::Drbg(err) => write!(f, "{err}"),
            CryptoError::InvalidKey => write!(f, "invalid key material"),
            CryptoError::InvalidSignature => write!(f, "invalid signature"),
            CryptoError::SigningFailed => write!(f, "signature generation failed"),
            CryptoError::BadDigestLen { expected, got } => {
                write!(f, "bad digest length: expected {expected}, got {got}")
            }
            CryptoError::PublicKey(msg) => write!(f, "public key error: {msg}"),
        }
    }
}

impl From<DrbgError> for CryptoError {
    fn from(err: DrbgError) -> Self {
        CryptoError::Drbg(err)
    }
}

#[cfg(feature = "std")]
impl std::error::Error for CryptoError {}

/// Generate an ML-DSA-87 keypair using the provided DRBG.
pub fn keypair_mldsa87(drbg: &mut HmacSha512Drbg) -> Result<Keypair, CryptoError> {
    let mut rng = DrbgRng::new(drbg.inner_mut());
    let kp = MlDsa87::key_gen(&mut rng);
    let sk = kp.signing_key().encode().to_vec();
    let pk = kp.verifying_key().encode().to_vec();
    Ok(Keypair {
        secret: sk,
        public: pk,
    })
}

/// Produce an ML-DSA-87 signature over `message` (usually a digest) with explicit context.
pub fn sign_mldsa87(
    drbg: &mut HmacSha512Drbg,
    secret_key: &[u8],
    message_digest: &[u8],
    digest_alg: DigestAlg,
    policy_hash: Option<&[u8]>,
) -> Result<Vec<u8>, CryptoError> {
    if message_digest.len() != digest_alg.output_len() {
        return Err(CryptoError::BadDigestLen {
            expected: digest_alg.output_len(),
            got: message_digest.len(),
        });
    }
    let enc =
        EncodedSigningKey::<MlDsa87>::try_from(secret_key).map_err(|_| CryptoError::InvalidKey)?;
    let sk = SigningKey::<MlDsa87>::decode(&enc);
    let transcript =
        transcript_digest("mldsa-87", digest_alg.as_str(), message_digest, policy_hash);
    drbg.inner_mut()
        .generate(&mut [], Some(&transcript))
        .map_err(DrbgError::from)?;
    let mut rng = DrbgRng::new(drbg.inner_mut());
    let sig = sk
        .sign_randomized(transcript.as_slice(), SIGNING_CONTEXT, &mut rng)
        .map_err(|_| CryptoError::SigningFailed)?;
    Ok(sig.encode().to_vec())
}

/// Verify an ML-DSA-87 signature over `message`.
pub fn verify_mldsa87(
    public_key: &[u8],
    message_digest: &[u8],
    digest_alg: DigestAlg,
    signature: &[u8],
    policy_hash: Option<&[u8]>,
) -> Result<(), CryptoError> {
    if message_digest.len() != digest_alg.output_len() {
        return Err(CryptoError::BadDigestLen {
            expected: digest_alg.output_len(),
            got: message_digest.len(),
        });
    }
    if public_key.len() != mldsa87::PUBLIC_KEY_LEN {
        return Err(CryptoError::InvalidSignature);
    }
    if signature.len() != mldsa87::SIGNATURE_LEN {
        return Err(CryptoError::InvalidSignature);
    }
    let enc_vk = EncodedVerifyingKey::<MlDsa87>::try_from(public_key)
        .map_err(|_| CryptoError::InvalidKey)?;
    let vk = VerifyingKey::<MlDsa87>::decode(&enc_vk);
    let enc_sig = EncodedSignature::<MlDsa87>::try_from(signature)
        .map_err(|_| CryptoError::InvalidSignature)?;
    let sig = MlSignature::<MlDsa87>::decode(&enc_sig).ok_or(CryptoError::InvalidSignature)?;
    let transcript =
        transcript_digest("mldsa-87", digest_alg.as_str(), message_digest, policy_hash);
    if vk.verify_with_context(transcript.as_slice(), SIGNING_CONTEXT, &sig) {
        Ok(())
    } else {
        Err(CryptoError::InvalidSignature)
    }
}

/// Compute the canonical key identifier for an ML-DSA-87 verifying key.
pub fn kid_from_public_key(public_key: &[u8]) -> Result<String, CryptoError> {
    let spki = public_key_to_spki(public_key)?;
    Ok(public::kid_from_spki_der(&spki))
}

/// Convert a raw ML-DSA-87 public key into canonical SPKI DER bytes.
pub fn public_key_to_spki(public_key: &[u8]) -> Result<Vec<u8>, CryptoError> {
    let enc_vk = EncodedVerifyingKey::<MlDsa87>::try_from(public_key)
        .map_err(|_| CryptoError::InvalidKey)?;
    let vk = VerifyingKey::<MlDsa87>::decode(&enc_vk);
    let spki = vk
        .to_public_key_der()
        .map_err(|_| CryptoError::InvalidKey)?;
    Ok(spki.as_bytes().to_vec())
}

/// Verify an ML-DSA-87 signature against an SPKI-encoded public key.
pub fn verify_mldsa87_spki(
    spki_der: &[u8],
    message_digest: &[u8],
    digest_alg: DigestAlg,
    signature: &[u8],
    policy_hash: Option<&[u8]>,
) -> Result<(), CryptoError> {
    let public_key =
        spki_subject_key_bytes(spki_der).map_err(|e| CryptoError::PublicKey(format!("{e}")))?;
    if public_key.len() != mldsa87::PUBLIC_KEY_LEN {
        return Err(CryptoError::InvalidSignature);
    }
    if signature.len() != mldsa87::SIGNATURE_LEN {
        return Err(CryptoError::InvalidSignature);
    }
    verify_mldsa87(
        &public_key,
        message_digest,
        digest_alg,
        signature,
        policy_hash,
    )
}
/// Return true when the supplied signature algorithm identifier satisfies Level-5 policy.
pub fn is_level5_sig_alg(alg: &str) -> bool {
    matches!(
        alg,
        "mldsa-87"
            | "slh-dsa-sha2-256s"
            | "slh-dsa-sha2-256f"
            | "slh-dsa-shake-256s"
            | "slh-dsa-shake-256f"
    )
}

#[cfg(test)]
mod tests {
    use super::*;
    use sha2::{Digest, Sha512};

    fn entropy(seed: u8) -> [u8; 48] {
        let mut out = [0u8; 48];
        for (i, byte) in out.iter_mut().enumerate() {
            *byte = seed.wrapping_add(i as u8);
        }
        out
    }

    fn nonce(seed: u8) -> [u8; 16] {
        let mut out = [0u8; 16];
        for (i, byte) in out.iter_mut().enumerate() {
            *byte = seed.wrapping_add((i * 5) as u8);
        }
        out
    }

    #[test]
    fn keygen_and_sign_verify() {
        let mut drbg = HmacSha512Drbg::new(&entropy(1), &nonce(2), Some(b"keygen")).unwrap();
        let kp = keypair_mldsa87(&mut drbg).expect("keypair");

        let mut signer_drbg = HmacSha512Drbg::new(&entropy(3), &nonce(4), Some(b"sign")).unwrap();
        let mut hasher = Sha512::new();
        hasher.update(b"deterministic digest");
        let digest = hasher.finalize().to_vec();
        let sig = sign_mldsa87(
            &mut signer_drbg,
            &kp.secret,
            &digest,
            DigestAlg::Sha512,
            None,
        )
        .expect("sign");

        verify_mldsa87(&kp.public, &digest, DigestAlg::Sha512, &sig, None).expect("verify");

        let mut bad_hasher = Sha512::new();
        bad_hasher.update(b"different");
        let bad_digest = bad_hasher.finalize().to_vec();
        assert!(verify_mldsa87(&kp.public, &bad_digest, DigestAlg::Sha512, &sig, None).is_err());
    }

    #[test]
    fn reject_wrong_digest_length() {
        let mut drbg = HmacSha512Drbg::new(&entropy(7), &nonce(8), Some(b"keygen")).unwrap();
        let kp = keypair_mldsa87(&mut drbg).expect("keypair");
        let mut signer_drbg = HmacSha512Drbg::new(&entropy(9), &nonce(10), Some(b"sign")).unwrap();
        let short = vec![0u8; DigestAlg::Sha512.output_len() - 1];
        assert!(matches!(
            sign_mldsa87(
                &mut signer_drbg,
                &kp.secret,
                &short,
                DigestAlg::Sha512,
                None
            ),
            Err(CryptoError::BadDigestLen { .. })
        ));
    }

    #[test]
    fn kid_from_public_key_has_expected_length() {
        let mut drbg = HmacSha512Drbg::new(&entropy(11), &nonce(12), Some(b"keygen")).unwrap();
        let kp = keypair_mldsa87(&mut drbg).expect("keypair");
        let kid = kid_from_public_key(&kp.public).expect("kid from public");
        assert_eq!(kid.len(), 16);
    }

    #[test]
    fn verify_rejects_truncated_signature() {
        let mut drbg = HmacSha512Drbg::new(&entropy(21), &nonce(22), Some(b"keygen")).unwrap();
        let kp = keypair_mldsa87(&mut drbg).expect("keypair");
        let mut signer_drbg = HmacSha512Drbg::new(&entropy(23), &nonce(24), Some(b"sign")).unwrap();
        let digest = vec![0xAB; DigestAlg::Sha512.output_len()];
        let sig = sign_mldsa87(
            &mut signer_drbg,
            &kp.secret,
            &digest,
            DigestAlg::Sha512,
            None,
        )
        .expect("sign");
        assert_eq!(sig.len(), mldsa87::SIGNATURE_LEN);

        let mut truncated = sig.clone();
        truncated.truncate(truncated.len() - 1);
        assert!(verify_mldsa87(&kp.public, &digest, DigestAlg::Sha512, &sig, None).is_ok());
        assert!(verify_mldsa87(&kp.public, &digest, DigestAlg::Sha512, &truncated, None).is_err());

        // SPKI path
        let spki = public_key_to_spki(&kp.public).expect("spki");
        assert!(verify_mldsa87_spki(&spki, &digest, DigestAlg::Sha512, &truncated, None).is_err());
    }
}