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.

4 files changed, 118 insertions, 57 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,
 {
diff --git a/embassy-boot/boot/src/large_erase.rs b/embassy-boot/boot/src/large_erase.rs
new file mode 100644
index 000000000..d00d43599
--- /dev/null
+++ b/embassy-boot/boot/src/large_erase.rs
@@ -0,0 +1,76 @@
+#![allow(unused)]
+
+use embedded_storage::nor_flash::{ErrorType, NorFlash, ReadNorFlash};
+use embedded_storage_async::nor_flash::{NorFlash as AsyncNorFlash, ReadNorFlash as AsyncReadNorFlash};
+
+use crate::Flash;
+
+pub struct LargeErase<F, const ERASE_SIZE: usize>(pub F);
+
+impl<F, const ERASE_SIZE: usize> LargeErase<F, ERASE_SIZE> {
+    pub const fn new(flash: F) -> Self {
+        Self(flash)
+    }
+}
+
+impl<F: Flash, const ERASE_SIZE: usize> Flash for LargeErase<F, ERASE_SIZE> {
+    const ERASE_VALUE: u8 = F::ERASE_VALUE;
+}
+
+impl<F: ErrorType, const ERASE_SIZE: usize> ErrorType for LargeErase<F, ERASE_SIZE> {
+    type Error = F::Error;
+}
+
+impl<F: NorFlash, const ERASE_SIZE: usize> NorFlash for LargeErase<F, ERASE_SIZE> {
+    const WRITE_SIZE: usize = F::ERASE_SIZE;
+    const ERASE_SIZE: usize = ERASE_SIZE;
+
+    fn erase(&mut self, from: u32, to: u32) -> Result<(), Self::Error> {
+        assert!(ERASE_SIZE >= F::ERASE_SIZE);
+        assert_eq!(0, ERASE_SIZE % F::ERASE_SIZE);
+        self.0.erase(from, to)
+    }
+
+    fn write(&mut self, offset: u32, bytes: &[u8]) -> Result<(), Self::Error> {
+        self.0.write(offset, bytes)
+    }
+}
+
+impl<F: ReadNorFlash, const ERASE_SIZE: usize> ReadNorFlash for LargeErase<F, ERASE_SIZE> {
+    const READ_SIZE: usize = F::READ_SIZE;
+
+    fn read(&mut self, offset: u32, bytes: &mut [u8]) -> Result<(), Self::Error> {
+        self.0.read(offset, bytes)
+    }
+
+    fn capacity(&self) -> usize {
+        self.0.capacity()
+    }
+}
+
+impl<F: AsyncNorFlash, const ERASE_SIZE: usize> AsyncNorFlash for LargeErase<F, ERASE_SIZE> {
+    const WRITE_SIZE: usize = F::ERASE_SIZE;
+    const ERASE_SIZE: usize = ERASE_SIZE;
+
+    async fn erase(&mut self, from: u32, to: u32) -> Result<(), Self::Error> {
+        assert!(ERASE_SIZE >= F::ERASE_SIZE);
+        assert_eq!(0, ERASE_SIZE % F::ERASE_SIZE);
+        self.0.erase(from, to).await
+    }
+
+    async fn write(&mut self, offset: u32, bytes: &[u8]) -> Result<(), Self::Error> {
+        self.0.write(offset, bytes).await
+    }
+}
+
+impl<F: AsyncReadNorFlash, const ERASE_SIZE: usize> AsyncReadNorFlash for LargeErase<F, ERASE_SIZE> {
+    const READ_SIZE: usize = F::READ_SIZE;
+
+    async fn read(&mut self, offset: u32, bytes: &mut [u8]) -> Result<(), Self::Error> {
+        self.0.read(offset, bytes).await
+    }
+
+    fn capacity(&self) -> usize {
+        self.0.capacity()
+    }
+}
diff --git a/embassy-boot/boot/src/lib.rs b/embassy-boot/boot/src/lib.rs
index 896498c0b..79759124b 100644
--- a/embassy-boot/boot/src/lib.rs
+++ b/embassy-boot/boot/src/lib.rs
@@ -7,6 +7,7 @@ mod fmt;
 
 mod boot_loader;
 mod firmware_updater;
+mod large_erase;
 mod mem_flash;
 mod partition;
 
@@ -48,6 +49,7 @@ mod tests {
     use futures::executor::block_on;
 
     use super::*;
+    use crate::large_erase::LargeErase;
     use crate::mem_flash::MemFlash;
 
     /*
@@ -99,14 +101,10 @@ mod tests {
 
         let mut bootloader: BootLoader = BootLoader::new(ACTIVE, DFU, STATE);
         let mut updater = FirmwareUpdater::new(DFU, STATE);
-        let mut offset = 0;
-        for chunk in update.chunks(4096) {
-            block_on(updater.write_firmware(offset, chunk, &mut flash)).unwrap();
-            offset += chunk.len();
-        }
+        block_on(updater.write_firmware(0, &update, &mut flash)).unwrap();
         block_on(updater.mark_updated(&mut flash, &mut aligned)).unwrap();
 
-        let mut page = [0; 4096];
+        let mut page = [0; 1024];
         assert_eq!(
             State::Swap,
             bootloader
@@ -158,7 +156,7 @@ mod tests {
         const DFU: Partition = Partition::new(0, 16384);
 
         let mut active = MemFlash::<16384, 4096, 8>::random();
-        let mut dfu = MemFlash::<16384, 2048, 8>::random();
+        let mut dfu = LargeErase::<_, 4096>::new(MemFlash::<16384, 2048, 8>::random());
         let mut state = MemFlash::<4096, 128, 4>::random();
         let mut aligned = [0; 4];
 
@@ -171,11 +169,7 @@ mod tests {
 
         let mut updater = FirmwareUpdater::new(DFU, STATE);
 
-        let mut offset = 0;
-        for chunk in update.chunks(2048) {
-            block_on(updater.write_firmware(offset, chunk, &mut dfu)).unwrap();
-            offset += chunk.len();
-        }
+        block_on(updater.write_firmware(0, &update, &mut dfu)).unwrap();
         block_on(updater.mark_updated(&mut state, &mut aligned)).unwrap();
 
         let mut bootloader: BootLoader = BootLoader::new(ACTIVE, DFU, STATE);
@@ -194,7 +188,7 @@ mod tests {
 
         // First DFU page is untouched
         for i in DFU.from + 4096..DFU.to {
-            assert_eq!(dfu.mem[i], original[i - DFU.from - 4096], "Index {}", i);
+            assert_eq!(dfu.0.mem[i], original[i - DFU.from - 4096], "Index {}", i);
         }
     }
 
@@ -206,7 +200,7 @@ mod tests {
         const DFU: Partition = Partition::new(0, 16384);
 
         let mut aligned = [0; 4];
-        let mut active = MemFlash::<16384, 2048, 4>::random();
+        let mut active = LargeErase::<_, 4096>::new(MemFlash::<16384, 2048, 4>::random());
         let mut dfu = MemFlash::<16384, 4096, 8>::random();
         let mut state = MemFlash::<4096, 128, 4>::random();
 
@@ -214,16 +208,12 @@ mod tests {
         let update: [u8; DFU.len()] = [rand::random::<u8>(); DFU.len()];
 
         for i in ACTIVE.from..ACTIVE.to {
-            active.mem[i] = original[i - ACTIVE.from];
+            active.0.mem[i] = original[i - ACTIVE.from];
         }
 
         let mut updater = FirmwareUpdater::new(DFU, STATE);
 
-        let mut offset = 0;
-        for chunk in update.chunks(4096) {
-            block_on(updater.write_firmware(offset, chunk, &mut dfu)).unwrap();
-            offset += chunk.len();
-        }
+        block_on(updater.write_firmware(0, &update, &mut dfu)).unwrap();
         block_on(updater.mark_updated(&mut state, &mut aligned)).unwrap();
 
         let mut bootloader: BootLoader = BootLoader::new(ACTIVE, DFU, STATE);
@@ -239,7 +229,7 @@ mod tests {
         );
 
         for i in ACTIVE.from..ACTIVE.to {
-            assert_eq!(active.mem[i], update[i - ACTIVE.from], "Index {}", i);
+            assert_eq!(active.0.mem[i], update[i - ACTIVE.from], "Index {}", i);
         }
 
         // First DFU page is untouched
diff --git a/embassy-boot/boot/src/mem_flash.rs b/embassy-boot/boot/src/mem_flash.rs
index 828aad9d9..2598bf4de 100644
--- a/embassy-boot/boot/src/mem_flash.rs
+++ b/embassy-boot/boot/src/mem_flash.rs
@@ -47,7 +47,6 @@ impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> Defaul
 impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> Flash
     for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
 {
-    const BLOCK_SIZE: usize = ERASE_SIZE;
     const ERASE_VALUE: u8 = 0xFF;
 }