libsoliton 0.1.3

//! ML-DSA-65 (FIPS 204) digital signatures via the `ml-dsa` crate.
//!
//! Uses the seed-based API with entropy from `getrandom`. Key generation
//! is cryptographically equivalent to the RNG-based API — `key_gen(rng)`
//! internally just draws 32 random bytes then calls `from_seed` (FIPS 204 §6.1).
//!
//! Secret keys are stored as the 32-byte seed (FIPS 204 §6.1 `ξ`), not the
//! ~4032-byte expanded form. The signing key is deterministically re-expanded
//! from the seed on each `sign()` call. This minimises key material in memory
//! and on disk, and matches the canonical representation in the standard.
//!
//! ## Hedged Signing via `sign_internal`
//!
//! Signing uses the `sign_internal` API with 32 bytes of fresh `getrandom`
//! entropy as the `rnd` parameter (hedged mode). This provides fault-injection
//! resistance: with deterministic signing, a single induced fault during signing
//! (voltage glitch, Rowhammer) can extract the secret key via differential
//! analysis. Hedged signing mixes fresh randomness into the nonce derivation,
//! making each signing run independent even on the same message.
//!
//! `sign_internal` / `verify_internal` skip the FIPS 204 domain separator and
//! context string, making soliton ML-DSA signatures incompatible with
//! standalone FIPS 204 verifiers. This is intentional — lo signatures are
//! always verified by soliton (or a reimplementation built from Specification.md),
//! never by a generic FIPS 204 implementation. The same approach is used for
//! the X-Wing combiner (§4).

use crate::error::{Error, Result};
use zeroize::{Zeroize, ZeroizeOnDrop, Zeroizing};

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

/// ML-DSA-65 secret key (32-byte seed, FIPS 204 §6.1 `ξ`).
///
/// The seed is sufficient to deterministically reconstruct the full signing key
/// via `SigningKey::from_seed`. Storing only the seed reduces key material
/// exposure compared to the ~4032-byte expanded form.
#[derive(Zeroize, ZeroizeOnDrop)]
pub struct SecretKey(pub(crate) Vec<u8>);

/// ML-DSA-65 signature (3309 bytes, FIPS 204).
#[derive(Clone, PartialEq, Eq)]
pub struct Signature(pub(crate) Vec<u8>);

/// ML-DSA-65 public key length (bytes).
const PK_LEN: usize = 1952;
/// ML-DSA-65 secret key length (bytes, seed form per FIPS 204 §6.1).
const SK_LEN: usize = 32;
/// ML-DSA-65 signature length (bytes).
const SIG_LEN: usize = 3309;

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 (internal trusted use).
    pub(crate) fn from_bytes_unchecked(bytes: Vec<u8>) -> Self {
        assert_eq!(
            bytes.len(),
            PK_LEN,
            "mldsa::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 with size validation.
    pub fn from_bytes(bytes: Vec<u8>) -> Result<Self> {
        // Wrap in Zeroizing so the raw Vec is zeroized on the error path
        // (SecretKey derives ZeroizeOnDrop, but that only fires on success).
        let mut bytes = Zeroizing::new(bytes);
        if bytes.len() != SK_LEN {
            return Err(Error::InvalidLength {
                expected: SK_LEN,
                got: bytes.len(),
            });
        }
        Ok(Self(std::mem::take(&mut *bytes)))
    }
}

impl Signature {
    /// 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() != SIG_LEN {
            return Err(Error::InvalidLength {
                expected: SIG_LEN,
                got: bytes.len(),
            });
        }
        Ok(Self(bytes))
    }

    /// Construct from raw bytes without size validation (internal trusted use).
    pub(crate) fn from_bytes_unchecked(bytes: Vec<u8>) -> Self {
        assert_eq!(
            bytes.len(),
            SIG_LEN,
            "mldsa::Signature::from_bytes_unchecked: wrong size"
        );
        Self(bytes)
    }
}

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

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

/// Return the expected signature length in bytes.
pub const fn sig_len() -> usize {
    SIG_LEN
}

