From 8d64460b89cca22d6fe9f9ce359b8e2aa5dab8e2 Mon Sep 17 00:00:00 2001
From: Vishal Mhatre <vishal.mhatre@gmail.com>
Date: Tue, 8 Aug 2023 12:15:37 -0700
Subject: [PATCH] [fix] Fail ROM boot flow on receving an invalid Mailbox
 command

---
 error/src/lib.rs                            |  1 +
 hw-model/src/lib.rs                         | 27 +++++++++++++++++++++
 rom/dev/doc/test-coverage/test-coverage.md  |  2 +-
 rom/dev/src/flow/cold_reset/fw_processor.rs |  2 +-
 rom/dev/tests/test_mailbox_errors.rs        | 20 +++++++--------
 5 files changed, 40 insertions(+), 12 deletions(-)

diff --git a/error/src/lib.rs b/error/src/lib.rs
index 53dbaf39b1..3fb6a408db 100644
--- a/error/src/lib.rs
+++ b/error/src/lib.rs
@@ -327,6 +327,7 @@ impl CaliptraError {
     pub const FW_PROC_INVALID_IMAGE_SIZE: CaliptraError = CaliptraError::new_const(0x01020002);
     pub const FW_PROC_MAILBOX_STATE_INCONSISTENT: CaliptraError =
         CaliptraError::new_const(0x01020003);
+    pub const FW_PROC_MAILBOX_INVALID_COMMAND: CaliptraError = CaliptraError::new_const(0x01020004);
 
     /// FMC Alias Layer : Certificate Verification Failure.
     pub const FMC_ALIAS_CERT_VERIFY: CaliptraError = CaliptraError::new_const(0x01030001);
diff --git a/hw-model/src/lib.rs b/hw-model/src/lib.rs
index 05fe53ac7c..ec9d914c9e 100644
--- a/hw-model/src/lib.rs
+++ b/hw-model/src/lib.rs
@@ -554,6 +554,33 @@ pub trait HwModel {
         }
     }
 
+    fn step_until_fatal_error(&mut self, expected_error: u32, max_wait_cycles: u32) {
+        let mut cycle_count = 0u32;
+        let initial_error = self.soc_ifc().cptra_fw_error_fatal().read();
+        loop {
+            let actual_error = self.soc_ifc().cptra_fw_error_fatal().read();
+            if actual_error == expected_error {
+                break;
+            }
+
+            if actual_error != initial_error {
+                panic!(
+                    "Expected the fatal error to be  \
+                    ({expected_error}), but error changed from \
+                    {initial_error} to {actual_error})"
+                );
+            }
+            self.step();
+            cycle_count += 1;
+            if cycle_count >= max_wait_cycles {
+                panic!(
+                    "Expected fatal error to be  \
+                    ({expected_error}), but was stuck at ({initial_error})"
+                );
+            }
+        }
+    }
+
     /// A register block that can be used to manipulate the soc_ifc peripheral
     /// over the simulated SoC->Caliptra APB bus.
     fn soc_ifc(&mut self) -> caliptra_registers::soc_ifc::RegisterBlock<BusMmio<Self::TBus<'_>>> {
diff --git a/rom/dev/doc/test-coverage/test-coverage.md b/rom/dev/doc/test-coverage/test-coverage.md
index 59006f7e46..181b8a53c9 100644
--- a/rom/dev/doc/test-coverage/test-coverage.md
+++ b/rom/dev/doc/test-coverage/test-coverage.md
@@ -103,7 +103,7 @@ Test Name |  Description | ROM Error Code
 # **Mailbox Tests**
 Test Name | ROM Stage | Description | ROM Error Code
 ---|---|---|---
-**test_unknown_command_is_not_fatal** | FW Downloader | Checks for sending invalid commands to the mailbox | N/A
+**test_unknown_command_is_fatal** | FW Downloader | Checks for sending invalid commands to the mailbox | N/A
 **test_mailbox_command_aborted_after_handle_fatal_error** | FW Downloader | Checks for sending an invalid fw image followed by a valid one | FW_PROC_INVALID_IMAGE_SIZE
 <br><br>
 # **General Integration Tests**
diff --git a/rom/dev/src/flow/cold_reset/fw_processor.rs b/rom/dev/src/flow/cold_reset/fw_processor.rs
index 398e7ca59b..ec5cf7f0d4 100644
--- a/rom/dev/src/flow/cold_reset/fw_processor.rs
+++ b/rom/dev/src/flow/cold_reset/fw_processor.rs
@@ -132,7 +132,7 @@ impl FirmwareProcessor {
                 if txn.cmd() != Self::MBOX_DOWNLOAD_FIRMWARE_CMD_ID {
                     cprintln!("Invalid command 0x{:08x} received", txn.cmd());
                     txn.start_txn().complete(false)?;
-                    continue;
+                    return Err(CaliptraError::FW_PROC_MAILBOX_INVALID_COMMAND);
                 }
 
                 // Re-borrow mailbox to work around https://github.com/rust-lang/rust/issues/54663
diff --git a/rom/dev/tests/test_mailbox_errors.rs b/rom/dev/tests/test_mailbox_errors.rs
index 9f8d7608a4..ee651130ca 100644
--- a/rom/dev/tests/test_mailbox_errors.rs
+++ b/rom/dev/tests/test_mailbox_errors.rs
@@ -6,9 +6,13 @@ use caliptra_hw_model::{Fuses, HwModel, ModelError};
 
 pub mod helpers;
 
+// Since the boot takes less than 30M cycles, we know something is wrong if
+// we're stuck at the same state for that duration.
+const MAX_WAIT_CYCLES: u32 = 30_000_000;
+
 #[test]
-fn test_unknown_command_is_not_fatal() {
-    let (mut hw, image_bundle) =
+fn test_unknown_command_is_fatal() {
+    let (mut hw, _image_bundle) =
         helpers::build_hw_model_and_image_bundle(Fuses::default(), ImageOptions::default());
 
     // This command does not exist
@@ -17,14 +21,10 @@ fn test_unknown_command_is_not_fatal() {
         Err(ModelError::MailboxCmdFailed(0))
     );
 
-    // The ROM does not currently report an error for this
-    // TODO: Is this right?
-    assert_eq!(hw.soc_ifc().cptra_fw_error_non_fatal().read(), 0);
-
-    // Make sure we can still upload new firmware after the unknown
-    // command.
-    hw.upload_firmware(&image_bundle.to_bytes().unwrap())
-        .unwrap();
+    hw.step_until_fatal_error(
+        CaliptraError::FW_PROC_MAILBOX_INVALID_COMMAND.into(),
+        MAX_WAIT_CYCLES,
+    );
 }
 
 #[test]