Merge #971

971: (embassy-boot): add blocking API to FirmwareUpdater r=lulf a=MathiasKoch

Also add a split `prepare_update` + `write_firmware` API, to allow for an optimized update API at the exchange of added complexity. 
The old API is left in place to allow users to choose the complexity level that fits their needs. 

Co-authored-by: Mathias <[email protected]>

1 files changed, 209 insertions, 7 deletions

diff --git a/embassy-boot/boot/src/lib.rs b/embassy-boot/boot/src/lib.rs
index 96878ace9..8286601ec 100644
--- a/embassy-boot/boot/src/lib.rs
+++ b/embassy-boot/boot/src/lib.rs
@@ -604,6 +604,21 @@ impl FirmwareUpdater {
         self.dfu.len()
     }
 
+    /// Obtain the current state.
+    ///
+    /// This is useful to check if the bootloader has just done a swap, in order
+    /// to do verifications and self-tests of the new image before calling
+    /// `mark_booted`.
+    pub async fn get_state<F: AsyncNorFlash>(&mut self, flash: &mut F, aligned: &mut [u8]) -> Result<State, F::Error> {
+        flash.read(self.state.from as u32, aligned).await?;
+
+        if !aligned.iter().any(|&b| b != SWAP_MAGIC) {
+            Ok(State::Swap)
+        } else {
+            Ok(State::Boot)
+        }
+    }
+
     /// Mark to trigger firmware swap on next boot.
     ///
     /// # Safety
