ml-dsa 0.1.1

#![no_std]
#![doc = include_str!("../README.md")]
#![doc(
    html_logo_url = "https://raw.githubusercontent.com/RustCrypto/meta/master/logo.svg",
    html_favicon_url = "https://raw.githubusercontent.com/RustCrypto/meta/master/logo.svg"
)]
#![cfg_attr(docsrs, feature(doc_cfg))]
#![allow(non_snake_case)] // Allow notation matching the spec
#![allow(clippy::similar_names)] // Allow notation matching the spec
#![allow(clippy::many_single_char_names)] // Allow notation matching the spec
#![allow(clippy::clone_on_copy)] // Be explicit about moving data

//! # Usage
//!
//! The following types provide the core functionality of this crate, and are all generic around the
//! [`MlDsaParams`] trait which defines the security level and is one of [`MlDsa44`], [`MlDsa65`],
//! or [`MlDsa87`]  (with `MlDsa65` recommended as providing the best balance of security and
//! performance):
//!
//! - [`SigningKey`]: secret key capable of generating signatures. Implements the [`KeyInit`],
//!   [`KeyExport`], [`Keypair`], and [`Signer`] traits, as well as [`Generate`] when the
//!   `rand_core` feature of this crate is enabled.
//! - [`VerifyingKey`]: public key associated with a given [`SigningKey`]. Implements the
//!   [`KeyInit`], [`KeyExport`], and [`Verifier`] traits.
//! - [`Signature`]: ML-DSA signature generated by a `SigningKey`, and verifiable by a
//!   `VerifyingKey`. Implements the [`SignatureEncoding`] trait.
//!
#![cfg_attr(feature = "getrandom", doc = "```")]
#![cfg_attr(not(feature = "getrandom"), doc = "```ignore")]
//! # fn main() -> Result<(), signature::Error> {
//! // NOTE: requires the `getrandom` feature is enabled
//! use ml_dsa::{MlDsa65, Generate, Keypair, SigningKey, Signer, Verifier};
//!
//! let sk = SigningKey::<MlDsa65>::generate();
//!
//! let msg = b"Hello world";
//! let sig = sk.sign(msg);
//!
//! sk.verifying_key().verify(msg, &sig)?;
//! # Ok(()) }
//! ```

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

#[cfg(feature = "pkcs8")]
pub mod pkcs8;

mod algebra;
mod crypto;
mod encode;
mod hint;
mod ntt;
mod param;
mod sampling;
mod signing;
mod verifying;

pub use crate::{
    param::{EncodedSignature, EncodedVerifyingKey, ExpandedSigningKeyBytes, MlDsaParams},
    signing::{ExpandedSigningKey, SigningKey},
    verifying::VerifyingKey,
};
pub use common::{self, KeyExport, KeyInit, KeySizeUser};
pub use signature::{self, Error, Keypair, SignatureEncoding, Signer, Verifier};

#[cfg(feature = "rand_core")]
pub use common::Generate;

use crate::{
    algebra::{AlgebraExt, Vector},
    crypto::H,
    hint::Hint,
    param::{ParameterSet, QMinus1},
};
use core::convert::{TryFrom, TryInto};
use hybrid_array::{
    Array,
    sizes::{U1, U2, U4, U5, U6, U7, U8, U17, U19, U32, U48, U55, U64, U75, U80, U88},
    typenum::{Diff, Length, Prod, Quot, Shleft},
};
use module_lattice::{MaybeBox, Truncate};
use shake::Shake256;

/// A 32-byte array, defined here for brevity because it is used several times
pub type B32 = Array<u8, U32>;

/// A 64-byte array, defined here for brevity because it is used several times
pub(crate) type B64 = Array<u8, U64>;

/// ML-DSA seeds are signing (private) keys, which are consistently 32-bytes across all security
/// levels, and are the preferred serialization for representing such keys.
pub type Seed = B32;

