9 files changed, 157 insertions, 196 deletions

diff --git a/embassy-boot/boot/src/boot_loader.rs b/embassy-boot/boot/src/boot_loader.rs
index 9d047f778..2412427c0 100644
--- a/embassy-boot/boot/src/boot_loader.rs
+++ b/embassy-boot/boot/src/boot_loader.rs
@@ -30,23 +30,16 @@ 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.
 pub trait FlashConfig {
+    /// The erase value of the state flash. Typically the default of 0xFF is used, but some flashes use a different value.
+    const STATE_ERASE_VALUE: u8 = 0xFF;
     /// Flash type used for the state partition.
-    type STATE: Flash;
+    type STATE: NorFlash;
     /// Flash type used for the active partition.
-    type ACTIVE: Flash;
+    type ACTIVE: NorFlash;
     /// Flash type used for the dfu partition.
-    type DFU: Flash;
+    type DFU: NorFlash;
 
     /// Return flash instance used to write/read to/from active partition.
     fn active(&mut self) -> &mut Self::ACTIVE;
@@ -56,8 +49,18 @@ pub trait FlashConfig {
     fn state(&mut self) -> &mut Self::STATE;
 }
 
-/// BootLoader works with any flash implementing embedded_storage and can also work with
-/// different page sizes and flash write sizes.
+trait FlashConfigEx {
+    fn page_size() -> usize;
+}
+
+impl<T: FlashConfig> FlashConfigEx for T {
+    /// Get the page size which is the "unit of operation" within the bootloader.
+    fn page_size() -> usize {
+        core::cmp::max(T::ACTIVE::ERASE_SIZE, T::DFU::ERASE_SIZE)
+    }
+}
+
+/// 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.
@@ -91,6 +94,9 @@ impl BootLoader {
     /// The DFU partition is assumed to be 1 page bigger than the active partition for the swap
     /// algorithm to work correctly.
     ///
+    /// The provided aligned_buf argument must satisfy any alignment requirements
+    /// given by the partition flashes. All flash operations will use this buffer.
+    ///
     /// SWAPPING
     ///
     /// Assume a flash size of 3 pages for the active partition, and 4 pages for the DFU partition.
@@ -169,87 +175,95 @@ impl BootLoader {
     /// |       DFU |            3 |      3 |      2 |      1 |      3 |
     /// +-----------+--------------+--------+--------+--------+--------+
     ///
-    pub fn prepare_boot<P: FlashConfig>(
-        &mut self,
-        p: &mut P,
-        magic: &mut [u8],
-        page: &mut [u8],
-    ) -> Result<State, BootError> {
+    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_partitions(self.active, self.dfu, self.state, page.len(), P::STATE::WRITE_SIZE);
-        assert_eq!(magic.len(), P::STATE::WRITE_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::ACTIVE::ERASE_SIZE);
+        assert_eq!(0, P::page_size() % P::DFU::WRITE_SIZE);
+        assert_eq!(0, P::page_size() % P::DFU::ERASE_SIZE);
+        assert!(aligned_buf.len() >= P::STATE::WRITE_SIZE);
+        assert_eq!(0, aligned_buf.len() % P::ACTIVE::WRITE_SIZE);
+        assert_eq!(0, aligned_buf.len() % P::DFU::WRITE_SIZE);
+        assert_partitions(self.active, self.dfu, self.state, P::page_size(), P::STATE::WRITE_SIZE);
 
         // Copy contents from partition N to active
-        let state = self.read_state(p, magic)?;
+        let state = self.read_state(p, aligned_buf)?;
         if state == State::Swap {
             //
             // Check if we already swapped. If we're in the swap state, this means we should revert
             // since the app has failed to mark boot as successful
             //
-            if !self.is_swapped(p, magic, page)? {
+            if !self.is_swapped(p, aligned_buf)? {
                 trace!("Swapping");
-                self.swap(p, magic, page)?;
+                self.swap(p, aligned_buf)?;
                 trace!("Swapping done");
             } else {
                 trace!("Reverting");
-                self.revert(p, magic, page)?;
+                self.revert(p, aligned_buf)?;
 
                 let state_flash = p.state();
+                let state_word = &mut aligned_buf[..P::STATE::WRITE_SIZE];
 
                 // Invalidate progress
-                magic.fill(!P::STATE::ERASE_VALUE);
+                state_word.fill(!P::STATE_ERASE_VALUE);
                 self.state
-                    .write_blocking(state_flash, P::STATE::WRITE_SIZE as u32, magic)?;
+                    .write_blocking(state_flash, P::STATE::WRITE_SIZE as u32, state_word)?;
 
                 // Clear magic and progress
                 self.state.wipe_blocking(state_flash)?;
 
                 // Set magic
-                magic.fill(BOOT_MAGIC);
-                self.state.write_blocking(state_flash, 0, magic)?;
+                state_word.fill(BOOT_MAGIC);
+                self.state.write_blocking(state_flash, 0, state_word)?;
             }
         }
         Ok(state)
     }
 
-    fn is_swapped<P: FlashConfig>(&mut self, p: &mut P, magic: &mut [u8], page: &mut [u8]) -> Result<bool, BootError> {
-        let page_size = page.len();
-        let page_count = self.active.len() / page_size;
-        let progress = self.current_progress(p, magic)?;
+    fn is_swapped<P: FlashConfig>(&mut self, p: &mut P, aligned_buf: &mut [u8]) -> Result<bool, BootError> {
+        let page_count = self.active.len() / P::page_size();
+        let progress = self.current_progress(p, aligned_buf)?;
 
         Ok(progress >= page_count * 2)
     }
 
-    fn current_progress<P: FlashConfig>(&mut self, config: &mut P, aligned: &mut [u8]) -> Result<usize, BootError> {
-        let write_size = aligned.len();
-        let max_index = ((self.state.len() - write_size) / write_size) - 2;
-        aligned.fill(!P::STATE::ERASE_VALUE);
-
+    fn current_progress<P: FlashConfig>(&mut self, config: &mut P, aligned_buf: &mut [u8]) -> Result<usize, BootError> {
+        let max_index = ((self.state.len() - P::STATE::WRITE_SIZE) / P::STATE::WRITE_SIZE) - 2;
         let state_flash = config.state();
+        let state_word = &mut aligned_buf[..P::STATE::WRITE_SIZE];
 
         self.state
-            .read_blocking(state_flash, P::STATE::WRITE_SIZE as u32, aligned)?;
-        if aligned.iter().any(|&b| b != P::STATE::ERASE_VALUE) {
+            .read_blocking(state_flash, P::STATE::WRITE_SIZE as u32, state_word)?;
+        if state_word.iter().any(|&b| b != P::STATE_ERASE_VALUE) {
             // Progress is invalid
             return Ok(max_index);
         }
 
         for index in 0..max_index {
-            self.state
-                .read_blocking(state_flash, (2 + index) as u32 * P::STATE::WRITE_SIZE as u32, aligned)?;
+            self.state.read_blocking(
+                state_flash,
+                (2 + index) as u32 * P::STATE::WRITE_SIZE as u32,
+                state_word,
+            )?;
 
-            if aligned.iter().any(|&b| b == P::STATE::ERASE_VALUE) {
+            if state_word.iter().any(|&b| b == P::STATE_ERASE_VALUE) {
                 return Ok(index);
             }
         }
         Ok(max_index)
     }
 
-    fn update_progress<P: FlashConfig>(&mut self, index: usize, p: &mut P, magic: &mut [u8]) -> Result<(), BootError> {
-        let aligned = magic;
-        aligned.fill(!P::STATE::ERASE_VALUE);
+    fn update_progress<P: FlashConfig>(
+        &mut self,
+        index: usize,
+        p: &mut P,
+        aligned_buf: &mut [u8],
+    ) -> Result<(), BootError> {
+        let state_word = &mut aligned_buf[..P::STATE::WRITE_SIZE];
+        state_word.fill(!P::STATE_ERASE_VALUE);
         self.state
-            .write_blocking(p.state(), (2 + index) as u32 * P::STATE::WRITE_SIZE as u32, aligned)?;
+            .write_blocking(p.state(), (2 + index) as u32 * P::STATE::WRITE_SIZE as u32, state_word)?;
         Ok(())
     }
 
@@ -259,26 +273,22 @@ impl BootLoader {
         from_offset: u32,
         to_offset: u32,
         p: &mut P,
-        magic: &mut [u8],
-        page: &mut [u8],
+        aligned_buf: &mut [u8],
     ) -> Result<(), BootError> {
-        let buf = page;
-        if self.current_progress(p, magic)? <= idx {
-            let mut offset = from_offset;
-            for chunk in buf.chunks_mut(P::DFU::BLOCK_SIZE) {
-                self.dfu.read_blocking(p.dfu(), offset, chunk)?;
-                offset += chunk.len() as u32;
-            }
+        if self.current_progress(p, aligned_buf)? <= idx {
+            let page_size = P::page_size() as u32;
 
             self.active
-                .erase_blocking(p.active(), to_offset, to_offset + buf.len() as u32)?;
+                .erase_blocking(p.active(), to_offset, to_offset + page_size)?;
 
-            let mut offset = to_offset;
-            for chunk in 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, magic)?;
+
+            self.update_progress(idx, p, aligned_buf)?;
         }
         Ok(())
     }
@@ -289,32 +299,28 @@ impl BootLoader {
         from_offset: u32,
         to_offset: u32,
         p: &mut P,
-        magic: &mut [u8],
-        page: &mut [u8],
+        aligned_buf: &mut [u8],
     ) -> Result<(), BootError> {
-        let buf = page;
-        if self.current_progress(p, magic)? <= idx {
-            let mut offset = from_offset;
-            for chunk in buf.chunks_mut(P::ACTIVE::BLOCK_SIZE) {
-                self.active.read_blocking(p.active(), offset, chunk)?;
-                offset += chunk.len() as u32;
-            }
+        if self.current_progress(p, aligned_buf)? <= idx {
+            let page_size = P::page_size() as u32;
 
             self.dfu
-                .erase_blocking(p.dfu(), to_offset as u32, to_offset + buf.len() as u32)?;
+                .erase_blocking(p.dfu(), to_offset as u32, to_offset + page_size)?;
 
-            let mut offset = to_offset;
-            for chunk in 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, magic)?;
+
+            self.update_progress(idx, p, aligned_buf)?;
         }
         Ok(())
     }
 
-    fn swap<P: FlashConfig>(&mut self, p: &mut P, magic: &mut [u8], page: &mut [u8]) -> Result<(), BootError> {
-        let page_size = page.len();
+    fn swap<P: FlashConfig>(&mut self, p: &mut P, aligned_buf: &mut [u8]) -> Result<(), BootError> {
+        let page_size = P::page_size();
         let page_count = self.active.len() / page_size;
         trace!("Page count: {}", page_count);
         for page_num in 0..page_count {
@@ -326,20 +332,20 @@ impl BootLoader {
             let active_from_offset = ((page_count - 1 - page_num) * page_size) as u32;
             let dfu_to_offset = ((page_count - page_num) * page_size) as u32;
             //trace!("Copy active {} to dfu {}", active_from_offset, dfu_to_offset);
-            self.copy_page_once_to_dfu(idx, active_from_offset, dfu_to_offset, p, magic, page)?;
+            self.copy_page_once_to_dfu(idx, active_from_offset, dfu_to_offset, p, aligned_buf)?;
 
             // Copy DFU page to the active page
             let active_to_offset = ((page_count - 1 - page_num) * page_size) as u32;
             let dfu_from_offset = ((page_count - 1 - page_num) * page_size) as u32;
             //trace!("Copy dfy {} to active {}", dfu_from_offset, active_to_offset);
-            self.copy_page_once_to_active(idx + 1, dfu_from_offset, active_to_offset, p, magic, page)?;
+            self.copy_page_once_to_active(idx + 1, dfu_from_offset, active_to_offset, p, aligned_buf)?;
         }
 
         Ok(())
     }
 
-    fn revert<P: FlashConfig>(&mut self, p: &mut P, magic: &mut [u8], page: &mut [u8]) -> Result<(), BootError> {
-        let page_size = page.len();
+    fn revert<P: FlashConfig>(&mut self, p: &mut P, aligned_buf: &mut [u8]) -> Result<(), BootError> {
+        let page_size = P::page_size();
         let page_count = self.active.len() / page_size;
         for page_num in 0..page_count {
             let idx = page_count * 2 + page_num * 2;
@@ -347,21 +353,22 @@ impl BootLoader {
             // Copy the bad active page to the DFU page
             let active_from_offset = (page_num * page_size) as u32;
             let dfu_to_offset = (page_num * page_size) as u32;
-            self.copy_page_once_to_dfu(idx, active_from_offset, dfu_to_offset, p, magic, page)?;
+            self.copy_page_once_to_dfu(idx, active_from_offset, dfu_to_offset, p, aligned_buf)?;
 
             // Copy the DFU page back to the active page
             let active_to_offset = (page_num * page_size) as u32;
             let dfu_from_offset = ((page_num + 1) * page_size) as u32;
-            self.copy_page_once_to_active(idx + 1, dfu_from_offset, active_to_offset, p, magic, page)?;
+            self.copy_page_once_to_active(idx + 1, dfu_from_offset, active_to_offset, p, aligned_buf)?;
         }
 
         Ok(())
     }
 
-    fn read_state<P: FlashConfig>(&mut self, config: &mut P, magic: &mut [u8]) -> Result<State, BootError> {
-        self.state.read_blocking(config.state(), 0, magic)?;
+    fn read_state<P: FlashConfig>(&mut self, config: &mut P, aligned_buf: &mut [u8]) -> Result<State, BootError> {
+        let state_word = &mut aligned_buf[..P::STATE::WRITE_SIZE];
+        self.state.read_blocking(config.state(), 0, state_word)?;
 
-        if !magic.iter().any(|&b| b != SWAP_MAGIC) {
+        if !state_word.iter().any(|&b| b != SWAP_MAGIC) {
             Ok(State::Swap)
         } else {
             Ok(State::Boot)
@@ -377,16 +384,16 @@ 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>
 where
-    F: NorFlash + ReadNorFlash,
+    F: NorFlash,
 {
     flash: F,
 }
 
-impl<F, const BLOCK_SIZE: usize, const ERASE_VALUE: u8> BootFlash<F, BLOCK_SIZE, ERASE_VALUE>
+impl<F> BootFlash<F>
 where
-    F: NorFlash + ReadNorFlash,
+    F: NorFlash,
 {
     /// Create a new instance of a bootable flash
     pub fn new(flash: F) -> Self {
@@ -394,24 +401,16 @@ where
     }
 }
 
-impl<F, const BLOCK_SIZE: usize, const ERASE_VALUE: u8> Flash for BootFlash<F, BLOCK_SIZE, 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> ErrorType for BootFlash<F>
 where
-    F: ReadNorFlash + NorFlash,
+    F: 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> NorFlash for BootFlash<F>
 where
-    F: ReadNorFlash + NorFlash,
+    F: NorFlash,
 {
     const WRITE_SIZE: usize = F::WRITE_SIZE;
     const ERASE_SIZE: usize = F::ERASE_SIZE;
@@ -425,9 +424,9 @@ where
     }
 }
 
-impl<F, const BLOCK_SIZE: usize, const ERASE_VALUE: u8> ReadNorFlash for BootFlash<F, BLOCK_SIZE, ERASE_VALUE>
+impl<F> ReadNorFlash for BootFlash<F>
 where
-    F: ReadNorFlash + NorFlash,
+    F: NorFlash,
 {
     const READ_SIZE: usize = F::READ_SIZE;
 
@@ -443,14 +442,14 @@ where
 /// Convenience provider that uses a single flash for all partitions.
 pub struct SingleFlashConfig<'a, F>
 where
-    F: Flash,
+    F: NorFlash,
 {
     flash: &'a mut F,
 }
 
 impl<'a, F> SingleFlashConfig<'a, F>
 where
-    F: Flash,
+    F: NorFlash,
 {
     /// Create a provider for a single flash.
     pub fn new(flash: &'a mut F) -> Self {
@@ -460,7 +459,7 @@ where
 
 impl<'a, F> FlashConfig for SingleFlashConfig<'a, F>
 where
-    F: Flash,
+    F: NorFlash,
 {
     type STATE = F;
     type ACTIVE = F;
@@ -480,9 +479,9 @@ where
 /// Convenience flash provider that uses separate flash instances for each partition.
 pub struct MultiFlashConfig<'a, ACTIVE, STATE, DFU>
 where
-    ACTIVE: Flash,
-    STATE: Flash,
-    DFU: Flash,
+    ACTIVE: NorFlash,
+    STATE: NorFlash,
+    DFU: NorFlash,
 {
     active: &'a mut ACTIVE,
     state: &'a mut STATE,
@@ -491,9 +490,9 @@ where
 
 impl<'a, ACTIVE, STATE, DFU> MultiFlashConfig<'a, ACTIVE, STATE, DFU>
 where
-    ACTIVE: Flash,
-    STATE: Flash,
-    DFU: Flash,
+    ACTIVE: NorFlash,
+    STATE: NorFlash,
+    DFU: NorFlash,
 {
     /// Create a new flash provider with separate configuration for all three partitions.
     pub fn new(active: &'a mut ACTIVE, state: &'a mut STATE, dfu: &'a mut DFU) -> Self {
@@ -503,9 +502,9 @@ where
 
 impl<'a, ACTIVE, STATE, DFU> FlashConfig for MultiFlashConfig<'a, ACTIVE, STATE, DFU>
 where
-    ACTIVE: Flash,
-    STATE: Flash,
-    DFU: Flash,
+    ACTIVE: NorFlash,
+    STATE: NorFlash,
+    DFU: NorFlash,
 {
     type STATE = STATE;
     type ACTIVE = ACTIVE;
diff --git a/embassy-boot/boot/src/lib.rs b/embassy-boot/boot/src/lib.rs
index d4078f1cb..605e5253c 100644
--- a/embassy-boot/boot/src/lib.rs
+++ b/embassy-boot/boot/src/lib.rs
@@ -11,7 +11,7 @@ mod firmware_updater;
 mod mem_flash;
 mod partition;
 
-pub use boot_loader::{BootError, BootFlash, BootLoader, Flash, FlashConfig, MultiFlashConfig, SingleFlashConfig};
+pub use boot_loader::{BootError, BootFlash, BootLoader, FlashConfig, MultiFlashConfig, SingleFlashConfig};
 pub use firmware_updater::{FirmwareUpdater, FirmwareUpdaterError};
 pub use partition::Partition;
 
@@ -78,12 +78,8 @@ mod tests {
 
         let mut bootloader: BootLoader = BootLoader::new(ACTIVE, DFU, STATE);
 
-        let mut magic = [0; 4];
         let mut page = [0; 4096];
-        assert_eq!(
-            State::Boot,
-            bootloader.prepare_boot(&mut flash, &mut magic, &mut page).unwrap()
-        );
+        assert_eq!(State::Boot, bootloader.prepare_boot(&mut flash, &mut page).unwrap());
     }
 
     #[test]
@@ -104,19 +100,14 @@ 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 magic = [0; 4];
-        let mut page = [0; 4096];
+        let mut page = [0; 1024];
         assert_eq!(
             State::Swap,
             bootloader
-                .prepare_boot(&mut SingleFlashConfig::new(&mut flash), &mut magic, &mut page)
+                .prepare_boot(&mut SingleFlashConfig::new(&mut flash), &mut page)
                 .unwrap()
         );
 
@@ -133,7 +124,7 @@ mod tests {
         assert_eq!(
             State::Swap,
             bootloader
-                .prepare_boot(&mut SingleFlashConfig::new(&mut flash), &mut magic, &mut page)
+                .prepare_boot(&mut SingleFlashConfig::new(&mut flash), &mut page)
                 .unwrap()
         );
 
@@ -151,7 +142,7 @@ mod tests {
         assert_eq!(
             State::Boot,
             bootloader
-                .prepare_boot(&mut SingleFlashConfig::new(&mut flash), &mut magic, &mut page)
+                .prepare_boot(&mut SingleFlashConfig::new(&mut flash), &mut page)
                 .unwrap()
         );
     }
@@ -177,25 +168,16 @@ 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);
-        let mut magic = [0; 4];
         let mut page = [0; 4096];
 
         assert_eq!(
             State::Swap,
             bootloader
-                .prepare_boot(
-                    &mut MultiFlashConfig::new(&mut active, &mut state, &mut dfu),
-                    &mut magic,
-                    &mut page
-                )
+                .prepare_boot(&mut MultiFlashConfig::new(&mut active, &mut state, &mut dfu), &mut page)
                 .unwrap()
         );
 
@@ -230,22 +212,16 @@ mod tests {
 
         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);
-        let mut magic = [0; 4];
         let mut page = [0; 4096];
         assert_eq!(
             State::Swap,
             bootloader
                 .prepare_boot(
                     &mut MultiFlashConfig::new(&mut active, &mut state, &mut dfu,),
-                    &mut magic,
                     &mut page
                 )
                 .unwrap()
diff --git a/embassy-boot/boot/src/mem_flash.rs b/embassy-boot/boot/src/mem_flash.rs
index 828aad9d9..dd85405c8 100644
--- a/embassy-boot/boot/src/mem_flash.rs
+++ b/embassy-boot/boot/src/mem_flash.rs
@@ -5,8 +5,6 @@ use core::ops::{Bound, Range, RangeBounds};
 use embedded_storage::nor_flash::{ErrorType, NorFlash, NorFlashError, NorFlashErrorKind, ReadNorFlash};
 use embedded_storage_async::nor_flash::{NorFlash as AsyncNorFlash, ReadNorFlash as AsyncReadNorFlash};
 
-use crate::Flash;
-
 pub struct MemFlash<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> {
     pub mem: [u8; SIZE],
     pub pending_write_successes: Option<usize>,
@@ -44,13 +42,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;
-}
-
 impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> ErrorType
     for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
 {
diff --git a/embassy-boot/nrf/src/lib.rs b/embassy-boot/nrf/src/lib.rs
index f40ae62d6..a2176f609 100644
--- a/embassy-boot/nrf/src/lib.rs
+++ b/embassy-boot/nrf/src/lib.rs
@@ -11,13 +11,12 @@ use embassy_nrf::wdt;
 use embedded_storage::nor_flash::{ErrorType, NorFlash, ReadNorFlash};
 
 /// A bootloader for nRF devices.
-pub struct BootLoader {
+pub struct BootLoader<const BUFFER_SIZE: usize = PAGE_SIZE> {
     boot: embassy_boot::BootLoader,
-    magic: AlignedBuffer<4>,
-    page: AlignedBuffer<PAGE_SIZE>,
+    aligned_buf: AlignedBuffer<BUFFER_SIZE>,
 }
 
-impl Default for BootLoader {
+impl Default for BootLoader<PAGE_SIZE> {
     /// Create a new bootloader instance using parameters from linker script
     fn default() -> Self {
         extern "C" {
@@ -56,20 +55,19 @@ impl Default for BootLoader {
     }
 }
 
-impl BootLoader {
+impl<const BUFFER_SIZE: usize> BootLoader<BUFFER_SIZE> {
     /// Create a new bootloader instance using the supplied partitions for active, dfu and state.
     pub fn new(active: Partition, dfu: Partition, state: Partition) -> Self {
         Self {
             boot: embassy_boot::BootLoader::new(active, dfu, state),
-            magic: AlignedBuffer([0; 4]),
-            page: AlignedBuffer([0; PAGE_SIZE]),
+            aligned_buf: AlignedBuffer([0; BUFFER_SIZE]),
         }
     }
 
     /// Inspect the bootloader state and perform actions required before booting, such as swapping
     /// firmware.
     pub fn prepare<F: FlashConfig>(&mut self, flash: &mut F) -> usize {
-        match self.boot.prepare_boot(flash, &mut self.magic.0, &mut self.page.0) {
+        match self.boot.prepare_boot(flash, &mut self.aligned_buf.0) {
             Ok(_) => self.boot.boot_address(),
             Err(_) => panic!("boot prepare error!"),
         }
diff --git a/embassy-boot/rp/src/lib.rs b/embassy-boot/rp/src/lib.rs
index 6df34133e..0031efa63 100644
--- a/embassy-boot/rp/src/lib.rs
+++ b/embassy-boot/rp/src/lib.rs
@@ -5,33 +5,31 @@
 mod fmt;
 
 pub use embassy_boot::{AlignedBuffer, BootFlash, FirmwareUpdater, FlashConfig, Partition, SingleFlashConfig, State};
-use embassy_rp::flash::{Flash, ERASE_SIZE, WRITE_SIZE};
+use embassy_rp::flash::{Flash, ERASE_SIZE};
 use embassy_rp::peripherals::{FLASH, WATCHDOG};
 use embassy_rp::watchdog::Watchdog;
 use embassy_time::Duration;
 use embedded_storage::nor_flash::{ErrorType, NorFlash, ReadNorFlash};
 
 /// A bootloader for RP2040 devices.
-pub struct BootLoader {
+pub struct BootLoader<const BUFFER_SIZE: usize = ERASE_SIZE> {
     boot: embassy_boot::BootLoader,
-    magic: AlignedBuffer<WRITE_SIZE>,
-    page: AlignedBuffer<ERASE_SIZE>,
+    aligned_buf: AlignedBuffer<BUFFER_SIZE>,
 }
 
-impl BootLoader {
+impl<const BUFFER_SIZE: usize> BootLoader<BUFFER_SIZE> {
     /// Create a new bootloader instance using the supplied partitions for active, dfu and state.
     pub fn new(active: Partition, dfu: Partition, state: Partition) -> Self {
         Self {
             boot: embassy_boot::BootLoader::new(active, dfu, state),
-            magic: AlignedBuffer([0; WRITE_SIZE]),
-            page: AlignedBuffer([0; ERASE_SIZE]),
+            aligned_buf: AlignedBuffer([0; BUFFER_SIZE]),
         }
     }
 
     /// Inspect the bootloader state and perform actions required before booting, such as swapping
     /// firmware.
     pub fn prepare<F: FlashConfig>(&mut self, flash: &mut F) -> usize {
-        match self.boot.prepare_boot(flash, self.magic.as_mut(), self.page.as_mut()) {
+        match self.boot.prepare_boot(flash, self.aligned_buf.as_mut()) {
             Ok(_) => embassy_rp::flash::FLASH_BASE + self.boot.boot_address(),
             Err(_) => panic!("boot prepare error!"),
         }
@@ -54,7 +52,7 @@ impl BootLoader {
     }
 }
 
-impl Default for BootLoader {
+impl Default for BootLoader<ERASE_SIZE> {
     /// Create a new bootloader instance using parameters from linker script
     fn default() -> Self {
         extern "C" {
diff --git a/embassy-boot/stm32/src/lib.rs b/embassy-boot/stm32/src/lib.rs
index 82f712c4d..1f63fcd63 100644
--- a/embassy-boot/stm32/src/lib.rs
+++ b/embassy-boot/stm32/src/lib.rs
@@ -7,26 +7,24 @@ mod fmt;
 pub use embassy_boot::{AlignedBuffer, BootFlash, FirmwareUpdater, FlashConfig, Partition, SingleFlashConfig, State};
 
 /// A bootloader for STM32 devices.
-pub struct BootLoader<const PAGE_SIZE: usize, const WRITE_SIZE: usize> {
+pub struct BootLoader<const BUFFER_SIZE: usize> {
     boot: embassy_boot::BootLoader,
-    magic: AlignedBuffer<WRITE_SIZE>,
-    page: AlignedBuffer<PAGE_SIZE>,
+    aligned_buf: AlignedBuffer<BUFFER_SIZE>,
 }
 
-impl<const PAGE_SIZE: usize, const WRITE_SIZE: usize> BootLoader<PAGE_SIZE, WRITE_SIZE> {
+impl<const BUFFER_SIZE: usize> BootLoader<BUFFER_SIZE> {
     /// Create a new bootloader instance using the supplied partitions for active, dfu and state.
     pub fn new(active: Partition, dfu: Partition, state: Partition) -> Self {
         Self {
             boot: embassy_boot::BootLoader::new(active, dfu, state),
-            magic: AlignedBuffer([0; WRITE_SIZE]),
-            page: AlignedBuffer([0; PAGE_SIZE]),
+            aligned_buf: AlignedBuffer([0; BUFFER_SIZE]),
         }
     }
 
     /// Inspect the bootloader state and perform actions required before booting, such as swapping
     /// firmware.
     pub fn prepare<F: FlashConfig>(&mut self, flash: &mut F) -> usize {
-        match self.boot.prepare_boot(flash, self.magic.as_mut(), self.page.as_mut()) {
+        match self.boot.prepare_boot(flash, self.aligned_buf.as_mut()) {
             Ok(_) => embassy_stm32::flash::FLASH_BASE + self.boot.boot_address(),
             Err(_) => panic!("boot prepare error!"),
         }
@@ -49,7 +47,7 @@ impl<const PAGE_SIZE: usize, const WRITE_SIZE: usize> BootLoader<PAGE_SIZE, WRIT
     }
 }
 
-impl<const PAGE_SIZE: usize, const WRITE_SIZE: usize> Default for BootLoader<PAGE_SIZE, WRITE_SIZE> {
+impl<const BUFFER_SIZE: usize> Default for BootLoader<BUFFER_SIZE> {
     /// Create a new bootloader instance using parameters from linker script
     fn default() -> Self {
         extern "C" {
diff --git a/examples/boot/bootloader/nrf/src/main.rs b/examples/boot/bootloader/nrf/src/main.rs
index aca3b857a..8818a23b8 100644
--- a/examples/boot/bootloader/nrf/src/main.rs
+++ b/examples/boot/bootloader/nrf/src/main.rs
@@ -27,9 +27,11 @@ fn main() -> ! {
     wdt_config.run_during_sleep = true;
     wdt_config.run_during_debug_halt = false;
 
-    let start = bl.prepare(&mut SingleFlashConfig::new(&mut BootFlash::<_, 4096>::new(
-        WatchdogFlash::start(Nvmc::new(p.NVMC), p.WDT, wdt_config),
-    )));
+    let start = bl.prepare(&mut SingleFlashConfig::new(&mut BootFlash::new(WatchdogFlash::start(
+        Nvmc::new(p.NVMC),
+        p.WDT,
+        wdt_config,
+    ))));
     unsafe { bl.load(start) }
 }
 
diff --git a/examples/boot/bootloader/rp/src/main.rs b/examples/boot/bootloader/rp/src/main.rs
index fb7f0522b..8129591fa 100644
--- a/examples/boot/bootloader/rp/src/main.rs
+++ b/examples/boot/bootloader/rp/src/main.rs
@@ -5,7 +5,6 @@ use cortex_m_rt::{entry, exception};
 #[cfg(feature = "defmt")]
 use defmt_rtt as _;
 use embassy_boot_rp::*;
-use embassy_rp::flash::ERASE_SIZE;
 use embassy_time::Duration;
 
 const FLASH_SIZE: usize = 2 * 1024 * 1024;
@@ -24,7 +23,7 @@ fn main() -> ! {
 
     let mut bl: BootLoader = BootLoader::default();
     let flash = WatchdogFlash::<FLASH_SIZE>::start(p.FLASH, p.WATCHDOG, Duration::from_secs(8));
-    let mut flash = BootFlash::<_, ERASE_SIZE>::new(flash);
+    let mut flash = BootFlash::new(flash);
     let start = bl.prepare(&mut SingleFlashConfig::new(&mut flash));
     core::mem::drop(flash);
 
diff --git a/examples/boot/bootloader/stm32/src/main.rs b/examples/boot/bootloader/stm32/src/main.rs
index 4b17cd799..b8027d19a 100644
--- a/examples/boot/bootloader/stm32/src/main.rs
+++ b/examples/boot/bootloader/stm32/src/main.rs
@@ -5,7 +5,7 @@ use cortex_m_rt::{entry, exception};
 #[cfg(feature = "defmt")]
 use defmt_rtt as _;
 use embassy_boot_stm32::*;
-use embassy_stm32::flash::{Flash, ERASE_SIZE, ERASE_VALUE, WRITE_SIZE};
+use embassy_stm32::flash::{Flash, ERASE_SIZE};
 
 #[entry]
 fn main() -> ! {
@@ -19,9 +19,9 @@ fn main() -> ! {
         }
     */
 
-    let mut bl: BootLoader<ERASE_SIZE, WRITE_SIZE> = BootLoader::default();
+    let mut bl: BootLoader<ERASE_SIZE> = BootLoader::default();
     let flash = Flash::new(p.FLASH);
-    let mut flash = BootFlash::<_, ERASE_SIZE, ERASE_VALUE>::new(flash);
+    let mut flash = BootFlash::new(flash);
     let start = bl.prepare(&mut SingleFlashConfig::new(&mut flash));
     core::mem::drop(flash);
     unsafe { bl.load(start) }