[feat] ML-DSA87 driver (#1778)

This change adds a driver for interacting with the MLDSA peripheral.

builder/src/firmware.rs

@@ -220,6 +220,11 @@ pub mod driver_tests {
         ..BASE_FWID
     };
 
+    pub const ML_DSA87: FwId = FwId {
+        bin_name: "ml_dsa87",
+        ..BASE_FWID
+    };
+
     pub const PCRBANK: FwId = FwId {
         bin_name: "pcrbank",
         ..BASE_FWID
@@ -421,6 +426,7 @@ pub const REGISTERED_FW: &[&FwId] = &[
     &driver_tests::MAILBOX_DRIVER_SENDER,
     &driver_tests::MAILBOX_DRIVER_NEGATIVE_TESTS,
     &driver_tests::MBOX_SEND_TXN_DROP,
+    &driver_tests::ML_DSA87,
     &driver_tests::PCRBANK,
     &driver_tests::SHA1,
     &driver_tests::SHA256,

drivers/src/lib.rs

@@ -85,7 +85,9 @@ pub use lms::{
     Sha256Digest, D_INTR, D_LEAF, D_MESG, D_PBLC,
 };
 pub use mailbox::{Mailbox, MailboxRecvTxn, MailboxSendTxn};
-pub use mldsa87::{MlDsa87, MlDsa87PubKey, MlDsa87Reg, MlDsa87Signature};
+pub use mldsa87::{
+    Mldsa87, Mldsa87Msg, Mldsa87PubKey, Mldsa87Result, Mldsa87SignRnd, Mldsa87Signature,
+};
 pub use okref::okmutref;
 pub use okref::okref;
 pub use pcr_bank::{PcrBank, PcrId};

drivers/src/memory_layout.rs

@@ -19,7 +19,7 @@ use crate::FirmwareHandoffTable;
 use caliptra_image_types::ImageManifest;
 
 #[cfg(test)]
-use crate::MlDsa87PubKey;
+use crate::Mldsa87PubKey;
 
 //
 // Memory Addresses
@@ -112,7 +112,7 @@ fn mem_layout_test_fht() {
 #[test]
 #[allow(clippy::assertions_on_constants)]
 fn mem_layout_test_idevid_mldsa_pub_key() {
-    assert!(IDEVID_MLDSA_PUB_KEY_MAX_SIZE as usize >= core::mem::size_of::<MlDsa87PubKey>());
+    assert!(IDEVID_MLDSA_PUB_KEY_MAX_SIZE as usize >= core::mem::size_of::<Mldsa87PubKey>());
     assert_eq!(
         (LDEVID_TBS_ORG - IDEVID_MLDSA_PUB_KEY_ORG),
         IDEVID_MLDSA_PUB_KEY_MAX_SIZE

drivers/src/mldsa87.rs

@@ -4,7 +4,7 @@ Licensed under the Apache-2.0 license.
 
 File Name:
 
-    mldsa87.rs
+    Mldsa87.rs
 
 Abstract:
 
@@ -14,56 +14,85 @@ Abstract:
 #![allow(dead_code)]
 
 use crate::{
-    array::{Array4x1157, Array4x648},
-    Array4x16,
+    array::{Array4x1157, Array4x16, Array4x648, Array4x8},
+    kv_access::{KvAccess, KvAccessErr},
+    wait, CaliptraError, CaliptraResult, KeyReadArgs, Trng,
 };
-use crate::{CaliptraResult, KeyReadArgs, Trng};
+#[cfg(not(feature = "no-cfi"))]
+use caliptra_cfi_derive::cfi_impl_fn;
+use caliptra_registers::mldsa::{MldsaReg, RegisterBlock};
 
 #[must_use]
 #[repr(u32)]
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
-pub enum MlDsa87Result {
+pub enum Mldsa87Result {
     Success = 0xAAAAAAAA,
     SigVerifyFailed = 0x55555555,
 }
 
 /// MLDSA-87 Public Key
-pub type MlDsa87PubKey = Array4x648;
+pub type Mldsa87PubKey = Array4x648;
 
 /// MLDSA-87 Signature
-pub type MlDsa87Signature = Array4x1157;
+pub type Mldsa87Signature = Array4x1157;
 
 /// MLDSA-87 Message (64 Bytes)
-pub type MlDsa87MsgScalar = Array4x16;
+pub type Mldsa87Msg = Array4x16;
 
-/// TEMP: Placeholder for MlDsa87Reg
-pub struct MlDsa87Reg {
-    _priv: (),
-}
-impl MlDsa87Reg {
-    /// # Safety
-    ///
-    /// Caller must ensure that all concurrent use of this
-    /// peripheral in the firmware is done so in a compatible
-    /// way. The simplest way to enforce this is to only call
-    /// this function once.
-    #[inline(always)]
-    pub unsafe fn new() -> Self {
-        Self { _priv: () }
-    }
-}
-/// END - TEMP: Placeholder for MlDsa87Reg
+/// MLDSA-87 Signature RND
+pub type Mldsa87SignRnd = Array4x8;
+
+type Mldsa87VerifyRes = Array4x16;
 
 /// MLDSA-87  API
-pub struct MlDsa87 {
-    mldsa87: MlDsa87Reg,
+pub struct Mldsa87 {
+    mldsa87: MldsaReg,
 }
 
-impl MlDsa87 {
-    pub fn new(mldsa87: MlDsa87Reg) -> Self {
+impl Mldsa87 {
+    pub fn new(mldsa87: MldsaReg) -> Self {
         Self { mldsa87 }
     }
 
+    // The trng only generates 12 dwords
+    fn generate_iv(trng: &mut Trng) -> CaliptraResult<Array4x16> {
+        let iv = {
+            let mut iv = [0; 16];
+            let iv1 = trng.generate()?;
+            let iv2 = trng.generate()?;
+            iv[..12].copy_from_slice(&iv1.0);
+            iv[12..16].copy_from_slice(&iv2.0[0..4]);
+            Array4x16::from(iv)
+        };
+        Ok(iv)
+    }
+
+    // Wait on the provided condition OR the error condition defined in this function
+    // In the event of the error condition being set, clear the error bits and return an error
+    fn wait<F>(regs: RegisterBlock<ureg::RealMmioMut>, condition: F) -> CaliptraResult<()>
+    where
+        F: Fn() -> bool,
+    {
+        let err_condition = || {
+            (u32::from(regs.intr_block_rf().error_global_intr_r().read()) != 0)
+                || (u32::from(regs.intr_block_rf().error_internal_intr_r().read()) != 0)
+        };
+
+        // Wait for either the given condition or the error condition
+        wait::until(|| (condition() || err_condition()));
+
+        if err_condition() {
+            // Clear the errors
+            // error_global_intr_r is RO
+            regs.intr_block_rf()
+                .error_internal_intr_r()
+                .write(|_| u32::from(regs.intr_block_rf().error_internal_intr_r().read()).into());
+            return Err(CaliptraError::DRIVER_MLDSA87_HW_ERROR);
+        }
+
+        Ok(())
+    }
+
     /// Generate MLDSA-87 Key Pair
     ///
     /// # Arguments
@@ -73,13 +102,42 @@ impl MlDsa87 {
     ///
     /// # Returns
     ///
-    /// * `MlDsa87PubKey` - Generated MLDSA-87 Public Key
+    /// * `Mldsa87PubKey` - Generated MLDSA-87 Public Key
     pub fn key_pair(
         &mut self,
-        _seed: &KeyReadArgs,
-        _trng: &mut Trng,
-    ) -> CaliptraResult<MlDsa87PubKey> {
-        Ok(MlDsa87PubKey::default())
+        seed: &KeyReadArgs,
+        trng: &mut Trng,
+    ) -> CaliptraResult<Mldsa87PubKey> {
+        let mldsa = self.mldsa87.regs_mut();
+
+        // Wait for hardware ready
+        Mldsa87::wait(mldsa, || mldsa.status().read().ready())?;
+
+        // Clear the hardware before start
+        mldsa.ctrl().write(|w| w.zeroize(true));
+
+        // Copy seed from keyvault
+        KvAccess::copy_from_kv(*seed, mldsa.kv_rd_seed_status(), mldsa.kv_rd_seed_ctrl())
+            .map_err(|err| err.into_read_seed_err())?;
+
+        // Generate an IV.
+        let iv = Self::generate_iv(trng)?;
+        KvAccess::copy_from_arr(&iv, mldsa.entropy())?;
+
+        // Program the command register for key generation
+        mldsa.ctrl().write(|w| w.ctrl(|w| w.keygen()));
+
+        // Wait for hardware ready
+        Mldsa87::wait(mldsa, || mldsa.status().read().valid())?;
+
+        // Copy pubkey
+        let pubkey = Mldsa87PubKey::read_from_reg(mldsa.pubkey());
+
+        // Clear the hardware when done
+        mldsa.ctrl().write(|w| w.zeroize(true));
+
+        Ok(pubkey)
+        // TODO check that pubkey is valid?
     }
 
     /// Sign the digest with specified private key. To defend against glitching
@@ -88,22 +146,75 @@ impl MlDsa87 {
     ///
     /// # Arguments
     ///
-    /// * `priv_key_in` - Key Vault slot containing the seed for the private key generation.
+    /// * `seed` - Key Vault slot containing the seed for deterministic MLDSA Key Pair generation.
     /// * `pub_key` - Public key to verify the signature with.
     /// * `msg` - Message to sign.
+    /// * `sign_rnd` - Signature RND input
     /// * `trng` - TRNG driver instance.
     ///
     /// # Returns
     ///
-    /// * `MlDsa87Signature` - Generated signature
+    /// * `Mldsa87Signature` - Generated signature
     pub fn sign(
         &mut self,
-        _priv_key_in: &KeyReadArgs,
-        _pub_key: &MlDsa87PubKey,
-        _msg: &MlDsa87MsgScalar,
-        _trng: &mut Trng,
-    ) -> CaliptraResult<MlDsa87Signature> {
-        Ok(MlDsa87Signature::default())
+        seed: &KeyReadArgs,
+        pub_key: &Mldsa87PubKey,
+        msg: &Mldsa87Msg,
+        sign_rnd: &Mldsa87SignRnd,
+        trng: &mut Trng,
+    ) -> CaliptraResult<Mldsa87Signature> {
+        let mldsa = self.mldsa87.regs_mut();
+
+        // Wait for hardware ready
+        Mldsa87::wait(mldsa, || mldsa.status().read().ready())?;
+
+        // Clear the hardware before start
+        mldsa.ctrl().write(|w| w.zeroize(true));
+
+        // Copy seed from keyvault
+        KvAccess::copy_from_kv(*seed, mldsa.kv_rd_seed_status(), mldsa.kv_rd_seed_ctrl())
+            .map_err(|err| err.into_read_seed_err())?;
+
+        // Copy digest
+        KvAccess::copy_from_arr(msg, mldsa.msg())?;
+
+        // Sign RND, TODO do we want deterministic?
+        KvAccess::copy_from_arr(sign_rnd, mldsa.sign_rnd())?;
+
+        // Generate an IV.
+        let iv = Self::generate_iv(trng)?;
+        KvAccess::copy_from_arr(&iv, mldsa.entropy())?;
+
+        // Program the command register for key generation
+        mldsa.ctrl().write(|w| w.ctrl(|w| w.keygen_sign()));
+
+        // Wait for hardware ready
+        Mldsa87::wait(mldsa, || mldsa.status().read().valid())?;
+
+        // Copy signature
+        let signature = Mldsa87Signature::read_from_reg(mldsa.signature());
+
+        // Clear the hardware when done
+        mldsa.ctrl().write(|w| w.zeroize(true));
+
+        let verify_res = self.verify_res(pub_key, msg, &signature)?;
+
+        let truncated_signature = &signature.0[..16];
+
+        if verify_res.0 == truncated_signature {
+            // We only have a 6, 8 and 12 dword cfi assert
+            caliptra_cfi_lib::cfi_assert_eq_12_words(
+                &verify_res.0[..12].try_into().unwrap(),
+                &truncated_signature[..12].try_into().unwrap(),
+            );
+            caliptra_cfi_lib::cfi_assert_eq_8_words(
+                &verify_res.0[8..].try_into().unwrap(),
+                &truncated_signature[8..].try_into().unwrap(),
+            );
+            Ok(signature)
+        } else {
+            Err(CaliptraError::DRIVER_MLDSA87_SIGN_VALIDATION_FAILED)
+        }
     }
 
     /// Verify the signature with specified public key and message.
@@ -116,13 +227,88 @@ impl MlDsa87 {
     ///
     /// # Result
     ///
-    /// *  `MlDsa87Result` - MlDsa87Result::Success if the signature verification passed else an error code.
+    /// *  `Mldsa87Result` - Mldsa87Result::Success if the signature verification passed else an error code.
+    fn verify_res(
+        &mut self,
+        pub_key: &Mldsa87PubKey,
+        msg: &Mldsa87Msg,
+        signature: &Mldsa87Signature,
+    ) -> CaliptraResult<Mldsa87VerifyRes> {
+        let mldsa = self.mldsa87.regs_mut();
+
+        // Wait for hardware ready
+        Mldsa87::wait(mldsa, || mldsa.status().read().ready())?;
+
+        // Clear the hardware before start
+        mldsa.ctrl().write(|w| w.zeroize(true));
+
+        // Copy digest
+        msg.write_to_reg(mldsa.msg());
+
+        // Copy pubkey
+        pub_key.write_to_reg(mldsa.pubkey());
+
+        // Copy signature
+        signature.write_to_reg(mldsa.signature());
+
+        // Program the command register for signature verification
+        mldsa.ctrl().write(|w| w.ctrl(|w| w.verifying()));
+
+        // Wait for hardware ready
+        Mldsa87::wait(mldsa, || mldsa.status().read().valid())?;
+
+        // Copy the random value
+        let verify_res = Array4x16::read_from_reg(mldsa.verify_res());
+
+        // Clear the hardware when done
+        mldsa.ctrl().write(|w| w.zeroize(true));
+
+        Ok(verify_res)
+    }
+
+    #[cfg_attr(not(feature = "no-cfi"), cfi_impl_fn)]
     pub fn verify(
         &mut self,
-        _pub_key: &MlDsa87PubKey,
-        _msg: &MlDsa87MsgScalar,
-        _signature: &MlDsa87Signature,
-    ) -> CaliptraResult<MlDsa87Result> {
-        Ok(MlDsa87Result::Success)
+        pub_key: &Mldsa87PubKey,
+        msg: &Mldsa87Msg,
+        signature: &Mldsa87Signature,
+    ) -> CaliptraResult<Mldsa87Result> {
+        let verify_res = self.verify_res(pub_key, msg, signature)?;
+
+        let truncated_signature = &signature.0[..16];
+
+        let result = if verify_res.0 == truncated_signature {
+            // We only have a 6, 8 and 12 dword cfi assert
+            caliptra_cfi_lib::cfi_assert_eq_12_words(
+                &verify_res.0[..12].try_into().unwrap(),
+                &truncated_signature[..12].try_into().unwrap(),
+            );
+            caliptra_cfi_lib::cfi_assert_eq_8_words(
+                &verify_res.0[8..].try_into().unwrap(),
+                &truncated_signature[8..].try_into().unwrap(),
+            );
+            Mldsa87Result::Success
+        } else {
+            Mldsa87Result::SigVerifyFailed
+        };
+
+        Ok(result)
+    }
+}
+
+/// Mldsa87 key access error trait
+trait MlDsaKeyAccessErr {
+    /// Convert to read seed operation error
+    fn into_read_seed_err(self) -> CaliptraError;
+}
+
+impl MlDsaKeyAccessErr for KvAccessErr {
+    /// Convert to read seed operation error
+    fn into_read_seed_err(self) -> CaliptraError {
+        match self {
+            KvAccessErr::KeyRead => CaliptraError::DRIVER_MLDSA87_READ_SEED_KV_READ,
+            KvAccessErr::KeyWrite => CaliptraError::DRIVER_MLDSA87_READ_SEED_KV_WRITE,
+            KvAccessErr::Generic => CaliptraError::DRIVER_MLDSA87_READ_SEED_KV_UNKNOWN,
+        }
     }
 }