krypteia-quantica 0.1.0

//! # ML-KEM — Module-Lattice-Based Key-Encapsulation Mechanism
//!
//! A pure-Rust implementation of **FIPS 203** (ML-KEM) providing three
//! parameter sets: [`MlKem512`], [`MlKem768`], and [`MlKem1024`].
//!
//! Uses only the Rust standard library. No external dependencies.
//!
//! ## Quick start
//!
//! ```no_run
//! use quantica::ml_kem::*;
//!
//! let mut rng = OsRng;
//!
//! // Key generation
//! let (ek, dk) = MlKem::<MlKem768>::keygen(&mut rng).unwrap();
//!
//! // Encapsulation (sender)
//! let (shared_secret_s, ciphertext) = MlKem::<MlKem768>::encaps(&ek, &mut rng).unwrap();
//!
//! // Decapsulation (receiver)
//! let shared_secret_r = MlKem::<MlKem768>::decaps(&dk, &ciphertext, &mut rng).unwrap();
//!
//! assert_eq!(shared_secret_s, shared_secret_r);
//! ```
//!
//! ## Side-channel countermeasures
//!
//! This implementation includes multiple layers of protection against
//! physical side-channel attacks:
//!
//! - **Constant-time**: no secret-dependent branches or memory accesses
//! - **Zeroization**: all secret intermediates erased via volatile writes
//! - **First-order masking**: secret polynomials split into additive shares (DPA/template)
//! - **Double decaps**: fault detection on FO comparison (DFA)
//! - **dk integrity**: H(ek) verification at decaps time (DFA)
//! - **NTT shuffling**: randomized butterfly order (SPA)
//!
//! ## Module overview
//!
//! | Module       | Description |
//! |--------------|-------------|
//! | [`params`]   | Parameter sets and the [`Params`] trait |
//! | [`kem`]      | Top-level ML-KEM algorithms (Algorithms 16-21) |
//! | [`kpke`]     | K-PKE component scheme (Algorithms 13-15) |
//! | [`ntt`]      | Number-Theoretic Transform and polynomial arithmetic |
//! | [`encode`]   | Encoding, decoding, compression, decompression |
//! | [`sample`]   | Polynomial sampling (NTT domain and CBD) |
//! | [`sha3`]     | SHA-3 / SHAKE primitives (FIPS 202) |
//! | [`rng`]      | Cryptographic RNG trait and OS-backed implementation |
//! | [`masked`]   | First-order arithmetic masking for polynomials |
//! | [`shuffle`]  | Fisher-Yates shuffle for NTT butterfly randomization |

/// Byte encoding/decoding and compression/decompression (Algorithms 3-6).
pub mod encode;
/// ML-KEM key encapsulation: keygen, encaps, decaps (Algorithms 16-21).
pub mod kem;
/// K-PKE component scheme: key generation, encryption, decryption (Algorithms 13-15).
pub mod kpke;
/// First-order arithmetic masking for DPA/template attack protection.
pub mod masked;
/// Number-Theoretic Transform and modular polynomial arithmetic.
pub mod ntt;
/// ML-KEM parameter sets and the [`Params`] trait.
pub mod params;
/// Cryptographic random number generation trait and OS-backed implementation.
pub mod rng;
/// Polynomial sampling algorithms: [`sample::sample_ntt`] and [`sample::sample_poly_cbd`].
pub mod sample;
/// SHA-3 and SHAKE hash function primitives (FIPS 202).
pub mod sha3;
/// Fisher-Yates shuffle for NTT butterfly index randomization (SPA protection).
pub mod shuffle;

pub use params::{MlKem512, MlKem768, MlKem1024, Params};
pub use rng::CryptoRng;
#[cfg(feature = "std")]
pub use rng::OsRng;

use crate::secret::{SecretArray, SecretBytes};
use alloc::vec::Vec;
use core::marker::PhantomData;

// =====================================================================
// Typed key / ciphertext / shared-secret wrappers
// =====================================================================

/// ML-KEM **encapsulation key** (the public half of a key pair).
///
/// Type-tagged with the parameter set `P` so the type system can
/// catch mismatches between security levels at compile time. The
/// underlying bytes are a plain `Vec<u8>` because the encapsulation
/// key is public material — no zeroization is performed on drop.
pub struct EncapsulationKey<P: Params> {
    bytes: Vec<u8>,
    _marker: PhantomData<P>,
}