@@ -660,12 +675,6 @@ impl FirmwareUpdater {
     ) -> Result<(), F::Error> {
         assert!(data.len() >= F::ERASE_SIZE);
 
-        trace!(
-            "Writing firmware at offset 0x{:x} len {}",
-            self.dfu.from + offset,
-            data.len()
-        );
-
         flash
             .erase(
                 (self.dfu.from + offset) as u32,
@@ -679,7 +688,156 @@ impl FirmwareUpdater {
             self.dfu.from + offset + data.len()
         );
 
-        let mut write_offset = self.dfu.from + offset;
+        FirmwareWriter(self.dfu)
+            .write_block(offset, data, flash, block_size)
+            .await?;
+
+        Ok(())
+    }
+
+    /// Prepare for an incoming DFU update by erasing the entire DFU area and
+    /// returning a `FirmwareWriter`.
+    ///
+    /// Using this instead of `write_firmware` allows for an optimized API in
+    /// exchange for added complexity.
+    pub async fn prepare_update<F: AsyncNorFlash>(&mut self, flash: &mut F) -> Result<FirmwareWriter, F::Error> {
+        flash.erase((self.dfu.from) as u32, (self.dfu.to) as u32).await?;
+
+        trace!("Erased from {} to {}", self.dfu.from, self.dfu.to);
+
+        Ok(FirmwareWriter(self.dfu))
+    }
+
+    //
+    // Blocking API
+    //
+
+    /// Obtain the current state.
+    ///
+    /// This is useful to check if the bootloader has just done a swap, in order
+    /// to do verifications and self-tests of the new image before calling
+    /// `mark_booted`.
+    pub fn get_state_blocking<F: NorFlash>(&mut self, flash: &mut F, aligned: &mut [u8]) -> Result<State, F::Error> {
+        flash.read(self.state.from as u32, aligned)?;
+
+        if !aligned.iter().any(|&b| b != SWAP_MAGIC) {
+            Ok(State::Swap)
+        } else {
+            Ok(State::Boot)
+        }
+    }
+
+    /// Mark to trigger firmware swap on next boot.
+    ///
+    /// # Safety
+    ///
+    /// The `aligned` buffer must have a size of F::WRITE_SIZE, and follow the alignment rules for the flash being written to.
+    pub fn mark_updated_blocking<F: NorFlash>(&mut self, flash: &mut F, aligned: &mut [u8]) -> Result<(), F::Error> {
+        assert_eq!(aligned.len(), F::WRITE_SIZE);
+        self.set_magic_blocking(aligned, SWAP_MAGIC, flash)
+    }
+
+    /// Mark firmware boot successful and stop rollback on reset.
+    ///
+    /// # Safety
+    ///
+    /// The `aligned` buffer must have a size of F::WRITE_SIZE, and follow the alignment rules for the flash being written to.
+    pub fn mark_booted_blocking<F: NorFlash>(&mut self, flash: &mut F, aligned: &mut [u8]) -> Result<(), F::Error> {
+        assert_eq!(aligned.len(), F::WRITE_SIZE);
+        self.set_magic_blocking(aligned, BOOT_MAGIC, flash)
+    }
+
+    fn set_magic_blocking<F: NorFlash>(
+        &mut self,
+        aligned: &mut [u8],
+        magic: u8,
+        flash: &mut F,
+    ) -> Result<(), F::Error> {
+        flash.read(self.state.from as u32, aligned)?;
+
+        if aligned.iter().any(|&b| b != magic) {
+            aligned.fill(0);
+
+            flash.write(self.state.from as u32, aligned)?;
+            flash.erase(self.state.from as u32, self.state.to as u32)?;
+
+            aligned.fill(magic);
+            flash.write(self.state.from as u32, aligned)?;
+        }
+        Ok(())
+    }
+
+    /// Write data to a flash page.
+    ///
+    /// The buffer must follow alignment requirements of the target flash and a multiple of page size big.
+    ///
+    /// # Safety
+    ///
+    /// Failing to meet alignment and size requirements may result in a panic.
+    pub fn write_firmware_blocking<F: NorFlash>(
+        &mut self,
+        offset: usize,
+        data: &[u8],
+        flash: &mut F,
+        block_size: usize,
+    ) -> Result<(), F::Error> {
+        assert!(data.len() >= F::ERASE_SIZE);
+
+        flash.erase(
+            (self.dfu.from + offset) as u32,
+            (self.dfu.from + offset + data.len()) as u32,
+        )?;
+
+        trace!(
+            "Erased from {} to {}",
+            self.dfu.from + offset,
+            self.dfu.from + offset + data.len()
+        );
+
+        FirmwareWriter(self.dfu).write_block_blocking(offset, data, flash, block_size)?;
+
+        Ok(())
+    }
+
+    /// Prepare for an incoming DFU update by erasing the entire DFU area and
+    /// returning a `FirmwareWriter`.
+    ///
+    /// Using this instead of `write_firmware_blocking` allows for an optimized
+    /// API in exchange for added complexity.
+    pub fn prepare_update_blocking<F: NorFlash>(&mut self, flash: &mut F) -> Result<FirmwareWriter, F::Error> {
+        flash.erase((self.dfu.from) as u32, (self.dfu.to) as u32)?;
+
+        trace!("Erased from {} to {}", self.dfu.from, self.dfu.to);
+
+        Ok(FirmwareWriter(self.dfu))
+    }
+}
+
+/// FirmwareWriter allows writing blocks to an already erased flash.
+pub struct FirmwareWriter(Partition);
+
+impl FirmwareWriter {
+    /// Write data to a flash page.
+    ///
+    /// The buffer must follow alignment requirements of the target flash and a multiple of page size big.
+    ///
+    /// # Safety
+    ///
+    /// Failing to meet alignment and size requirements may result in a panic.
+    pub async fn write_block<F: AsyncNorFlash>(
+        &mut self,
+        offset: usize,
+        data: &[u8],
+        flash: &mut F,
+        block_size: usize,
+    ) -> Result<(), F::Error> {
+        trace!(
+            "Writing firmware at offset 0x{:x} len {}",
+            self.0.from + offset,
+            data.len()
+        );
+
+        let mut write_offset = self.0.from + offset;
         for chunk in data.chunks(block_size) {
             trace!("Wrote chunk at {}: {:?}", write_offset, chunk);
             flash.write(write_offset as u32, chunk).await?;
@@ -702,6 +860,50 @@ impl FirmwareUpdater {
 
         Ok(())
     }
+
+    /// Write data to a flash page.
+    ///
+    /// The buffer must follow alignment requirements of the target flash and a multiple of page size big.
+    ///
+    /// # Safety
+    ///
+    /// Failing to meet alignment and size requirements may result in a panic.
+    pub fn write_block_blocking<F: NorFlash>(
+        &mut self,
+        offset: usize,
+        data: &[u8],
+        flash: &mut F,
+        block_size: usize,
+    ) -> Result<(), F::Error> {
+        trace!(
+            "Writing firmware at offset 0x{:x} len {}",
+            self.0.from + offset,
+            data.len()
+        );
+
+        let mut write_offset = self.0.from + offset;
+        for chunk in data.chunks(block_size) {
+            trace!("Wrote chunk at {}: {:?}", write_offset, chunk);
+            flash.write(write_offset as u32, chunk)?;
+            write_offset += chunk.len();
+        }
+        /*
+        trace!("Wrote data, reading back for verification");
+
+        let mut buf: [u8; 4096] = [0; 4096];
+        let mut data_offset = 0;
+        let mut read_offset = self.dfu.from + offset;
+        for chunk in buf.chunks_mut(block_size) {
+            flash.read(read_offset as u32, chunk).await?;
+            trace!("Read chunk at {}: {:?}", read_offset, chunk);
+            assert_eq!(&data[data_offset..data_offset + block_size], chunk);
+            read_offset += chunk.len();
+            data_offset += chunk.len();
+        }
+        */
+
+        Ok(())
+    }
 }
 
 #[cfg(test)]