krypteia-quantica 0.1.0

//! ML-DSA: Module-Lattice-Based Digital Signature Standard (FIPS 204).
//!
//! This crate implements the ML-DSA (formerly CRYSTALS-Dilithium) digital signature
//! scheme as specified in FIPS 204. ML-DSA is a post-quantum lattice-based signature
//! scheme built on the hardness of the Module Learning With Errors (M-LWE) and
//! Module Short Integer Solution (M-SIS) problems.
//!
//! Three parameter sets are provided, corresponding to NIST security levels 2, 3, and 5:
//!
//! - [`MlDsa44Scheme`] -- ML-DSA-44 (security level 2, ~128-bit classical security)
//! - [`MlDsa65Scheme`] -- ML-DSA-65 (security level 3, ~192-bit classical security)
//! - [`MlDsa87Scheme`] -- ML-DSA-87 (security level 5, ~256-bit classical security)
//!
//! # Examples
//!
//! ```rust
//! use quantica::ml_dsa::{MlDsa44Scheme, OsRng, MlDsa};
//!
//! let mut rng = OsRng;
//! let (pk, sk) = MlDsa44Scheme::keygen(&mut rng).unwrap();
//! let msg = b"Hello, post-quantum world!";
//! let sig = MlDsa44Scheme::sign(&sk, msg, b"", &mut rng).unwrap();
//! let valid = MlDsa44Scheme::verify(&pk, msg, b"", &sig).unwrap();
//! assert!(valid);
//! ```

/// Decomposition, rounding, and hint functions for signatures.
pub mod decompose;
/// Core ML-DSA key generation, signing, and verification algorithms.
pub mod dsa;
/// Encoding and decoding of keys, signatures, and polynomials.
pub mod encode;
/// Number Theoretic Transform for polynomial arithmetic.
pub mod ntt;
/// ML-DSA parameter sets and constants (FIPS 204, Table 1).
pub mod params;
/// Cryptographic random number generation trait and OS-backed implementation.
pub mod rng;
/// Sampling algorithms for matrix, secret, and masking generation.
pub mod sample;
/// Keccak/SHA-3/SHAKE hash function implementations (FIPS 202).
pub mod sha3;

/// First-order arithmetic masking for ML-DSA secret polynomials
/// (DPA / template-attack countermeasure). Available with the
/// `sca-protected` Cargo feature.
#[cfg(feature = "sca-protected")]
pub mod masked;

/// Fisher-Yates shuffled NTT for ML-DSA secret polynomials
/// (SPA / trace-alignment countermeasure). Available with the
/// `sca-protected` Cargo feature.
#[cfg(feature = "sca-protected")]
pub mod shuffle;

/// Small polynomial representation (i16, 512 B/poly) for secret vectors.
/// Uses the ML-KEM NTT (q=3329) for multiplications. Available with
/// the `small-secret` Cargo feature.
#[cfg(feature = "small-secret")]
pub mod smallpoly;

/// Compressed polynomial storage (3 bytes/coeff, 768 B/poly) and
/// compressed challenge (68 bytes + schoolbook multiply).
/// Available with `compressed-poly` or `compressed-challenge`.
#[cfg(any(feature = "compressed-poly", feature = "compressed-challenge"))]
pub mod compressed;

use alloc::vec::Vec;
use core::marker::PhantomData;

pub use params::{MlDsa44, MlDsa65, MlDsa87, Params};
pub use rng::CryptoRng;
#[cfg(feature = "std")]
pub use rng::OsRng;

use crate::secret::SecretBytes;

// =====================================================================
// Typed key / signature wrappers
// =====================================================================

/// ML-DSA **verifying key** (the public half of a key pair).
///
/// Type-tagged with the parameter set `P`. Public material — no
/// zeroization is performed on drop.
pub struct VerifyingKey<P: Params> {
    bytes: Vec<u8>,
    _marker: PhantomData<P>,
}

impl<P: Params> VerifyingKey<P> {
    /// Wrap a raw byte slice. Length is validated against
    /// [`Params::PK_LEN`].
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, MlDsaError> {
        if bytes.len() != P::PK_LEN {
            return Err(MlDsaError::InvalidPublicKey);
        }
        Ok(Self {
            bytes: bytes.to_vec(),
            _marker: PhantomData,
        })
    }

    /// Borrow the encoded verifying key as a byte slice.
    pub fn as_bytes(&self) -> &[u8] {
        &self.bytes
    }

    /// Length in bytes (always [`Params::PK_LEN`]).
    pub fn len(&self) -> usize {
        self.bytes.len()
    }
}