impl<P: Params> EncapsulationKey<P> {
    /// Wrap a raw byte vector. Length is validated against
    /// [`Params::EK_LEN`] for the parameter set `P`.
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, MlKemError> {
        if bytes.len() != P::EK_LEN {
            return Err(MlKemError::InvalidEncapsulationKey);
        }
        Ok(Self {
            bytes: bytes.to_vec(),
            _marker: PhantomData,
        })
    }

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

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

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

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

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

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

impl<P: Params> DecapsulationKey<P> {
    /// Wrap a raw byte slice into a zeroizing decapsulation key.
    /// Length is validated against [`Params::DK_LEN`].
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, MlKemError> {
        if bytes.len() != P::DK_LEN {
            return Err(MlKemError::InvalidDecapsulationKey);
        }
        Ok(Self {
            bytes: SecretBytes::from_slice(bytes),
            _marker: PhantomData,
        })
    }

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

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

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

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

/// ML-KEM **ciphertext** wrapping the encapsulated shared secret.
///
/// Type-tagged with `P`. Ciphertexts are not secret (they travel on
/// the wire), so no zeroization is performed.
pub struct Ciphertext<P: Params> {
    bytes: Vec<u8>,
    _marker: PhantomData<P>,
}

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

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

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

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

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

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

/// 32-byte ML-KEM shared secret.
///
/// Backed by a [`SecretArray<32>`] which wipes itself on [`Drop`].
/// Equality is constant-time via `silentops::ct_eq`.
pub type SharedSecret = SecretArray<32>;

/// Errors returned by ML-KEM operations.
///
/// All error variants are designed to avoid leaking secret information;
/// timing is independent of the specific failure path.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MlKemError {
    /// The cryptographic random number generator failed to produce bytes.
    ///
    /// This typically indicates a system-level failure (e.g., `/dev/urandom`
    /// is unavailable).
    RngFailure,
    /// The provided encapsulation (public) key has an invalid length or
    /// fails the modulus check required by FIPS 203 Section 7.2.
    InvalidEncapsulationKey,
    /// The provided decapsulation (private) key has an invalid length or
    /// fails the `H(ek)` integrity check embedded in the key.
    ///
    /// This may indicate storage corruption or fault injection.
    InvalidDecapsulationKey,
    /// The provided ciphertext has an invalid length for the parameter set.
    InvalidCiphertext,
}

impl core::fmt::Display for MlKemError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            Self::RngFailure => write!(f, "Random bit generation failed"),
            Self::InvalidEncapsulationKey => write!(f, "Invalid encapsulation key"),
            Self::InvalidDecapsulationKey => write!(f, "Invalid decapsulation key"),
            Self::InvalidCiphertext => write!(f, "Invalid ciphertext"),
        }
    }
}

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

/// Main ML-KEM interface, generic over a [`Params`] parameter set.
///
/// This is the primary API for ML-KEM key encapsulation. It wraps the
/// lower-level functions in [`kem`] and provides a convenient, type-safe
/// interface parameterized by security level.
///
/// # Type parameters
///
/// * `P` - One of [`MlKem512`], [`MlKem768`], or [`MlKem1024`], selecting
///   the security category (1, 3, or 5 respectively).
///
/// # Example
///
/// ```no_run
/// use quantica::ml_kem::*;
///
/// let mut rng = OsRng;
/// let (ek, dk) = MlKem::<MlKem768>::keygen(&mut rng).unwrap();
/// let (ss, ct) = MlKem::<MlKem768>::encaps(&ek, &mut rng).unwrap();
/// let ss2 = MlKem::<MlKem768>::decaps(&dk, &ct, &mut rng).unwrap();
/// assert_eq!(ss, ss2);
/// ```
pub struct MlKem<P: Params>(core::marker::PhantomData<P>);