/// An ML-DSA signature
#[derive(Clone, Debug, PartialEq)]
pub struct Signature<P: MlDsaParams> {
    c_tilde: Array<u8, P::Lambda>,
    z: MaybeBox<Vector<P::L>>,
    h: Hint<P>,
}

impl<P: MlDsaParams> Signature<P> {
    /// Encode the signature in a fixed-size byte array.
    // Algorithm 26 sigEncode
    pub fn encode(&self) -> EncodedSignature<P> {
        let c_tilde = self.c_tilde.clone();
        let z = P::encode_z(&self.z);
        let h = self.h.bit_pack();
        P::concat_sig(c_tilde, z, h)
    }

    /// Decode the signature from an appropriately sized byte array.
    // Algorithm 27 sigDecode
    pub fn decode(enc: &EncodedSignature<P>) -> Option<Self> {
        let (c_tilde, z, h) = P::split_sig(enc);

        let c_tilde = c_tilde.clone();
        let z = MaybeBox::new(P::decode_z(z));
        let h = Hint::bit_unpack(h)?;

        if z.infinity_norm() >= P::GAMMA1_MINUS_BETA {
            return None;
        }

        Some(Self { c_tilde, z, h })
    }
}

impl<'a, P: MlDsaParams> TryFrom<&'a [u8]> for Signature<P> {
    type Error = Error;

    fn try_from(value: &'a [u8]) -> Result<Self, Self::Error> {
        let enc = EncodedSignature::<P>::try_from(value).map_err(|_| Error::new())?;
        Self::decode(&enc).ok_or(Error::new())
    }
}

impl<P: MlDsaParams> TryInto<EncodedSignature<P>> for Signature<P> {
    type Error = Error;

    fn try_into(self) -> Result<EncodedSignature<P>, Self::Error> {
        Ok(self.encode())
    }
}

impl<P: MlDsaParams> SignatureEncoding for Signature<P> {
    type Repr = EncodedSignature<P>;
}

impl<P: MlDsaParams> core::hash::Hash for Signature<P> {
    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
        self.encode().hash(state);
    }
}

struct MuBuilder(H);

impl MuBuilder {
    fn new(tr: &[u8], ctx: &[u8]) -> Self {
        let mut h = H::default();
        h = h.absorb(tr);
        h = h.absorb(&[0]);
        h = h.absorb(&[Truncate::truncate(ctx.len())]);
        h = h.absorb(ctx);

        Self(h)
    }

    fn internal(tr: &[u8], Mp: &[&[u8]]) -> B64 {
        let mut h = H::default().absorb(tr);

        for m in Mp {
            h = h.absorb(m);
        }

        h.squeeze_new()
    }

    fn message(mut self, M: &[&[u8]]) -> B64 {
        for m in M {
            self.0 = self.0.absorb(m);
        }

        self.0.squeeze_new()
    }

    fn finish(mut self) -> B64 {
        self.0.squeeze_new()
    }
}

impl AsMut<Shake256> for MuBuilder {
    fn as_mut(&mut self) -> &mut Shake256 {
        self.0.updatable()
    }
}

/// `MlDsa44` is the parameter set for security category 2, providing the equivalent of 128-bit
/// symmetric security.
#[derive(Clone, Copy, Debug, Default, PartialEq)]
pub struct MlDsa44;

impl ParameterSet for MlDsa44 {
    type K = U4;
    type L = U4;
    type Eta = U2;
    type Gamma1 = Shleft<U1, U17>;
    type Gamma2 = Quot<QMinus1, U88>;
    type TwoGamma2 = Prod<U2, Self::Gamma2>;
    type W1Bits = Length<Diff<Quot<U88, U2>, U1>>;
    type Lambda = U32;
    type Omega = U80;
    const TAU: usize = 39;
}

/// `MlDsa65` is the parameter set for security category 3, providing the equivalent of 192-bit
/// symmetric security, and is the recommended parameter set.
///
/// This set provides the best balance between performance and security.
#[derive(Clone, Copy, Debug, Default, PartialEq)]
pub struct MlDsa65;

