From 85c4b01d9fbd0e59d99c22d1d454c2462beaecca Mon Sep 17 00:00:00 2001
From: Sree Revoori <srrevoori@google.com>
Date: Thu, 4 Apr 2024 21:04:11 +0000
Subject: [PATCH] Prevent mailbox buffer overflow vulnerability

---
 runtime/src/mailbox.rs                          | 9 +++++++--
 runtime/src/main.rs                             | 2 +-
 runtime/test-fw/src/mock_rt_test_interactive.rs | 5 +++--
 3 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/runtime/src/mailbox.rs b/runtime/src/mailbox.rs
index 09c79163c9..2a4620b67a 100644
--- a/runtime/src/mailbox.rs
+++ b/runtime/src/mailbox.rs
@@ -53,9 +53,14 @@ impl Mailbox {
     }
 
     /// Set the length of the current mailbox data in bytes
-    pub fn set_dlen(&mut self, len: u32) {
+    pub fn set_dlen(&mut self, len: u32) -> CaliptraResult<()> {
+        if len > memory_layout::MBOX_SIZE {
+            return Err(CaliptraError::RUNTIME_MAILBOX_INVALID_PARAMS);
+        }
+
         let mbox = self.mbox.regs_mut();
         mbox.dlen().write(|_| len);
+        Ok(())
     }
 
     /// Get the length of the current mailbox data in words
@@ -141,7 +146,7 @@ impl Mailbox {
 
     /// Write a word-aligned `buf` to the mailbox
     pub fn write_response(&mut self, buf: &[u8]) -> CaliptraResult<()> {
-        self.set_dlen(buf.len() as u32);
+        self.set_dlen(buf.len() as u32)?;
         self.copy_bytes_to_mbox(buf);
         Ok(())
     }
diff --git a/runtime/src/main.rs b/runtime/src/main.rs
index 7771e4793b..6c844d30e6 100644
--- a/runtime/src/main.rs
+++ b/runtime/src/main.rs
@@ -91,7 +91,7 @@ pub extern "C" fn entry_point() -> ! {
     if let Err(e) = caliptra_runtime::handle_mailbox_commands(&mut drivers) {
         handle_fatal_error(e.into());
     }
-    loop {}
+    caliptra_drivers::ExitCtrl::exit(0xff);
 }
 
 #[no_mangle]
diff --git a/runtime/test-fw/src/mock_rt_test_interactive.rs b/runtime/test-fw/src/mock_rt_test_interactive.rs
index 5d8f525a60..ead4f66871 100644
--- a/runtime/test-fw/src/mock_rt_test_interactive.rs
+++ b/runtime/test-fw/src/mock_rt_test_interactive.rs
@@ -116,7 +116,8 @@ fn read_pcr_log(persistent_data: &PersistentDataAccessor, mbox: &mut Mailbox) {
         (core::mem::size_of::<PcrLogEntry>() * pcr_entry_count)
             .try_into()
             .unwrap(),
-    );
+    )
+    .unwrap();
     mbox.set_status(MboxStatusE::DataReady);
 }
 
@@ -136,7 +137,7 @@ fn read_pcrs(mbox: &mut Mailbox) {
         swap_word_bytes_inplace(&mut pcr_bytes);
         mbox.copy_bytes_to_mbox(pcr.as_bytes()).unwrap();
     }
-    mbox.set_dlen((48 * PCR_COUNT).try_into().unwrap());
+    mbox.set_dlen((48 * PCR_COUNT).try_into().unwrap()).unwrap();
     mbox.set_status(MboxStatusE::DataReady);
 }