merge embassy/master

12 files changed, 1523 insertions, 1422 deletions

diff --git a/embassy-boot/boot/Cargo.toml b/embassy-boot/boot/Cargo.toml
index 04409cdc7..39f501570 100644
--- a/embassy-boot/boot/Cargo.toml
+++ b/embassy-boot/boot/Cargo.toml
@@ -24,6 +24,7 @@ features = ["defmt"]
 
 [dependencies]
 defmt = { version = "0.3", optional = true }
+digest = "0.10"
 log = { version = "0.4", optional = true  }
 ed25519-dalek = { version = "1.0.1", default_features = false, features = ["u32_backend"], optional = true }
 embassy-sync = { version = "0.1.0", path = "../../embassy-sync" }
@@ -37,6 +38,7 @@ log = "0.4"
 env_logger = "0.9"
 rand = "0.7" # ed25519-dalek v1.0.1 depends on this exact version
 futures = { version = "0.3", features = ["executor"] }
+sha1 = "0.10.5"
 
 [dev-dependencies.ed25519-dalek]
 default_features = false
@@ -50,4 +52,4 @@ ed25519-salty = ["dep:salty", "_verify"]
 nightly = ["dep:embedded-storage-async"]
 
 #Internal features
-_verify = []
\ No newline at end of file
+_verify = []
diff --git a/embassy-boot/boot/src/boot_loader.rs b/embassy-boot/boot/src/boot_loader.rs
new file mode 100644
index 000000000..b959de2c4
--- /dev/null
+++ b/embassy-boot/boot/src/boot_loader.rs
@@ -0,0 +1,533 @@
+use embedded_storage::nor_flash::{ErrorType, NorFlash, NorFlashError, NorFlashErrorKind, ReadNorFlash};
+
+use crate::{Partition, State, BOOT_MAGIC, SWAP_MAGIC};
+
+/// Errors returned by bootloader
+#[derive(PartialEq, Eq, Debug)]
+pub enum BootError {
+    /// Error from flash.
+    Flash(NorFlashErrorKind),
+    /// Invalid bootloader magic
+    BadMagic,
+}
+
+#[cfg(feature = "defmt")]
+impl defmt::Format for BootError {
+    fn format(&self, fmt: defmt::Formatter) {
+        match self {
+            BootError::Flash(_) => defmt::write!(fmt, "BootError::Flash(_)"),
+            BootError::BadMagic => defmt::write!(fmt, "BootError::BadMagic"),
+        }
+    }
+}
+
+impl<E> From<E> for BootError
+where
+    E: NorFlashError,
+{
+    fn from(error: E) -> Self {
+        BootError::Flash(error.kind())
+    }
+}
+
+/// 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: NorFlash;
+    /// Flash type used for the active partition.
+    type ACTIVE: NorFlash;
+    /// Flash type used for the dfu partition.
+    type DFU: NorFlash;
+
+    /// Return flash instance used to write/read to/from active partition.
+    fn active(&mut self) -> &mut Self::ACTIVE;
+    /// Return flash instance used to write/read to/from dfu partition.
+    fn dfu(&mut self) -> &mut Self::DFU;
+    /// Return flash instance used to write/read to/from bootloader state.
+    fn state(&mut self) -> &mut Self::STATE;
+}
+
+trait FlashConfigEx {
+    fn page_size() -> u32;
+}
+
+impl<T: FlashConfig> FlashConfigEx for T {
+    /// Get the page size which is the "unit of operation" within the bootloader.
+    fn page_size() -> u32 {
+        core::cmp::max(T::ACTIVE::ERASE_SIZE, T::DFU::ERASE_SIZE) as u32
+    }
+}
+
+/// 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.
+    // | Range    | Description                                                                      |
+    // | 0..1     | Magic indicating bootloader state. BOOT_MAGIC means boot, SWAP_MAGIC means swap. |
+    // | 1..2     | Progress validity. ERASE_VALUE means valid, !ERASE_VALUE means invalid.          |
+    // | 2..2 + N | Progress index used while swapping or reverting                                  |
+    state: Partition,
+    // Location of the partition which will be booted from
+    active: Partition,
+    // Location of the partition which will be swapped in when requested
+    dfu: Partition,
+}
+
+impl BootLoader {
+    /// Create a new instance of a bootloader with the given partitions.
+    ///
+    /// - All partitions must be aligned with the PAGE_SIZE const generic parameter.
+    /// - The dfu partition must be at least PAGE_SIZE bigger than the active partition.
+    pub fn new(active: Partition, dfu: Partition, state: Partition) -> Self {
+        Self { active, dfu, state }
+    }
+
+    /// Return the offset of the active partition into the active flash.
+    pub fn boot_address(&self) -> usize {
+        self.active.from as usize
+    }
+
+    /// Perform necessary boot preparations like swapping images.
+    ///
+    /// 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.
+    /// The swap index contains the copy progress, as to allow continuation of the copy process on
+    /// power failure. The index counter is represented within 1 or more pages (depending on total
+    /// flash size), where a page X is considered swapped if index at location (X + WRITE_SIZE)
+    /// contains a zero value. This ensures that index updates can be performed atomically and
+    /// avoid a situation where the wrong index value is set (page write size is "atomic").
+    ///
+    /// +-----------+------------+--------+--------+--------+--------+
+    /// | Partition | Swap Index | Page 0 | Page 1 | Page 3 | Page 4 |
+    /// +-----------+------------+--------+--------+--------+--------+
+    /// |    Active |          0 |      1 |      2 |      3 |      - |
+    /// |       DFU |          0 |      3 |      2 |      1 |      X |
+    /// +-----------+------------+--------+--------+--------+--------+
+    ///
+    /// The algorithm starts by copying 'backwards', and after the first step, the layout is
+    /// as follows:
+    ///
+    /// +-----------+------------+--------+--------+--------+--------+
+    /// | Partition | Swap Index | Page 0 | Page 1 | Page 3 | Page 4 |
+    /// +-----------+------------+--------+--------+--------+--------+
+    /// |    Active |          1 |      1 |      2 |      1 |      - |
+    /// |       DFU |          1 |      3 |      2 |      1 |      3 |
+    /// +-----------+------------+--------+--------+--------+--------+
+    ///
+    /// The next iteration performs the same steps
+    ///
+    /// +-----------+------------+--------+--------+--------+--------+
+    /// | Partition | Swap Index | Page 0 | Page 1 | Page 3 | Page 4 |
+    /// +-----------+------------+--------+--------+--------+--------+
+    /// |    Active |          2 |      1 |      2 |      1 |      - |
+    /// |       DFU |          2 |      3 |      2 |      2 |      3 |
+    /// +-----------+------------+--------+--------+--------+--------+
+    ///
+    /// And again until we're done
+    ///
+    /// +-----------+------------+--------+--------+--------+--------+
+    /// | Partition | Swap Index | Page 0 | Page 1 | Page 3 | Page 4 |
+    /// +-----------+------------+--------+--------+--------+--------+
+    /// |    Active |          3 |      3 |      2 |      1 |      - |
+    /// |       DFU |          3 |      3 |      1 |      2 |      3 |
+    /// +-----------+------------+--------+--------+--------+--------+
+    ///
+    /// REVERTING
+    ///
+    /// The reverting algorithm uses the swap index to discover that images were swapped, but that
+    /// the application failed to mark the boot successful. In this case, the revert algorithm will
+    /// run.
+    ///
+    /// The revert index is located separately from the swap index, to ensure that revert can continue
+    /// on power failure.
+    ///
+    /// The revert algorithm works forwards, by starting copying into the 'unused' DFU page at the start.
+    ///
+    /// +-----------+--------------+--------+--------+--------+--------+
+    /// | Partition | Revert Index | Page 0 | Page 1 | Page 3 | Page 4 |
+    //*/
+    /// +-----------+--------------+--------+--------+--------+--------+
+    /// |    Active |            3 |      1 |      2 |      1 |      - |
+    /// |       DFU |            3 |      3 |      1 |      2 |      3 |
+    /// +-----------+--------------+--------+--------+--------+--------+
+    ///
+    ///
+    /// +-----------+--------------+--------+--------+--------+--------+
+    /// | Partition | Revert Index | Page 0 | Page 1 | Page 3 | Page 4 |
+    /// +-----------+--------------+--------+--------+--------+--------+
+    /// |    Active |            3 |      1 |      2 |      1 |      - |
+    /// |       DFU |            3 |      3 |      2 |      2 |      3 |
+    /// +-----------+--------------+--------+--------+--------+--------+
+    ///
+    /// +-----------+--------------+--------+--------+--------+--------+
+    /// | Partition | Revert Index | Page 0 | Page 1 | Page 3 | Page 4 |
+    /// +-----------+--------------+--------+--------+--------+--------+
+    /// |    Active |            3 |      1 |      2 |      3 |      - |
+    /// |       DFU |            3 |      3 |      2 |      1 |      3 |
+    /// +-----------+--------------+--------+--------+--------+--------+
+    ///
+    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!(0, P::page_size() % aligned_buf.len() as u32);
+        assert_eq!(0, P::page_size() % P::ACTIVE::WRITE_SIZE as u32);
+        assert_eq!(0, P::page_size() % P::ACTIVE::ERASE_SIZE as u32);
+        assert_eq!(0, P::page_size() % P::DFU::WRITE_SIZE as u32);
+        assert_eq!(0, P::page_size() % P::DFU::ERASE_SIZE as u32);
+        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, 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, aligned_buf)? {
+                trace!("Swapping");
+                self.swap(p, aligned_buf)?;
+                trace!("Swapping done");
+            } else {
+                trace!("Reverting");
+                self.revert(p, aligned_buf)?;
+
+                let state_flash = p.state();
+                let state_word = &mut aligned_buf[..P::STATE::WRITE_SIZE];
+
+                // Invalidate progress
+                state_word.fill(!P::STATE_ERASE_VALUE);
+                self.state
+                    .write_blocking(state_flash, P::STATE::WRITE_SIZE as u32, state_word)?;
+
+                // Clear magic and progress
+                self.state.wipe_blocking(state_flash)?;
+
+                // Set 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, aligned_buf: &mut [u8]) -> Result<bool, BootError> {
+        let page_count = (self.active.size() / P::page_size()) as usize;
+        let progress = self.current_progress(p, aligned_buf)?;
+
+        Ok(progress >= page_count * 2)
+    }
+
+    fn current_progress<P: FlashConfig>(&mut self, config: &mut P, aligned_buf: &mut [u8]) -> Result<usize, BootError> {
+        let write_size = P::STATE::WRITE_SIZE as u32;
+        let max_index = (((self.state.size() - write_size) / write_size) - 2) as usize;
+        let state_flash = config.state();
+        let state_word = &mut aligned_buf[..write_size as usize];
+
+        self.state.read_blocking(state_flash, write_size, 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 * write_size, state_word)?;
+
+            if state_word.iter().any(|&b| b == P::STATE_ERASE_VALUE) {
+                return Ok(index);
+            }
+        }
+        Ok(max_index)
+    }
+
+    fn update_progress<P: FlashConfig>(
+        &mut self,
+        progress_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 + progress_index) as u32 * P::STATE::WRITE_SIZE as u32,
+            state_word,
+        )?;
+        Ok(())
+    }
+
+    fn copy_page_once_to_active<P: FlashConfig>(
+        &mut self,
+        progress_index: usize,
+        from_offset: u32,
+        to_offset: u32,
+        p: &mut P,
+        aligned_buf: &mut [u8],
+    ) -> Result<(), BootError> {
+        if self.current_progress(p, aligned_buf)? <= progress_index {
+            let page_size = P::page_size() as u32;
+
+            self.active
+                .erase_blocking(p.active(), to_offset, to_offset + page_size)?;
+
+            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(progress_index, p, aligned_buf)?;
+        }
+        Ok(())
+    }
+
+    fn copy_page_once_to_dfu<P: FlashConfig>(
+        &mut self,
+        progress_index: usize,
+        from_offset: u32,
+        to_offset: u32,
+        p: &mut P,
+        aligned_buf: &mut [u8],
+    ) -> Result<(), BootError> {
+        if self.current_progress(p, aligned_buf)? <= progress_index {
+            let page_size = P::page_size() as u32;
+
+            self.dfu
+                .erase_blocking(p.dfu(), to_offset as u32, to_offset + page_size)?;
+
+            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(progress_index, p, aligned_buf)?;
+        }
+        Ok(())
+    }
+
+    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.size() / page_size;
+        for page_num in 0..page_count {
+            let progress_index = (page_num * 2) as usize;
+
+            // Copy active page to the 'next' DFU page.
+            let active_from_offset = (page_count - 1 - page_num) * page_size;
+            let dfu_to_offset = (page_count - page_num) * page_size;
+            //trace!("Copy active {} to dfu {}", active_from_offset, dfu_to_offset);
+            self.copy_page_once_to_dfu(progress_index, 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;
+            let dfu_from_offset = (page_count - 1 - page_num) * page_size;
+            //trace!("Copy dfy {} to active {}", dfu_from_offset, active_to_offset);
+            self.copy_page_once_to_active(progress_index + 1, dfu_from_offset, active_to_offset, p, aligned_buf)?;
+        }
+
+        Ok(())
+    }
+
+    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.size() / page_size;
+        for page_num in 0..page_count {
+            let progress_index = (page_count * 2 + page_num * 2) as usize;
+
+            // Copy the bad active page to the DFU page
+            let active_from_offset = page_num * page_size;
+            let dfu_to_offset = page_num * page_size;
+            self.copy_page_once_to_dfu(progress_index, 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;
+            let dfu_from_offset = (page_num + 1) * page_size;
+            self.copy_page_once_to_active(progress_index + 1, dfu_from_offset, active_to_offset, p, aligned_buf)?;
+        }
+
+        Ok(())
+    }
+
+    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 !state_word.iter().any(|&b| b != SWAP_MAGIC) {
+            Ok(State::Swap)
+        } else {
+            Ok(State::Boot)
+        }
+    }
+}
+
+fn assert_partitions(active: Partition, dfu: Partition, state: Partition, page_size: u32, state_write_size: usize) {
+    assert_eq!(active.size() % page_size, 0);
+    assert_eq!(dfu.size() % page_size, 0);
+    assert!(dfu.size() - active.size() >= page_size);
+    assert!(2 + 2 * (active.size() / page_size) <= state.size() / state_write_size as u32);
+}
+
+/// A flash wrapper implementing the Flash and embedded_storage traits.
+pub struct BootFlash<F>
+where
+    F: NorFlash,
+{
+    flash: F,
+}
+
+impl<F> BootFlash<F>
+where
+    F: NorFlash,
+{
+    /// Create a new instance of a bootable flash
+    pub fn new(flash: F) -> Self {
+        Self { flash }
+    }
+}
+
+impl<F> ErrorType for BootFlash<F>
+where
+    F: NorFlash,
+{
+    type Error = F::Error;
+}
+
+impl<F> NorFlash for BootFlash<F>
+where
+    F: NorFlash,
+{
+    const WRITE_SIZE: usize = F::WRITE_SIZE;
+    const ERASE_SIZE: usize = F::ERASE_SIZE;
+
+    fn erase(&mut self, from: u32, to: u32) -> Result<(), Self::Error> {
+        F::erase(&mut self.flash, from, to)
+    }
+
+    fn write(&mut self, offset: u32, bytes: &[u8]) -> Result<(), Self::Error> {
+        F::write(&mut self.flash, offset, bytes)
+    }
+}
+
+impl<F> ReadNorFlash for BootFlash<F>
+where
+    F: NorFlash,
+{
+    const READ_SIZE: usize = F::READ_SIZE;
+
+    fn read(&mut self, offset: u32, bytes: &mut [u8]) -> Result<(), Self::Error> {
+        F::read(&mut self.flash, offset, bytes)
+    }
+
+    fn capacity(&self) -> usize {
+        F::capacity(&self.flash)
+    }
+}
+
+/// Convenience provider that uses a single flash for all partitions.
+pub struct SingleFlashConfig<'a, F>
+where
+    F: NorFlash,
+{
+    flash: &'a mut F,
+}
+
+impl<'a, F> SingleFlashConfig<'a, F>
+where
+    F: NorFlash,
+{
+    /// Create a provider for a single flash.
+    pub fn new(flash: &'a mut F) -> Self {
+        Self { flash }
+    }
+}
+
+impl<'a, F> FlashConfig for SingleFlashConfig<'a, F>
+where
+    F: NorFlash,
+{
+    type STATE = F;
+    type ACTIVE = F;
+    type DFU = F;
+
+    fn active(&mut self) -> &mut Self::STATE {
+        self.flash
+    }
+    fn dfu(&mut self) -> &mut Self::ACTIVE {
+        self.flash
+    }
+    fn state(&mut self) -> &mut Self::DFU {
+        self.flash
+    }
+}
+
+/// Convenience flash provider that uses separate flash instances for each partition.
+pub struct MultiFlashConfig<'a, ACTIVE, STATE, DFU>
+where
+    ACTIVE: NorFlash,
+    STATE: NorFlash,
+    DFU: NorFlash,
+{
+    active: &'a mut ACTIVE,
+    state: &'a mut STATE,
+    dfu: &'a mut DFU,
+}
+
+impl<'a, ACTIVE, STATE, DFU> MultiFlashConfig<'a, ACTIVE, STATE, DFU>
+where
+    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 {
+        Self { active, state, dfu }
+    }
+}
+
+impl<'a, ACTIVE, STATE, DFU> FlashConfig for MultiFlashConfig<'a, ACTIVE, STATE, DFU>
+where
+    ACTIVE: NorFlash,
+    STATE: NorFlash,
+    DFU: NorFlash,
+{
+    type STATE = STATE;
+    type ACTIVE = ACTIVE;
+    type DFU = DFU;
+
+    fn active(&mut self) -> &mut Self::ACTIVE {
+        self.active
+    }
+    fn dfu(&mut self) -> &mut Self::DFU {
+        self.dfu
+    }
+    fn state(&mut self) -> &mut Self::STATE {
+        self.state
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    #[should_panic]
+    fn test_range_asserts() {
+        const ACTIVE: Partition = Partition::new(4096, 4194304);
+        const DFU: Partition = Partition::new(4194304, 2 * 4194304);
+        const STATE: Partition = Partition::new(0, 4096);
+        assert_partitions(ACTIVE, DFU, STATE, 4096, 4);
+    }
+}
diff --git a/embassy-boot/boot/src/digest_adapters/ed25519_dalek.rs b/embassy-boot/boot/src/digest_adapters/ed25519_dalek.rs
new file mode 100644
index 000000000..a184d1c51
--- /dev/null
+++ b/embassy-boot/boot/src/digest_adapters/ed25519_dalek.rs
@@ -0,0 +1,30 @@
+use digest::typenum::U64;
+use digest::{FixedOutput, HashMarker, OutputSizeUser, Update};
+use ed25519_dalek::Digest as _;
+
+pub struct Sha512(ed25519_dalek::Sha512);
+
+impl Default for Sha512 {
+    fn default() -> Self {
+        Self(ed25519_dalek::Sha512::new())
+    }
+}
+
+impl Update for Sha512 {
+    fn update(&mut self, data: &[u8]) {
+        self.0.update(data)
+    }
+}
+
+impl FixedOutput for Sha512 {
+    fn finalize_into(self, out: &mut digest::Output<Self>) {
+        let result = self.0.finalize();
+        out.as_mut_slice().copy_from_slice(result.as_slice())
+    }
+}
+
+impl OutputSizeUser for Sha512 {
+    type OutputSize = U64;
+}
+
+impl HashMarker for Sha512 {}
diff --git a/embassy-boot/boot/src/digest_adapters/mod.rs b/embassy-boot/boot/src/digest_adapters/mod.rs
new file mode 100644
index 000000000..9b4b4b60c
--- /dev/null
+++ b/embassy-boot/boot/src/digest_adapters/mod.rs
@@ -0,0 +1,5 @@
+#[cfg(feature = "ed25519-dalek")]
+pub(crate) mod ed25519_dalek;
+
+#[cfg(feature = "ed25519-salty")]
+pub(crate) mod salty;
diff --git a/embassy-boot/boot/src/digest_adapters/salty.rs b/embassy-boot/boot/src/digest_adapters/salty.rs
new file mode 100644
index 000000000..2b5dcf3af
--- /dev/null
+++ b/embassy-boot/boot/src/digest_adapters/salty.rs
@@ -0,0 +1,29 @@
+use digest::typenum::U64;
+use digest::{FixedOutput, HashMarker, OutputSizeUser, Update};
+
+pub struct Sha512(salty::Sha512);
+
+impl Default for Sha512 {
+    fn default() -> Self {
+        Self(salty::Sha512::new())
+    }
+}
+
+impl Update for Sha512 {
+    fn update(&mut self, data: &[u8]) {
+        self.0.update(data)
+    }
+}
+
+impl FixedOutput for Sha512 {
+    fn finalize_into(self, out: &mut digest::Output<Self>) {
+        let result = self.0.finalize();
+        out.as_mut_slice().copy_from_slice(result.as_slice())
+    }
+}
+
+impl OutputSizeUser for Sha512 {
+    type OutputSize = U64;
+}
+
+impl HashMarker for Sha512 {}
diff --git a/embassy-boot/boot/src/firmware_updater.rs b/embassy-boot/boot/src/firmware_updater.rs
new file mode 100644
index 000000000..a2f822f4a
--- /dev/null
+++ b/embassy-boot/boot/src/firmware_updater.rs
@@ -0,0 +1,534 @@
+use digest::Digest;
+use embedded_storage::nor_flash::{NorFlash, NorFlashError, NorFlashErrorKind};
+use embedded_storage_async::nor_flash::NorFlash as AsyncNorFlash;
+
+use crate::{Partition, State, BOOT_MAGIC, SWAP_MAGIC};
+
+/// Errors returned by FirmwareUpdater
+#[derive(Debug)]
+pub enum FirmwareUpdaterError {
+    /// Error from flash.
+    Flash(NorFlashErrorKind),
+    /// Signature errors.
+    Signature(signature::Error),
+}
+
+#[cfg(feature = "defmt")]
+impl defmt::Format for FirmwareUpdaterError {
+    fn format(&self, fmt: defmt::Formatter) {
+        match self {
+            FirmwareUpdaterError::Flash(_) => defmt::write!(fmt, "FirmwareUpdaterError::Flash(_)"),
+            FirmwareUpdaterError::Signature(_) => defmt::write!(fmt, "FirmwareUpdaterError::Signature(_)"),
+        }
+    }
+}
+
+impl<E> From<E> for FirmwareUpdaterError
+where
+    E: NorFlashError,
+{
+    fn from(error: E) -> Self {
+        FirmwareUpdaterError::Flash(error.kind())
+    }
+}
+
+/// FirmwareUpdater is an application API for interacting with the BootLoader without the ability to
+/// 'mess up' the internal bootloader state
+pub struct FirmwareUpdater {
+    state: Partition,
+    dfu: Partition,
+}
+
+impl Default for FirmwareUpdater {
+    fn default() -> Self {
+        extern "C" {
+            static __bootloader_state_start: u32;
+            static __bootloader_state_end: u32;
+            static __bootloader_dfu_start: u32;
+            static __bootloader_dfu_end: u32;
+        }
+
+        let dfu = unsafe {
+            Partition::new(
+                &__bootloader_dfu_start as *const u32 as u32,
+                &__bootloader_dfu_end as *const u32 as u32,
+            )
+        };
+        let state = unsafe {
+            Partition::new(
+                &__bootloader_state_start as *const u32 as u32,
+                &__bootloader_state_end as *const u32 as u32,
+            )
+        };
+
+        trace!("DFU: 0x{:x} - 0x{:x}", dfu.from, dfu.to);
+        trace!("STATE: 0x{:x} - 0x{:x}", state.from, state.to);
+        FirmwareUpdater::new(dfu, state)
+    }
+}
+
+impl FirmwareUpdater {
+    /// Create a firmware updater instance with partition ranges for the update and state partitions.
+    pub const fn new(dfu: Partition, state: Partition) -> Self {
+        Self { dfu, state }
+    }
+
+    /// 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,
+        state_flash: &mut F,
+        aligned: &mut [u8],
+    ) -> Result<State, FirmwareUpdaterError> {
+        self.state.read(state_flash, 0, aligned).await?;
+
+        if !aligned.iter().any(|&b| b != SWAP_MAGIC) {
+            Ok(State::Swap)
+        } else {
+            Ok(State::Boot)
+        }
+    }
+
+    /// Verify the DFU given a public key. If there is an error then DO NOT
+    /// proceed with updating the firmware as it must be signed with a
+    /// corresponding private key (otherwise it could be malicious firmware).
+    ///
+    /// Mark to trigger firmware swap on next boot if verify suceeds.
+    ///
+    /// If the "ed25519-salty" feature is set (or another similar feature) then the signature is expected to have
+    /// been generated from a SHA-512 digest of the firmware bytes.
+    ///
+    /// If no signature feature is set then this method will always return a
+    /// signature error.
+    ///
+    /// # Safety
+    ///
+    /// The `_aligned` buffer must have a size of F::WRITE_SIZE, and follow the alignment rules for the flash being read from
+    /// and written to.
+    #[cfg(feature = "_verify")]
+    pub async fn verify_and_mark_updated<F: AsyncNorFlash>(
+        &mut self,
+        _state_and_dfu_flash: &mut F,
+        _public_key: &[u8],
+        _signature: &[u8],
+        _update_len: u32,
+        _aligned: &mut [u8],
+    ) -> Result<(), FirmwareUpdaterError> {
+        assert_eq!(_aligned.len(), F::WRITE_SIZE);
+        assert!(_update_len <= self.dfu.size());
+
+        #[cfg(feature = "ed25519-dalek")]
+        {
+            use ed25519_dalek::{PublicKey, Signature, SignatureError, Verifier};
+
+            use crate::digest_adapters::ed25519_dalek::Sha512;
+
+            let into_signature_error = |e: SignatureError| FirmwareUpdaterError::Signature(e.into());
+
+            let public_key = PublicKey::from_bytes(_public_key).map_err(into_signature_error)?;
+            let signature = Signature::from_bytes(_signature).map_err(into_signature_error)?;
+
+            let mut message = [0; 64];
+            self.hash::<_, Sha512>(_state_and_dfu_flash, _update_len, _aligned, &mut message)
+                .await?;
+
+            public_key.verify(&message, &signature).map_err(into_signature_error)?
+        }
+        #[cfg(feature = "ed25519-salty")]
+        {
+            use salty::constants::{PUBLICKEY_SERIALIZED_LENGTH, SIGNATURE_SERIALIZED_LENGTH};
+            use salty::{PublicKey, Signature};
+
+            use crate::digest_adapters::salty::Sha512;
+
+            fn into_signature_error<E>(_: E) -> FirmwareUpdaterError {
+                FirmwareUpdaterError::Signature(signature::Error::default())
+            }
+
+            let public_key: [u8; PUBLICKEY_SERIALIZED_LENGTH] = _public_key.try_into().map_err(into_signature_error)?;
+            let public_key = PublicKey::try_from(&public_key).map_err(into_signature_error)?;
+            let signature: [u8; SIGNATURE_SERIALIZED_LENGTH] = _signature.try_into().map_err(into_signature_error)?;
+            let signature = Signature::try_from(&signature).map_err(into_signature_error)?;
+
+            let mut message = [0; 64];
+            self.hash::<_, Sha512>(_state_and_dfu_flash, _update_len, _aligned, &mut message)
+                .await?;
+
+            let r = public_key.verify(&message, &signature);
+            trace!(
+                "Verifying with public key {}, signature {} and message {} yields ok: {}",
+                public_key.to_bytes(),
+                signature.to_bytes(),
+                message,
+                r.is_ok()
+            );
+            r.map_err(into_signature_error)?
+        }
+
+        self.set_magic(_aligned, SWAP_MAGIC, _state_and_dfu_flash).await
+    }
+
+    /// Verify the update in DFU with any digest.
+    pub async fn hash<F: AsyncNorFlash, D: Digest>(
+        &mut self,
+        dfu_flash: &mut F,
+        update_len: u32,
+        chunk_buf: &mut [u8],
+        output: &mut [u8],
+    ) -> Result<(), FirmwareUpdaterError> {
+        let mut digest = D::new();
+        for offset in (0..update_len).step_by(chunk_buf.len()) {
+            self.dfu.read(dfu_flash, offset, chunk_buf).await?;
+            let len = core::cmp::min((update_len - offset) as usize, chunk_buf.len());
+            digest.update(&chunk_buf[..len]);
+        }
+        output.copy_from_slice(digest.finalize().as_slice());
+        Ok(())
+    }
+
+    /// 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.
+    #[cfg(not(feature = "_verify"))]
+    pub async fn mark_updated<F: AsyncNorFlash>(
+        &mut self,
+        state_flash: &mut F,
+        aligned: &mut [u8],
+    ) -> Result<(), FirmwareUpdaterError> {
+        assert_eq!(aligned.len(), F::WRITE_SIZE);
+        self.set_magic(aligned, SWAP_MAGIC, state_flash).await
+    }
+
+    /// 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 async fn mark_booted<F: AsyncNorFlash>(
+        &mut self,
+        state_flash: &mut F,
+        aligned: &mut [u8],
+    ) -> Result<(), FirmwareUpdaterError> {
+        assert_eq!(aligned.len(), F::WRITE_SIZE);
+        self.set_magic(aligned, BOOT_MAGIC, state_flash).await
+    }
+
+    async fn set_magic<F: AsyncNorFlash>(
+        &mut self,
+        aligned: &mut [u8],
+        magic: u8,
+        state_flash: &mut F,
+    ) -> Result<(), FirmwareUpdaterError> {
+        self.state.read(state_flash, 0, aligned).await?;
+
+        if aligned.iter().any(|&b| b != magic) {
+            // Read progress validity
+            self.state.read(state_flash, F::WRITE_SIZE as u32, aligned).await?;
+
+            // FIXME: Do not make this assumption.
+            const STATE_ERASE_VALUE: u8 = 0xFF;
+
+            if aligned.iter().any(|&b| b != STATE_ERASE_VALUE) {
+                // The current progress validity marker is invalid
+            } else {
+                // Invalidate progress
+                aligned.fill(!STATE_ERASE_VALUE);
+                self.state.write(state_flash, F::WRITE_SIZE as u32, aligned).await?;
+            }
+
+            // Clear magic and progress
+            self.state.wipe(state_flash).await?;
+
+            // Set magic
+            aligned.fill(magic);
+            self.state.write(state_flash, 0, aligned).await?;
+        }
+        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 async fn write_firmware<F: AsyncNorFlash>(
+        &mut self,
+        offset: usize,
+        data: &[u8],
+        dfu_flash: &mut F,
+    ) -> Result<(), FirmwareUpdaterError> {
+        assert!(data.len() >= F::ERASE_SIZE);
+
+        self.dfu
+            .erase(dfu_flash, offset as u32, (offset + data.len()) as u32)
+            .await?;
+
+        self.dfu.write(dfu_flash, offset as u32, data).await?;
+
+        Ok(())
+    }
+
+    /// Prepare for an incoming DFU update by erasing the entire DFU area and
+    /// returning its `Partition`.
+    ///
+    /// 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,
+        dfu_flash: &mut F,
+    ) -> Result<Partition, FirmwareUpdaterError> {
+        self.dfu.wipe(dfu_flash).await?;
+
+        Ok(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,
+        state_flash: &mut F,
+        aligned: &mut [u8],
+    ) -> Result<State, FirmwareUpdaterError> {
+        self.state.read_blocking(state_flash, 0, aligned)?;
+
+        if !aligned.iter().any(|&b| b != SWAP_MAGIC) {
+            Ok(State::Swap)
+        } else {
+            Ok(State::Boot)
+        }
+    }
+
+    /// Verify the DFU given a public key. If there is an error then DO NOT
+    /// proceed with updating the firmware as it must be signed with a
+    /// corresponding private key (otherwise it could be malicious firmware).
+    ///
+    /// Mark to trigger firmware swap on next boot if verify suceeds.
+    ///
+    /// If the "ed25519-salty" feature is set (or another similar feature) then the signature is expected to have
+    /// been generated from a SHA-512 digest of the firmware bytes.
+    ///
+    /// If no signature feature is set then this method will always return a
+    /// signature error.
+    ///
+    /// # Safety
+    ///
+    /// The `_aligned` buffer must have a size of F::WRITE_SIZE, and follow the alignment rules for the flash being read from
+    /// and written to.
+    #[cfg(feature = "_verify")]
+    pub fn verify_and_mark_updated_blocking<F: NorFlash>(
+        &mut self,
+        _state_and_dfu_flash: &mut F,
+        _public_key: &[u8],
+        _signature: &[u8],
+        _update_len: u32,
+        _aligned: &mut [u8],
+    ) -> Result<(), FirmwareUpdaterError> {
+        assert_eq!(_aligned.len(), F::WRITE_SIZE);
+        assert!(_update_len <= self.dfu.size());
+
+        #[cfg(feature = "ed25519-dalek")]
+        {
+            use ed25519_dalek::{PublicKey, Signature, SignatureError, Verifier};
+
+            use crate::digest_adapters::ed25519_dalek::Sha512;
+
+            let into_signature_error = |e: SignatureError| FirmwareUpdaterError::Signature(e.into());
+
+            let public_key = PublicKey::from_bytes(_public_key).map_err(into_signature_error)?;
+            let signature = Signature::from_bytes(_signature).map_err(into_signature_error)?;
+
+            let mut message = [0; 64];
+            self.hash_blocking::<_, Sha512>(_state_and_dfu_flash, _update_len, _aligned, &mut message)?;
+
+            public_key.verify(&message, &signature).map_err(into_signature_error)?
+        }
+        #[cfg(feature = "ed25519-salty")]
+        {
+            use salty::constants::{PUBLICKEY_SERIALIZED_LENGTH, SIGNATURE_SERIALIZED_LENGTH};
+            use salty::{PublicKey, Signature};
+
+            use crate::digest_adapters::salty::Sha512;
+
+            fn into_signature_error<E>(_: E) -> FirmwareUpdaterError {
+                FirmwareUpdaterError::Signature(signature::Error::default())
+            }
+
+            let public_key: [u8; PUBLICKEY_SERIALIZED_LENGTH] = _public_key.try_into().map_err(into_signature_error)?;
+            let public_key = PublicKey::try_from(&public_key).map_err(into_signature_error)?;
+            let signature: [u8; SIGNATURE_SERIALIZED_LENGTH] = _signature.try_into().map_err(into_signature_error)?;
+            let signature = Signature::try_from(&signature).map_err(into_signature_error)?;
+
+            let mut message = [0; 64];
+            self.hash_blocking::<_, Sha512>(_state_and_dfu_flash, _update_len, _aligned, &mut message)?;
+
+            let r = public_key.verify(&message, &signature);
+            trace!(
+                "Verifying with public key {}, signature {} and message {} yields ok: {}",
+                public_key.to_bytes(),
+                signature.to_bytes(),
+                message,
+                r.is_ok()
+            );
+            r.map_err(into_signature_error)?
+        }
+
+        self.set_magic_blocking(_aligned, SWAP_MAGIC, _state_and_dfu_flash)
+    }
+
+    /// Verify the update in DFU with any digest.
+    pub fn hash_blocking<F: NorFlash, D: Digest>(
+        &mut self,
+        dfu_flash: &mut F,
+        update_len: u32,
+        chunk_buf: &mut [u8],
+        output: &mut [u8],
+    ) -> Result<(), FirmwareUpdaterError> {
+        let mut digest = D::new();
+        for offset in (0..update_len).step_by(chunk_buf.len()) {
+            self.dfu.read_blocking(dfu_flash, offset, chunk_buf)?;
+            let len = core::cmp::min((update_len - offset) as usize, chunk_buf.len());
+            digest.update(&chunk_buf[..len]);
+        }
+        output.copy_from_slice(digest.finalize().as_slice());
+        Ok(())
+    }
+
+    /// 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.
+    #[cfg(not(feature = "_verify"))]
+    pub fn mark_updated_blocking<F: NorFlash>(
+        &mut self,
+        state_flash: &mut F,
+        aligned: &mut [u8],
+    ) -> Result<(), FirmwareUpdaterError> {
+        assert_eq!(aligned.len(), F::WRITE_SIZE);
+        self.set_magic_blocking(aligned, SWAP_MAGIC, state_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,
+        state_flash: &mut F,
+        aligned: &mut [u8],
+    ) -> Result<(), FirmwareUpdaterError> {
+        assert_eq!(aligned.len(), F::WRITE_SIZE);
+        self.set_magic_blocking(aligned, BOOT_MAGIC, state_flash)
+    }
+
+    fn set_magic_blocking<F: NorFlash>(
+        &mut self,
+        aligned: &mut [u8],
+        magic: u8,
+        state_flash: &mut F,
+    ) -> Result<(), FirmwareUpdaterError> {
+        self.state.read_blocking(state_flash, 0, aligned)?;
+
+        if aligned.iter().any(|&b| b != magic) {
+            // Read progress validity
+            self.state.read_blocking(state_flash, F::WRITE_SIZE as u32, aligned)?;
+
+            // FIXME: Do not make this assumption.
+            const STATE_ERASE_VALUE: u8 = 0xFF;
+
+            if aligned.iter().any(|&b| b != STATE_ERASE_VALUE) {
+                // The current progress validity marker is invalid
+            } else {
+                // Invalidate progress
+                aligned.fill(!STATE_ERASE_VALUE);
+                self.state.write_blocking(state_flash, F::WRITE_SIZE as u32, aligned)?;
+            }
+
+            // Clear magic and progress
+            self.state.wipe_blocking(state_flash)?;
+
+            // Set magic
+            aligned.fill(magic);
+            self.state.write_blocking(state_flash, 0, 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],
+        dfu_flash: &mut F,
+    ) -> Result<(), FirmwareUpdaterError> {
+        assert!(data.len() >= F::ERASE_SIZE);
+
+        self.dfu
+            .erase_blocking(dfu_flash, offset as u32, (offset + data.len()) as u32)?;
+
+        self.dfu.write_blocking(dfu_flash, offset as u32, data)?;
+
+        Ok(())
+    }
+
+    /// Prepare for an incoming DFU update by erasing the entire DFU area and
+    /// returning its `Partition`.
+    ///
+    /// 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<Partition, FirmwareUpdaterError> {
+        self.dfu.wipe_blocking(flash)?;
+
+        Ok(self.dfu)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use futures::executor::block_on;
+    use sha1::{Digest, Sha1};
+
+    use super::*;
+    use crate::mem_flash::MemFlash;
+
+    #[test]
+    fn can_verify_sha1() {
+        const STATE: Partition = Partition::new(0, 4096);
+        const DFU: Partition = Partition::new(65536, 131072);
+
+        let mut flash = MemFlash::<131072, 4096, 8>::default();
+
+        let update = [0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66];
+        let mut to_write = [0; 4096];
+        to_write[..7].copy_from_slice(update.as_slice());
+
+        let mut updater = FirmwareUpdater::new(DFU, STATE);
+        block_on(updater.write_firmware(0, to_write.as_slice(), &mut flash)).unwrap();
+        let mut chunk_buf = [0; 2];
+        let mut hash = [0; 20];
+        block_on(updater.hash::<_, Sha1>(&mut flash, update.len() as u32, &mut chunk_buf, &mut hash)).unwrap();
+
+        assert_eq!(Sha1::digest(update).as_slice(), hash);
+    }
+}
diff --git a/embassy-boot/boot/src/lib.rs b/embassy-boot/boot/src/lib.rs
index 7ce0c664a..e268d8883 100644
--- a/embassy-boot/boot/src/lib.rs
+++ b/embassy-boot/boot/src/lib.rs
@@ -5,36 +5,18 @@
 #![doc = include_str!("../README.md")]
 mod fmt;
 
-use embedded_storage::nor_flash::{ErrorType, NorFlash, NorFlashError, NorFlashErrorKind, ReadNorFlash};
+mod boot_loader;
+mod digest_adapters;
+mod firmware_updater;
+mod mem_flash;
+mod partition;
 
-#[cfg(feature = "nightly")]
-use embedded_storage_async::nor_flash::NorFlash as AsyncNorFlash;
+pub use boot_loader::{BootError, BootFlash, BootLoader, FlashConfig, MultiFlashConfig, SingleFlashConfig};
+pub use firmware_updater::{FirmwareUpdater, FirmwareUpdaterError};
+pub use partition::Partition;
 
-const BOOT_MAGIC: u8 = 0xD0;
-const SWAP_MAGIC: u8 = 0xF0;
-
-/// A region in flash used by the bootloader.
-#[derive(Copy, Clone, Debug)]
-#[cfg_attr(feature = "defmt", derive(defmt::Format))]
-pub struct Partition {
-    /// Start of the flash region.
-    pub from: usize,
-    /// End of the flash region.
-    pub to: usize,
-}
-
-impl Partition {
-    /// Create a new partition with the provided range
-    pub const fn new(from: usize, to: usize) -> Self {
-        Self { from, to }
-    }
-
-    /// Return the length of the partition
-    #[allow(clippy::len_without_is_empty)]
-    pub const fn len(&self) -> usize {
-        self.to - self.from
-    }
-}
+pub(crate) const BOOT_MAGIC: u8 = 0xD0;
+pub(crate) const SWAP_MAGIC: u8 = 0xF0;
 
 /// The state of the bootloader after running prepare.
 #[derive(PartialEq, Eq, Debug)]
@@ -46,34 +28,6 @@ pub enum State {
     Swap,
 }
 
-/// Errors returned by bootloader
-#[derive(PartialEq, Eq, Debug)]
-pub enum BootError {
-    /// Error from flash.
-    Flash(NorFlashErrorKind),
-    /// Invalid bootloader magic
-    BadMagic,
-}
-
-#[cfg(feature = "defmt")]
-impl defmt::Format for BootError {
-    fn format(&self, fmt: defmt::Formatter) {
-        match self {
-            BootError::Flash(_) => defmt::write!(fmt, "BootError::Flash(_)"),
-            BootError::BadMagic => defmt::write!(fmt, "BootError::BadMagic"),
-        }
-    }
-}
-
-impl<E> From<E> for BootError
-where
-    E: NorFlashError,
-{
-    fn from(error: E) -> Self {
-        BootError::Flash(error.kind())
-    }
-}
-
 /// Buffer aligned to 32 byte boundary, largest known alignment requirement for embassy-boot.
 #[repr(align(32))]
 pub struct AlignedBuffer<const N: usize>(pub [u8; N]);
@@ -90,1128 +44,12 @@ impl<const N: usize> AsMut<[u8]> for AlignedBuffer<N> {
     }
 }
 
-/// Extension of the embedded-storage flash type information with block size and erase value.
-pub trait Flash: NorFlash + ReadNorFlash {
-    /// 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
-pub trait FlashConfig {
-    /// Flash type used for the state partition.
-    type STATE: Flash;
-    /// Flash type used for the active partition.
-    type ACTIVE: Flash;
-    /// Flash type used for the dfu partition.
-    type DFU: Flash;
-
-    /// Return flash instance used to write/read to/from active partition.
-    fn active(&mut self) -> &mut Self::ACTIVE;
-    /// Return flash instance used to write/read to/from dfu partition.
-    fn dfu(&mut self) -> &mut Self::DFU;
-    /// Return flash instance used to write/read to/from bootloader state.
-    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.
-pub struct BootLoader {
-    // Page with current state of bootloader. The state partition has the following format:
-    // | Range          | Description                                                                      |
-    // | 0 - WRITE_SIZE | Magic indicating bootloader state. BOOT_MAGIC means boot, SWAP_MAGIC means swap. |
-    // | WRITE_SIZE - N | Progress index used while swapping or reverting                                  |
-    state: Partition,
-    // Location of the partition which will be booted from
-    active: Partition,
-    // Location of the partition which will be swapped in when requested
-    dfu: Partition,
-}
-
-impl BootLoader {
-    /// Create a new instance of a bootloader with the given partitions.
-    ///
-    /// - All partitions must be aligned with the PAGE_SIZE const generic parameter.
-    /// - The dfu partition must be at least PAGE_SIZE bigger than the active partition.
-    pub fn new(active: Partition, dfu: Partition, state: Partition) -> Self {
-        Self { active, dfu, state }
-    }
-
-    /// Return the boot address for the active partition.
-    pub fn boot_address(&self) -> usize {
-        self.active.from
-    }
-
-    /// Perform necessary boot preparations like swapping images.
-    ///
-    /// The DFU partition is assumed to be 1 page bigger than the active partition for the swap
-    /// algorithm to work correctly.
-    ///
-    /// SWAPPING
-    ///
-    /// Assume a flash size of 3 pages for the active partition, and 4 pages for the DFU partition.
-    /// The swap index contains the copy progress, as to allow continuation of the copy process on
-    /// power failure. The index counter is represented within 1 or more pages (depending on total
-    /// flash size), where a page X is considered swapped if index at location (X + WRITE_SIZE)
-    /// contains a zero value. This ensures that index updates can be performed atomically and
-    /// avoid a situation where the wrong index value is set (page write size is "atomic").
-    ///
-    /// +-----------+------------+--------+--------+--------+--------+
-    /// | Partition | Swap Index | Page 0 | Page 1 | Page 3 | Page 4 |
-    /// +-----------+------------+--------+--------+--------+--------+
-    /// |    Active |          0 |      1 |      2 |      3 |      - |
-    /// |       DFU |          0 |      3 |      2 |      1 |      X |
-    /// +-----------+------------+--------+--------+--------+--------+
-    ///
-    /// The algorithm starts by copying 'backwards', and after the first step, the layout is
-    /// as follows:
-    ///
-    /// +-----------+------------+--------+--------+--------+--------+
-    /// | Partition | Swap Index | Page 0 | Page 1 | Page 3 | Page 4 |
-    /// +-----------+------------+--------+--------+--------+--------+
-    /// |    Active |          1 |      1 |      2 |      1 |      - |
-    /// |       DFU |          1 |      3 |      2 |      1 |      3 |
-    /// +-----------+------------+--------+--------+--------+--------+
-    ///
-    /// The next iteration performs the same steps
-    ///
-    /// +-----------+------------+--------+--------+--------+--------+
-    /// | Partition | Swap Index | Page 0 | Page 1 | Page 3 | Page 4 |
-    /// +-----------+------------+--------+--------+--------+--------+
-    /// |    Active |          2 |      1 |      2 |      1 |      - |
-    /// |       DFU |          2 |      3 |      2 |      2 |      3 |
-    /// +-----------+------------+--------+--------+--------+--------+
-    ///
-    /// And again until we're done
-    ///
-    /// +-----------+------------+--------+--------+--------+--------+
-    /// | Partition | Swap Index | Page 0 | Page 1 | Page 3 | Page 4 |
-    /// +-----------+------------+--------+--------+--------+--------+
-    /// |    Active |          3 |      3 |      2 |      1 |      - |
-    /// |       DFU |          3 |      3 |      1 |      2 |      3 |
-    /// +-----------+------------+--------+--------+--------+--------+
-    ///
-    /// REVERTING
-    ///
-    /// The reverting algorithm uses the swap index to discover that images were swapped, but that
-    /// the application failed to mark the boot successful. In this case, the revert algorithm will
-    /// run.
-    ///
-    /// The revert index is located separately from the swap index, to ensure that revert can continue
-    /// on power failure.
-    ///
-    /// The revert algorithm works forwards, by starting copying into the 'unused' DFU page at the start.
-    ///
-    /// +-----------+--------------+--------+--------+--------+--------+
-    /// | Partition | Revert Index | Page 0 | Page 1 | Page 3 | Page 4 |
-    //*/
-    /// +-----------+--------------+--------+--------+--------+--------+
-    /// |    Active |            3 |      1 |      2 |      1 |      - |
-    /// |       DFU |            3 |      3 |      1 |      2 |      3 |
-    /// +-----------+--------------+--------+--------+--------+--------+
-    ///
-    ///
-    /// +-----------+--------------+--------+--------+--------+--------+
-    /// | Partition | Revert Index | Page 0 | Page 1 | Page 3 | Page 4 |
-    /// +-----------+--------------+--------+--------+--------+--------+
-    /// |    Active |            3 |      1 |      2 |      1 |      - |
-    /// |       DFU |            3 |      3 |      2 |      2 |      3 |
-    /// +-----------+--------------+--------+--------+--------+--------+
-    ///
-    /// +-----------+--------------+--------+--------+--------+--------+
-    /// | Partition | Revert Index | Page 0 | Page 1 | Page 3 | Page 4 |
-    /// +-----------+--------------+--------+--------+--------+--------+
-    /// |    Active |            3 |      1 |      2 |      3 |      - |
-    /// |       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> {
-        // 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);
-
-        // Copy contents from partition N to active
-        let state = self.read_state(p, magic)?;
-        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)? {
-                trace!("Swapping");
-                self.swap(p, magic, page)?;
-                trace!("Swapping done");
-            } else {
-                trace!("Reverting");
-                self.revert(p, magic, page)?;
-
-                // Overwrite magic and reset progress
-                let fstate = p.state();
-                magic.fill(!P::STATE::ERASE_VALUE);
-                fstate.write(self.state.from as u32, magic)?;
-                fstate.erase(self.state.from as u32, self.state.to as u32)?;
-
-                magic.fill(BOOT_MAGIC);
-                fstate.write(self.state.from as u32, magic)?;
-            }
-        }
-        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)?;
-
-        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) - 1;
-        aligned.fill(!P::STATE::ERASE_VALUE);
-
-        let flash = config.state();
-        for i in 0..max_index {
-            flash.read((self.state.from + write_size + i * write_size) as u32, aligned)?;
-
-            if aligned.iter().any(|&b| b == P::STATE::ERASE_VALUE) {
-                return Ok(i);
-            }
-        }
-        Ok(max_index)
-    }
-
-    fn update_progress<P: FlashConfig>(&mut self, idx: usize, p: &mut P, magic: &mut [u8]) -> Result<(), BootError> {
-        let flash = p.state();
-        let write_size = magic.len();
-        let w = self.state.from + write_size + idx * write_size;
-
-        let aligned = magic;
-        aligned.fill(!P::STATE::ERASE_VALUE);
-        flash.write(w as u32, aligned)?;
-        Ok(())
-    }
-
-    fn active_addr(&self, n: usize, page_size: usize) -> usize {
-        self.active.from + n * page_size
-    }
-
-    fn dfu_addr(&self, n: usize, page_size: usize) -> usize {
-        self.dfu.from + n * page_size
-    }
-
-    fn copy_page_once_to_active<P: FlashConfig>(
-        &mut self,
-        idx: usize,
-        from_page: usize,
-        to_page: usize,
-        p: &mut P,
-        magic: &mut [u8],
-        page: &mut [u8],
-    ) -> Result<(), BootError> {
-        let buf = page;
-        if self.current_progress(p, magic)? <= idx {
-            let mut offset = from_page;
-            for chunk in buf.chunks_mut(P::DFU::BLOCK_SIZE) {
-                p.dfu().read(offset as u32, chunk)?;
-                offset += chunk.len();
-            }
-
-            p.active().erase(to_page as u32, (to_page + buf.len()) as u32)?;
-
-            let mut offset = to_page;
-            for chunk in buf.chunks(P::ACTIVE::BLOCK_SIZE) {
-                p.active().write(offset as u32, chunk)?;
-                offset += chunk.len();
-            }
-            self.update_progress(idx, p, magic)?;
-        }
-        Ok(())
-    }
-
-    fn copy_page_once_to_dfu<P: FlashConfig>(
-        &mut self,
-        idx: usize,
-        from_page: usize,
-        to_page: usize,
-        p: &mut P,
-        magic: &mut [u8],
-        page: &mut [u8],
-    ) -> Result<(), BootError> {
-        let buf = page;
-        if self.current_progress(p, magic)? <= idx {
-            let mut offset = from_page;
-            for chunk in buf.chunks_mut(P::ACTIVE::BLOCK_SIZE) {
-                p.active().read(offset as u32, chunk)?;
-                offset += chunk.len();
-            }
-
-            p.dfu().erase(to_page as u32, (to_page + buf.len()) as u32)?;
-
-            let mut offset = to_page;
-            for chunk in buf.chunks(P::DFU::BLOCK_SIZE) {
-                p.dfu().write(offset as u32, chunk)?;
-                offset += chunk.len();
-            }
-            self.update_progress(idx, p, magic)?;
-        }
-        Ok(())
-    }
-
-    fn swap<P: FlashConfig>(&mut self, p: &mut P, magic: &mut [u8], page: &mut [u8]) -> Result<(), BootError> {
-        let page_size = page.len();
-        let page_count = self.active.len() / page_size;
-        trace!("Page count: {}", page_count);
-        for page_num in 0..page_count {
-            trace!("COPY PAGE {}", page_num);
-            // Copy active page to the 'next' DFU page.
-            let active_page = self.active_addr(page_count - 1 - page_num, page_size);
-            let dfu_page = self.dfu_addr(page_count - page_num, page_size);
-            //trace!("Copy active {} to dfu {}", active_page, dfu_page);
-            self.copy_page_once_to_dfu(page_num * 2, active_page, dfu_page, p, magic, page)?;
-
-            // Copy DFU page to the active page
-            let active_page = self.active_addr(page_count - 1 - page_num, page_size);
-            let dfu_page = self.dfu_addr(page_count - 1 - page_num, page_size);
-            //trace!("Copy dfy {} to active {}", dfu_page, active_page);
-            self.copy_page_once_to_active(page_num * 2 + 1, dfu_page, active_page, p, magic, page)?;
-        }
-
-        Ok(())
-    }
-
-    fn revert<P: FlashConfig>(&mut self, p: &mut P, magic: &mut [u8], page: &mut [u8]) -> Result<(), BootError> {
-        let page_size = page.len();
-        let page_count = self.active.len() / page_size;
-        for page_num in 0..page_count {
-            // Copy the bad active page to the DFU page
-            let active_page = self.active_addr(page_num, page_size);
-            let dfu_page = self.dfu_addr(page_num, page_size);
-            self.copy_page_once_to_dfu(page_count * 2 + page_num * 2, active_page, dfu_page, p, magic, page)?;
-
-            // Copy the DFU page back to the active page
-            let active_page = self.active_addr(page_num, page_size);
-            let dfu_page = self.dfu_addr(page_num + 1, page_size);
-            self.copy_page_once_to_active(page_count * 2 + page_num * 2 + 1, dfu_page, active_page, p, magic, page)?;
-        }
-
-        Ok(())
-    }
-
-    fn read_state<P: FlashConfig>(&mut self, config: &mut P, magic: &mut [u8]) -> Result<State, BootError> {
-        let flash = config.state();
-        flash.read(self.state.from as u32, magic)?;
-
-        if !magic.iter().any(|&b| b != SWAP_MAGIC) {
-            Ok(State::Swap)
-        } else {
-            Ok(State::Boot)
-        }
-    }
-}
-
-fn assert_partitions(active: Partition, dfu: Partition, state: Partition, page_size: usize, write_size: usize) {
-    assert_eq!(active.len() % page_size, 0);
-    assert_eq!(dfu.len() % page_size, 0);
-    assert!(dfu.len() - active.len() >= page_size);
-    assert!(2 * (active.len() / page_size) <= (state.len() - write_size) / write_size);
-}
-
-/// Convenience provider that uses a single flash for all partitions.
-pub struct SingleFlashConfig<'a, F>
-where
-    F: Flash,
-{
-    flash: &'a mut F,
-}
-
-impl<'a, F> SingleFlashConfig<'a, F>
-where
-    F: Flash,
-{
-    /// Create a provider for a single flash.
-    pub fn new(flash: &'a mut F) -> Self {
-        Self { flash }
-    }
-}
-
-impl<'a, F> FlashConfig for SingleFlashConfig<'a, F>
-where
-    F: Flash,
-{
-    type STATE = F;
-    type ACTIVE = F;
-    type DFU = F;
-
-    fn active(&mut self) -> &mut Self::STATE {
-        self.flash
-    }
-    fn dfu(&mut self) -> &mut Self::ACTIVE {
-        self.flash
-    }
-    fn state(&mut self) -> &mut Self::DFU {
-        self.flash
-    }
-}
-
-/// A flash wrapper implementing the Flash and embedded_storage traits.
-pub struct BootFlash<F, const BLOCK_SIZE: usize, 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>
-where
-    F: NorFlash + ReadNorFlash,
-{
-    /// Create a new instance of a bootable flash
-    pub fn new(flash: F) -> Self {
-        Self { flash }
-    }
-}
-
-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>
-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>
-where
-    F: ReadNorFlash + NorFlash,
-{
-    const WRITE_SIZE: usize = F::WRITE_SIZE;
-    const ERASE_SIZE: usize = F::ERASE_SIZE;
-
-    fn erase(&mut self, from: u32, to: u32) -> Result<(), Self::Error> {
-        F::erase(&mut self.flash, from, to)
-    }
-
-    fn write(&mut self, offset: u32, bytes: &[u8]) -> Result<(), Self::Error> {
-        F::write(&mut self.flash, offset, bytes)
-    }
-}
-
-impl<F, const BLOCK_SIZE: usize, const ERASE_VALUE: u8> ReadNorFlash for BootFlash<F, BLOCK_SIZE, ERASE_VALUE>
-where
-    F: ReadNorFlash + NorFlash,
-{
-    const READ_SIZE: usize = F::READ_SIZE;
-
-    fn read(&mut self, offset: u32, bytes: &mut [u8]) -> Result<(), Self::Error> {
-        F::read(&mut self.flash, offset, bytes)
-    }
-
-    fn capacity(&self) -> usize {
-        F::capacity(&self.flash)
-    }
-}
-
-/// 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: &'a mut ACTIVE,
-    state: &'a mut STATE,
-    dfu: &'a mut DFU,
-}
-
-impl<'a, ACTIVE, STATE, DFU> MultiFlashConfig<'a, ACTIVE, STATE, DFU>
-where
-    ACTIVE: Flash,
-    STATE: Flash,
-    DFU: Flash,
-{
-    /// 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 {
-        Self { active, state, dfu }
-    }
-}
-
-impl<'a, ACTIVE, STATE, DFU> FlashConfig for MultiFlashConfig<'a, ACTIVE, STATE, DFU>
-where
-    ACTIVE: Flash,
-    STATE: Flash,
-    DFU: Flash,
-{
-    type STATE = STATE;
-    type ACTIVE = ACTIVE;
-    type DFU = DFU;
-
-    fn active(&mut self) -> &mut Self::ACTIVE {
-        self.active
-    }
-    fn dfu(&mut self) -> &mut Self::DFU {
-        self.dfu
-    }
-    fn state(&mut self) -> &mut Self::STATE {
-        self.state
-    }
-}
-/// Errors returned by FirmwareUpdater
-#[derive(Debug)]
-pub enum FirmwareUpdaterError {
-    /// Error from flash.
-    Flash(NorFlashErrorKind),
-    /// Signature errors.
-    Signature(signature::Error),
-}
-
-#[cfg(feature = "defmt")]
-impl defmt::Format for FirmwareUpdaterError {
-    fn format(&self, fmt: defmt::Formatter) {
-        match self {
-            FirmwareUpdaterError::Flash(_) => defmt::write!(fmt, "FirmwareUpdaterError::Flash(_)"),
-            FirmwareUpdaterError::Signature(_) => defmt::write!(fmt, "FirmwareUpdaterError::Signature(_)"),
-        }
-    }
-}
-
-impl<E> From<E> for FirmwareUpdaterError
-where
-    E: NorFlashError,
-{
-    fn from(error: E) -> Self {
-        FirmwareUpdaterError::Flash(error.kind())
-    }
-}
-
-/// FirmwareUpdater is an application API for interacting with the BootLoader without the ability to
-/// 'mess up' the internal bootloader state
-pub struct FirmwareUpdater {
-    state: Partition,
-    dfu: Partition,
-}
-
-impl Default for FirmwareUpdater {
-    fn default() -> Self {
-        extern "C" {
-            static __bootloader_state_start: u32;
-            static __bootloader_state_end: u32;
-            static __bootloader_dfu_start: u32;
-            static __bootloader_dfu_end: u32;
-        }
-
-        let dfu = unsafe {
-            Partition::new(
-                &__bootloader_dfu_start as *const u32 as usize,
-                &__bootloader_dfu_end as *const u32 as usize,
-            )
-        };
-        let state = unsafe {
-            Partition::new(
-                &__bootloader_state_start as *const u32 as usize,
-                &__bootloader_state_end as *const u32 as usize,
-            )
-        };
-
-        trace!("DFU: 0x{:x} - 0x{:x}", dfu.from, dfu.to);
-        trace!("STATE: 0x{:x} - 0x{:x}", state.from, state.to);
-        FirmwareUpdater::new(dfu, state)
-    }
-}
-
-impl FirmwareUpdater {
-    /// Create a firmware updater instance with partition ranges for the update and state partitions.
-    pub const fn new(dfu: Partition, state: Partition) -> Self {
-        Self { dfu, state }
-    }
-
-    /// Return the length of the DFU area
-    pub fn firmware_len(&self) -> usize {
-        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`.
-    #[cfg(feature = "nightly")]
-    pub async fn get_state<F: AsyncNorFlash>(
-        &mut self,
-        flash: &mut F,
-        aligned: &mut [u8],
-    ) -> Result<State, FirmwareUpdaterError> {
-        flash.read(self.state.from as u32, aligned).await?;
-
-        if !aligned.iter().any(|&b| b != SWAP_MAGIC) {
-            Ok(State::Swap)
-        } else {
-            Ok(State::Boot)
-        }
-    }
-
-    /// Verify the DFU given a public key. If there is an error then DO NOT
-    /// proceed with updating the firmware as it must be signed with a
-    /// corresponding private key (otherwise it could be malicious firmware).
-    ///
-    /// Mark to trigger firmware swap on next boot if verify suceeds.
-    ///
-    /// If the "ed25519-salty" feature is set (or another similar feature) then the signature is expected to have
-    /// been generated from a SHA-512 digest of the firmware bytes.
-    ///
-    /// If no signature feature is set then this method will always return a
-    /// signature error.
-    ///
-    /// # Safety
-    ///
-    /// The `_aligned` buffer must have a size of F::WRITE_SIZE, and follow the alignment rules for the flash being read from
-    /// and written to.
-    #[cfg(feature = "_verify")]
-    pub async fn verify_and_mark_updated<F: AsyncNorFlash>(
-        &mut self,
-        _flash: &mut F,
-        _public_key: &[u8],
-        _signature: &[u8],
-        _update_len: usize,
-        _aligned: &mut [u8],
-    ) -> Result<(), FirmwareUpdaterError> {
-        let _end = self.dfu.from + _update_len;
-        let _read_size = _aligned.len();
-
-        assert_eq!(_aligned.len(), F::WRITE_SIZE);
-        assert!(_end <= self.dfu.to);
-
-        #[cfg(feature = "ed25519-dalek")]
-        {
-            use ed25519_dalek::{Digest, PublicKey, Sha512, Signature, SignatureError, Verifier};
-
-            let into_signature_error = |e: SignatureError| FirmwareUpdaterError::Signature(e.into());
-
-            let public_key = PublicKey::from_bytes(_public_key).map_err(into_signature_error)?;
-            let signature = Signature::from_bytes(_signature).map_err(into_signature_error)?;
-
-            let mut digest = Sha512::new();
-
-            let mut offset = self.dfu.from;
-            let last_offset = _end / _read_size * _read_size;
-
-            while offset < last_offset {
-                _flash.read(offset as u32, _aligned).await?;
-                digest.update(&_aligned);
-                offset += _read_size;
-            }
-
-            let remaining = _end % _read_size;
-
-            if remaining > 0 {
-                _flash.read(last_offset as u32, _aligned).await?;
-                digest.update(&_aligned[0..remaining]);
-            }
-
-            public_key
-                .verify(&digest.finalize(), &signature)
-                .map_err(into_signature_error)?
-        }
-        #[cfg(feature = "ed25519-salty")]
-        {
-            use salty::constants::{PUBLICKEY_SERIALIZED_LENGTH, SIGNATURE_SERIALIZED_LENGTH};
-            use salty::{PublicKey, Sha512, Signature};
-
-            fn into_signature_error<E>(_: E) -> FirmwareUpdaterError {
-                FirmwareUpdaterError::Signature(signature::Error::default())
-            }
-
-            let public_key: [u8; PUBLICKEY_SERIALIZED_LENGTH] = _public_key.try_into().map_err(into_signature_error)?;
-            let public_key = PublicKey::try_from(&public_key).map_err(into_signature_error)?;
-            let signature: [u8; SIGNATURE_SERIALIZED_LENGTH] = _signature.try_into().map_err(into_signature_error)?;
-            let signature = Signature::try_from(&signature).map_err(into_signature_error)?;
-
-            let mut digest = Sha512::new();
-
-            let mut offset = self.dfu.from;
-            let last_offset = _end / _read_size * _read_size;
-
-            while offset < last_offset {
-                _flash.read(offset as u32, _aligned).await?;
-                digest.update(&_aligned);
-                offset += _read_size;
-            }
-
-            let remaining = _end % _read_size;
-
-            if remaining > 0 {
-                _flash.read(last_offset as u32, _aligned).await?;
-                digest.update(&_aligned[0..remaining]);
-            }
-
-            let message = digest.finalize();
-            let r = public_key.verify(&message, &signature);
-            trace!(
-                "Verifying with public key {}, signature {} and message {} yields ok: {}",
-                public_key.to_bytes(),
-                signature.to_bytes(),
-                message,
-                r.is_ok()
-            );
-            r.map_err(into_signature_error)?
-        }
-
-        self.set_magic(_aligned, SWAP_MAGIC, _flash).await
-    }
-
-    /// 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.
-    #[cfg(not(feature = "_verify"))]
-    #[cfg(feature = "nightly")]
-    pub async fn mark_updated<F: AsyncNorFlash>(
-        &mut self,
-        flash: &mut F,
-        aligned: &mut [u8],
-    ) -> Result<(), FirmwareUpdaterError> {
-        assert_eq!(aligned.len(), F::WRITE_SIZE);
-        self.set_magic(aligned, SWAP_MAGIC, flash).await
-    }
-
-    /// 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.
-    #[cfg(feature = "nightly")]
-    pub async fn mark_booted<F: AsyncNorFlash>(
-        &mut self,
-        flash: &mut F,
-        aligned: &mut [u8],
-    ) -> Result<(), FirmwareUpdaterError> {
-        assert_eq!(aligned.len(), F::WRITE_SIZE);
-        self.set_magic(aligned, BOOT_MAGIC, flash).await
-    }
-
-    #[cfg(feature = "nightly")]
-    async fn set_magic<F: AsyncNorFlash>(
-        &mut self,
-        aligned: &mut [u8],
-        magic: u8,
-        flash: &mut F,
-    ) -> Result<(), FirmwareUpdaterError> {
-        flash.read(self.state.from as u32, aligned).await?;
-
-        if aligned.iter().any(|&b| b != magic) {
-            aligned.fill(0);
-
-            flash.write(self.state.from as u32, aligned).await?;
-            flash.erase(self.state.from as u32, self.state.to as u32).await?;
-
-            aligned.fill(magic);
-            flash.write(self.state.from as u32, aligned).await?;
-        }
-        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.
-    #[cfg(feature = "nightly")]
-    pub async fn write_firmware<F: AsyncNorFlash>(
-        &mut self,
-        offset: usize,
-        data: &[u8],
-        flash: &mut F,
-        block_size: usize,
-    ) -> Result<(), FirmwareUpdaterError> {
-        assert!(data.len() >= F::ERASE_SIZE);
-
-        flash
-            .erase(
-                (self.dfu.from + offset) as u32,
-                (self.dfu.from + offset + data.len()) as u32,
-            )
-            .await?;
-
-        trace!(
-            "Erased from {} to {}",
-            self.dfu.from + offset,
-            self.dfu.from + offset + data.len()
-        );
-
-        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.
-    #[cfg(feature = "nightly")]
-    pub async fn prepare_update<F: AsyncNorFlash>(
-        &mut self,
-        flash: &mut F,
-    ) -> Result<FirmwareWriter, FirmwareUpdaterError> {
-        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, FirmwareUpdaterError> {
-        flash.read(self.state.from as u32, aligned)?;
-
-        if !aligned.iter().any(|&b| b != SWAP_MAGIC) {
-            Ok(State::Swap)
-        } else {
-            Ok(State::Boot)
-        }
-    }
-
-    /// Verify the DFU given a public key. If there is an error then DO NOT
-    /// proceed with updating the firmware as it must be signed with a
-    /// corresponding private key (otherwise it could be malicious firmware).
-    ///
-    /// Mark to trigger firmware swap on next boot if verify suceeds.
-    ///
-    /// If the "ed25519-salty" feature is set (or another similar feature) then the signature is expected to have
-    /// been generated from a SHA-512 digest of the firmware bytes.
-    ///
-    /// If no signature feature is set then this method will always return a
-    /// signature error.
-    ///
-    /// # Safety
-    ///
-    /// The `_aligned` buffer must have a size of F::WRITE_SIZE, and follow the alignment rules for the flash being read from
-    /// and written to.
-    #[cfg(feature = "_verify")]
-    pub fn verify_and_mark_updated_blocking<F: NorFlash>(
-        &mut self,
-        _flash: &mut F,
-        _public_key: &[u8],
-        _signature: &[u8],
-        _update_len: usize,
-        _aligned: &mut [u8],
-    ) -> Result<(), FirmwareUpdaterError> {
-        let _end = self.dfu.from + _update_len;
-        let _read_size = _aligned.len();
-
-        assert_eq!(_aligned.len(), F::WRITE_SIZE);
-        assert!(_end <= self.dfu.to);
-
-        #[cfg(feature = "ed25519-dalek")]
-        {
-            use ed25519_dalek::{Digest, PublicKey, Sha512, Signature, SignatureError, Verifier};
-
-            let into_signature_error = |e: SignatureError| FirmwareUpdaterError::Signature(e.into());
-
-            let public_key = PublicKey::from_bytes(_public_key).map_err(into_signature_error)?;
-            let signature = Signature::from_bytes(_signature).map_err(into_signature_error)?;
-
-            let mut digest = Sha512::new();
-
-            let mut offset = self.dfu.from;
-            let last_offset = _end / _read_size * _read_size;
-
-            while offset < last_offset {
-                _flash.read(offset as u32, _aligned)?;
-                digest.update(&_aligned);
-                offset += _read_size;
-            }
-
-            let remaining = _end % _read_size;
-
-            if remaining > 0 {
-                _flash.read(last_offset as u32, _aligned)?;
-                digest.update(&_aligned[0..remaining]);
-            }
-
-            public_key
-                .verify(&digest.finalize(), &signature)
-                .map_err(into_signature_error)?
-        }
-        #[cfg(feature = "ed25519-salty")]
-        {
-            use salty::constants::{PUBLICKEY_SERIALIZED_LENGTH, SIGNATURE_SERIALIZED_LENGTH};
-            use salty::{PublicKey, Sha512, Signature};
-
-            fn into_signature_error<E>(_: E) -> FirmwareUpdaterError {
-                FirmwareUpdaterError::Signature(signature::Error::default())
-            }
-
-            let public_key: [u8; PUBLICKEY_SERIALIZED_LENGTH] = _public_key.try_into().map_err(into_signature_error)?;
-            let public_key = PublicKey::try_from(&public_key).map_err(into_signature_error)?;
-            let signature: [u8; SIGNATURE_SERIALIZED_LENGTH] = _signature.try_into().map_err(into_signature_error)?;
-            let signature = Signature::try_from(&signature).map_err(into_signature_error)?;
-
-            let mut digest = Sha512::new();
-
-            let mut offset = self.dfu.from;
-            let last_offset = _end / _read_size * _read_size;
-
-            while offset < last_offset {
-                _flash.read(offset as u32, _aligned)?;
-                digest.update(&_aligned);
-                offset += _read_size;
-            }
-
-            let remaining = _end % _read_size;
-
-            if remaining > 0 {
-                _flash.read(last_offset as u32, _aligned)?;
-                digest.update(&_aligned[0..remaining]);
-            }
-
-            let message = digest.finalize();
-            let r = public_key.verify(&message, &signature);
-            trace!(
-                "Verifying with public key {}, signature {} and message {} yields ok: {}",
-                public_key.to_bytes(),
-                signature.to_bytes(),
-                message,
-                r.is_ok()
-            );
-            r.map_err(into_signature_error)?
-        }
-
-        self.set_magic_blocking(_aligned, SWAP_MAGIC, _flash)
-    }
-
-    /// 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.
-    #[cfg(not(feature = "_verify"))]
-    pub fn mark_updated_blocking<F: NorFlash>(
-        &mut self,
-        flash: &mut F,
-        aligned: &mut [u8],
-    ) -> Result<(), FirmwareUpdaterError> {
-        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<(), FirmwareUpdaterError> {
-        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<(), FirmwareUpdaterError> {
-        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<(), FirmwareUpdaterError> {
-        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, FirmwareUpdaterError> {
-        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.
-    #[cfg(feature = "nightly")]
-    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?;
-            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(())
-    }
-
-    /// 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)]
 mod tests {
-    use core::convert::Infallible;
-
-    use embedded_storage::nor_flash::ErrorType;
-    use embedded_storage_async::nor_flash::ReadNorFlash as AsyncReadNorFlash;
     use futures::executor::block_on;
 
     use super::*;
+    use crate::mem_flash::MemFlash;
 
     /*
     #[test]
@@ -1234,18 +72,14 @@ mod tests {
         const ACTIVE: Partition = Partition::new(4096, 61440);
         const DFU: Partition = Partition::new(61440, 122880);
 
-        let mut flash = MemFlash::<131072, 4096, 4>([0xff; 131072]);
-        flash.0[0..4].copy_from_slice(&[BOOT_MAGIC; 4]);
+        let mut flash = MemFlash::<131072, 4096, 4>::default();
+        flash.mem[0..4].copy_from_slice(&[BOOT_MAGIC; 4]);
         let mut flash = SingleFlashConfig::new(&mut flash);
 
         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]
@@ -1254,66 +88,49 @@ mod tests {
         const STATE: Partition = Partition::new(0, 4096);
         const ACTIVE: Partition = Partition::new(4096, 61440);
         const DFU: Partition = Partition::new(61440, 122880);
-        let mut flash = MemFlash::<131072, 4096, 4>([0xff; 131072]);
+        let mut flash = MemFlash::<131072, 4096, 4>::random();
 
-        let original: [u8; ACTIVE.len()] = [rand::random::<u8>(); ACTIVE.len()];
-        let update: [u8; DFU.len()] = [rand::random::<u8>(); DFU.len()];
+        let original = [rand::random::<u8>(); ACTIVE.size() as usize];
+        let update = [rand::random::<u8>(); ACTIVE.size() as usize];
         let mut aligned = [0; 4];
 
-        for i in ACTIVE.from..ACTIVE.to {
-            flash.0[i] = original[i - ACTIVE.from];
-        }
+        flash.program(ACTIVE.from, &original).unwrap();
 
         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, 4096)).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()
         );
 
-        for i in ACTIVE.from..ACTIVE.to {
-            assert_eq!(flash.0[i], update[i - ACTIVE.from], "Index {}", i);
-        }
-
+        flash.assert_eq(ACTIVE.from, &update);
         // First DFU page is untouched
-        for i in DFU.from + 4096..DFU.to {
-            assert_eq!(flash.0[i], original[i - DFU.from - 4096], "Index {}", i);
-        }
+        flash.assert_eq(DFU.from + 4096, &original);
 
         // Running again should cause a revert
         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()
         );
 
-        for i in ACTIVE.from..ACTIVE.to {
-            assert_eq!(flash.0[i], original[i - ACTIVE.from], "Index {}", i);
-        }
-
+        flash.assert_eq(ACTIVE.from, &original);
         // Last page is untouched
-        for i in DFU.from..DFU.to - 4096 {
-            assert_eq!(flash.0[i], update[i - DFU.from], "Index {}", i);
-        }
+        flash.assert_eq(DFU.from, &update);
 
         // Mark as booted
         block_on(updater.mark_booted(&mut flash, &mut aligned)).unwrap();
         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()
         );
     }
@@ -1325,50 +142,34 @@ mod tests {
         const ACTIVE: Partition = Partition::new(4096, 16384);
         const DFU: Partition = Partition::new(0, 16384);
 
-        let mut active = MemFlash::<16384, 4096, 8>([0xff; 16384]);
-        let mut dfu = MemFlash::<16384, 2048, 8>([0xff; 16384]);
-        let mut state = MemFlash::<4096, 128, 4>([0xff; 4096]);
+        let mut active = MemFlash::<16384, 4096, 8>::random();
+        let mut dfu = MemFlash::<16384, 2048, 8>::random();
+        let mut state = MemFlash::<4096, 128, 4>::random();
         let mut aligned = [0; 4];
 
-        let original: [u8; ACTIVE.len()] = [rand::random::<u8>(); ACTIVE.len()];
-        let update: [u8; DFU.len()] = [rand::random::<u8>(); DFU.len()];
+        let original = [rand::random::<u8>(); ACTIVE.size() as usize];
+        let update = [rand::random::<u8>(); ACTIVE.size() as usize];
 
-        for i in ACTIVE.from..ACTIVE.to {
-            active.0[i] = original[i - ACTIVE.from];
-        }
+        active.program(ACTIVE.from, &original).unwrap();
 
         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, chunk.len())).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()
         );
 
-        for i in ACTIVE.from..ACTIVE.to {
-            assert_eq!(active.0[i], update[i - ACTIVE.from], "Index {}", i);
-        }
-
+        active.assert_eq(ACTIVE.from, &update);
         // First DFU page is untouched
-        for i in DFU.from + 4096..DFU.to {
-            assert_eq!(dfu.0[i], original[i - DFU.from - 4096], "Index {}", i);
-        }
+        dfu.assert_eq(DFU.from + 4096, &original);
     }
 
     #[test]
@@ -1379,57 +180,35 @@ mod tests {
         const DFU: Partition = Partition::new(0, 16384);
 
         let mut aligned = [0; 4];
-        let mut active = MemFlash::<16384, 2048, 4>([0xff; 16384]);
-        let mut dfu = MemFlash::<16384, 4096, 8>([0xff; 16384]);
-        let mut state = MemFlash::<4096, 128, 4>([0xff; 4096]);
+        let mut active = MemFlash::<16384, 2048, 4>::random();
+        let mut dfu = MemFlash::<16384, 4096, 8>::random();
+        let mut state = MemFlash::<4096, 128, 4>::random();
 
-        let original: [u8; ACTIVE.len()] = [rand::random::<u8>(); ACTIVE.len()];
-        let update: [u8; DFU.len()] = [rand::random::<u8>(); DFU.len()];
+        let original = [rand::random::<u8>(); ACTIVE.size() as usize];
+        let update = [rand::random::<u8>(); ACTIVE.size() as usize];
 
-        for i in ACTIVE.from..ACTIVE.to {
-            active.0[i] = original[i - ACTIVE.from];
-        }
+        active.program(ACTIVE.from, &original).unwrap();
 
         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, chunk.len())).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()
         );
 
-        for i in ACTIVE.from..ACTIVE.to {
-            assert_eq!(active.0[i], update[i - ACTIVE.from], "Index {}", i);
-        }
-
+        active.assert_eq(ACTIVE.from, &update);
         // First DFU page is untouched
-        for i in DFU.from + 4096..DFU.to {
-            assert_eq!(dfu.0[i], original[i - DFU.from - 4096], "Index {}", i);
-        }
-    }
-
-    #[test]
-    #[should_panic]
-    fn test_range_asserts() {
-        const ACTIVE: Partition = Partition::new(4096, 4194304);
-        const DFU: Partition = Partition::new(4194304, 2 * 4194304);
-        const STATE: Partition = Partition::new(0, 4096);
-        assert_partitions(ACTIVE, DFU, STATE, 4096, 4);
+        dfu.assert_eq(DFU.from + 4096, &original);
     }
 
     #[test]
@@ -1458,13 +237,13 @@ mod tests {
 
         const STATE: Partition = Partition::new(0, 4096);
         const DFU: Partition = Partition::new(4096, 8192);
-        let mut flash = MemFlash::<8192, 4096, 4>([0xff; 8192]);
+        let mut flash = MemFlash::<8192, 4096, 4>::default();
 
         let firmware_len = firmware.len();
 
         let mut write_buf = [0; 4096];
         write_buf[0..firmware_len].copy_from_slice(firmware);
-        NorFlash::write(&mut flash, DFU.from as u32, &write_buf).unwrap();
+        DFU.write_blocking(&mut flash, 0, &write_buf).unwrap();
 
         // On with the test
 
@@ -1476,117 +255,9 @@ mod tests {
             &mut flash,
             &public_key.to_bytes(),
             &signature.to_bytes(),
-            firmware_len,
+            firmware_len as u32,
             &mut aligned,
         ))
         .is_ok());
     }
-    struct MemFlash<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize>([u8; SIZE]);
-
-    impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> NorFlash
-        for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
-    {
-        const WRITE_SIZE: usize = WRITE_SIZE;
-        const ERASE_SIZE: usize = ERASE_SIZE;
-        fn erase(&mut self, from: u32, to: u32) -> Result<(), Self::Error> {
-            let from = from as usize;
-            let to = to as usize;
-            assert!(from % ERASE_SIZE == 0);
-            assert!(to % ERASE_SIZE == 0, "To: {}, erase size: {}", to, ERASE_SIZE);
-            for i in from..to {
-                self.0[i] = 0xFF;
-            }
-            Ok(())
-        }
-
-        fn write(&mut self, offset: u32, data: &[u8]) -> Result<(), Self::Error> {
-            assert!(data.len() % WRITE_SIZE == 0);
-            assert!(offset as usize % WRITE_SIZE == 0);
-            assert!(offset as usize + data.len() <= SIZE);
-
-            self.0[offset as usize..offset as usize + data.len()].copy_from_slice(data);
-
-            Ok(())
-        }
-    }
-
-    impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> ErrorType
-        for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
-    {
-        type Error = Infallible;
-    }
-
-    impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> ReadNorFlash
-        for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
-    {
-        const READ_SIZE: usize = 1;
-
-        fn read(&mut self, offset: u32, buf: &mut [u8]) -> Result<(), Self::Error> {
-            let len = buf.len();
-            buf[..].copy_from_slice(&self.0[offset as usize..offset as usize + len]);
-            Ok(())
-        }
-
-        fn capacity(&self) -> usize {
-            SIZE
-        }
-    }
-
-    impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> super::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> AsyncReadNorFlash
-        for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
-    {
-        const READ_SIZE: usize = 1;
-
-        async fn read(&mut self, offset: u32, buf: &mut [u8]) -> Result<(), Self::Error> {
-            let len = buf.len();
-            buf[..].copy_from_slice(&self.0[offset as usize..offset as usize + len]);
-            Ok(())
-        }
-
-        fn capacity(&self) -> usize {
-            SIZE
-        }
-    }
-
-    impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> AsyncNorFlash
-        for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
-    {
-        const WRITE_SIZE: usize = WRITE_SIZE;
-        const ERASE_SIZE: usize = ERASE_SIZE;
-
-        async fn erase(&mut self, from: u32, to: u32) -> Result<(), Self::Error> {
-            let from = from as usize;
-            let to = to as usize;
-            assert!(from % ERASE_SIZE == 0);
-            assert!(to % ERASE_SIZE == 0);
-            for i in from..to {
-                self.0[i] = 0xFF;
-            }
-            Ok(())
-        }
-
-        async fn write(&mut self, offset: u32, data: &[u8]) -> Result<(), Self::Error> {
-            info!("Writing {} bytes to 0x{:x}", data.len(), offset);
-            assert!(data.len() % WRITE_SIZE == 0);
-            assert!(offset as usize % WRITE_SIZE == 0);
-            assert!(
-                offset as usize + data.len() <= SIZE,
-                "OFFSET: {}, LEN: {}, FLASH SIZE: {}",
-                offset,
-                data.len(),
-                SIZE
-            );
-
-            self.0[offset as usize..offset as usize + data.len()].copy_from_slice(data);
-
-            Ok(())
-        }
-    }
 }
diff --git a/embassy-boot/boot/src/mem_flash.rs b/embassy-boot/boot/src/mem_flash.rs
new file mode 100644
index 000000000..c62379b24
--- /dev/null
+++ b/embassy-boot/boot/src/mem_flash.rs
@@ -0,0 +1,164 @@
+#![allow(unused)]
+
+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};
+
+pub struct MemFlash<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> {
+    pub mem: [u8; SIZE],
+    pub pending_write_successes: Option<usize>,
+}
+
+#[derive(Debug)]
+pub struct MemFlashError;
+
+impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE> {
+    pub const fn new(fill: u8) -> Self {
+        Self {
+            mem: [fill; SIZE],
+            pending_write_successes: None,
+        }
+    }
+
+    #[cfg(test)]
+    pub fn random() -> Self {
+        let mut mem = [0; SIZE];
+        for byte in mem.iter_mut() {
+            *byte = rand::random::<u8>();
+        }
+        Self {
+            mem,
+            pending_write_successes: None,
+        }
+    }
+
+    pub fn program(&mut self, offset: u32, bytes: &[u8]) -> Result<(), MemFlashError> {
+        let offset = offset as usize;
+        assert!(bytes.len() % WRITE_SIZE == 0);
+        assert!(offset % WRITE_SIZE == 0);
+        assert!(offset + bytes.len() <= SIZE);
+
+        self.mem[offset..offset + bytes.len()].copy_from_slice(bytes);
+
+        Ok(())
+    }
+
+    pub fn assert_eq(&self, offset: u32, expectation: &[u8]) {
+        for i in 0..expectation.len() {
+            assert_eq!(self.mem[offset as usize + i], expectation[i], "Index {}", i);
+        }
+    }
+}
+
+impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> Default
+    for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
+{
+    fn default() -> Self {
+        Self::new(0xFF)
+    }
+}
+
+impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> ErrorType
+    for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
+{
+    type Error = MemFlashError;
+}
+
+impl NorFlashError for MemFlashError {
+    fn kind(&self) -> NorFlashErrorKind {
+        NorFlashErrorKind::Other
+    }
+}
+
+impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> ReadNorFlash
+    for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
+{
+    const READ_SIZE: usize = 1;
+
+    fn read(&mut self, offset: u32, bytes: &mut [u8]) -> Result<(), Self::Error> {
+        let len = bytes.len();
+        bytes.copy_from_slice(&self.mem[offset as usize..offset as usize + len]);
+        Ok(())
+    }
+
+    fn capacity(&self) -> usize {
+        SIZE
+    }
+}
+
+impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> NorFlash
+    for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
+{
+    const WRITE_SIZE: usize = WRITE_SIZE;
+    const ERASE_SIZE: usize = ERASE_SIZE;
+
+    fn erase(&mut self, from: u32, to: u32) -> Result<(), Self::Error> {
+        let from = from as usize;
+        let to = to as usize;
+        assert!(from % ERASE_SIZE == 0);
+        assert!(to % ERASE_SIZE == 0, "To: {}, erase size: {}", to, ERASE_SIZE);
+        for i in from..to {
+            self.mem[i] = 0xFF;
+        }
+        Ok(())
+    }
+
+    fn write(&mut self, offset: u32, bytes: &[u8]) -> Result<(), Self::Error> {
+        let offset = offset as usize;
+        assert!(bytes.len() % WRITE_SIZE == 0);
+        assert!(offset % WRITE_SIZE == 0);
+        assert!(offset + bytes.len() <= SIZE);
+
+        if let Some(pending_successes) = self.pending_write_successes {
+            if pending_successes > 0 {
+                self.pending_write_successes = Some(pending_successes - 1);
+            } else {
+                return Err(MemFlashError);
+            }
+        }
+
+        for ((offset, mem_byte), new_byte) in self
+            .mem
+            .iter_mut()
+            .enumerate()
+            .skip(offset)
+            .take(bytes.len())
+            .zip(bytes)
+        {
+            assert_eq!(0xFF, *mem_byte, "Offset {} is not erased", offset);
+            *mem_byte = *new_byte;
+        }
+
+        Ok(())
+    }
+}
+
+impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> AsyncReadNorFlash
+    for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
+{
+    const READ_SIZE: usize = 1;
+
+    async fn read(&mut self, offset: u32, bytes: &mut [u8]) -> Result<(), Self::Error> {
+        <Self as ReadNorFlash>::read(self, offset, bytes)
+    }
+
+    fn capacity(&self) -> usize {
+        <Self as ReadNorFlash>::capacity(self)
+    }
+}
+
+impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> AsyncNorFlash
+    for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
+{
+    const WRITE_SIZE: usize = WRITE_SIZE;
+    const ERASE_SIZE: usize = ERASE_SIZE;
+
+    async fn erase(&mut self, from: u32, to: u32) -> Result<(), Self::Error> {
+        <Self as NorFlash>::erase(self, from, to)
+    }
+
+    async fn write(&mut self, offset: u32, bytes: &[u8]) -> Result<(), Self::Error> {
+        <Self as NorFlash>::write(self, offset, bytes)
+    }
+}
diff --git a/embassy-boot/boot/src/partition.rs b/embassy-boot/boot/src/partition.rs
new file mode 100644
index 000000000..7529059b6
--- /dev/null
+++ b/embassy-boot/boot/src/partition.rs
@@ -0,0 +1,139 @@
+use embedded_storage::nor_flash::{NorFlash, ReadNorFlash};
+use embedded_storage_async::nor_flash::{NorFlash as AsyncNorFlash, ReadNorFlash as AsyncReadNorFlash};
+
+/// A region in flash used by the bootloader.
+#[derive(Copy, Clone, Debug)]
+#[cfg_attr(feature = "defmt", derive(defmt::Format))]
+pub struct Partition {
+    /// The offset into the flash where the partition starts.
+    pub from: u32,
+    /// The offset into the flash where the partition ends.
+    pub to: u32,
+}
+
+impl Partition {
+    /// Create a new partition with the provided range
+    pub const fn new(from: u32, to: u32) -> Self {
+        Self { from, to }
+    }
+
+    /// Return the size of the partition
+    pub const fn size(&self) -> u32 {
+        self.to - self.from
+    }
+
+    /// Read from the partition on the provided flash
+    pub async fn read<F: AsyncReadNorFlash>(
+        &self,
+        flash: &mut F,
+        offset: u32,
+        bytes: &mut [u8],
+    ) -> Result<(), F::Error> {
+        let offset = self.from as u32 + offset;
+        flash.read(offset, bytes).await
+    }
+
+    /// Write to the partition on the provided flash
+    pub async fn write<F: AsyncNorFlash>(&self, flash: &mut F, offset: u32, bytes: &[u8]) -> Result<(), F::Error> {
+        let offset = self.from as u32 + offset;
+        flash.write(offset, bytes).await?;
+        trace!("Wrote from 0x{:x} len {}", offset, bytes.len());
+        Ok(())
+    }
+
+    /// Erase part of the partition on the provided flash
+    pub async fn erase<F: AsyncNorFlash>(&self, flash: &mut F, from: u32, to: u32) -> Result<(), F::Error> {
+        let from = self.from as u32 + from;
+        let to = self.from as u32 + to;
+        flash.erase(from, to).await?;
+        trace!("Erased from 0x{:x} to 0x{:x}", from, to);
+        Ok(())
+    }
+
+    /// Erase the entire partition
+    pub(crate) async fn wipe<F: AsyncNorFlash>(&self, flash: &mut F) -> Result<(), F::Error> {
+        let from = self.from as u32;
+        let to = self.to as u32;
+        flash.erase(from, to).await?;
+        trace!("Wiped from 0x{:x} to 0x{:x}", from, to);
+        Ok(())
+    }
+
+    /// Read from the partition on the provided flash
+    pub fn read_blocking<F: ReadNorFlash>(&self, flash: &mut F, offset: u32, bytes: &mut [u8]) -> Result<(), F::Error> {
+        let offset = self.from as u32 + offset;
+        flash.read(offset, bytes)
+    }
+
+    /// Write to the partition on the provided flash
+    pub fn write_blocking<F: NorFlash>(&self, flash: &mut F, offset: u32, bytes: &[u8]) -> Result<(), F::Error> {
+        let offset = self.from as u32 + offset;
+        flash.write(offset, bytes)?;
+        trace!("Wrote from 0x{:x} len {}", offset, bytes.len());
+        Ok(())
+    }
+
+    /// Erase part of the partition on the provided flash
+    pub fn erase_blocking<F: NorFlash>(&self, flash: &mut F, from: u32, to: u32) -> Result<(), F::Error> {
+        let from = self.from as u32 + from;
+        let to = self.from as u32 + to;
+        flash.erase(from, to)?;
+        trace!("Erased from 0x{:x} to 0x{:x}", from, to);
+        Ok(())
+    }
+
+    /// Erase the entire partition
+    pub(crate) fn wipe_blocking<F: NorFlash>(&self, flash: &mut F) -> Result<(), F::Error> {
+        let from = self.from as u32;
+        let to = self.to as u32;
+        flash.erase(from, to)?;
+        trace!("Wiped from 0x{:x} to 0x{:x}", from, to);
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::mem_flash::MemFlash;
+    use crate::Partition;
+
+    #[test]
+    fn can_erase() {
+        let mut flash = MemFlash::<1024, 64, 4>::new(0x00);
+        let partition = Partition::new(256, 512);
+
+        partition.erase_blocking(&mut flash, 64, 192).unwrap();
+
+        for (index, byte) in flash.mem.iter().copied().enumerate().take(256 + 64) {
+            assert_eq!(0x00, byte, "Index {}", index);
+        }
+
+        for (index, byte) in flash.mem.iter().copied().enumerate().skip(256 + 64).take(128) {
+            assert_eq!(0xFF, byte, "Index {}", index);
+        }
+
+        for (index, byte) in flash.mem.iter().copied().enumerate().skip(256 + 64 + 128) {
+            assert_eq!(0x00, byte, "Index {}", index);
+        }
+    }
+
+    #[test]
+    fn can_wipe() {
+        let mut flash = MemFlash::<1024, 64, 4>::new(0x00);
+        let partition = Partition::new(256, 512);
+
+        partition.wipe_blocking(&mut flash).unwrap();
+
+        for (index, byte) in flash.mem.iter().copied().enumerate().take(256) {
+            assert_eq!(0x00, byte, "Index {}", index);
+        }
+
+        for (index, byte) in flash.mem.iter().copied().enumerate().skip(256).take(256) {
+            assert_eq!(0xFF, byte, "Index {}", index);
+        }
+
+        for (index, byte) in flash.mem.iter().copied().enumerate().skip(512) {
+            assert_eq!(0x00, byte, "Index {}", index);
+        }
+    }
+}
diff --git a/embassy-boot/nrf/src/lib.rs b/embassy-boot/nrf/src/lib.rs
index 5cc6ba448..48bbd7e2a 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" {
@@ -31,20 +30,20 @@ impl Default for BootLoader {
 
         let active = unsafe {
             Partition::new(
-                &__bootloader_active_start as *const u32 as usize,
-                &__bootloader_active_end as *const u32 as usize,
+                &__bootloader_active_start as *const u32 as u32,
+                &__bootloader_active_end as *const u32 as u32,
             )
         };
         let dfu = unsafe {
             Partition::new(
-                &__bootloader_dfu_start as *const u32 as usize,
-                &__bootloader_dfu_end as *const u32 as usize,
+                &__bootloader_dfu_start as *const u32 as u32,
+                &__bootloader_dfu_end as *const u32 as u32,
             )
         };
         let state = unsafe {
             Partition::new(
-                &__bootloader_state_start as *const u32 as usize,
-                &__bootloader_state_end as *const u32 as usize,
+                &__bootloader_state_start as *const u32 as u32,
+                &__bootloader_state_end as *const u32 as u32,
             )
         };
 
@@ -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..c3cb22299 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" {
@@ -68,20 +66,20 @@ impl Default for BootLoader {
 
         let active = unsafe {
             Partition::new(
-                &__bootloader_active_start as *const u32 as usize,
-                &__bootloader_active_end as *const u32 as usize,
+                &__bootloader_active_start as *const u32 as u32,
+                &__bootloader_active_end as *const u32 as u32,
             )
         };
         let dfu = unsafe {
             Partition::new(
-                &__bootloader_dfu_start as *const u32 as usize,
-                &__bootloader_dfu_end as *const u32 as usize,
+                &__bootloader_dfu_start as *const u32 as u32,
+                &__bootloader_dfu_end as *const u32 as u32,
             )
         };
         let state = unsafe {
             Partition::new(
-                &__bootloader_state_start as *const u32 as usize,
-                &__bootloader_state_end as *const u32 as usize,
+                &__bootloader_state_start as *const u32 as u32,
+                &__bootloader_state_end as *const u32 as u32,
             )
         };
 
diff --git a/embassy-boot/stm32/src/lib.rs b/embassy-boot/stm32/src/lib.rs
index 82f712c4d..94404697f 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" {
@@ -63,20 +61,20 @@ impl<const PAGE_SIZE: usize, const WRITE_SIZE: usize> Default for BootLoader<PAG
 
         let active = unsafe {
             Partition::new(
-                &__bootloader_active_start as *const u32 as usize,
-                &__bootloader_active_end as *const u32 as usize,
+                &__bootloader_active_start as *const u32 as u32,
+                &__bootloader_active_end as *const u32 as u32,
             )
         };
         let dfu = unsafe {
             Partition::new(
-                &__bootloader_dfu_start as *const u32 as usize,
-                &__bootloader_dfu_end as *const u32 as usize,
+                &__bootloader_dfu_start as *const u32 as u32,
+                &__bootloader_dfu_end as *const u32 as u32,
             )
         };
         let state = unsafe {
             Partition::new(
-                &__bootloader_state_start as *const u32 as usize,
-                &__bootloader_state_end as *const u32 as usize,
+                &__bootloader_state_start as *const u32 as u32,
+                &__bootloader_state_end as *const u32 as u32,
             )
         };