impl<P: Params> AsRef<[u8]> for VerifyingKey<P> {
    fn as_ref(&self) -> &[u8] {
        &self.bytes
    }
}

impl<P: Params> core::ops::Deref for VerifyingKey<P> {
    type Target = [u8];
    fn deref(&self) -> &[u8] {
        &self.bytes
    }
}

impl<P: Params> Clone for VerifyingKey<P> {
    fn clone(&self) -> Self {
        Self {
            bytes: self.bytes.clone(),
            _marker: PhantomData,
        }
    }
}

/// ML-DSA **signing key** (the private half of a key pair).
///
/// Backed by [`SecretBytes`] — wipes its memory on [`Drop`] via
/// `silentops::ct_zeroize`. Type-tagged with `P` to prevent
/// cross-parameter-set use.
pub struct SigningKey<P: Params> {
    bytes: SecretBytes,
    _marker: PhantomData<P>,
}

impl<P: Params> SigningKey<P> {
    /// Wrap a raw byte slice. Length is validated against
    /// [`Params::SK_LEN`].
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, MlDsaError> {
        if bytes.len() != P::SK_LEN {
            return Err(MlDsaError::InvalidSecretKey);
        }
        Ok(Self {
            bytes: SecretBytes::from_slice(bytes),
            _marker: PhantomData,
        })
    }

    /// Borrow the encoded signing key as a byte slice.
    pub fn as_bytes(&self) -> &[u8] {
        self.bytes.as_bytes()
    }

    /// Length in bytes (always [`Params::SK_LEN`]).
    pub fn len(&self) -> usize {
        self.bytes.len()
    }
}

impl<P: Params> AsRef<[u8]> for SigningKey<P> {
    fn as_ref(&self) -> &[u8] {
        self.bytes.as_bytes()
    }
}

impl<P: Params> core::ops::Deref for SigningKey<P> {
    type Target = [u8];
    fn deref(&self) -> &[u8] {
        self.bytes.as_bytes()
    }
}

/// ML-DSA **signature**. Type-tagged with the parameter set `P`.
///
/// Signatures are public material (transmitted alongside the message)
/// and are not zeroized.
pub struct Signature<P: Params> {
    bytes: Vec<u8>,
    _marker: PhantomData<P>,
}

impl<P: Params> Signature<P> {
    /// Wrap a raw byte slice. Length is validated against
    /// [`Params::SIG_LEN`].
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, MlDsaError> {
        if bytes.len() != P::SIG_LEN {
            return Err(MlDsaError::InvalidSignature);
        }
        Ok(Self {
            bytes: bytes.to_vec(),
            _marker: PhantomData,
        })
    }

    /// Borrow the encoded signature as a byte slice.
    pub fn as_bytes(&self) -> &[u8] {
        &self.bytes
    }

    /// Length in bytes (always [`Params::SIG_LEN`]).
    pub fn len(&self) -> usize {
        self.bytes.len()
    }
}

impl<P: Params> AsRef<[u8]> for Signature<P> {
    fn as_ref(&self) -> &[u8] {
        &self.bytes
    }
}

impl<P: Params> core::ops::Deref for Signature<P> {
    type Target = [u8];
    fn deref(&self) -> &[u8] {
        &self.bytes
    }
}

impl<P: Params> Clone for Signature<P> {
    fn clone(&self) -> Self {
        Self {
            bytes: self.bytes.clone(),
            _marker: PhantomData,
        }
    }
}

/// Error types for ML-DSA operations.
///
/// Each variant corresponds to a specific failure mode that can occur during
/// key generation, signing, or verification.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum MlDsaError {
    /// Random number generation failed.
    ///
    /// Returned when the underlying RNG (e.g., `/dev/urandom`) cannot provide bytes.
    RngFailure,
    /// Invalid public key (wrong length or format).
    ///
    /// The provided public key does not have the expected byte length for the
    /// chosen parameter set.
    InvalidPublicKey,
    /// Invalid secret key (wrong length or format).
    ///
    /// The provided secret key does not have the expected byte length for the
    /// chosen parameter set.
    InvalidSecretKey,
    /// Invalid signature (wrong length or format).
    ///
    /// The provided signature does not have the expected byte length for the
    /// chosen parameter set, or its internal encoding is malformed.
    InvalidSignature,
    /// Context string too long (> 255 bytes).
    ///
    /// FIPS 204 limits the optional context string to at most 255 bytes.
    ContextTooLong,
    /// Signature verification failed.
    VerificationFailed,
}

