Repository: chipsalliance/caliptra-sw
Pull request #176 (2023-04-13 21:01:12)

ROM: Don't indicate failure after successful firmware download.
Also, have HwModel::download_firmware() check for mailbox success, and extract general purpose code into HwModel::mailbox_execute() function.

More tests are in the pipeline, but they are currently blocked on incorrect behavior of the emulator's mailbox peripheral.


Patch:
diff --git a/hw-model/src/lib.rs b/hw-model/src/lib.rs
index e5c87d35..207c04f8 100644
--- a/hw-model/src/lib.rs
+++ b/hw-model/src/lib.rs
@@ -26,8 +26,6 @@ pub use model_emulated::ModelEmulated;
 #[cfg(feature = "verilator")]
 pub use model_verilated::ModelVerilated;
 
-use caliptra_emu_types::RvSize;
-
 /// Ideally, general-purpose functions would return `impl HwModel` instead of
 /// `DefaultHwModel` to prevent users from calling functions that aren't
 /// available on all HwModel implementations.  Unfortunately, rust-analyzer
@@ -95,9 +93,13 @@ pub struct BootParams<'a> {
     pub fw_image: Option<&'a [u8]>,
 }
 
-#[derive(Debug)]
+#[derive(Debug, Eq, PartialEq)]
 pub enum ModelError {
-    MailboxErr,
+    MailboxCmdFailed,
+    UnableToLockMailbox,
+    BufferTooLargeForMailbox,
+    UploadFirmwareUnexpectedResponse,
+    UnknownCommandStatus(u32),
     NotReadyForFwErr,
     ReadyForFirmwareTimeout { cycles: u32 },
 }
