Add support for Bonus Data in Cosmo host flash (#1991)

In Cosmo and Grapefruit, we divide a 128 MiB flash chip into two (now
three) sections:

- 32 MiB of slot 0 (host boot data)
- 32 MiB of slot 1 (host boot data)
- 64 MiB of Bonus Data (previously unused)

The host boot slots use the lowest 64 KiB for "persistent data", which
right now encodes which slot to select at boot time. This behavior is
identical to Gimlet, and is unchanged in the PR.

This PR consists of three standalone commits, which do the following:

- Check flash addresses in the `cosmo-hf` driver, to avoid writing into
  other sections by accident
- Use a strong type for an absolute address (versus `u32` for a
  section-relative address)
- Add new Hiffy APIs to read and write the upper 64 MiB of flash
- These APIs are added to the `HostFlash` interface, so they're also
  added to Gimlet's host flash driver, which just returns an error if you
  try to use them.

Under the hood, the W25Q01JV uses two 64 MiB dies and exposes manual
APIs to select which die is active (e.g. checking BUSY status only reads
from the active die). However, all of the QSPI commands that we use take
a 4-byte address and automatically switch between dies, so everything
Just Works™ (feel free to double-check this,
[here's the docs](https://drive.google.com/file/d/1qcurvqkNYha0cWxEeEhD_J2qGgi_r4AL/))

drv/cosmo-hf/src/hf.rs

@@ -15,15 +15,17 @@ use ringbuf::ringbuf_entry_root as ringbuf_entry;
 use userlib::{task_slot, RecvMessage, UnwrapLite};
 use zerocopy::{AsBytes, FromBytes};
 
-use crate::{FlashDriver, Trace, PAGE_SIZE_BYTES, SECTOR_SIZE_BYTES};
+use crate::{
+    FlashAddr, FlashDriver, Trace, PAGE_SIZE_BYTES, SECTOR_SIZE_BYTES,
+};
 
 task_slot!(HASH, hash_driver);
 
 /// We break the 128 MiB flash chip into 2x 32 MiB slots, to match Gimlet
 ///
-/// The upper 64 MiB are unused (which is good, because it's a separate die and
-/// requires special handling).
+/// The upper 64 MiB are used for Bonus Data.
 const SLOT_SIZE_BYTES: u32 = 1024 * 1024 * 32;
+const BONUS_SIZE_BYTES: u32 = 1024 * 1024 * 64;
 
 pub struct ServerImpl {
     pub drv: FlashDriver,
@@ -51,6 +53,7 @@ impl ServerImpl {
         out
     }
 
+    /// Checks whether the given (relative) address is writable
     fn check_addr_writable(
         &self,
         addr: u32,
@@ -65,18 +68,51 @@ impl ServerImpl {
         }
     }
 
-    fn flash_base(&self) -> u32 {
-        Self::flash_base_for(self.dev)
+    /// Returns the current device's absolute base address
+    fn flash_base(&self) -> FlashAddr {
+        // This is always valid, so we can unwrap it here
+        FlashAddr::new(Self::flash_base_for(self.dev)).unwrap_lite()
     }
 
-    fn flash_addr(&self, offset: u32) -> u32 {
-        Self::flash_addr_for(offset, self.dev)
+    /// Converts a relative address to an absolute address in out current device
+    fn flash_addr(&self, offset: u32, size: u32) -> Result<FlashAddr, HfError> {
+        if offset
+            .checked_add(size)
+            .is_some_and(|a| a < SLOT_SIZE_BYTES)
+        {
+            Self::flash_addr_for(offset, self.dev)
+        } else {
+            Err(HfError::BadAddress)
+        }
     }
 
-    fn flash_addr_for(offset: u32, dev: HfDevSelect) -> u32 {
-        offset + Self::flash_base_for(dev)
+    /// Converts a relative address to an absolute address in bonus space
+    fn bonus_addr(offset: u32, len: u32) -> Result<FlashAddr, HfError> {
+        if offset
+            .checked_add(len)
+            .is_some_and(|a| a < BONUS_SIZE_BYTES)
+        {
+            let addr = offset
+                .checked_add(2 * SLOT_SIZE_BYTES)
+                .ok_or(HfError::BadAddress)?;
+            Ok(FlashAddr(addr))
+        } else {
+            Err(HfError::BadAddress)
+        }
     }
 
+    /// Converts a relative address to an absolute address in a device slot
+    fn flash_addr_for(
+        offset: u32,
+        dev: HfDevSelect,
+    ) -> Result<FlashAddr, HfError> {
+        let addr = offset
+            .checked_add(Self::flash_base_for(dev))
+            .ok_or(HfError::BadAddress)?;
+        FlashAddr::new(addr).ok_or(HfError::BadAddress)
+    }
+
+    /// Return the absolute flash address base for the given virtual device
     fn flash_base_for(dev: HfDevSelect) -> u32 {
         match dev {
             HfDevSelect::Flash0 => 0,
@@ -100,7 +136,7 @@ impl ServerImpl {
             let mut data = HfRawPersistentData::new_zeroed();
             self.drv
                 .flash_read(
-                    Self::flash_addr_for(addr, dev),
+                    Self::flash_addr_for(addr, dev).unwrap_lite(),
                     &mut data.as_bytes_mut(),
                 )
                 .unwrap_lite(); // flash_read is infallible when using a slice
@@ -146,24 +182,27 @@ impl ServerImpl {
     ///
     /// If `addr` is `None`, then we're out of available space; erase all of
     /// sector 0 and write to address 0 upon success.
+    ///
+    /// # Panics
+    /// If `addr` points outside the slot
     fn write_raw_persistent_data_to_addr(
         &mut self,
         addr: Option<u32>,
         raw_data: &HfRawPersistentData,
         dev: HfDevSelect,
     ) {
         let addr = match addr {
-            Some(a) => Self::flash_addr_for(a, dev),
+            Some(a) => Self::flash_addr_for(a, dev).unwrap_lite(),
             None => {
-                let addr = Self::flash_addr_for(0, dev);
+                let addr = Self::flash_addr_for(0, dev).unwrap_lite();
                 self.drv.flash_sector_erase(addr);
                 addr
             }
         };
         // flash_write is infallible when given a slice
         self.drv
             .flash_write(addr, &mut raw_data.as_bytes())
-            .unwrap_lite()
+            .unwrap_lite();
     }
 
     /// Ensures that the persistent data is consistent between the virtual devs
@@ -240,7 +279,9 @@ impl idl::InOrderHostFlashImpl for ServerImpl {
         // chip.  Instead, use the sector erase to erase the currently-active
         // virtual device.
         for offset in (0..SLOT_SIZE_BYTES).step_by(SECTOR_SIZE_BYTES as usize) {
-            self.drv.flash_sector_erase(self.flash_addr(offset));
+            self.drv.flash_sector_erase(
+                self.flash_addr(offset, SECTOR_SIZE_BYTES)?,
+            );
         }
         Ok(())
     }
@@ -255,9 +296,10 @@ impl idl::InOrderHostFlashImpl for ServerImpl {
     ) -> Result<(), RequestError<HfError>> {
         self.check_addr_writable(addr, protect)?;
         self.drv.check_flash_mux_state()?;
+        let addr = self.flash_addr(addr, data.len() as u32)?;
         self.drv
             .flash_write(
-                self.flash_addr(addr),
+                addr,
                 &mut LeaseBufReader::<_, 32>::from(data.into_inner()),
             )
             .map_err(|()| RequestError::went_away())
@@ -273,7 +315,7 @@ impl idl::InOrderHostFlashImpl for ServerImpl {
         self.drv.check_flash_mux_state()?;
         self.drv
             .flash_read(
-                self.flash_addr(addr),
+                self.flash_addr(addr, dest.len() as u32)?,
                 &mut LeaseBufWriter::<_, 32>::from(dest.into_inner()),
             )
             .map_err(|_| RequestError::went_away())
@@ -293,7 +335,50 @@ impl idl::InOrderHostFlashImpl for ServerImpl {
     ) -> Result<(), RequestError<HfError>> {
         self.drv.check_flash_mux_state()?;
         self.check_addr_writable(addr, protect)?;
-        self.drv.flash_sector_erase(self.flash_addr(addr));
+        self.drv.flash_sector_erase(self.flash_addr(addr, 0)?);
+        Ok(())
+    }
+
+    /// Writes a page to the bonus region of flash
+    fn bonus_page_program(
+        &mut self,
+        _: &RecvMessage,
+        addr: u32,
+        data: LenLimit<Leased<R, [u8]>, PAGE_SIZE_BYTES>,
+    ) -> Result<(), RequestError<HfError>> {
+        // TODO check mux state?
+        self.drv
+            .flash_write(
+                Self::bonus_addr(addr, data.len() as u32)?,
+                &mut LeaseBufReader::<_, 32>::from(data.into_inner()),
+            )
+            .map_err(|()| RequestError::went_away())
+    }
+
+    /// Reads a page from the bonus region of flash
+    fn bonus_read(
+        &mut self,
+        _: &RecvMessage,
+        addr: u32,
+        dest: LenLimit<Leased<W, [u8]>, PAGE_SIZE_BYTES>,
+    ) -> Result<(), RequestError<HfError>> {
+        // TODO check mux state?
+        self.drv
+            .flash_read(
+                Self::bonus_addr(addr, dest.len() as u32)?,
+                &mut LeaseBufWriter::<_, 32>::from(dest.into_inner()),
+            )
+            .map_err(|_| RequestError::went_away())
+    }
+
+    /// Erases a 64 KiB sector in the bonus flash device
+    fn bonus_sector_erase(
+        &mut self,
+        _: &RecvMessage,
+        addr: u32,
+    ) -> Result<(), RequestError<HfError>> {
+        // TODO check mux state?
+        self.drv.flash_sector_erase(Self::bonus_addr(addr, 0)?);
         Ok(())
     }
 
@@ -344,12 +429,8 @@ impl idl::InOrderHostFlashImpl for ServerImpl {
             ringbuf_entry!(Trace::HashInitError(e));
             return Err(HfError::HashError.into());
         }
-        let begin = self.flash_addr(addr) as usize;
-        // TODO: Begin may be an address beyond physical end of
-        // flash part and may wrap around.
-        let end = begin
-            .checked_add(len as usize)
-            .ok_or(HfError::HashBadRange)?;
+        let begin = self.flash_addr(addr, len)?.get() as usize;
+        let end = begin + len as usize;
 
         let mut buf = [0u8; PAGE_SIZE_BYTES];
         for addr in (begin..end).step_by(buf.len()) {
@@ -358,7 +439,12 @@ impl idl::InOrderHostFlashImpl for ServerImpl {
             // a lease (where writing into the lease fails if the client goes
             // away).  Giving it a buffer is infallible.
             self.drv
-                .flash_read(addr as u32, &mut &mut buf[..size])
+                .flash_read(
+                    // This unwrap is safe because we already checked the range
+                    // when building the initial `begin` address
+                    FlashAddr::new(addr as u32).unwrap_lite(),
+                    &mut &mut buf[..size],
+                )
                 .unwrap_lite();
             if let Err(e) = hash_driver.update(size as u32, &buf[..size]) {
                 ringbuf_entry!(Trace::HashUpdateError(e));
@@ -497,6 +583,32 @@ impl idl::InOrderHostFlashImpl for FailServer {
         Err(self.err.into())
     }
 
+    fn bonus_page_program(
+        &mut self,
+        _: &RecvMessage,
+        _addr: u32,
+        _data: LenLimit<Leased<R, [u8]>, PAGE_SIZE_BYTES>,
+    ) -> Result<(), RequestError<HfError>> {
+        Err(self.err.into())
+    }
+
+    fn bonus_read(
+        &mut self,
+        _: &RecvMessage,
+        _offset: u32,
+        _dest: LenLimit<Leased<W, [u8]>, PAGE_SIZE_BYTES>,
+    ) -> Result<(), RequestError<HfError>> {
+        Err(self.err.into())
+    }
+
+    fn bonus_sector_erase(
+        &mut self,
+        _: &RecvMessage,
+        _addr: u32,
+    ) -> Result<(), RequestError<HfError>> {
+        Err(self.err.into())
+    }
+
     fn get_mux(
         &mut self,
         _: &RecvMessage,

drv/cosmo-hf/src/main.rs

@@ -4,7 +4,7 @@
 
 //! Minimal driver for FMC-attached NOR flash, implementing the `hf` API
 //!
-//! The NOR flash chip is a W25Q01JVZEIQ, which is a 1 GiB NOR flash.  It is
+//! The NOR flash chip is a W25Q01JVZEIQ, which is a 1 GBit NOR flash.  It is
 //! connected to the FPGA over SPI / QSPI.
 //!
 //! # References
@@ -58,6 +58,9 @@ pub const PAGE_SIZE_BYTES: usize = 256;
 /// behavior of the Gimlet host flash driver.
 pub const SECTOR_SIZE_BYTES: u32 = 65_536;
 
+/// Total flash size is 128 MiB
+pub const FLASH_SIZE_BYTES: u32 = 128 * 1024 * 1024;
+
 #[export_name = "main"]
 fn main() -> ! {
     // Wait for the FPGA to be configured; the sequencer task only starts its
@@ -86,6 +89,23 @@ fn main() -> ! {
     }
 }
 
+/// Absolute memory address
+#[derive(Copy, Clone)]
+struct FlashAddr(u32);
+
+impl FlashAddr {
+    fn new(v: u32) -> Option<Self> {
+        if v < FLASH_SIZE_BYTES {
+            Some(FlashAddr(v))
+        } else {
+            None
+        }
+    }
+    fn get(&self) -> u32 {
+        self.0
+    }
+}
+
 /// Driver for a QSPI NOR flash controlled by an FPGA over FMC
 struct FlashDriver;
 
@@ -227,10 +247,10 @@ impl FlashDriver {
     }
 
     /// Erases the 64KiB flash sector containing the given address
-    fn flash_sector_erase(&mut self, addr: u32) {
+    fn flash_sector_erase(&mut self, addr: FlashAddr) {
         self.flash_write_enable();
         self.write_reg(reg::DATA_BYTES, 0);
-        self.write_reg(reg::ADDR, addr);
+        self.write_reg(reg::ADDR, addr.0);
         self.write_reg(reg::DUMMY_CYCLES, 0);
         self.write_reg(reg::INSTR, instr::BLOCK_ERASE_64KB_4B);
         self.wait_fpga_busy();
@@ -245,7 +265,7 @@ impl FlashDriver {
     /// provided lease; when given a slice, it is infallible.
     fn flash_read(
         &mut self,
-        offset: u32,
+        offset: FlashAddr,
         dest: &mut dyn idol_runtime::BufWriter<'_>,
     ) -> Result<(), ()> {
         loop {
@@ -255,7 +275,7 @@ impl FlashDriver {
             }
             self.clear_fifos();
             self.write_reg(reg::DATA_BYTES, len as u32);
-            self.write_reg(reg::ADDR, offset);
+            self.write_reg(reg::ADDR, offset.0);
             self.write_reg(reg::DUMMY_CYCLES, 8);
             self.write_reg(reg::INSTR, instr::FAST_READ_QUAD_OUTPUT_4B);
             for i in 0..len.div_ceil(4) {
@@ -279,7 +299,7 @@ impl FlashDriver {
     /// provided lease; when given a slice, it is infallible.
     fn flash_write(
         &mut self,
-        addr: u32,
+        addr: FlashAddr,
         data: &mut dyn idol_runtime::BufReader<'_>,
     ) -> Result<(), ()> {
         loop {
@@ -290,7 +310,7 @@ impl FlashDriver {
 
             self.flash_write_enable();
             self.write_reg(reg::DATA_BYTES, len as u32);
-            self.write_reg(reg::ADDR, addr);
+            self.write_reg(reg::ADDR, addr.0);
             self.write_reg(reg::DUMMY_CYCLES, 0);
             for i in 0..len.div_ceil(4) {
                 let mut v = [0u8; 4];
@@ -368,8 +388,8 @@ impl FlashDriver {
         });
     }
 
-    fn set_espi_addr_offset(&self, v: u32) {
-        self.write_reg(reg::SP5_FLASH_OFFSET, v);
+    fn set_espi_addr_offset(&self, v: FlashAddr) {
+        self.write_reg(reg::SP5_FLASH_OFFSET, v.0);
     }
 }