libsoliton 0.1.3

//! ML-KEM-768 (FIPS 203) via the `ml-kem` crate.
//!
//! Uses the deterministic API (`generate_deterministic`, `encapsulate_deterministic`)
//! with entropy from `getrandom`. This is cryptographically equivalent to the
//! RNG-based API — the non-deterministic functions internally just draw random
//! bytes and pass them to the deterministic functions (see FIPS 203 §7.1-7.3).

use crate::error::{Error, Result};
use ml_kem::array::Array;
use ml_kem::{B32, EncapsulateDeterministic, EncodedSizeUser, KemCore, MlKem768};
use zeroize::{Zeroize, ZeroizeOnDrop};

/// ML-KEM-768 public key (1184 bytes).
#[derive(Clone, PartialEq, Eq)]
pub struct PublicKey(pub(crate) Vec<u8>);

/// ML-KEM-768 secret key (2400 bytes, expanded form).
///
/// # Security
///
/// Must not be resized after construction — `ZeroizeOnDrop` only zeroizes
/// the current allocation; a prior allocation freed by `Vec` resize would not
/// be zeroized.
#[derive(Zeroize, ZeroizeOnDrop)]
pub struct SecretKey(pub(crate) Vec<u8>);

/// ML-KEM-768 ciphertext (1088 bytes).
#[derive(Clone, PartialEq, Eq)]
pub struct Ciphertext(pub(crate) Vec<u8>);

/// ML-KEM-768 shared secret (32 bytes, FIPS 203 §7.3).
#[derive(Zeroize, ZeroizeOnDrop)]
pub struct SharedSecret(pub(crate) [u8; 32]);

impl PublicKey {
    /// Return the raw byte representation.
    pub fn as_bytes(&self) -> &[u8] {
        &self.0
    }

    /// Construct from raw bytes with size validation.
    pub fn from_bytes(bytes: Vec<u8>) -> Result<Self> {
        if bytes.len() != PK_LEN {
            return Err(Error::InvalidLength {
                expected: PK_LEN,
                got: bytes.len(),
            });
        }
        Ok(Self(bytes))
    }

    /// Construct from raw bytes without size validation.
    ///
    /// # Invariants
    ///
    /// Caller must ensure `bytes.len()` equals `PK_LEN` (1184 bytes).
    pub(crate) fn from_bytes_unchecked(bytes: Vec<u8>) -> Self {
        assert_eq!(
            bytes.len(),
            PK_LEN,
            "mlkem::PublicKey::from_bytes_unchecked: wrong size"
        );
        Self(bytes)
    }
}

impl SecretKey {
    /// Return the raw byte representation.
    pub(crate) fn as_bytes(&self) -> &[u8] {
        &self.0
    }

    /// Construct from raw bytes without size validation.
    ///
    /// # Invariants
    ///
    /// Caller must ensure `bytes.len()` equals `SK_LEN` (2400 bytes).
    /// The `Vec` must not be resized after construction.
    pub(crate) fn from_bytes_unchecked(bytes: Vec<u8>) -> Self {
        assert_eq!(
            bytes.len(),
            SK_LEN,
            "mlkem::SecretKey::from_bytes_unchecked: wrong size"
        );
        Self(bytes)
    }
}

impl Ciphertext {
    /// Return the raw byte representation.
    pub fn as_bytes(&self) -> &[u8] {
        &self.0
    }

    /// Construct from raw bytes with size validation.
    pub fn from_bytes(bytes: Vec<u8>) -> Result<Self> {
        if bytes.len() != CT_LEN {
            return Err(Error::InvalidLength {
                expected: CT_LEN,
                got: bytes.len(),
            });
        }
        Ok(Self(bytes))
    }

    /// Construct from raw bytes without size validation.
    ///
    /// # Invariants
    ///
    /// Caller must ensure `bytes.len()` equals `CT_LEN` (1088 bytes).
    pub(crate) fn from_bytes_unchecked(bytes: Vec<u8>) -> Self {
        assert_eq!(
            bytes.len(),
            CT_LEN,
            "mlkem::Ciphertext::from_bytes_unchecked: wrong size"
        );
        Self(bytes)
    }
}

impl SharedSecret {
    /// Return the raw 32-byte shared secret.
    pub fn as_bytes(&self) -> &[u8; 32] {
        &self.0
    }
}

/// ML-KEM-768 public key length (bytes).
const PK_LEN: usize = 1184;
/// ML-KEM-768 secret key length (bytes, expanded form).
const SK_LEN: usize = 2400;
/// ML-KEM-768 ciphertext length (bytes).
const CT_LEN: usize = 1088;

/// Return the expected public key length in bytes.
pub const fn pk_len() -> usize {
    PK_LEN
}