/// Generate an ML-DSA-65 keypair.
///
/// Uses `getrandom` as the entropy source, passed to the deterministic
/// `from_seed` API. Equivalent to ML-DSA.KeyGen() from FIPS 204 §6.1.
///
/// The secret key is the 32-byte seed `ξ`, not the expanded form.
///
/// # Security
///
/// The 32-byte seed and its `B32` copy are zeroized immediately after key
/// generation. The extracted seed is wrapped in `Zeroizing` before being
/// moved into `SecretKey` (which derives `ZeroizeOnDrop`).
pub fn keygen() -> Result<(PublicKey, SecretKey)> {
    use ml_dsa::{B32, KeyGen, MlDsa65};

    // Generate 32-byte seed from OS CSPRNG.
    let mut seed_bytes = [0u8; 32];
    super::random::random_bytes(&mut seed_bytes);
    let mut seed: B32 = seed_bytes.into();
    // [u8; 32] is Copy — B32::from() received a bitwise copy, so the
    // original stack value must be explicitly zeroized.
    seed_bytes.zeroize();

    let kp = MlDsa65::from_seed(&seed);
    seed.zeroize();

    let pk_bytes = kp.verifying_key().encode().to_vec();
    // B32 (hybrid_array::Array<u8, U32>) implements Zeroize but not ZeroizeOnDrop.
    // Bind the return value to explicitly zeroize the seed copy.
    let mut seed_out = kp.to_seed();
    let mut sk_bytes = zeroize::Zeroizing::new(seed_out.to_vec());
    seed_out.zeroize();

    // 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-DSA-65 VK size mismatch: got {}, expected {PK_LEN}",
        pk_bytes.len()
    );
    assert_eq!(
        sk_bytes.len(),
        SK_LEN,
        "ML-DSA-65 seed size mismatch: got {}, expected {SK_LEN}",
        sk_bytes.len()
    );

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

/// Sign a message with ML-DSA-65 (hedged mode).
///
/// Re-expands the signing key from the stored 32-byte seed, generates 32
/// bytes of fresh entropy from `getrandom`, and signs via `sign_internal`
/// with the entropy as the `rnd` parameter. See module doc for rationale.
///
/// # Security
///
/// The expanded `SigningKey` is zeroized on drop (`ml-dsa` feature `zeroize`).
/// The 32-byte seed is borrowed from `SecretKey`, not copied. The `rnd`
/// buffer is zeroized after use.
pub fn sign(sk: &SecretKey, message: &[u8]) -> Result<Signature> {
    use ml_dsa::{B32, MlDsa65, Seed, SigningKey};

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

    // ml-dsa requires its own typed key — raw &[u8] cannot be passed directly.
    // ZeroizeOnDrop guaranteed by ml-dsa features = ["zeroize"].
    let seed: &Seed = sk.0.as_slice().try_into().map_err(|_| Error::Internal)?;
    let signing_key = SigningKey::<MlDsa65>::from_seed(seed);

    // Hedged signing: 32 bytes of fresh getrandom entropy mixed into the
    // nonce derivation (FIPS 204 Algorithm 2, `rnd` parameter). Provides
    // fault-injection resistance — see module doc.
    let mut rnd_bytes = [0u8; 32];
    super::random::random_bytes(&mut rnd_bytes);
    let mut rnd: B32 = rnd_bytes.into();
    // [u8; 32] is Copy — .into() received a bitwise copy.
    rnd_bytes.zeroize();

    // sign_internal skips the FIPS 204 domain separator / context string.
    // Paired with verify_internal in verify() — both use MuBuilder::internal,
    // so signatures are self-consistent within soliton.
    let sig = signing_key.sign_internal(&[message], &rnd);
    // B32 (hybrid_array::Array<u8, U32>) implements Zeroize but not ZeroizeOnDrop.
    // Explicitly zeroize the hedged randomness after signing — same pattern as
    // seed zeroization in keygen().
    rnd.zeroize();

    let sig_bytes = sig.encode().to_vec();
    // Runtime assert (not debug_assert) because SIG_LEN comes from a foreign
    // crate's type-level constant — a crate update could silently change it.
    assert_eq!(
        sig_bytes.len(),
        SIG_LEN,
        "ML-DSA-65 sig size mismatch: got {}, expected {SIG_LEN}",
        sig_bytes.len()
    );

    Ok(Signature(sig_bytes))
}

/// Verify an ML-DSA-65 signature.
///
/// Returns `Ok(())` on valid signature, `Err(VerificationFailed)` otherwise.
pub fn verify(pk: &PublicKey, message: &[u8], signature: &Signature) -> Result<()> {
    use ml_dsa::{EncodedSignature, EncodedVerifyingKey, MlDsa65, VerifyingKey};

    if pk.0.len() != PK_LEN {
        return Err(Error::InvalidLength {
            expected: PK_LEN,
            got: pk.0.len(),
        });
    }
    if signature.0.len() != SIG_LEN {
        return Err(Error::InvalidLength {
            expected: SIG_LEN,
            got: signature.0.len(),
        });
    }

    // ml-dsa requires its own typed key — raw &[u8] cannot be passed directly.
    let vk_enc: &EncodedVerifyingKey<MlDsa65> =
        pk.0.as_slice().try_into().map_err(|_| Error::Internal)?;
    // Infallible: decode() from a correctly-sized EncodedVerifyingKey always
    // succeeds — the encoding is a fixed-length byte array, not a parsed format.
    let vk = VerifyingKey::<MlDsa65>::decode(vk_enc);

    // ml-dsa requires its own typed signature — raw &[u8] cannot be passed directly.
    let sig_enc: &EncodedSignature<MlDsa65> = signature
        .0
        .as_slice()
        .try_into()
        .map_err(|_| Error::Internal)?;
    // VerificationFailed (not InvalidData) — avoids leaking whether the
    // signature was malformed vs mathematically incorrect to the caller.
    let sig = ml_dsa::Signature::<MlDsa65>::decode(sig_enc).ok_or(Error::VerificationFailed)?;

    // verify_internal skips the FIPS 204 domain separator / context string.
    // Paired with sign_internal in sign() — both use MuBuilder::internal.
    // sign_internal takes &[&[u8]] (slice-of-slices); verify_internal takes
    // &[u8] (flat slice). The asymmetry is safe: SHA3's streaming absorb
    // produces identical hashes regardless of chunking boundaries.
    if vk.verify_internal(message, &sig) {
        Ok(())
    } else {
        Err(Error::VerificationFailed)
    }
}

