Assert active and dfu have same erase size and copy in smaller chunks

The copy from active to dfu (and vice versa) is now done in smaller portions depending on aligned_buf, which now does not need to be erase_size big.

1 files changed, 31 insertions, 35 deletions

diff --git a/embassy-boot/boot/src/boot_loader.rs b/embassy-boot/boot/src/boot_loader.rs
index 698075599..db067da5b 100644
--- a/embassy-boot/boot/src/boot_loader.rs
+++ b/embassy-boot/boot/src/boot_loader.rs
@@ -32,14 +32,13 @@ where
 
 /// Extension of the embedded-storage flash type information with block size and erase value.
 pub trait Flash: NorFlash {
-    /// The block size that should be used when writing to flash. For most builtin flashes, this is the same as the erase
-    /// size of the flash, but for external QSPI flash modules, this can be lower.
-    const BLOCK_SIZE: usize;
     /// The erase value of the flash. Typically the default of 0xFF is used, but some flashes use a different value.
     const ERASE_VALUE: u8 = 0xFF;
 }
 
-/// Trait defining the flash handles used for active and DFU partition
+/// Trait defining the flash handles used for active and DFU partition.
+/// The ACTIVE and DFU erase sizes must be equal. If this is not the case, then consider adding an adapter for the
+/// smallest flash to increase its erase size such that they match. See e.g. [`crate::large_erase::LargeErase`].
 pub trait FlashConfig {
     /// Flash type used for the state partition.
     type STATE: Flash;
@@ -62,12 +61,12 @@ trait FlashConfigEx {
 
 impl<T: FlashConfig> FlashConfigEx for T {
     fn page_size() -> usize {
-        core::cmp::max(T::ACTIVE::ERASE_SIZE, T::DFU::ERASE_SIZE)
+        assert_eq!(T::ACTIVE::ERASE_SIZE, T::DFU::ERASE_SIZE);
+        T::ACTIVE::ERASE_SIZE
     }
 }
 
-/// BootLoader works with any flash implementing embedded_storage and can also work with
-/// different page sizes and flash write sizes.
+/// BootLoader works with any flash implementing embedded_storage.
 pub struct BootLoader {
     // Page with current state of bootloader. The state partition has the following format:
     // All ranges are in multiples of WRITE_SIZE bytes.
@@ -184,7 +183,9 @@ impl BootLoader {
     ///
     pub fn prepare_boot<P: FlashConfig>(&mut self, p: &mut P, aligned_buf: &mut [u8]) -> Result<State, BootError> {
         // Ensure we have enough progress pages to store copy progress
-        assert_eq!(aligned_buf.len(), P::page_size());
+        assert_eq!(0, P::page_size() % aligned_buf.len());
+        assert_eq!(0, P::page_size() % P::ACTIVE::WRITE_SIZE);
+        assert_eq!(0, P::page_size() % P::DFU::WRITE_SIZE);
         assert!(aligned_buf.len() >= P::STATE::WRITE_SIZE);
         assert_partitions(self.active, self.dfu, self.state, P::page_size(), P::STATE::WRITE_SIZE);
 
@@ -277,20 +278,18 @@ impl BootLoader {
         aligned_buf: &mut [u8],
     ) -> Result<(), BootError> {
         if self.current_progress(p, aligned_buf)? <= idx {
-            let mut offset = from_offset;
-            for chunk in aligned_buf.chunks_mut(P::DFU::BLOCK_SIZE) {
-                self.dfu.read_blocking(p.dfu(), offset, chunk)?;
-                offset += chunk.len() as u32;
-            }
+            let page_size = P::page_size() as u32;
 
             self.active
-                .erase_blocking(p.active(), to_offset, to_offset + P::page_size() as u32)?;
+                .erase_blocking(p.active(), to_offset, to_offset + page_size)?;
 
-            let mut offset = to_offset;
-            for chunk in aligned_buf.chunks(P::ACTIVE::BLOCK_SIZE) {
-                self.active.write_blocking(p.active(), offset, chunk)?;
-                offset += chunk.len() as u32;
+            for offset_in_page in (0..page_size).step_by(aligned_buf.len()) {
+                self.dfu
+                    .read_blocking(p.dfu(), from_offset + offset_in_page as u32, aligned_buf)?;
+                self.active
+                    .write_blocking(p.active(), to_offset + offset_in_page as u32, aligned_buf)?;
             }
+
             self.update_progress(idx, p, aligned_buf)?;
         }
         Ok(())
@@ -305,20 +304,18 @@ impl BootLoader {
         aligned_buf: &mut [u8],
     ) -> Result<(), BootError> {
         if self.current_progress(p, aligned_buf)? <= idx {
-            let mut offset = from_offset;
-            for chunk in aligned_buf.chunks_mut(P::ACTIVE::BLOCK_SIZE) {
-                self.active.read_blocking(p.active(), offset, chunk)?;
-                offset += chunk.len() as u32;
-            }
+            let page_size = P::page_size() as u32;
 
             self.dfu
-                .erase_blocking(p.dfu(), to_offset as u32, to_offset + P::page_size() as u32)?;
+                .erase_blocking(p.dfu(), to_offset as u32, to_offset + page_size)?;
 
-            let mut offset = to_offset;
-            for chunk in aligned_buf.chunks(P::DFU::BLOCK_SIZE) {
-                self.dfu.write_blocking(p.dfu(), offset, chunk)?;
-                offset += chunk.len() as u32;
+            for offset_in_page in (0..page_size).step_by(aligned_buf.len()) {
+                self.active
+                    .read_blocking(p.active(), from_offset + offset_in_page as u32, aligned_buf)?;
+                self.dfu
+                    .write_blocking(p.dfu(), to_offset + offset_in_page as u32, aligned_buf)?;
             }
+
             self.update_progress(idx, p, aligned_buf)?;
         }
         Ok(())
@@ -389,14 +386,14 @@ fn assert_partitions(active: Partition, dfu: Partition, state: Partition, page_s
 }
 
 /// A flash wrapper implementing the Flash and embedded_storage traits.
-pub struct BootFlash<F, const BLOCK_SIZE: usize, const ERASE_VALUE: u8 = 0xFF>
+pub struct BootFlash<F, const ERASE_VALUE: u8 = 0xFF>
 where
     F: NorFlash + ReadNorFlash,
 {
     flash: F,
 }
 
-impl<F, const BLOCK_SIZE: usize, const ERASE_VALUE: u8> BootFlash<F, BLOCK_SIZE, ERASE_VALUE>
+impl<F, const ERASE_VALUE: u8> BootFlash<F, ERASE_VALUE>
 where
     F: NorFlash + ReadNorFlash,
 {
@@ -406,22 +403,21 @@ where
     }
 }
 
-impl<F, const BLOCK_SIZE: usize, const ERASE_VALUE: u8> Flash for BootFlash<F, BLOCK_SIZE, ERASE_VALUE>
+impl<F, const ERASE_VALUE: u8> Flash for BootFlash<F, ERASE_VALUE>
 where
     F: NorFlash + ReadNorFlash,
 {
-    const BLOCK_SIZE: usize = BLOCK_SIZE;
     const ERASE_VALUE: u8 = ERASE_VALUE;
 }
 
-impl<F, const BLOCK_SIZE: usize, const ERASE_VALUE: u8> ErrorType for BootFlash<F, BLOCK_SIZE, ERASE_VALUE>
+impl<F, const ERASE_VALUE: u8> ErrorType for BootFlash<F, ERASE_VALUE>
 where
     F: ReadNorFlash + NorFlash,
 {
     type Error = F::Error;
 }
 
-impl<F, const BLOCK_SIZE: usize, const ERASE_VALUE: u8> NorFlash for BootFlash<F, BLOCK_SIZE, ERASE_VALUE>
+impl<F, const ERASE_VALUE: u8> NorFlash for BootFlash<F, ERASE_VALUE>
 where
     F: ReadNorFlash + NorFlash,
 {
@@ -437,7 +433,7 @@ where
     }
 }
 
-impl<F, const BLOCK_SIZE: usize, const ERASE_VALUE: u8> ReadNorFlash for BootFlash<F, BLOCK_SIZE, ERASE_VALUE>
+impl<F, const ERASE_VALUE: u8> ReadNorFlash for BootFlash<F, ERASE_VALUE>
 where
     F: ReadNorFlash + NorFlash,
 {