[update] MLDSA driver update for KAT

drivers/src/array.rs

@@ -163,6 +163,7 @@ pub type Array4x16 = Array4xN<16, 64>;
 pub type Array4x32 = Array4xN<32, 128>;
 pub type Array4x648 = Array4xN<648, 2592>;
 pub type Array4x1157 = Array4xN<1157, 4628>;
+pub type Array4x1224 = Array4xN<1224, 4896>;
 
 #[cfg(test)]
 mod tests {

drivers/src/ecc384.rs

@@ -59,7 +59,7 @@ impl From<KeyReadArgs> for Ecc384Seed<'_> {
     }
 }
 
-/// ECC-384 Public Key output
+/// ECC-384 Private Key output
 #[derive(Debug)]
 pub enum Ecc384PrivKeyOut<'a> {
     /// Array

drivers/src/lib.rs

@@ -84,7 +84,8 @@ pub use lms::{
 };
 pub use mailbox::{Mailbox, MailboxRecvTxn, MailboxSendTxn};
 pub use mldsa87::{
-    Mldsa87, Mldsa87Msg, Mldsa87PubKey, Mldsa87Result, Mldsa87SignRnd, Mldsa87Signature,
+    Mldsa87, Mldsa87Msg, Mldsa87PubKey, Mldsa87Result, Mldsa87Seed, Mldsa87SignRnd,
+    Mldsa87Signature,
 };
 pub use okref::okmutref;
 pub use okref::okref;

drivers/src/mldsa87.rs

@@ -14,7 +14,7 @@ Abstract:
 #![allow(dead_code)]
 
 use crate::{
-    array::{Array4x1157, Array4x16, Array4x648, Array4x8},
+    array::{Array4x1157, Array4x1224, Array4x16, Array4x648, Array4x8},
     kv_access::{KvAccess, KvAccessErr},
     wait, CaliptraError, CaliptraResult, KeyReadArgs, Trng,
 };