#[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(), 1952);
        assert_eq!(sk.as_bytes().len(), 32);
    }

    #[test]
    fn sign_verify_round_trip() {
        let (pk, sk) = keygen().unwrap();
        let sig = sign(&sk, b"hello mldsa").unwrap();
        assert!(verify(&pk, b"hello mldsa", &sig).is_ok());
    }

    #[test]
    fn sign_size() {
        let (_, sk) = keygen().unwrap();
        let sig = sign(&sk, b"test").unwrap();
        assert_eq!(sig.as_bytes().len(), 3309);
    }

    #[test]
    fn sign_verify_empty_message() {
        let (pk, sk) = keygen().unwrap();
        let sig = sign(&sk, b"").unwrap();
        assert!(verify(&pk, b"", &sig).is_ok());
    }

    #[test]
    fn verify_wrong_message() {
        let (pk, sk) = keygen().unwrap();
        let sig = sign(&sk, b"message one").unwrap();
        assert!(matches!(
            verify(&pk, b"message two", &sig),
            Err(Error::VerificationFailed)
        ));
    }

    #[test]
    fn verify_wrong_key() {
        let (_, sk) = keygen().unwrap();
        let (pk2, _) = keygen().unwrap();
        let sig = sign(&sk, b"test").unwrap();
        assert!(matches!(
            verify(&pk2, b"test", &sig),
            Err(Error::VerificationFailed)
        ));
    }

    #[test]
    fn verify_tampered_signature() {
        let (pk, sk) = keygen().unwrap();
        let sig = sign(&sk, b"test").unwrap();
        let mut bad = sig.as_bytes().to_vec();
        bad[0] ^= 0xFF;
        let bad_sig = Signature::from_bytes(bad).unwrap();
        assert!(matches!(
            verify(&pk, b"test", &bad_sig),
            Err(Error::VerificationFailed)
        ));
    }

    #[test]
    fn verify_all_zeros_signature() {
        let (pk, _) = keygen().unwrap();
        let zeros_sig = Signature::from_bytes(vec![0u8; 3309]).unwrap();
        assert!(matches!(
            verify(&pk, b"test", &zeros_sig),
            Err(Error::VerificationFailed)
        ));
    }

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

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

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

    #[test]
    fn sign_hedged_both_verify() {
        let (pk, sk) = keygen().unwrap();
        let msg = b"hedged signing test";
        let sig1 = sign(&sk, msg).unwrap();
        let sig2 = sign(&sk, msg).unwrap();
        // Hedged signing: same message produces different signatures (fresh
        // rnd each time), but both must verify.
        assert_ne!(
            sig1.as_bytes(),
            sig2.as_bytes(),
            "hedged signatures on the same message must differ"
        );
        assert!(verify(&pk, msg, &sig1).is_ok());
        assert!(verify(&pk, msg, &sig2).is_ok());
    }

    proptest::proptest! {
        // Default case count (256) — ML-DSA keygen is computationally expensive;
        // 256 iterations provide adequate coverage without excessive test time.
        #[test]
        fn proptest_round_trip(
            msg in proptest::collection::vec(proptest::prelude::any::<u8>(), 0..1024),
            flip_byte in 0..3309usize,
        ) {
            let kg = keygen();
            proptest::prop_assert!(kg.is_ok());
            let (pk, sk) = kg.unwrap();
            let sign_result = sign(&sk, &msg);
            proptest::prop_assert!(sign_result.is_ok());
            let sig = sign_result.unwrap();
            proptest::prop_assert!(verify(&pk, &msg, &sig).is_ok());

            // Forgery rejection: mutated signature must fail verification.
            let mut bad_sig_bytes = sig.as_bytes().to_vec();
            bad_sig_bytes[flip_byte] ^= 0x01;
            let bad_sig = Signature::from_bytes(bad_sig_bytes).unwrap();
            proptest::prop_assert!(verify(&pk, &msg, &bad_sig).is_err());

            // Forgery rejection: mutated message must fail verification.
            if !msg.is_empty() {
                let mut bad_msg = msg.clone();
                bad_msg[0] ^= 0x01;
                proptest::prop_assert!(verify(&pk, &bad_msg, &sig).is_err());
            }
        }
    }
}