impl core::fmt::Display for MlDsaError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            MlDsaError::RngFailure => write!(f, "RNG failure"),
            MlDsaError::InvalidPublicKey => write!(f, "Invalid public key"),
            MlDsaError::InvalidSecretKey => write!(f, "Invalid secret key"),
            MlDsaError::InvalidSignature => write!(f, "Invalid signature"),
            MlDsaError::ContextTooLong => write!(f, "Context string too long"),
            MlDsaError::VerificationFailed => write!(f, "Signature verification failed"),
        }
    }
}

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

/// Generic ML-DSA interface parameterized by security level.
///
/// This struct provides the high-level API for ML-DSA key generation, signing,
/// and verification. The type parameter `P` selects the parameter set
/// ([`MlDsa44`], [`MlDsa65`], or [`MlDsa87`]).
///
/// All methods are stateless; [`MlDsa`] carries no runtime data and exists only
/// to bind the parameter set at the type level.
pub struct MlDsa<P: Params> {
    _marker: PhantomData<P>,
}

impl<P: Params> MlDsa<P> {
    /// Generate a new ML-DSA key pair.
    ///
    /// Implements Algorithm 1 of FIPS 204 (ML-DSA.KeyGen). Draws 32 random
    /// bytes from `rng` and derives a public key / secret key pair.
    ///
    /// Returns `(pk, sk)` where `pk` has length [`Self::PK_LEN`] and `sk` has
    /// length [`Self::SK_LEN`].
    ///
    /// # Errors
    ///
    /// Returns [`MlDsaError::RngFailure`] if the RNG cannot provide bytes.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use quantica::ml_dsa::{MlDsa, MlDsa44, OsRng};
    ///
    /// let mut rng = OsRng;
    /// let (pk, sk) = MlDsa::<MlDsa44>::keygen(&mut rng).unwrap();
    /// assert_eq!(pk.len(), MlDsa::<MlDsa44>::PK_LEN);
    /// assert_eq!(sk.len(), MlDsa::<MlDsa44>::SK_LEN);
    /// ```
    pub fn keygen(rng: &mut dyn CryptoRng) -> Result<(VerifyingKey<P>, SigningKey<P>), MlDsaError> {
        let (pk_v, sk_v) = dsa::keygen::<P>(rng)?;
        Ok((
            VerifyingKey {
                bytes: pk_v,
                _marker: PhantomData,
            },
            SigningKey {
                bytes: SecretBytes::from_vec(sk_v),
                _marker: PhantomData,
            },
        ))
    }

    /// Deterministic key generation from a 32-byte seed.
    ///
    /// Implements Algorithm 6 of FIPS 204 (ML-DSA.KeyGen_internal). This is
    /// primarily useful for testing with known-answer vectors.
    ///
    /// - `xi`: 32-byte random seed.
    ///
    /// Returns `(pk, sk)`.
    pub fn keygen_internal(xi: &[u8; 32]) -> (Vec<u8>, Vec<u8>) {
        dsa::keygen_internal::<P>(xi)
    }

    /// Sign a message with an optional context string.
    ///
    /// Implements Algorithm 2 of FIPS 204 (ML-DSA.Sign). Uses **hedged signing**:
    /// 32 random bytes are drawn from `rng` and mixed with the secret key material
    /// to produce the per-signature nonce. This provides resilience against
    /// fault attacks compared to purely deterministic signing.
    ///
    /// The signing algorithm uses a rejection sampling loop internally: candidate
    /// signatures are generated until one passes all norm checks, so execution
    /// time may vary.
    ///
    /// - `sk`: secret key (must be [`Self::SK_LEN`] bytes).
    /// - `msg`: message to sign (arbitrary length).
    /// - `ctx`: optional context string (at most 255 bytes).
    /// - `rng`: source of randomness for hedged signing.
    ///
    /// Returns a signature of length [`Self::SIG_LEN`].
    ///
    /// # Errors
    ///
    /// - [`MlDsaError::InvalidSecretKey`] if `sk` has the wrong length.
    /// - [`MlDsaError::ContextTooLong`] if `ctx` exceeds 255 bytes.
    /// - [`MlDsaError::RngFailure`] if the RNG cannot provide bytes.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use quantica::ml_dsa::{MlDsa, MlDsa44, OsRng};
    ///
    /// let mut rng = OsRng;
    /// let (pk, sk) = MlDsa::<MlDsa44>::keygen(&mut rng).unwrap();
    /// let sig = MlDsa::<MlDsa44>::sign(&sk, b"message", b"", &mut rng).unwrap();
    /// assert_eq!(sig.len(), MlDsa::<MlDsa44>::SIG_LEN);
    /// ```
    pub fn sign(
        sk: &SigningKey<P>,
        msg: &[u8],
        ctx: &[u8],
        rng: &mut dyn CryptoRng,
    ) -> Result<Signature<P>, MlDsaError> {
        let sig_v = dsa::sign::<P>(sk.as_bytes(), msg, ctx, rng)?;
        Ok(Signature {
            bytes: sig_v,
            _marker: PhantomData,
        })
    }