impl<P: Params> MlKem<P> {
    /// Generate an ML-KEM key pair.
    ///
    /// Produces an encapsulation key (public) and a decapsulation key (private)
    /// using randomness from the provided RNG. Implements Algorithm 19 of FIPS 203.
    ///
    /// The decapsulation key includes an integrity hash `H(ek)` that is verified
    /// during decapsulation to detect storage corruption (DFA protection).
    ///
    /// # Arguments
    ///
    /// * `rng` - A cryptographic random number generator implementing [`CryptoRng`].
    ///
    /// # Returns
    ///
    /// A tuple `(encapsulation_key, decapsulation_key)` of typed
    /// wrappers. The decapsulation key auto-zeroizes on `Drop`.
    ///
    /// # Errors
    ///
    /// Returns [`MlKemError::RngFailure`] if the RNG fails to produce bytes.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use quantica::ml_kem::*;
    /// let mut rng = OsRng;
    /// let (ek, dk) = MlKem::<MlKem768>::keygen(&mut rng).unwrap();
    /// assert_eq!(ek.len(), MlKem768::EK_LEN);
    /// assert_eq!(dk.len(), MlKem768::DK_LEN);
    /// ```
    pub fn keygen(rng: &mut impl CryptoRng) -> Result<(EncapsulationKey<P>, DecapsulationKey<P>), MlKemError> {
        let (ek_v, dk_v) = kem::keygen::<P>(rng)?;
        Ok((
            EncapsulationKey {
                bytes: ek_v,
                _marker: PhantomData,
            },
            DecapsulationKey {
                bytes: SecretBytes::from_vec(dk_v),
                _marker: PhantomData,
            },
        ))
    }

    /// Encapsulate a shared secret against an encapsulation key.
    ///
    /// Given a public encapsulation key, generates a fresh 32-byte shared
    /// secret and the corresponding ciphertext. Implements Algorithm 20 of
    /// FIPS 203 with input validation (length and modulus checks).
    ///
    /// # Arguments
    ///
    /// * `ek` - The encapsulation (public) key, exactly [`Params::EK_LEN`] bytes.
    /// * `rng` - A cryptographic random number generator implementing [`CryptoRng`].
    ///
    /// # Returns
    ///
    /// A tuple `(shared_secret, ciphertext)` where the shared secret is 32 bytes.
    ///
    /// # Errors
    ///
    /// * [`MlKemError::InvalidEncapsulationKey`] if `ek` has wrong length or fails
    ///   the modulus check.
    /// * [`MlKemError::RngFailure`] if the RNG fails.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use quantica::ml_kem::*;
    /// let mut rng = OsRng;
    /// let (ek, _dk) = MlKem::<MlKem768>::keygen(&mut rng).unwrap();
    /// let (shared_secret, ciphertext) = MlKem::<MlKem768>::encaps(&ek, &mut rng).unwrap();
    /// ```
    pub fn encaps(
        ek: &EncapsulationKey<P>,
        rng: &mut impl CryptoRng,
    ) -> Result<(SharedSecret, Ciphertext<P>), MlKemError> {
        let (ss, ct) = kem::encaps::<P>(ek.as_bytes(), rng)?;
        Ok((
            SharedSecret::new(ss),
            Ciphertext {
                bytes: ct,
                _marker: PhantomData,
            },
        ))
    }

    /// Decapsulate a ciphertext with full DFA protection.
    ///
    /// Recovers the 32-byte shared secret from a ciphertext using the
    /// decapsulation (private) key. Implements Algorithm 21 of FIPS 203
    /// with two additional DFA countermeasures:
    ///
    /// 1. **dk integrity check** -- verifies `H(ek)` stored in `dk` to detect
    ///    fault injection on key material in memory.
    /// 2. **Double computation** -- runs the internal decapsulation twice and
    ///    compares results. A single-fault attack can only corrupt one execution,
    ///    so divergent results indicate fault injection.
    ///
    /// Recommended for embedded and high-security contexts where physical
    /// fault attacks are in the threat model.
    ///
    /// # Arguments
    ///
    /// * `dk` - The decapsulation (private) key, exactly [`Params::DK_LEN`] bytes.
    /// * `ct` - The ciphertext, exactly [`Params::CT_LEN`] bytes.
    /// * `rng` - A cryptographic random number generator (reserved for future use).
    ///
    /// # Returns
    ///
    /// The 32-byte shared secret.
    ///
    /// # Errors
    ///
    /// * [`MlKemError::InvalidDecapsulationKey`] if `dk` has wrong length or
    ///   fails the integrity check.
    /// * [`MlKemError::InvalidCiphertext`] if `ct` has wrong length.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use quantica::ml_kem::*;
    /// let mut rng = OsRng;
    /// let (ek, dk) = MlKem::<MlKem768>::keygen(&mut rng).unwrap();
    /// let (ss, ct) = MlKem::<MlKem768>::encaps(&ek, &mut rng).unwrap();
    /// let ss2 = MlKem::<MlKem768>::decaps(&dk, &ct, &mut rng).unwrap();
    /// assert_eq!(ss, ss2);
    /// ```
    pub fn decaps(
        dk: &DecapsulationKey<P>,
        ct: &Ciphertext<P>,
        rng: &mut impl CryptoRng,
    ) -> Result<SharedSecret, MlKemError> {
        kem::decaps::<P>(dk.as_bytes(), ct.as_bytes(), rng).map(SharedSecret::new)
    }