@@ -35,6 +35,9 @@ pub enum Mldsa87Result {
 /// MLDSA-87 Public Key
 pub type Mldsa87PubKey = Array4x648;
 
+/// MLDSA-87 Private Key
+pub type Mldsa87PrivKey = Array4x1224;
+
 /// MLDSA-87 Signature
 pub type Mldsa87Signature = Array4x1157;
 
@@ -46,6 +49,40 @@ pub type Mldsa87SignRnd = Array4x8;
 
 type Mldsa87VerifyRes = Array4x16;
 
+/// MLDSA-87 Seed
+#[derive(Debug, Copy, Clone)]
+pub enum Mldsa87Seed<'a> {
+    /// Array
+    Array4x8(&'a Array4x8),
+
+    /// Key Vault Key
+    Key(KeyReadArgs),
+
+    /// Private Key
+    PrivKey(&'a Mldsa87PrivKey),
+}
+
+impl<'a> From<&'a Array4x8> for Mldsa87Seed<'a> {
+    /// Converts to this type from the input type.
+    fn from(value: &'a Array4x8) -> Self {
+        Self::Array4x8(value)
+    }
+}
+
+impl From<KeyReadArgs> for Mldsa87Seed<'_> {
+    /// Converts to this type from the input type.
+    fn from(value: KeyReadArgs) -> Self {
+        Self::Key(value)
+    }
+}
+
+impl<'a> From<&'a Mldsa87PrivKey> for Mldsa87Seed<'a> {
+    /// Converts to this type from the input type.
+    fn from(value: &'a Mldsa87PrivKey) -> Self {
+        Self::PrivKey(value)
+    }
+}
+
 /// MLDSA-87  API
 pub struct Mldsa87 {
     mldsa87: MldsaReg,
@@ -99,16 +136,18 @@ impl Mldsa87 {
     ///
     /// # Arguments
     ///
-    /// * `seed` - Key Vault slot containing the seed for deterministic MLDSA Key Pair generation.
+    /// * `seed` - Either an array of 4x8 bytes or a key vault key to use as seed.
     /// * `trng` - TRNG driver instance.
+    /// * `priv_key_out` - Optional output parameter to store the private key.
     ///
     /// # Returns
     ///
     /// * `Mldsa87PubKey` - Generated MLDSA-87 Public Key
     pub fn key_pair(
         &mut self,
-        seed: &KeyReadArgs,
+        seed: &Mldsa87Seed,
         trng: &mut Trng,
+        priv_key_out: Option<&mut Mldsa87PrivKey>,
     ) -> CaliptraResult<Mldsa87PubKey> {
         let mldsa = self.mldsa87.regs_mut();
 
@@ -121,9 +160,15 @@ impl Mldsa87 {
         // Wait for hardware ready
         Mldsa87::wait(mldsa, || mldsa.status().read().ready())?;
 
-        // 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 seed to the hardware
+        match seed {
+            Mldsa87Seed::Array4x8(arr) => KvAccess::copy_from_arr(arr, mldsa.seed())?,
+            Mldsa87Seed::Key(key) => {
+                KvAccess::copy_from_kv(*key, mldsa.kv_rd_seed_status(), mldsa.kv_rd_seed_ctrl())
+                    .map_err(|err| err.into_read_seed_err())?
+            }
+            Mldsa87Seed::PrivKey(_) => Err(CaliptraError::DRIVER_MLDSA87_KEY_GEN_SEED_BAD_USAGE)?,
+        }
 
         // Generate an IV.
         let iv = Self::generate_iv(trng)?;
@@ -138,6 +183,11 @@ impl Mldsa87 {
         // Copy pubkey
         let pubkey = Mldsa87PubKey::read_from_reg(mldsa.pubkey());
 
+        // Copy private key if requested.
+        if let Some(priv_key) = priv_key_out {
+            *priv_key = Mldsa87PrivKey::read_from_reg(mldsa.privkey_out());
+        }
+
         // Clear the hardware when done
         mldsa.ctrl().write(|w| w.zeroize(true));
 
@@ -162,7 +212,7 @@ impl Mldsa87 {
     /// * `Mldsa87Signature` - Generated signature
     pub fn sign(
         &mut self,
-        seed: &KeyReadArgs,
+        seed: &Mldsa87Seed,
         pub_key: &Mldsa87PubKey,
         msg: &Mldsa87Msg,
         sign_rnd: &Mldsa87SignRnd,
@@ -179,9 +229,17 @@ impl Mldsa87 {
         // Wait for hardware ready
         Mldsa87::wait(mldsa, || mldsa.status().read().ready())?;
 
-        // 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 seed or the private key to the hardware
+        match seed {
+            Mldsa87Seed::Array4x8(arr) => KvAccess::copy_from_arr(arr, mldsa.seed())?,
+            Mldsa87Seed::Key(key) => {
+                KvAccess::copy_from_kv(*key, mldsa.kv_rd_seed_status(), mldsa.kv_rd_seed_ctrl())
+                    .map_err(|err| err.into_read_seed_err())?
+            }
+            Mldsa87Seed::PrivKey(priv_key) => {
+                KvAccess::copy_from_arr(priv_key, mldsa.privkey_in())?
+            }
+        }
 
         // Copy digest
         KvAccess::copy_from_arr(msg, mldsa.msg())?;

drivers/test-fw/src/bin/mldsa87_tests.rs

@@ -18,7 +18,7 @@ Abstract:
 use caliptra_cfi_lib::CfiCounter;
 use caliptra_drivers::{
     Array4x12, Hmac, HmacData, HmacKey, HmacMode, HmacTag, KeyId, KeyReadArgs, KeyUsage,
-    KeyWriteArgs, Mldsa87, Mldsa87Msg, Mldsa87PubKey, Mldsa87Result, Mldsa87SignRnd,
+    KeyWriteArgs, Mldsa87, Mldsa87Msg, Mldsa87PubKey, Mldsa87Result, Mldsa87Seed, Mldsa87SignRnd,
     Mldsa87Signature, Trng,
 };
 use caliptra_registers::csrng::CsrngReg;
@@ -302,8 +302,9 @@ fn test_gen_key_pair() {
     )
     .unwrap();
 
-    let seed = KeyReadArgs::new(KEY_ID);
-    let public_key = ml_dsa87.key_pair(&seed, &mut trng).unwrap();
+    let public_key = ml_dsa87
+        .key_pair(&Mldsa87Seed::Key(KeyReadArgs::new(KEY_ID)), &mut trng, None)
+        .unwrap();
     assert_eq!(public_key, Mldsa87PubKey::from(PUBKEY));
 }
 
@@ -321,12 +322,11 @@ fn test_sign() {
     };
 
     let sign_rnd = Mldsa87SignRnd::default(); // Deterministic signing
-    let seed = KeyReadArgs::new(KEY_ID); // Reuse SEED
 
     let signature = ml_dsa87
         .sign(
-            &seed,
-            &Mldsa87PubKey::from(PUBKEY),
+            &Mldsa87Seed::Key(KeyReadArgs::new(KEY_ID)),
+            &Mldsa87PubKey::from(PUBKEY), // Reuse SEED
             &MESSAGE.into(),
             &sign_rnd,
             &mut trng,

error/src/lib.rs

@@ -117,6 +117,8 @@ impl CaliptraError {
     pub const DRIVER_MLDSA87_HW_ERROR: CaliptraError = CaliptraError::new_const(0x00058003);
     pub const DRIVER_MLDSA87_SIGN_VALIDATION_FAILED: CaliptraError =
         CaliptraError::new_const(0x00058004);
+    pub const DRIVER_MLDSA87_KEY_GEN_SEED_BAD_USAGE: CaliptraError =
+        CaliptraError::new_const(0x00058005);
 
     pub const DRIVER_KV_ERASE_USE_LOCK_SET_FAILURE: CaliptraError =
         CaliptraError::new_const(0x00060001);

fmc/src/flow/crypto.rs

@@ -14,6 +14,7 @@ use caliptra_common::{
 use caliptra_drivers::{
     okref, Array4x12, CaliptraResult, Ecc384PrivKeyIn, Ecc384PrivKeyOut, Ecc384PubKey,
     Ecc384Result, Ecc384Signature, HmacMode, KeyId, KeyReadArgs, KeyUsage, KeyWriteArgs,
+    Mldsa87Seed,
 };
 
 pub enum Crypto {}
@@ -157,9 +158,11 @@ impl Crypto {
         Crypto::env_hmac_kdf(env, cdi, label, None, key_pair_seed, HmacMode::Hmac512)?;
 
         // Generate the public key.
-        let pub_key = env
-            .mldsa
-            .key_pair(&KeyReadArgs::new(key_pair_seed), &mut env.trng)?;
+        let pub_key = env.mldsa.key_pair(
+            &Mldsa87Seed::Key(KeyReadArgs::new(key_pair_seed)),
+            &mut env.trng,
+            None,
+        )?;
 
         Ok(MlDsaKeyPair {
             key_pair_seed,

rom/dev/src/crypto.rs

@@ -226,9 +226,11 @@ impl Crypto {
         Crypto::env_hmac_kdf(env, cdi, label, None, key_pair_seed, HmacMode::Hmac512)?;
 
         // Generate the public key.
-        let pub_key = env
-            .mldsa87
-            .key_pair(&KeyReadArgs::new(key_pair_seed), &mut env.trng)?;
+        let pub_key = env.mldsa87.key_pair(
+            &Mldsa87Seed::Key(KeyReadArgs::new(key_pair_seed)),
+            &mut env.trng,
+            None,
+        )?;
 
         Ok(MlDsaKeyPair {
             key_pair_seed,
@@ -260,9 +262,8 @@ impl Crypto {
     ) -> CaliptraResult<Mldsa87Signature> {
         let mut digest = env.sha2_512_384.sha512_digest(data);
         let digest = okmutref(&mut digest)?;
-        let priv_key_args = KeyReadArgs::new(priv_key);
         let result = env.mldsa87.sign(
-            &priv_key_args,
+            &Mldsa87Seed::Key(KeyReadArgs::new(priv_key)),
             pub_key,
             digest,
             &Mldsa87SignRnd::default(),