impl ParameterSet for MlDsa65 {
    type K = U6;
    type L = U5;
    type Eta = U4;
    type Gamma1 = Shleft<U1, U19>;
    type Gamma2 = Quot<QMinus1, U32>;
    type TwoGamma2 = Prod<U2, Self::Gamma2>;
    type W1Bits = Length<Diff<Quot<U32, U2>, U1>>;
    type Lambda = U48;
    type Omega = U55;
    const TAU: usize = 49;
}

/// `MlDsa87` is the parameter set for security category 5, providing the equivalent of 256-bit
/// symmetric security.
#[derive(Clone, Copy, Debug, Default, PartialEq)]
pub struct MlDsa87;

impl ParameterSet for MlDsa87 {
    type K = U8;
    type L = U7;
    type Eta = U2;
    type Gamma1 = Shleft<U1, U19>;
    type Gamma2 = Quot<QMinus1, U32>;
    type TwoGamma2 = Prod<U2, Self::Gamma2>;
    type W1Bits = Length<Diff<Quot<U32, U2>, U1>>;
    type Lambda = U64;
    type Omega = U75;
    const TAU: usize = 60;
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::param::*;
    use hybrid_array::typenum::Unsigned;
    use signature::Keypair;

    #[test]
    fn output_sizes() {
        //           priv pub  sig
        // ML-DSA-44 2560 1312 2420
        // ML-DSA-65 4032 1952 3309
        // ML-DSA-87 4896 2592 4627
        assert_eq!(SigningKeySize::<MlDsa44>::USIZE, 2560);
        assert_eq!(VerifyingKeySize::<MlDsa44>::USIZE, 1312);
        assert_eq!(SignatureSize::<MlDsa44>::USIZE, 2420);

        assert_eq!(SigningKeySize::<MlDsa65>::USIZE, 4032);
        assert_eq!(VerifyingKeySize::<MlDsa65>::USIZE, 1952);
        assert_eq!(SignatureSize::<MlDsa65>::USIZE, 3309);

        assert_eq!(SigningKeySize::<MlDsa87>::USIZE, 4896);
        assert_eq!(VerifyingKeySize::<MlDsa87>::USIZE, 2592);
        assert_eq!(SignatureSize::<MlDsa87>::USIZE, 4627);
    }

    fn encode_decode_round_trip_test<P>()
    where
        P: MlDsaParams + PartialEq,
    {
        let seed = Array::default();
        let ssk = SigningKey::from_seed(&seed);
        assert_eq!(ssk.to_seed(), seed);

        let esk = ssk.expanded_key();
        let vk = ssk.verifying_key();

        let vk_bytes = vk.encode();
        let vk2 = VerifyingKey::<P>::decode(&vk_bytes);
        assert!(vk == vk2);

        #[allow(deprecated)]
        {
            let sk_bytes = esk.to_expanded();
            let sk2 = ExpandedSigningKey::<P>::from_expanded(&sk_bytes);
            assert!(esk == &sk2);

            let M = b"Hello world";
            let rnd = Array([0u8; 32]);
            let sig = esk.sign_internal(&[M], &rnd);
            let sig_bytes = sig.encode();
            let sig2 = Signature::<P>::decode(&sig_bytes).unwrap();
            assert!(sig == sig2);
        }
    }

    #[test]
    fn encode_decode_round_trip() {
        encode_decode_round_trip_test::<MlDsa44>();
        encode_decode_round_trip_test::<MlDsa65>();
        encode_decode_round_trip_test::<MlDsa87>();
    }

    fn public_from_private_test<P>()
    where
        P: MlDsaParams + PartialEq,
    {
        let ssk = SigningKey::<P>::from_seed(&Array::default());
        let esk = ssk.expanded_key();
        let vk = ssk.verifying_key();
        let vk_derived = esk.verifying_key();

        assert!(vk == vk_derived);
    }

    #[test]
    fn public_from_private() {
        public_from_private_test::<MlDsa44>();
        public_from_private_test::<MlDsa65>();
        public_from_private_test::<MlDsa87>();
    }

