Add support for ML-DSA

Signed-off-by: Jakub Jelen <[email protected]>

Author: Jakuje

cryptoki/src/mechanism/mldsa.rs

@@ -0,0 +1,152 @@
+//! ML-DSA mechanism types
+
+use crate::mechanism::{Mechanism, MechanismType};
+
+use cryptoki_sys::*;
+use std::{convert::TryInto, marker::PhantomData, ptr::null_mut};
+
+/// The hedge type for ML-DSA signature
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub enum HedgeType {
+    /// CKH_HEDGE_PREFERRED: Default
+    /// token may create either a hedged signature or a deterministic signature
+    Preferred,
+    /// CKH_HEDGE_REQUIRED
+    /// token must produce a hedged signature or fail
+    Required,
+    /// CKH_DETERMINISTIC_REQUIRED
+    /// token must produce a deterministic signature or fail
+    DeterministicRequired,
+}
+
+impl HedgeType {
+    /// Convert the enum to Cryptoki value
+    pub fn into_ck(&self) -> CK_HEDGE_TYPE {
+        match self {
+            HedgeType::Preferred => CKH_HEDGE_PREFERRED,
+            HedgeType::Required => CKH_HEDGE_REQUIRED,
+            HedgeType::DeterministicRequired => CKH_DETERMINISTIC_REQUIRED,
+        }
+    }
+}
+
+/// The ML-DSA additional context for signatures
+///
+/// This structure wraps `CK_SIGN_ADDITIONAL_CONTEXT` structure.
+#[derive(Debug, Clone, Copy)]
+pub struct SignAdditionalContext<'a> {
+    inner: Option<CK_SIGN_ADDITIONAL_CONTEXT>,
+    /// Marker type to ensure we don't outlive the data
+    _marker: PhantomData<&'a [u8]>,
+}
+
+impl SignAdditionalContext<'_> {
+    /// Construct ML-DSA signature parameters.
+    ///
+    /// # Arguments
+    ///
+    /// * `hedge` - The HedgeType.
+    /// * `context` - The context.
+    ///
+    /// # Returns
+    ///
+    /// A new SignAdditionalContext struct.
+    pub fn new(hedge: HedgeType, context: Option<&[u8]>) -> Self {
+        if hedge == HedgeType::Preferred && context.is_none() {
+            return Self {
+                inner: None,
+                _marker: PhantomData,
+            };
+        }
+
+        let (p_context, ul_context_len) = match context {
+            Some(c) => (
+                c.as_ptr() as *mut _,
+                c.len().try_into().expect("usize can not fit in CK_ULONG"),
+            ),
+            None => (null_mut() as *mut _, 0),
+        };
+        Self {
+            inner: Some(CK_SIGN_ADDITIONAL_CONTEXT {
+                hedgeVariant: hedge.into_ck(),
+                pContext: p_context,
+                ulContextLen: ul_context_len,
+            }),
+            _marker: PhantomData,
+        }
+    }
+
+    /// Retrieve the inner `CK_SIGN_ADDITIONAL_CONTEXT` struct, if present.
+    ///
+    /// This method provides a reference to the `CK_SIGN_ADDITIONAL_CONTEXT`
+    /// struct encapsulated within the `SignAdditionalContext`, if the signature
+    /// scheme requires additional parameters.
+    ///
+    /// # Returns
+    ///
+    /// `Some(&CK_SIGN_ADDITIONAL_CONTEXT)` if the signature scheme has associated
+    /// parameters, otherwise `None`.
+    pub fn inner(&self) -> Option<&CK_SIGN_ADDITIONAL_CONTEXT> {
+        self.inner.as_ref()
+    }
+}
+
+/// The ML-DSA additional context for signatures with hashing information
+///
+/// This structure wraps `CK_HASH_SIGN_ADDITIONAL_CONTEXT` structure.
+#[derive(Debug, Clone, Copy)]
+pub struct HashSignAdditionalContext<'a> {
+    inner: CK_HASH_SIGN_ADDITIONAL_CONTEXT,
+    /// Marker type to ensure we don't outlive the data
+    _marker: PhantomData<&'a [u8]>,
+}
+
+impl HashSignAdditionalContext<'_> {
+    /// Construct HashML-DSA Signature parameters.
+    ///
+    /// # Arguments
+    ///
+    /// * `hedge` - The HedgeType.
+    /// * `context` - The context
+    /// * `hash` - The hash type
+    ///
+    /// # Returns
+    ///
+    /// A new SignAdditionalContext struct.
+    pub fn new(hedge: HedgeType, context: Option<&[u8]>, hash: MechanismType) -> Self {
+        let (p_context, ul_context_len) = match context {
+            Some(c) => (
+                c.as_ptr() as *mut _,
+                c.len().try_into().expect("usize can not fit in CK_ULONG"),
+            ),
+            None => (null_mut(), 0),
+        };
+        Self {
+            inner: CK_HASH_SIGN_ADDITIONAL_CONTEXT {
+                hedgeVariant: hedge.into_ck(),
+                pContext: p_context,
+                ulContextLen: ul_context_len,
+                hash: hash.into(),
+            },
+            _marker: PhantomData,
+        }
+    }
+
+    /// Retrieve the inner `CK_HASH_SIGN_ADDITIONAL_CONTEXT` struct.
+    ///
+    /// This method provides a reference to the `CK_HASH_SIGN_ADDITIONAL_CONTEXT`
+    /// struct encapsulated within the `HashSignAdditionalContext`.
+    ///
+    /// # Returns
+    ///
+    /// `&CK_HASH_SIGN_ADDITIONAL_CONTEXT`.
+    pub fn inner(&self) -> &CK_HASH_SIGN_ADDITIONAL_CONTEXT {
+        &self.inner
+    }
+}
+
+impl<'a> From<HashSignAdditionalContext<'a>> for Mechanism<'a> {
+    fn from(params: HashSignAdditionalContext<'a>) -> Self {
+        Mechanism::HashMlDsa(params)
+    }
+}

