RustCrypto-signatures/ml-dsa/src/lib.rs

#![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
//!
#![cfg_attr(feature = "rand_core", doc = "```")]
#![cfg_attr(not(feature = "rand_core"), doc = "```ignore")]
//! # fn main() -> Result<(), signature::Error> {
//! use ml_dsa::{
//!     signature::{Keypair, Signer, Verifier},
//!     MlDsa65, KeyGen,
//! };
//! use getrandom::{SysRng, rand_core::UnwrapErr};
//!
//! let mut rng = UnwrapErr(SysRng);
//! let sk = MlDsa65::key_gen(&mut rng);
//!
//! 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 signature::{self, Error};

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

#[cfg(feature = "rand_core")]
use rand_core::CryptoRng;

/// 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> signature::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.
#[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.
#[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.
#[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;
}

/// A parameter set that knows how to generate key pairs
pub trait KeyGen: MlDsaParams {
    /// The type that is returned by key generation
    type KeyPair: signature::Keypair;

    /// Generate a signing key pair from the specified RNG
    #[cfg(feature = "rand_core")]
    fn key_gen<R: CryptoRng + ?Sized>(rng: &mut R) -> Self::KeyPair;

    /// Deterministically generate a signing key pair from the specified seed
    ///
    /// This method reflects the ML-DSA.KeyGen_internal algorithm from FIPS 204.
    fn from_seed(xi: &B32) -> Self::KeyPair;
}

impl<P> KeyGen for P
where
    P: MlDsaParams,
{
    type KeyPair = SigningKey<P>;

    /// Generate a signing key pair from the specified RNG
    // Algorithm 1 ML-DSA.KeyGen()
    #[cfg(feature = "rand_core")]
    fn key_gen<R: CryptoRng + ?Sized>(rng: &mut R) -> SigningKey<P> {
        let mut xi = B32::default();
        rng.fill_bytes(&mut xi);
        Self::from_seed(&xi)
    }

    /// Deterministically generate a signing key pair from the specified seed
    // Algorithm 6 ML-DSA.KeyGen_internal
    fn from_seed(seed: &Seed) -> SigningKey<P> {
        SigningKey::from_seed(seed)
    }
}

/// Type which opportunistically uses `Box` when the `alloc` feature is available but falls back to
/// a stack-allocated type when it's unavailable.
#[derive(Clone, Debug, PartialEq)]
pub(crate) struct MaybeBox<T> {
    #[cfg(not(feature = "alloc"))]
    inner: T,
    #[cfg(feature = "alloc")]
    inner: alloc::boxed::Box<T>,
}

impl<T> MaybeBox<T> {
    /// Create a new `MaybeBox`, using `Box` if `alloc` is available.
    #[inline]
    pub(crate) fn new(inner: T) -> Self {
        #[cfg(not(feature = "alloc"))]
        {
            Self { inner }
        }
        #[cfg(feature = "alloc")]
        Self {
            inner: alloc::boxed::Box::new(inner),
        }
    }
}

impl<T> Deref for MaybeBox<T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl<T> DerefMut for MaybeBox<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner
    }
}

#[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 = P::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 = 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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>();
    }
}