    fn sign_verify_round_trip_test<P>()
    where
        P: MlDsaParams,
    {
        let ssk = SigningKey::<P>::from_seed(&Array::default());
        let esk = ssk.expanded_key();
        let vk = ssk.verifying_key();

        let M = b"Hello world";
        let rnd = Array([0u8; 32]);
        let sig = esk.sign_internal(&[M], &rnd);

        assert!(vk.verify_internal(M, &sig));
    }

    #[test]
    fn sign_verify_round_trip() {
        sign_verify_round_trip_test::<MlDsa44>();
        sign_verify_round_trip_test::<MlDsa65>();
        sign_verify_round_trip_test::<MlDsa87>();
    }

    #[test]
    fn sign_mu_verify_mu_round_trip() {
        fn sign_mu_verify_mu<P>()
        where
            P: MlDsaParams,
        {
            let ssk = SigningKey::<P>::from_seed(&Array::default());
            let esk = ssk.expanded_key();
            let vk = ssk.verifying_key();

            let M = b"Hello world";
            let rnd = Array([0u8; 32]);
            let mu = MuBuilder::internal(&esk.tr, &[M]);
            let sig = esk.raw_sign_mu(&mu, &rnd);

            assert!(vk.raw_verify_mu(&mu, &sig));
        }
        sign_mu_verify_mu::<MlDsa44>();
        sign_mu_verify_mu::<MlDsa65>();
        sign_mu_verify_mu::<MlDsa87>();
    }

    #[test]
    fn sign_mu_verify_internal_round_trip() {
        fn sign_mu_verify_internal<P>()
        where
            P: MlDsaParams,
        {
            let ssk = SigningKey::<P>::from_seed(&Array::default());
            let esk = ssk.expanded_key();
            let vk = ssk.verifying_key();

            let M = b"Hello world";
            let rnd = Array([0u8; 32]);
            let mu = MuBuilder::internal(&esk.tr, &[M]);
            let sig = esk.raw_sign_mu(&mu, &rnd);

            assert!(vk.verify_internal(M, &sig));
        }
        sign_mu_verify_internal::<MlDsa44>();
        sign_mu_verify_internal::<MlDsa65>();
        sign_mu_verify_internal::<MlDsa87>();
    }

    #[test]
    fn sign_internal_verify_mu_round_trip() {
        fn sign_internal_verify_mu<P>()
        where
            P: MlDsaParams,
        {
            let ssk = SigningKey::<P>::from_seed(&Array::default());
            let esk = ssk.expanded_key();
            let vk = ssk.verifying_key();

            let M = b"Hello world";
            let rnd = Array([0u8; 32]);
            let mu = MuBuilder::internal(&esk.tr, &[M]);
            let sig = esk.sign_internal(&[M], &rnd);

            assert!(vk.raw_verify_mu(&mu, &sig));
        }
        sign_internal_verify_mu::<MlDsa44>();
        sign_internal_verify_mu::<MlDsa65>();
        sign_internal_verify_mu::<MlDsa87>();
    }

    #[test]
    fn from_seed_implementations_match() {
        fn assert_from_seed_equality<P>()
        where
            P: MlDsaParams,
        {
            let seed = Seed::default();
            let ssk = SigningKey::<P>::from_seed(&seed);
            let sk1 = ExpandedSigningKey::<P>::from_seed(&seed);
            assert_eq!(ssk.expanded_key(), &sk1);
        }
        assert_from_seed_equality::<MlDsa44>();
        assert_from_seed_equality::<MlDsa65>();
        assert_from_seed_equality::<MlDsa87>();
    }

    #[test]
    fn to_seed_returns_correct_seed() {
        fn test_to_seed<P: MlDsaParams>() {
            let seed = Array([
                1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23,
                24, 25, 26, 27, 28, 29, 30, 31, 32,
            ]);
            let kp = SigningKey::<P>::from_seed(&seed);
            assert_eq!(kp.to_seed(), seed);
        }
        test_to_seed::<MlDsa44>();
        test_to_seed::<MlDsa65>();
        test_to_seed::<MlDsa87>();
    }