/// Return the expected secret key length in bytes.
pub const fn sk_len() -> usize {
    SK_LEN
}

/// Return the expected ciphertext length in bytes.
pub const fn ct_len() -> usize {
    CT_LEN
}

/// Generate 32 random bytes as a `B32` (hybrid-array `Array<u8, U32>`).
///
/// KL1-class stack residue: `arr` is returned by value, so Rust's calling
/// convention copies 32 bytes into the caller's stack frame. The callee-frame
/// copy persists unzeroized until the stack is overwritten. `panic = "abort"`
/// prevents unwinding from extending this lifetime.
fn random_b32() -> B32 {
    let mut buf = [0u8; 32];
    super::random::random_bytes(&mut buf);
    let arr = Array::from(buf);
    // [u8; 32] is Copy — Array::from() received a bitwise copy, so the
    // original stack value must be explicitly zeroized.
    buf.zeroize();
    arr
}

/// Generate an ML-KEM-768 keypair.
///
/// Uses `getrandom` as the entropy source, passed to the deterministic
/// key generation API. This is cryptographically equivalent to ML-KEM.KeyGen()
/// from FIPS 203 §7.1 — the standard algorithm draws `d` and `z` from an RNG
/// then calls KeyGen_internal(d, z).
///
/// # Security
///
/// Seeds `d` and `z` are zeroized immediately after key generation. The secret
/// key bytes are wrapped in `Zeroizing` during construction and moved into
/// `SecretKey` (which derives `ZeroizeOnDrop`).
pub fn keygen() -> Result<(PublicKey, SecretKey)> {
    let mut d = random_b32();
    let mut z = random_b32();
    let (dk, ek) = MlKem768::generate_deterministic(&d, &z);
    // d and z together determine the entire ML-KEM-768 key pair.
    // B32 (hybrid-array::Array) implements Zeroize but not ZeroizeOnDrop.
    d.zeroize();
    z.zeroize();

    let pk_bytes = ek.as_bytes().to_vec();
    let mut sk_bytes = zeroize::Zeroizing::new(dk.as_bytes().to_vec());

    // Runtime asserts (not debug_assert) because these sizes come from a foreign
    // crate's type-level constants — a crate update could silently change them.
    assert_eq!(
        pk_bytes.len(),
        PK_LEN,
        "ML-KEM-768 EK size mismatch: got {}, expected {PK_LEN}",
        pk_bytes.len()
    );
    assert_eq!(
        sk_bytes.len(),
        SK_LEN,
        "ML-KEM-768 DK size mismatch: got {}, expected {SK_LEN}",
        sk_bytes.len()
    );

    Ok((
        PublicKey(pk_bytes),
        SecretKey(std::mem::take(&mut *sk_bytes)),
    ))
}

/// Encapsulate to an ML-KEM-768 public key.
///
/// Returns (ciphertext, shared_secret).
///
/// Uses `getrandom` as the entropy source, passed to the deterministic
/// encapsulation API. Equivalent to ML-KEM.Encaps() from FIPS 203 §7.2.
///
/// # Security
///
/// The random message `m` and raw shared secret `ss` are zeroized after use.
/// The returned `SharedSecret` implements `ZeroizeOnDrop`.
pub fn encapsulate(pk: &PublicKey) -> Result<(Ciphertext, SharedSecret)> {
    if pk.0.len() != PK_LEN {
        return Err(Error::InvalidLength {
            expected: PK_LEN,
            got: pk.0.len(),
        });
    }

    // ml-kem requires its own typed key — raw &[u8] cannot be passed directly.
    let ek_enc: &ml_kem::Encoded<<MlKem768 as KemCore>::EncapsulationKey> =
        pk.0.as_slice().try_into().map_err(|_| Error::Internal)?;
    let ek = <MlKem768 as KemCore>::EncapsulationKey::from_bytes(ek_enc);

    let mut m = random_b32();
    // Encapsulation is infallible for a well-formed EncapsulationKey — the only
    // error path in the ml-kem crate is a type-level length mismatch, which
    // cannot occur because `ek` was just reconstructed from validated bytes.
    let (ct, mut ss) = ek
        .encapsulate_deterministic(&m)
        .map_err(|_| Error::Internal)?;
    // m alone determines the shared secret for this encapsulation.
    m.zeroize();

    let ct_bytes: &[u8] = ct.as_ref();
    let mut ss_bytes = [0u8; 32];
    ss_bytes.copy_from_slice(ss.as_ref());
    // B32 (hybrid-array::Array) implements Zeroize but not ZeroizeOnDrop.
    ss.zeroize();
    let result = SharedSecret(ss_bytes);
    // [u8; 32] is Copy — SharedSecret() received a bitwise copy, so the
    // original stack value must be explicitly zeroized.
    ss_bytes.zeroize();

    Ok((Ciphertext(ct_bytes.to_vec()), result))
}

