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;
 }