cryptoki/src/mechanism/mod.rs

@@ -8,6 +8,7 @@ pub mod ekdf;
 pub mod elliptic_curve;
 pub mod hkdf;
 mod mechanism_info;
+pub mod mldsa;
 pub mod rsa;
 pub mod vendor_defined;
 
@@ -334,6 +335,58 @@ impl MechanismType {
     /// ML-KEM encapsulation and decapsulation mechanism
     pub const ML_KEM: MechanismType = MechanismType { val: CKM_ML_KEM };
 
+    // ML-DSA
+    /// ML-DSA key pair generation mechanism
+    pub const ML_DSA_KEY_PAIR_GEN: MechanismType = MechanismType {
+        val: CKM_ML_DSA_KEY_PAIR_GEN,
+    };
+    /// ML-DSA signature mechanism
+    pub const ML_DSA: MechanismType = MechanismType { val: CKM_ML_DSA };
+    /// HashML-DSA signature mechanism
+    pub const HASH_ML_DSA: MechanismType = MechanismType {
+        val: CKM_HASH_ML_DSA,
+    };
+    /// HashML-DSA signature mechanism with SHA224
+    pub const HASH_ML_DSA_SHA224: MechanismType = MechanismType {
+        val: CKM_HASH_ML_DSA_SHA224,
+    };
+    /// HashML-DSA signature mechanism with SHA256
+    pub const HASH_ML_DSA_SHA256: MechanismType = MechanismType {
+        val: CKM_HASH_ML_DSA_SHA256,
+    };
+    /// HashML-DSA signature mechanism with SHA384
+    pub const HASH_ML_DSA_SHA384: MechanismType = MechanismType {
+        val: CKM_HASH_ML_DSA_SHA384,
+    };
+    /// HashML-DSA signature mechanism with SHA512
+    pub const HASH_ML_DSA_SHA512: MechanismType = MechanismType {
+        val: CKM_HASH_ML_DSA_SHA512,
+    };
+    /// HashML-DSA signature mechanism with SHA3-224
+    pub const HASH_ML_DSA_SHA3_224: MechanismType = MechanismType {
+        val: CKM_HASH_ML_DSA_SHA3_224,
+    };
+    /// HashML-DSA signature mechanism with SHA3-256
+    pub const HASH_ML_DSA_SHA3_256: MechanismType = MechanismType {
+        val: CKM_HASH_ML_DSA_SHA3_256,
+    };
+    /// HashML-DSA signature mechanism with SHA3-384
+    pub const HASH_ML_DSA_SHA3_384: MechanismType = MechanismType {
+        val: CKM_HASH_ML_DSA_SHA3_384,
+    };
+    /// HashML-DSA signature mechanism with SHA3-512
+    pub const HASH_ML_DSA_SHA3_512: MechanismType = MechanismType {
+        val: CKM_HASH_ML_DSA_SHA3_512,
+    };
+    /// HashML-DSA signature mechanism with SHAKE128
+    pub const HASH_ML_DSA_SHAKE128: MechanismType = MechanismType {
+        val: CKM_HASH_ML_DSA_SHAKE128,
+    };
+    /// HashML-DSA signature mechanism with SHAKE256
+    pub const HASH_ML_DSA_SHAKE256: MechanismType = MechanismType {
+        val: CKM_HASH_ML_DSA_SHAKE256,
+    };
+
     /// Create vendor defined mechanism
     ///
     /// # Arguments
@@ -725,6 +778,19 @@ impl MechanismType {
             CKM_HKDF_DATA => String::from(stringify!(CKM_HKDF_DATA)),
             CKM_ML_KEM_KEY_PAIR_GEN => String::from(stringify!(CKM_ML_KEM_KEY_PAIR_GEN)),
             CKM_ML_KEM => String::from(stringify!(CKM_ML_KEM)),
+            CKM_ML_DSA_KEY_PAIR_GEN => String::from(stringify!(CKM_ML_DSA_KEY_PAIR_GEN)),
+            CKM_ML_DSA => String::from(stringify!(CKM_ML_DSA)),
+            CKM_HASH_ML_DSA => String::from(stringify!(CKM_HASH_ML_DSA)),
+            CKM_HASH_ML_DSA_SHA224 => String::from(stringify!(CKM_HASH_ML_DSA_SHA224)),
+            CKM_HASH_ML_DSA_SHA256 => String::from(stringify!(CKM_HASH_ML_DSA_SHA256)),
+            CKM_HASH_ML_DSA_SHA384 => String::from(stringify!(CKM_HASH_ML_DSA_SHA384)),
+            CKM_HASH_ML_DSA_SHA512 => String::from(stringify!(CKM_HASH_ML_DSA_SHA512)),
+            CKM_HASH_ML_DSA_SHA3_224 => String::from(stringify!(CKM_HASH_ML_DSA_SHA3_224)),
+            CKM_HASH_ML_DSA_SHA3_256 => String::from(stringify!(CKM_HASH_ML_DSA_SHA3_256)),
+            CKM_HASH_ML_DSA_SHA3_384 => String::from(stringify!(CKM_HASH_ML_DSA_SHA3_384)),
+            CKM_HASH_ML_DSA_SHA3_512 => String::from(stringify!(CKM_HASH_ML_DSA_SHA3_512)),
+            CKM_HASH_ML_DSA_SHAKE128 => String::from(stringify!(CKM_HASH_ML_DSA_SHAKE128)),
+            CKM_HASH_ML_DSA_SHAKE256 => String::from(stringify!(CKM_HASH_ML_DSA_SHAKE256)),
             _ => format!("unknown {mech:08x}"),
         }
     }