    /// Deterministic signing (internal / testing use).
    ///
    /// Implements Algorithm 7 of FIPS 204 (ML-DSA.Sign_internal). The caller
    /// supplies the pre-formatted message `m_prime` and a 32-byte `rnd` value
    /// directly. Setting `rnd` to all zeros produces fully deterministic
    /// signatures; using random bytes produces hedged signatures.
    ///
    /// - `sk`: encoded secret key.
    /// - `m_prime`: pre-formatted message (`0x00 || len(ctx) || ctx || msg`).
    /// - `rnd`: 32-byte randomness (all zeros for deterministic mode).
    ///
    /// # Errors
    ///
    /// Returns [`MlDsaError::InvalidSecretKey`] if `sk` has the wrong length.
    pub fn sign_internal(sk: &[u8], m_prime: &[u8], rnd: &[u8; 32]) -> Result<Vec<u8>, MlDsaError> {
        dsa::sign_internal::<P>(sk, m_prime, rnd)
    }

    /// Verify a signature on a message with an optional context string.
    ///
    /// Implements Algorithm 3 of FIPS 204 (ML-DSA.Verify).
    ///
    /// - `pk`: public key (must be [`Self::PK_LEN`] bytes).
    /// - `msg`: the signed message.
    /// - `ctx`: the context string used during signing (at most 255 bytes).
    /// - `sig`: the signature (must be [`Self::SIG_LEN`] bytes).
    ///
    /// Returns `Ok(true)` if the signature is valid, `Ok(false)` if verification
    /// fails (invalid signature content), or an `Err` for structural issues.
    ///
    /// # Errors
    ///
    /// - [`MlDsaError::InvalidPublicKey`] if `pk` has the wrong length.
    /// - [`MlDsaError::InvalidSignature`] if `sig` has the wrong length.
    /// - [`MlDsaError::ContextTooLong`] if `ctx` exceeds 255 bytes.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use quantica::ml_dsa::{MlDsa, MlDsa44, OsRng};
    ///
    /// let mut rng = OsRng;
    /// let (pk, sk) = MlDsa::<MlDsa44>::keygen(&mut rng).unwrap();
    /// let sig = MlDsa::<MlDsa44>::sign(&sk, b"msg", b"", &mut rng).unwrap();
    /// assert!(MlDsa::<MlDsa44>::verify(&pk, b"msg", b"", &sig).unwrap());
    /// ```
    pub fn verify(pk: &VerifyingKey<P>, msg: &[u8], ctx: &[u8], sig: &Signature<P>) -> Result<bool, MlDsaError> {
        dsa::verify::<P>(pk.as_bytes(), msg, ctx, sig.as_bytes())
    }

    /// Verify a signature (internal / testing use).
    ///
    /// Implements Algorithm 8 of FIPS 204 (ML-DSA.Verify_internal). The caller
    /// supplies the pre-formatted message `m_prime` directly.
    ///
    /// - `pk`: encoded public key.
    /// - `m_prime`: pre-formatted message.
    /// - `sig`: encoded signature.
    ///
    /// Returns `Ok(true)` on success or `Ok(false)` when the signature is invalid.
    pub fn verify_internal(pk: &[u8], m_prime: &[u8], sig: &[u8]) -> Result<bool, MlDsaError> {
        dsa::verify_internal::<P>(pk, m_prime, sig)
    }

    /// Public key length in bytes for this parameter set.
    pub const PK_LEN: usize = P::PK_LEN;

    /// Secret key length in bytes for this parameter set.
    pub const SK_LEN: usize = P::SK_LEN;

    /// Signature length in bytes for this parameter set.
    pub const SIG_LEN: usize = P::SIG_LEN;
}

/// Convenience alias for ML-DSA-44 (NIST security level 2).
pub type MlDsa44Scheme = MlDsa<MlDsa44>;
/// Convenience alias for ML-DSA-65 (NIST security level 3).
pub type MlDsa65Scheme = MlDsa<MlDsa65>;
/// Convenience alias for ML-DSA-87 (NIST security level 5).
pub type MlDsa87Scheme = MlDsa<MlDsa87>;