    /// Decapsulate a ciphertext without double computation (faster variant).
    ///
    /// Same as [`MlKem::decaps`] but omits the double-computation DFA
    /// countermeasure, making it roughly twice as fast. The `H(ek)` integrity
    /// check on the decapsulation key is still performed.
    ///
    /// Use this for software-only contexts where physical fault injection
    /// is not part of the threat model.
    ///
    /// # Arguments
    ///
    /// * `dk` - The decapsulation (private) key, exactly [`Params::DK_LEN`] bytes.
    /// * `ct` - The ciphertext, exactly [`Params::CT_LEN`] bytes.
    ///
    /// # Returns
    ///
    /// The 32-byte shared secret.
    ///
    /// # Errors
    ///
    /// * [`MlKemError::InvalidDecapsulationKey`] if `dk` has wrong length or
    ///   fails the integrity check.
    /// * [`MlKemError::InvalidCiphertext`] if `ct` has wrong length.
    pub fn decaps_fast(dk: &DecapsulationKey<P>, ct: &Ciphertext<P>) -> Result<SharedSecret, MlKemError> {
        kem::decaps_single::<P>(dk.as_bytes(), ct.as_bytes()).map(SharedSecret::new)
    }

    // --- Deterministic internal functions (for testing / CAVP) ---

    /// Deterministic key generation for testing and CAVP validation.
    ///
    /// Implements Algorithm 16 of FIPS 203 directly, using caller-supplied
    /// seeds `d` and `z` instead of drawing them from an RNG.
    ///
    /// # Arguments
    ///
    /// * `d` - 32-byte seed for K-PKE key generation.
    /// * `z` - 32-byte implicit rejection value stored in the decapsulation key.
    ///
    /// # Returns
    ///
    /// A tuple `(encapsulation_key, decapsulation_key)` as byte vectors.
    pub fn keygen_internal(d: &[u8; 32], z: &[u8; 32]) -> (Vec<u8>, Vec<u8>) {
        kem::keygen_internal::<P>(d, z)
    }

    /// Deterministic encapsulation for testing and CAVP validation.
    ///
    /// Implements Algorithm 17 of FIPS 203 directly, using a caller-supplied
    /// message `m` instead of drawing it from an RNG. No input validation
    /// is performed on `ek`.
    ///
    /// # Arguments
    ///
    /// * `ek` - The encapsulation (public) key.
    /// * `m` - 32-byte random message seed.
    ///
    /// # Returns
    ///
    /// A tuple `(shared_secret, ciphertext)`.
    pub fn encaps_internal(ek: &[u8], m: &[u8; 32]) -> ([u8; 32], Vec<u8>) {
        kem::encaps_internal::<P>(ek, m)
    }

    /// Deterministic decapsulation for testing and CAVP validation.
    ///
    /// Implements Algorithm 18 of FIPS 203 directly, with no input
    /// validation or DFA countermeasures. All comparisons are still
    /// constant-time.
    ///
    /// # Arguments
    ///
    /// * `dk` - The decapsulation (private) key.
    /// * `ct` - The ciphertext.
    ///
    /// # Returns
    ///
    /// The 32-byte shared secret.
    pub fn decaps_internal(dk: &[u8], ct: &[u8]) -> [u8; 32] {
        kem::decaps_internal::<P>(dk, ct)
    }
}