/// Decapsulate an ML-KEM-768 ciphertext.
///
/// Returns the 32-byte shared secret on success.
///
/// # Security
///
/// The raw shared secret from the `ml-kem` crate is zeroized after copying
/// into the returned `SharedSecret`. The returned type implements `ZeroizeOnDrop`.
pub fn decapsulate(sk: &SecretKey, ct: &Ciphertext) -> Result<SharedSecret> {
    use kem::Decapsulate;

    if sk.0.len() != SK_LEN {
        return Err(Error::InvalidLength {
            expected: SK_LEN,
            got: sk.0.len(),
        });
    }
    if ct.0.len() != CT_LEN {
        return Err(Error::InvalidLength {
            expected: CT_LEN,
            got: ct.0.len(),
        });
    }

    // ml-kem requires its own typed key — raw &[u8] cannot be passed directly.
    let dk_enc: &ml_kem::Encoded<<MlKem768 as KemCore>::DecapsulationKey> =
        sk.0.as_slice().try_into().map_err(|_| Error::Internal)?;
    let dk = <MlKem768 as KemCore>::DecapsulationKey::from_bytes(dk_enc);

    // ml-kem requires its own typed ciphertext — raw &[u8] cannot be passed directly.
    let ct_inner: ml_kem::Ciphertext<MlKem768> =
        ct.0.as_slice()
            .try_into()
            .map_err(|_| Error::InvalidLength {
                expected: CT_LEN,
                got: ct.0.len(),
            })?;

    // The ml-kem crate implements FIPS 203 implicit rejection: invalid
    // ciphertexts produce a pseudorandom shared secret via Ok(ss), never
    // Err. The map_err path covers only structural/type-level errors. If a
    // future crate version returns Err for ciphertext validity, this
    // early-return would break constant-time implicit rejection.
    let mut ss = dk
        .decapsulate(&ct_inner)
        .map_err(|_| Error::DecapsulationFailed)?;
    let mut ss_bytes = [0u8; 32];
    ss_bytes.copy_from_slice(ss.as_ref());
    // B32 (hybrid-array::Array) implements Zeroize but not ZeroizeOnDrop.
    ss.zeroize();
    let result = SharedSecret(ss_bytes);
    // [u8; 32] is Copy — SharedSecret() received a bitwise copy, so the
    // original stack value must be explicitly zeroized.
    ss_bytes.zeroize();
    Ok(result)
}

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

    #[test]
    fn keygen_sizes() {
        let (pk, sk) = keygen().unwrap();
        assert_eq!(pk.as_bytes().len(), 1184);
        assert_eq!(sk.as_bytes().len(), 2400);
    }

    #[test]
    fn round_trip() {
        let (pk, sk) = keygen().unwrap();
        let (ct, ss_enc) = encapsulate(&pk).unwrap();
        let ss_dec = decapsulate(&sk, &ct).unwrap();
        assert_eq!(ss_enc.as_bytes(), ss_dec.as_bytes());
    }

    #[test]
    fn wrong_pk_size() {
        assert!(matches!(
            PublicKey::from_bytes(vec![0u8; 100]),
            Err(Error::InvalidLength {
                expected: 1184,
                got: 100
            })
        ));
    }

    #[test]
    fn wrong_ct_size() {
        assert!(matches!(
            Ciphertext::from_bytes(vec![0u8; 100]),
            Err(Error::InvalidLength {
                expected: 1088,
                got: 100
            })
        ));
    }

    #[test]
    fn tampered_ct() {
        let (pk, sk) = keygen().unwrap();
        let (ct, ss_enc) = encapsulate(&pk).unwrap();
        let mut bad_ct_bytes = ct.as_bytes().to_vec();
        bad_ct_bytes[0] ^= 0xFF;
        let bad_ct = Ciphertext::from_bytes_unchecked(bad_ct_bytes);
        // ML-KEM implicit rejection: decaps succeeds but produces different SS.
        let ss_dec = decapsulate(&sk, &bad_ct).unwrap();
        assert_ne!(ss_enc.as_bytes(), ss_dec.as_bytes());
    }

    #[test]
    fn round_trip_repeated() {
        // ML-KEM keygen uses getrandom — proptest's RNG is irrelevant and
        // seed-reproducibility is impossible. A plain loop provides identical
        // coverage without misleading proptest shrinking semantics.
        for _ in 0..256 {
            let (pk, sk) = keygen().unwrap();
            let (ct, ss_enc) = encapsulate(&pk).unwrap();
            let ss_dec = decapsulate(&sk, &ct).unwrap();
            assert_eq!(ss_enc.as_bytes(), ss_dec.as_bytes());
        }
    }
}