@@ -105,7 +107,16 @@ impl Error for ModelError {}
 impl Display for ModelError {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
-            ModelError::MailboxErr => write!(f, "Mailbox error"),
+            ModelError::MailboxCmdFailed => write!(f, "Mailbox command failed"),
+            ModelError::UnableToLockMailbox => write!(f, "Unable to lock mailbox"),
+            ModelError::BufferTooLargeForMailbox => write!(f, "Buffer too large for mailbox"),
+            ModelError::UploadFirmwareUnexpectedResponse => {
+                write!(f, "Received unexpected response after uploading firmware")
+            }
+            ModelError::UnknownCommandStatus(status) => write!(
+                f,
+                "Received unknown command status from mailbox peripheral: 0x{status:x}"
+            ),
             ModelError::NotReadyForFwErr => write!(f, "Not ready for firmware"),
             ModelError::ReadyForFirmwareTimeout { cycles } => write!(
                 f,
@@ -330,44 +341,105 @@ pub trait HwModel {
 
     fn tracing_hint(&mut self, enable: bool);
 
-    /// Upload firmware to the mailbox.
-    fn upload_firmware(&mut self, firmware: &[u8]) -> Result<(), ModelError> {
+    /// Executes `cmd` with request data `buf`. Returns `Ok(Some(_))` if
+    /// the uC responded with data, `Ok(None)` if the uC indicated success
+    /// without data, Err(ModelError::MailboxCmdFailed) if the microcontroller
+    /// responded with an error, or other model errors if there was a problem
+    /// communicating with the mailbox.
+    fn mailbox_execute(
+        &mut self,
+        cmd: u32,
+        buf: &[u8],
+    ) -> std::result::Result<Option<Vec<u8>>, ModelError> {
+        const MAILBOX_SIZE: u32 = 128 * 1024;
+
         if self.soc_mbox().lock().read().lock() {
-            return Err(ModelError::MailboxErr);
-        }
-        if !self.soc_mbox().lock().read().lock() {
-            return Err(ModelError::MailboxErr);
+            return Err(ModelError::UnableToLockMailbox);
         }
-        #[cfg(feature = "verilator")]
-        if !self.soc_ifc().cptra_flow_status().read().ready_for_fw() {
-            return Err(ModelError::NotReadyForFwErr);
+        let Ok(input_len) = u32::try_from(buf.len()) else {
+            return Err(ModelError::BufferTooLargeForMailbox);
+        };
+        if input_len > MAILBOX_SIZE {
+            return Err(ModelError::BufferTooLargeForMailbox);
         }
 
-        self.soc_mbox().cmd().write(|_| FW_LOAD_CMD_OPCODE);
+        self.soc_mbox().cmd().write(|_| cmd);
+        self.soc_mbox().dlen().write(|_| input_len);
 
-        self.soc_mbox().dlen().write(|_| firmware.len() as u32);
+        writeln!(
+            self.output().logger(),
+            "<<< Executing mbox cmd 0x{cmd:08x} ({input_len} bytes) from SoC"
+        )
+        .unwrap();
+        let mut remaining = buf;
+        while remaining.len() >= 4 {
+            // Panic is impossible because the subslice is always 4 bytes
+            let word = u32::from_le_bytes(remaining[..4].try_into().unwrap());
+            self.soc_mbox().datain().write(|_| word);
+            remaining = &remaining[4..];
+        }
+        if !remaining.is_empty() {
+            let mut word_bytes = [0u8; 4];
+            word_bytes[..remaining.len()].copy_from_slice(remaining);
+            let word = u32::from_le_bytes(word_bytes);
+            self.soc_mbox().datain().write(|_| word);
+        }
 
-        let word_size = RvSize::Word as usize;
-        let remainder = firmware.len() % word_size;
-        let n = firmware.len() - remainder;
+        // Ask the microcontroller to execute this command
+        self.soc_mbox().execute().write(|w| w.execute(true));
 
-        for idx in (0..n).step_by(word_size) {
-            let val = u32::from_le_bytes(firmware[idx..idx + word_size].try_into().unwrap());
-            self.soc_mbox().datain().write(|_| val);
+        // Wait for the microcontroller to finish executing
+        while self.soc_mbox().status().read().status().cmd_busy() {
+            self.step();
+        }
+        let status = self.soc_mbox().status().read().status();
+        if status.cmd_failure() {
+            writeln!(self.output().logger(), ">>> mbox cmd response: failed").unwrap();
+            self.soc_mbox().execute().write(|w| w.execute(false));
+            return Err(ModelError::MailboxCmdFailed);
+        }
+        if status.cmd_complete() {
+            writeln!(self.output().logger(), ">>> mbox cmd response: success").unwrap();
+            self.soc_mbox().execute().write(|w| w.execute(false));
+            return Ok(None);
+        }
+        if !status.data_ready() {
+            return Err(ModelError::UnknownCommandStatus(status as u32));
         }
 
-        // Handle the remainder bytes.
-        if remainder > 0 {
-            let mut last_word = firmware[n] as u32;
-            for idx in 1..remainder {
-                last_word |= (firmware[n + idx] as u32) << (idx << 3);
-            }
-            self.soc_mbox().datain().write(|_| last_word);
+        let mut result = vec![];
+        let mut dlen = self.soc_mbox().dlen().read();
+
+        writeln!(
+            self.output().logger(),
+            ">>> mbox cmd response data ({dlen} bytes)"
+        )
+        .unwrap();
+        while dlen >= 4 {
+            result.extend_from_slice(&self.soc_mbox().dataout().read().to_le_bytes());
+            dlen -= 4;
+        }
+        if dlen > 0 {
+            // Unwrap cannot panic because dlen is less than 4
+            result.extend_from_slice(
+                &self.soc_mbox().dataout().read().to_le_bytes()[..usize::try_from(dlen).unwrap()],
+            );
         }
 
-        // Set Execute Bit
-        self.soc_mbox().execute().write(|w| w.execute(true));
+        self.soc_mbox().execute().write(|w| w.execute(false));
+        // mbox_fsm_ps isn't updated immediately after execute is cleared (!?),
+        // so step an extra clock cycle to wait for fm_ps to update
+        self.step();
+        assert!(self.soc_mbox().status().read().mbox_fsm_ps().mbox_idle());
+        Ok(Some(result))
+    }
 
+    /// Upload firmware to the mailbox.
+    fn upload_firmware(&mut self, firmware: &[u8]) -> Result<(), ModelError> {
+        let response = self.mailbox_execute(FW_LOAD_CMD_OPCODE, firmware)?;
+        if response.is_some() {
+            return Err(ModelError::UploadFirmwareUnexpectedResponse);
+        }
         Ok(())
     }
 }
diff --git a/rom/dev/src/flow/cold_reset/fmc_alias.rs b/rom/dev/src/flow/cold_reset/fmc_alias.rs
index e6fd6c8f..bb81d0b6 100644
--- a/rom/dev/src/flow/cold_reset/fmc_alias.rs
+++ b/rom/dev/src/flow/cold_reset/fmc_alias.rs
@@ -57,7 +57,7 @@ impl DiceLayer for FmcAliasLayer {
         );
 
         // Download the image
-        let txn = Self::download_image(env)?;
+        let mut txn = Self::download_image(env)?;
 
         // Load the manifest
         let manifest = Self::load_manifest(&txn)?;
@@ -72,7 +72,11 @@ impl DiceLayer for FmcAliasLayer {
         Self::extend_pcrs(env)?;
 
         // Load the image
-        Self::load_image(env, &manifest, txn)?;
+        Self::load_image(env, &manifest, &txn)?;
+
+        // Complete the mailbox transaction indicating success.
+        txn.complete(true)?;
+        drop(txn);
 
         // At this point PCR0 & PCR1 must have the same value. We use the value
         // of PCR1 as the UDS for deriving the CDI
@@ -191,7 +195,7 @@ impl FmcAliasLayer {
     fn load_image(
         _env: &RomEnv,
         manifest: &ImageManifest,
-        txn: MailboxRecvTxn,
+        txn: &MailboxRecvTxn,
     ) -> CaliptraResult<()> {
         cprintln!(
             "[afmc] Loading FMC at address 0x{:08x} len {}",
@@ -219,10 +223,6 @@ impl FmcAliasLayer {
 
         txn.copy_request(0, runtime_dest)?;
 
-        // Drop the tranaction and release the Mailbox lock after the image
-        // has been successfully verified and loaded in memory
-        drop(txn);
-
         Ok(())
     }
 


Test patch:
diff --git a/rom/dev/tests/test_fmcalias_derivation.rs b/rom/dev/tests/test_fmcalias_derivation.rs
index 52aa6997..c159a8b2 100644
--- a/rom/dev/tests/test_fmcalias_derivation.rs
+++ b/rom/dev/tests/test_fmcalias_derivation.rs
@@ -1,5 +1,5 @@
 // Licensed under the Apache-2.0 license
-use caliptra_hw_model::{Fuses, HwModel};
+use caliptra_hw_model::{Fuses, HwModel, ModelError};
 use caliptra_image_types::IMAGE_BYTE_SIZE;
 
 mod helpers;
@@ -13,18 +13,32 @@ fn test_zero_firmware_size() {
     let mut output = vec![];
 
     // Zero-sized firmware.
-    hw.upload_firmware(&[]).unwrap();
+    assert_eq!(
+        hw.upload_firmware(&[]).unwrap_err(),
+        ModelError::MailboxCmdFailed
+    );
     let result = hw.copy_output_until_non_fatal_error(INVALID_IMAGE_SIZE, &mut output);
     assert!(result.is_ok());
 }
 
 #[test]
 fn test_firmware_gt_max_size() {
+    const FW_LOAD_CMD_OPCODE: u32 = 0x4657_4C44;
+
+    // Firmware size > 128 KB.
+
     let (mut hw, _image_bundle) = helpers::build_hw_model_and_image_bundle(Fuses::default());
     let mut output = vec![];
 
-    // Firmware size > 128 KB.
-    hw.upload_firmware(&vec![0u8; IMAGE_BYTE_SIZE + 1]).unwrap();
+    // Manually put the oversize data in the mailbox because
+    // HwModel::upload_firmware won't let us.
+    assert!(!hw.soc_mbox().lock().read().lock());
+    hw.soc_mbox().cmd().write(|_| FW_LOAD_CMD_OPCODE);
+    hw.soc_mbox().dlen().write(|_| (IMAGE_BYTE_SIZE + 1) as u32);
+    for i in 0..((IMAGE_BYTE_SIZE + 1 + 3) / 4) {
+        hw.soc_mbox().datain().write(|_| i as u32);
+    }
+    hw.soc_mbox().execute().write(|w| w.execute(true));
     let result = hw.copy_output_until_non_fatal_error(INVALID_IMAGE_SIZE, &mut output);
     assert!(result.is_ok());
 }
diff --git a/rom/dev/tests/test_image_validation.rs b/rom/dev/tests/test_image_validation.rs
index 8e1ab539..7be85551 100644
--- a/rom/dev/tests/test_image_validation.rs
+++ b/rom/dev/tests/test_image_validation.rs
@@ -3,7 +3,7 @@
 use caliptra_builder::{ImageOptions, APP_WITH_UART, FMC_WITH_UART, ROM_WITH_UART};
 use caliptra_drivers::Array4x12;
 use caliptra_hw_model::{
-    BootParams, DeviceLifecycle, Fuses, HwModel, InitParams, SecurityState, U4,
+    BootParams, DeviceLifecycle, Fuses, HwModel, InitParams, ModelError, SecurityState, U4,
 };
 use caliptra_image_elf::ElfExecutable;
 use caliptra_image_fake_keys::{
@@ -65,8 +65,12 @@ const SIGNATURE_S: [u8; 48] = [
 fn test_invalid_manifest_marker() {
     let (mut hw, mut image_bundle) = helpers::build_hw_model_and_image_bundle(Fuses::default());
     image_bundle.manifest.marker = 0xDEADBEEF;
-    hw.upload_firmware(&image_bundle.to_bytes().unwrap())
-        .unwrap();
+
+    assert_eq!(
+        ModelError::MailboxCmdFailed,
+        hw.upload_firmware(&image_bundle.to_bytes().unwrap())
+            .unwrap_err()
+    );
     let mut output = vec![];
     let result = hw.copy_output_until_non_fatal_error(MANIFEST_MARKER_MISMATCH, &mut output);
     assert!(result.is_ok());
@@ -76,8 +80,12 @@ fn test_invalid_manifest_marker() {
 fn test_invalid_manifest_size() {
     let (mut hw, mut image_bundle) = helpers::build_hw_model_and_image_bundle(Fuses::default());
     image_bundle.manifest.size = (core::mem::size_of::<ImageManifest>() - 1) as u32;
-    hw.upload_firmware(&image_bundle.to_bytes().unwrap())
-        .unwrap();
+
+    assert_eq!(
+        ModelError::MailboxCmdFailed,
+        hw.upload_firmware(&image_bundle.to_bytes().unwrap())
+            .unwrap_err()
+    );
     let mut output = vec![];
     let result = hw.copy_output_until_non_fatal_error(MANIFEST_SIZE_MISMATCH, &mut output);
     assert!(result.is_ok());
@@ -91,8 +99,12 @@ fn test_preamble_zero_vendor_pubkey_digest() {
         ..Default::default()
     };
     let (mut hw, image_bundle) = helpers::build_hw_model_and_image_bundle(fuses);
-    hw.upload_firmware(&image_bundle.to_bytes().unwrap())
-        .unwrap();
+
+    assert_eq!(
+        ModelError::MailboxCmdFailed,
+        hw.upload_firmware(&image_bundle.to_bytes().unwrap())
+            .unwrap_err()
+    );
     let mut output = vec![];
     let result = hw.copy_output_until_non_fatal_error(VENDOR_PUB_KEY_DIGEST_INVALID, &mut output);
     assert!(result.is_ok());
@@ -107,8 +119,11 @@ fn test_preamble_vendor_pubkey_digest_mismatch() {
     };
 
     let (mut hw, image_bundle) = helpers::build_hw_model_and_image_bundle(fuses);
-    hw.upload_firmware(&image_bundle.to_bytes().unwrap())
-        .unwrap();
+    assert_eq!(
+        ModelError::MailboxCmdFailed,
+        hw.upload_firmware(&image_bundle.to_bytes().unwrap())
+            .unwrap_err()
+    );
     let mut output = vec![];
     let result = hw.copy_output_until_non_fatal_error(VENDOR_PUB_KEY_DIGEST_MISMATCH, &mut output);
     assert!(result.is_ok());
@@ -122,8 +137,12 @@ fn test_preamble_owner_pubkey_digest_mismatch() {
     };
 
     let (mut hw, image_bundle) = helpers::build_hw_model_and_image_bundle(fuses);
-    hw.upload_firmware(&image_bundle.to_bytes().unwrap())
-        .unwrap();
+
+    assert_eq!(
+        ModelError::MailboxCmdFailed,
+        hw.upload_firmware(&image_bundle.to_bytes().unwrap())
+            .unwrap_err()
+    );
     let mut output = vec![];
     let result = hw.copy_output_until_non_fatal_error(OWNER_PUB_KEY_DIGEST_MISMATCH, &mut output);
     assert!(result.is_ok());
@@ -168,13 +187,17 @@ fn test_preamble_vendor_pubkey_revocation() {
             caliptra_builder::build_and_sign_image(&FMC_WITH_UART, &APP_WITH_UART, image_options)
                 .unwrap();
 
-        hw.upload_firmware(&image_bundle.to_bytes().unwrap())
-            .unwrap();
-
         let result = if key_idx == LAST_KEY_IDX {
             // Last key is never revoked.
+            hw.upload_firmware(&image_bundle.to_bytes().unwrap())
+                .unwrap();
             hw.copy_output_until_exit_success(&mut output)
         } else {
+            assert_eq!(
+                ModelError::MailboxCmdFailed,
+                hw.upload_firmware(&image_bundle.to_bytes().unwrap())
+                    .unwrap_err()
+            );
             hw.copy_output_until_non_fatal_error(VENDOR_ECC_PUB_KEY_REVOKED, &mut output)
         };
 
@@ -187,8 +210,12 @@ fn test_preamble_vendor_pubkey_out_of_bounds() {
     let mut output = vec![];
     let (mut hw, mut image_bundle) = helpers::build_hw_model_and_image_bundle(Fuses::default());
     image_bundle.manifest.preamble.vendor_ecc_pub_key_idx = VENDOR_ECC_KEY_COUNT;
-    hw.upload_firmware(&image_bundle.to_bytes().unwrap())
-        .unwrap();
+
+    assert_eq!(
+        ModelError::MailboxCmdFailed,
+        hw.upload_firmware(&image_bundle.to_bytes().unwrap())
+            .unwrap_err()
+    );
     let result =
         hw.copy_output_until_non_fatal_error(VENDOR_ECC_PUB_KEY_INDEX_OUT_OF_BOUNDS, &mut output);
 
@@ -208,13 +235,18 @@ fn test_header_verify_vendor_sig_zero_pubkey() {
         .x
         .fill(0);
 
-    hw.upload_firmware(&image_bundle.to_bytes().unwrap())
-        .unwrap();
+    assert_eq!(
+        ModelError::MailboxCmdFailed,
+        hw.upload_firmware(&image_bundle.to_bytes().unwrap())
+            .unwrap_err()
+    );
 
     let result =
         hw.copy_output_until_non_fatal_error(VENDOR_PUB_KEY_DIGEST_INVALID_ARG, &mut output);
     assert!(result.is_ok());
 
+    let (mut hw, mut image_bundle) = helpers::build_hw_model_and_image_bundle(Fuses::default());
+
     // Set ecc_pub_key.y to zero.
     image_bundle.manifest.preamble.vendor_pub_keys.ecc_pub_keys[vendor_ecc_pub_key_idx].x =
         ecc_pub_key_x_backup;
@@ -222,8 +254,11 @@ fn test_header_verify_vendor_sig_zero_pubkey() {
         .y
         .fill(0);
 
-    hw.upload_firmware(&image_bundle.to_bytes().unwrap())
-        .unwrap();
+    assert_eq!(
+        ModelError::MailboxCmdFailed,
+        hw.upload_firmware(&image_bundle.to_bytes().unwrap())
+            .unwrap_err()
+    );
     let result =
         hw.copy_output_until_non_fatal_error(VENDOR_PUB_KEY_DIGEST_INVALID_ARG, &mut output);
     assert!(result.is_ok());
@@ -238,18 +273,26 @@ fn test_header_verify_vendor_sig_zero_signature() {
     let vendor_sig_r_backup = image_bundle.manifest.preamble.vendor_sigs.ecc_sig.r;
     image_bundle.manifest.preamble.vendor_sigs.ecc_sig.r.fill(0);
 
-    hw.upload_firmware(&image_bundle.to_bytes().unwrap())
-        .unwrap();
+    assert_eq!(
+        ModelError::MailboxCmdFailed,
+        hw.upload_firmware(&image_bundle.to_bytes().unwrap())
+            .unwrap_err()
+    );
     let result =
         hw.copy_output_until_non_fatal_error(VENDOR_ECC_SIGNATURE_INVALID_ARG, &mut output);
     assert!(result.is_ok());
 
+    let (mut hw, mut image_bundle) = helpers::build_hw_model_and_image_bundle(Fuses::default());
+
     // Set vendor_sig.s to zero.
     image_bundle.manifest.preamble.vendor_sigs.ecc_sig.r = vendor_sig_r_backup;
     image_bundle.manifest.preamble.vendor_sigs.ecc_sig.s.fill(0);
 
-    hw.upload_firmware(&image_bundle.to_bytes().unwrap())
-        .unwrap();
+    assert_eq!(
+        ModelError::MailboxCmdFailed,
+        hw.upload_firmware(&image_bundle.to_bytes().unwrap())
+            .unwrap_err()
+    );
     let result =
         hw.copy_output_until_non_fatal_error(VENDOR_ECC_SIGNATURE_INVALID_ARG, &mut output);
     assert!(result.is_ok());
@@ -272,11 +315,16 @@ fn test_header_verify_vendor_sig_mismatch() {
         .y
         .clone_from_slice(Array4x12::from(PUB_KEY_Y).0.as_slice());
 
-    hw.upload_firmware(&image_bundle.to_bytes().unwrap())
-        .unwrap();
+    assert_eq!(
+        ModelError::MailboxCmdFailed,
+        hw.upload_firmware(&image_bundle.to_bytes().unwrap())
+            .unwrap_err()
+    );
     let result = hw.copy_output_until_non_fatal_error(VENDOR_ECC_SIGNATURE_INVALID, &mut output);
     assert!(result.is_ok());
 
+    let (mut hw, mut image_bundle) = helpers::build_hw_model_and_image_bundle(Fuses::default());
+
     // Modify the vendor signature.
     image_bundle.manifest.preamble.vendor_pub_keys.ecc_pub_keys[vendor_ecc_pub_key_idx] =
         ecc_pub_key_backup;
@@ -295,8 +343,11 @@ fn test_header_verify_vendor_sig_mismatch() {
         .s
         .clone_from_slice(Array4x12::from(SIGNATURE_S).0.as_slice());
 
-    hw.upload_firmware(&image_bundle.to_bytes().unwrap())
-        .unwrap();
+    assert_eq!(
+        ModelError::MailboxCmdFailed,
+        hw.upload_firmware(&image_bundle.to_bytes().unwrap())
+            .unwrap_err()
+    );
     let result = hw.copy_output_until_non_fatal_error(VENDOR_ECC_SIGNATURE_INVALID, &mut output);
     assert!(result.is_ok());
 }
@@ -311,8 +362,11 @@ fn test_header_verify_vendor_pub_key_in_preamble_and_header() {
         image_bundle.manifest.preamble.vendor_ecc_pub_key_idx + 1;
     update_header(&mut image_bundle);
 
-    hw.upload_firmware(&image_bundle.to_bytes().unwrap())
-        .unwrap();
+    assert_eq!(
+        ModelError::MailboxCmdFailed,
+        hw.upload_firmware(&image_bundle.to_bytes().unwrap())
+            .unwrap_err()
+    );
     let result =
         hw.copy_output_until_non_fatal_error(VENDOR_ECC_PUB_KEY_INDEX_MISMATCH, &mut output);
     assert!(result.is_ok());
@@ -359,8 +413,11 @@ fn test_header_verify_owner_sig_zero_pubkey_x() {
     })
     .unwrap();
 
-    hw.upload_firmware(&image_bundle.to_bytes().unwrap())
-        .unwrap();
+    assert_eq!(
+        ModelError::MailboxCmdFailed,
+        hw.upload_firmware(&image_bundle.to_bytes().unwrap())
+            .unwrap_err()
+    );
 
     let result =
         hw.copy_output_until_non_fatal_error(OWNER_PUB_KEY_DIGEST_INVALID_ARG, &mut output);
@@ -408,8 +465,11 @@ fn test_header_verify_owner_sig_zero_pubkey_y() {
     })
     .unwrap();
 
-    hw.upload_firmware(&image_bundle.to_bytes().unwrap())
-        .unwrap();
+    assert_eq!(
+        ModelError::MailboxCmdFailed,
+        hw.upload_firmware(&image_bundle.to_bytes().unwrap())
+            .unwrap_err()
+    );
 
     let result =
         hw.copy_output_until_non_fatal_error(OWNER_PUB_KEY_DIGEST_INVALID_ARG, &mut output);
@@ -451,8 +511,11 @@ fn test_header_verify_owner_sig_zero_signature_r() {
     // Set owner_sig.r to zero.
     image_bundle.manifest.preamble.owner_sigs.ecc_sig.r.fill(0);
 
-    hw.upload_firmware(&image_bundle.to_bytes().unwrap())
-        .unwrap();
+    assert_eq!(
+        ModelError::MailboxCmdFailed,
+        hw.upload_firmware(&image_bundle.to_bytes().unwrap())
+            .unwrap_err()
+    );
 
     let result = hw.copy_output_until_non_fatal_error(OWNER_ECC_SIGNATURE_INVALID_ARG, &mut output);
     assert!(result.is_ok());
@@ -493,8 +556,11 @@ fn test_header_verify_owner_sig_zero_signature_s() {
     // Set owner_sig.s to zero.
     image_bundle.manifest.preamble.owner_sigs.ecc_sig.s.fill(0);
 
-    hw.upload_firmware(&image_bundle.to_bytes().unwrap())
-        .unwrap();
+    assert_eq!(
+        ModelError::MailboxCmdFailed,
+        hw.upload_firmware(&image_bundle.to_bytes().unwrap())
+            .unwrap_err()
+    );
 
     let result = hw.copy_output_until_non_fatal_error(OWNER_ECC_SIGNATURE_INVALID_ARG, &mut output);
     assert!(result.is_ok());
@@ -509,8 +575,11 @@ fn test_toc_invalid_entry_count() {
     image_bundle.manifest.header.toc_len = caliptra_image_types::MAX_TOC_ENTRY_COUNT + 1;
     update_header(&mut image_bundle);
 
-    hw.upload_firmware(&image_bundle.to_bytes().unwrap())
-        .unwrap();
+    assert_eq!(
+        ModelError::MailboxCmdFailed,
+        hw.upload_firmware(&image_bundle.to_bytes().unwrap())
+            .unwrap_err()
+    );
     let result = hw.copy_output_until_non_fatal_error(TOC_ENTRY_COUNT_INVALID, &mut output);
     assert!(result.is_ok());
 }
@@ -524,8 +593,11 @@ fn test_toc_invalid_toc_digest() {
     image_bundle.manifest.header.toc_digest[0] = 0xDEADBEEF;
     update_header(&mut image_bundle);
 
-    hw.upload_firmware(&image_bundle.to_bytes().unwrap())
-        .unwrap();
+    assert_eq!(
+        ModelError::MailboxCmdFailed,
+        hw.upload_firmware(&image_bundle.to_bytes().unwrap())
+            .unwrap_err()
+    );
     let result = hw.copy_output_until_non_fatal_error(TOC_DIGEST_MISMATCH, &mut output);
     assert!(result.is_ok());
 }
@@ -549,7 +621,10 @@ fn test_toc_fmc_range_overlap() {
         runtime_new_offset,
         runtime_new_size,
     );
-    hw.upload_firmware(&image).unwrap();
+    assert_eq!(
+        ModelError::MailboxCmdFailed,
+        hw.upload_firmware(&image).unwrap_err()
+    );
     let result = hw.copy_output_until_non_fatal_error(FMC_RUNTIME_OVERLAP, &mut output);
     assert!(result.is_ok());
 
@@ -567,7 +642,11 @@ fn test_toc_fmc_range_overlap() {
         runtime_new_offset,
         runtime_new_size,
     );
-    hw.upload_firmware(&image).unwrap();
+
+    assert_eq!(
+        ModelError::MailboxCmdFailed,
+        hw.upload_firmware(&image).unwrap_err()
+    );
     let result = hw.copy_output_until_non_fatal_error(FMC_RUNTIME_OVERLAP, &mut output);
     assert!(result.is_ok());
 
@@ -585,7 +664,11 @@ fn test_toc_fmc_range_overlap() {
         runtime_new_offset,
         runtime_new_size,
     );
-    hw.upload_firmware(&image).unwrap();
+
+    assert_eq!(
+        ModelError::MailboxCmdFailed,
+        hw.upload_firmware(&image).unwrap_err()
+    );
     let result = hw.copy_output_until_non_fatal_error(FMC_RUNTIME_OVERLAP, &mut output);
     assert!(result.is_ok());
 }
@@ -606,7 +689,10 @@ fn test_toc_fmc_range_incorrect_order() {
         runtime_new_offset,
         runtime_new_size,
     );
-    hw.upload_firmware(&image).unwrap();
+    assert_eq!(
+        ModelError::MailboxCmdFailed,
+        hw.upload_firmware(&image).unwrap_err()
+    );
     let result = hw.copy_output_until_non_fatal_error(FMC_RUNTIME_INCORRECT_ORDER, &mut output);
     assert!(result.is_ok());
 }
@@ -619,8 +705,11 @@ fn test_fmc_digest_mismatch() {
     // Change the FMC image.
     image_bundle.fmc[0..4].copy_from_slice(0xDEADBEEFu32.as_bytes());
 
-    hw.upload_firmware(&image_bundle.to_bytes().unwrap())
-        .unwrap();
+    assert_eq!(
+        ModelError::MailboxCmdFailed,
+        hw.upload_firmware(&image_bundle.to_bytes().unwrap())
+            .unwrap_err()
+    );
     let result = hw.copy_output_until_non_fatal_error(FMC_DIGEST_MISMATCH, &mut output);
     assert!(result.is_ok());
 }
@@ -633,8 +722,11 @@ fn test_runtime_digest_mismatch() {
     // Change the FMC image.
     image_bundle.runtime[0..4].copy_from_slice(0xDEADBEEFu32.as_bytes());
 
-    hw.upload_firmware(&image_bundle.to_bytes().unwrap())
-        .unwrap();
+    assert_eq!(
+        ModelError::MailboxCmdFailed,
+        hw.upload_firmware(&image_bundle.to_bytes().unwrap())
+            .unwrap_err()
+    );
     let result = hw.copy_output_until_non_fatal_error(RUNTIME_DIGEST_MISMATCH, &mut output);
     assert!(result.is_ok());
 }