From 65a65dab36790aa8faeedaba5f93683c23b7bdca Mon Sep 17 00:00:00 2001
From: Matt Keeter <matt@oxide.computer>
Date: Fri, 7 Feb 2025 10:05:13 -0500
Subject: [PATCH] Add APOB message to host_sp_comms

---
 lib/host-sp-messages/src/lib.rs |  5 ++
 task/host-sp-comms/src/main.rs  | 84 +++++++++++++++++++++++++++++++++
 2 files changed, 89 insertions(+)

diff --git a/lib/host-sp-messages/src/lib.rs b/lib/host-sp-messages/src/lib.rs
index 2f24f0f95..74738d008 100644
--- a/lib/host-sp-messages/src/lib.rs
+++ b/lib/host-sp-messages/src/lib.rs
@@ -123,6 +123,10 @@ pub enum HostToSp {
         // We use a raw `u8` here for the same reason as in `KeyLookup` above.
         key: u8,
     },
+    // APOB is followed by a binary data blob, which should be written to flash
+    APOB {
+        offset: u64,
+    },
 }
 
 /// The order of these cases is critical! We are relying on hubpack's encoding
@@ -185,6 +189,7 @@ pub enum SpToHost {
         name: [u8; 32],
     },
     KeySetResult(#[count(children)] KeySetResult),
+    APOBResult(u8),
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, num_derive::FromPrimitive)]
diff --git a/task/host-sp-comms/src/main.rs b/task/host-sp-comms/src/main.rs
index 88923d98f..9e29acd3c 100644
--- a/task/host-sp-comms/src/main.rs
+++ b/task/host-sp-comms/src/main.rs
@@ -137,6 +137,11 @@ enum Trace {
         #[count(children)]
         message: SpToHost,
     },
+    APOBWriteError {
+        offset: u64,
+        #[count(children)]
+        err: APOBError,
+    },
 }
 
 counted_ringbuf!(Trace, 20, Trace::None);
@@ -163,6 +168,31 @@ enum Timers {
     TxPeriodicZeroByte,
 }
 
+#[derive(Copy, Clone, Debug, Eq, PartialEq, counters::Count)]
+enum APOBError {
+    OffsetOverflow {
+        offset: u64,
+    },
+    NotErased {
+        offset: u32,
+    },
+    EraseFailed {
+        offset: u32,
+        #[count(children)]
+        err: drv_hf_api::HfError,
+    },
+    WriteFailed {
+        offset: u32,
+        #[count(children)]
+        err: drv_hf_api::HfError,
+    },
+    ReadFailed {
+        offset: u32,
+        #[count(children)]
+        err: drv_hf_api::HfError,
+    },
+}
+
 #[export_name = "main"]
 fn main() -> ! {
     let mut server = ServerImpl::claim_static_resources();
@@ -965,6 +995,15 @@ impl ServerImpl {
                     }),
                 }
             }
+            HostToSp::APOB { offset } => {
+                Some(match Self::apob_write(&self.hf, offset, data) {
+                    Ok(()) => SpToHost::APOBResult(0),
+                    Err(err) => {
+                        ringbuf_entry!(Trace::APOBWriteError { offset, err });
+                        SpToHost::APOBResult(1)
+                    }
+                })
+            }
         };
 
         if let Some(response) = response {
@@ -993,6 +1032,51 @@ impl ServerImpl {
         Ok(())
     }
 
+    /// Write data to the bonus region of flash
+    ///
+    /// This does not take `&self` because we need to force a split borrow
+    fn apob_write(
+        hf: &HostFlash,
+        mut offset: u64,
+        data: &[u8],
+    ) -> Result<(), APOBError> {
+        for chunk in data.chunks(drv_hf_api::PAGE_SIZE_BYTES) {
+            Self::apob_write_page(
+                hf,
+                offset
+                    .try_into()
+                    .map_err(|_| APOBError::OffsetOverflow { offset })?,
+                chunk,
+            )?;
+            offset += chunk.len() as u64;
+        }
+        Ok(())
+    }
+
+    /// Write a single page of data to the bonus region of flash
+    ///
+    /// This does not take `&self` because we need to force a split borrow
+    fn apob_write_page(
+        hf: &HostFlash,
+        offset: u32,
+        data: &[u8],
+    ) -> Result<(), APOBError> {
+        if offset as usize % drv_hf_api::SECTOR_SIZE_BYTES == 0 {
+            hf.bonus_sector_erase(offset)
+                .map_err(|err| APOBError::EraseFailed { offset, err })?;
+        } else {
+            // Read back the page and confirm that it's all empty
+            let mut scratch = [0u8; drv_hf_api::PAGE_SIZE_BYTES];
+            hf.bonus_read(offset, &mut scratch[..data.len()])
+                .map_err(|err| APOBError::ReadFailed { offset, err })?;
+            if !scratch[..data.len()].iter().all(|b| *b == 0xFF) {
+                return Err(APOBError::NotErased { offset });
+            }
+        }
+        hf.bonus_page_program(offset, data)
+            .map_err(|err| APOBError::WriteFailed { offset, err })
+    }
+
     fn handle_sprot(
         &mut self,
         sequence: u64,