@@ -811,6 +877,18 @@ impl TryFrom<CK_MECHANISM_TYPE> for MechanismType {
             CKM_HKDF_DATA => Ok(MechanismType::HKDF_DATA),
             CKM_ML_KEM_KEY_PAIR_GEN => Ok(MechanismType::ML_KEM_KEY_PAIR_GEN),
             CKM_ML_KEM => Ok(MechanismType::ML_KEM),
+            CKM_ML_DSA_KEY_PAIR_GEN => Ok(MechanismType::ML_DSA_KEY_PAIR_GEN),
+            CKM_ML_DSA => Ok(MechanismType::ML_DSA),
+            CKM_HASH_ML_DSA => Ok(MechanismType::HASH_ML_DSA),
+            CKM_HASH_ML_DSA_SHA224 => Ok(MechanismType::HASH_ML_DSA_SHA224),
+            CKM_HASH_ML_DSA_SHA256 => Ok(MechanismType::HASH_ML_DSA_SHA256),
+            CKM_HASH_ML_DSA_SHA384 => Ok(MechanismType::HASH_ML_DSA_SHA384),
+            CKM_HASH_ML_DSA_SHA512 => Ok(MechanismType::HASH_ML_DSA_SHA512),
+            CKM_HASH_ML_DSA_SHA3_224 => Ok(MechanismType::HASH_ML_DSA_SHA3_224),
+            CKM_HASH_ML_DSA_SHA3_256 => Ok(MechanismType::HASH_ML_DSA_SHA3_256),
+            CKM_HASH_ML_DSA_SHA3_384 => Ok(MechanismType::HASH_ML_DSA_SHA3_384),
+            CKM_HASH_ML_DSA_SHA3_512 => Ok(MechanismType::HASH_ML_DSA_SHA3_512),
+            CKM_HASH_ML_DSA_SHAKE128 => Ok(MechanismType::HASH_ML_DSA_SHAKE128),
             other => {
                 error!("Mechanism type {} is not supported.", other);
                 Err(Error::NotSupported)
@@ -1039,6 +1117,34 @@ pub enum Mechanism<'a> {
     /// ML-KEM key encacpsulation/decapsulation mechanism
     MlKem,
 
+    // ML-DSA
+    /// ML-DSA key pair generation mechanism
+    MlDsaKeyPairGen,
+    /// ML-DSA signature mechanism
+    MlDsa(mldsa::SignAdditionalContext<'a>),
+    /// HashML-DSA signature mechanism
+    HashMlDsa(mldsa::HashSignAdditionalContext<'a>),
+    /// HashML-DSA signature mechanism with SHA224
+    HashMlDsaSha224(mldsa::SignAdditionalContext<'a>),
+    /// HashML-DSA signature mechanism with SHA256
+    HashMlDsaSha256(mldsa::SignAdditionalContext<'a>),
+    /// HashML-DSA signature mechanism with SHA384
+    HashMlDsaSha384(mldsa::SignAdditionalContext<'a>),
+    /// HashML-DSA signature mechanism with SHA512
+    HashMlDsaSha512(mldsa::SignAdditionalContext<'a>),
+    /// HashML-DSA signature mechanism with SHA3-224
+    HashMlDsaSha3_224(mldsa::SignAdditionalContext<'a>),
+    /// HashML-DSA signature mechanism with SHA3-256
+    HashMlDsaSha3_256(mldsa::SignAdditionalContext<'a>),
+    /// HashML-DSA signature mechanism with SHA3-384
+    HashMlDsaSha3_384(mldsa::SignAdditionalContext<'a>),
+    /// HashML-DSA signature mechanism with SHA3-512
+    HashMlDsaSha3_512(mldsa::SignAdditionalContext<'a>),
+    /// HashML-DSA signature mechanism with SHAKE128
+    HashMlDsaShake128(mldsa::SignAdditionalContext<'a>),
+    /// HashML-DSA signature mechanism with SHAKE256
+    HashMlDsaShake256(mldsa::SignAdditionalContext<'a>),
+
     /// Vendor defined mechanism
     VendorDefined(VendorDefinedMechanism<'a>),
 }
@@ -1123,6 +1229,20 @@ impl Mechanism<'_> {
             Mechanism::MlKemKeyPairGen => MechanismType::ML_KEM_KEY_PAIR_GEN,
             Mechanism::MlKem => MechanismType::ML_KEM,
 
+            Mechanism::MlDsaKeyPairGen => MechanismType::ML_DSA_KEY_PAIR_GEN,
+            Mechanism::MlDsa(_) => MechanismType::ML_DSA,
+            Mechanism::HashMlDsa(_) => MechanismType::HASH_ML_DSA,
+            Mechanism::HashMlDsaSha224(_) => MechanismType::HASH_ML_DSA_SHA224,
+            Mechanism::HashMlDsaSha256(_) => MechanismType::HASH_ML_DSA_SHA256,
+            Mechanism::HashMlDsaSha384(_) => MechanismType::HASH_ML_DSA_SHA384,
+            Mechanism::HashMlDsaSha512(_) => MechanismType::HASH_ML_DSA_SHA512,
+            Mechanism::HashMlDsaSha3_224(_) => MechanismType::HASH_ML_DSA_SHA3_224,
+            Mechanism::HashMlDsaSha3_256(_) => MechanismType::HASH_ML_DSA_SHA3_256,
+            Mechanism::HashMlDsaSha3_384(_) => MechanismType::HASH_ML_DSA_SHA3_384,
+            Mechanism::HashMlDsaSha3_512(_) => MechanismType::HASH_ML_DSA_SHA3_512,
+            Mechanism::HashMlDsaShake128(_) => MechanismType::HASH_ML_DSA_SHAKE128,
+            Mechanism::HashMlDsaShake256(_) => MechanismType::HASH_ML_DSA_SHAKE256,
+
             Mechanism::VendorDefined(vm) => MechanismType {
                 val: vm.inner.mechanism,
             },
@@ -1175,6 +1295,25 @@ impl From<&Mechanism<'_>> for CK_MECHANISM {
             Mechanism::HkdfDerive(params) | Mechanism::HkdfData(params) => {
                 make_mechanism(mechanism, params)
             }
+            Mechanism::HashMlDsa(params) => make_mechanism(mechanism, params),
+            Mechanism::MlDsa(params)
+            | Mechanism::HashMlDsaSha224(params)
+            | Mechanism::HashMlDsaSha256(params)
+            | Mechanism::HashMlDsaSha384(params)
+            | Mechanism::HashMlDsaSha512(params)
+            | Mechanism::HashMlDsaSha3_224(params)
+            | Mechanism::HashMlDsaSha3_256(params)
+            | Mechanism::HashMlDsaSha3_384(params)
+            | Mechanism::HashMlDsaSha3_512(params)
+            | Mechanism::HashMlDsaShake128(params)
+            | Mechanism::HashMlDsaShake256(params) => match params.inner() {
+                None => CK_MECHANISM {
+                    mechanism,
+                    pParameter: null_mut(),
+                    ulParameterLen: 0,
+                },
+                Some(params) => make_mechanism(mechanism, params),
+            },
             // Mechanisms without parameters
             Mechanism::AesKeyGen
             | Mechanism::AesEcb
@@ -1221,7 +1360,8 @@ impl From<&Mechanism<'_>> for CK_MECHANISM {
             | Mechanism::GenericSecretKeyGen
             | Mechanism::HkdfKeyGen
             | Mechanism::MlKemKeyPairGen
-            | Mechanism::MlKem => CK_MECHANISM {
+            | Mechanism::MlKem
+            | Mechanism::MlDsaKeyPairGen => CK_MECHANISM {
                 mechanism,
                 pParameter: null_mut(),
                 ulParameterLen: 0,