    #[test]
    fn verification_rejects_invalid_signature() {
        fn test_invalid_sig<P: MlDsaParams>() {
            let kp = SigningKey::<P>::from_seed(&Array::default());
            let vk = kp.verifying_key();

            let msg = b"Hello world";
            let rnd = Array([0u8; 32]);
            let mut sig = kp.expanded_key().sign_internal(&[msg], &rnd);
            sig.c_tilde[0] ^= 0xFF;

            assert!(!vk.verify_with_context(msg, &[], &sig));
        }
        test_invalid_sig::<MlDsa44>();
        test_invalid_sig::<MlDsa65>();
        test_invalid_sig::<MlDsa87>();
    }

    #[test]
    fn verification_rejects_wrong_message() {
        fn test_wrong_msg<P: MlDsaParams>() {
            let kp = SigningKey::<P>::from_seed(&Array::default());
            let vk = kp.verifying_key();

            let msg1 = b"Hello world";
            let msg2 = b"Wrong message";
            let rnd = Array([0u8; 32]);
            let sig = kp.expanded_key().sign_internal(&[msg1], &rnd);

            assert!(!vk.verify_with_context(msg2, &[], &sig));
        }
        test_wrong_msg::<MlDsa44>();
        test_wrong_msg::<MlDsa65>();
        test_wrong_msg::<MlDsa87>();
    }

    #[test]
    fn context_length_validation() {
        fn test_ctx_length<P: MlDsaParams>() {
            let ssk = SigningKey::<P>::from_seed(&Array::default());
            let sk = ssk.expanded_key();
            let vk = ssk.verifying_key();

            let msg = b"Hello world";
            let long_ctx = [0u8; 256];
            let short_ctx = [0u8; 255];

            assert!(sk.sign_deterministic(msg, &long_ctx).is_err());

            let sig = sk.sign_deterministic(msg, &short_ctx).unwrap();
            assert!(!vk.verify_with_context(msg, &long_ctx, &sig));
            assert!(vk.verify_with_context(msg, &short_ctx, &sig));
        }
        test_ctx_length::<MlDsa44>();
        test_ctx_length::<MlDsa65>();
        test_ctx_length::<MlDsa87>();
    }

    #[test]
    fn derived_verifying_key_validates_signatures() {
        fn test_derived_vk<P: MlDsaParams>() {
            let seed = Array([42u8; 32]);
            let ssk = SigningKey::<P>::from_seed(&seed);
            let sk = ssk.expanded_key();
            let derived_vk = sk.verifying_key();

            let msg = b"Test message for derived key";
            let rnd = Array([0u8; 32]);
            let sig = sk.sign_internal(&[msg], &rnd);

            assert!(derived_vk.verify_internal(msg, &sig));
            assert_eq!(derived_vk.encode(), ssk.verifying_key().encode());
        }
        test_derived_vk::<MlDsa44>();
        test_derived_vk::<MlDsa65>();
        test_derived_vk::<MlDsa87>();
    }

    #[test]
    #[cfg(feature = "alloc")]
    fn debug_implementations() {
        extern crate alloc;
        use core::fmt::Write;

        fn test_debug<P: MlDsaParams>() {
            let kp = SigningKey::<P>::from_seed(&Array::default());

            let mut kp_debug = alloc::string::String::new();
            write!(&mut kp_debug, "{:?}", kp).unwrap();
            assert!(kp_debug.contains("SigningKey"));

            let mut sk_debug = alloc::string::String::new();
            write!(&mut sk_debug, "{:?}", kp.expanded_key()).unwrap();
            assert!(sk_debug.contains("ExpandedSigningKey"));
        }
        test_debug::<MlDsa44>();
        test_debug::<MlDsa65>();
        test_debug::<MlDsa87>();
    }
}