Merge remote-tracking branch 'origin/main' into nrf-pdm

25 files changed, 2331 insertions, 1037 deletions

diff --git a/embassy-boot/boot/Cargo.toml b/embassy-boot/boot/Cargo.toml
index a42f88688..415d7960f 100644
--- a/embassy-boot/boot/Cargo.toml
+++ b/embassy-boot/boot/Cargo.toml
@@ -1,25 +1,56 @@
 [package]
 edition = "2021"
 name = "embassy-boot"
-version = "0.1.0"
-description = "Bootloader using Embassy"
+version = "0.1.1"
+description = "A lightweight bootloader supporting firmware updates in a power-fail-safe way, with trial boots and rollbacks."
+license = "MIT OR Apache-2.0"
+repository = "https://github.com/embassy-rs/embassy"
+categories = [
+    "embedded",
+    "no-std",
+    "asynchronous",
+]
 
 [package.metadata.embassy_docs]
 src_base = "https://github.com/embassy-rs/embassy/blob/embassy-boot-v$VERSION/embassy-boot/boot/src/"
 src_base_git = "https://github.com/embassy-rs/embassy/blob/$COMMIT/embassy-boot/boot/src/"
 target = "thumbv7em-none-eabi"
+features = ["defmt"]
+
+[package.metadata.docs.rs]
+features = ["defmt"]
 
 [lib]
 
 [dependencies]
 defmt = { version = "0.3", optional = true }
+digest = "0.10"
 log = { version = "0.4", optional = true  }
-embassy-sync = { version = "0.1.0", path = "../../embassy-sync" }
+ed25519-dalek = { version = "1.0.1", default_features = false, features = ["u32_backend"], optional = true }
+embassy-embedded-hal = { version = "0.1.0", path = "../../embassy-embedded-hal" }
+embassy-sync = { version = "0.2.0", path = "../../embassy-sync" }
 embedded-storage = "0.3.0"
-embedded-storage-async = "0.3.0"
+embedded-storage-async = { version = "0.4.0", optional = true }
+salty = { git = "https://github.com/ycrypto/salty.git", rev = "a9f17911a5024698406b75c0fac56ab5ccf6a8c7", optional = true }
+signature = { version = "1.6.4", default-features = false }
 
 [dev-dependencies]
 log = "0.4"
 env_logger = "0.9"
-rand = "0.8"
+rand = "0.7" # ed25519-dalek v1.0.1 depends on this exact version
 futures = { version = "0.3", features = ["executor"] }
+sha1 = "0.10.5"
+critical-section = { version = "1.1.1", features = ["std"] }
+
+[dev-dependencies.ed25519-dalek]
+default_features = false
+features = ["rand", "std", "u32_backend"]
+
+[features]
+ed25519-dalek = ["dep:ed25519-dalek", "_verify"]
+ed25519-salty = ["dep:salty", "_verify"]
+
+nightly = ["dep:embedded-storage-async", "embassy-embedded-hal/nightly"]
+
+#Internal features
+_verify = []
diff --git a/embassy-boot/boot/README.md b/embassy-boot/boot/README.md
new file mode 100644
index 000000000..07755bc6c
--- /dev/null
+++ b/embassy-boot/boot/README.md
@@ -0,0 +1,31 @@
+# embassy-boot
+
+An [Embassy](https://embassy.dev) project.
+
+A lightweight bootloader supporting firmware updates in a power-fail-safe way, with trial boots and rollbacks.
+
+The bootloader can be used either as a library or be flashed directly with the default configuration derived from linker scripts.
+
+By design, the bootloader does not provide any network capabilities. Networking capabilities for fetching new firmware can be provided by the user application, using the bootloader as a library for updating the firmware, or by using the bootloader as a library and adding this capability yourself.
+
+## Hardware support
+
+The bootloader supports different hardware in separate crates:
+
+* `embassy-boot-nrf` - for the nRF microcontrollers.
+* `embassy-boot-rp` - for the RP2040 microcontrollers.
+* `embassy-boot-stm32` - for the STM32 microcontrollers.
+
+## Minimum supported Rust version (MSRV)
+
+`embassy-boot` is guaranteed to compile on the latest stable Rust version at the time of release. It might compile with older versions but that may change in any new patch release.
+
+## License
+
+This work is licensed under either of
+
+- Apache License, Version 2.0 ([LICENSE-APACHE](LICENSE-APACHE) or
+  <http://www.apache.org/licenses/LICENSE-2.0>)
+- MIT license ([LICENSE-MIT](LICENSE-MIT) or <http://opensource.org/licenses/MIT>)
+
+at your option.
diff --git a/embassy-boot/boot/src/boot_loader.rs b/embassy-boot/boot/src/boot_loader.rs
new file mode 100644
index 000000000..a8c19197b
--- /dev/null
+++ b/embassy-boot/boot/src/boot_loader.rs
@@ -0,0 +1,421 @@
+use core::cell::RefCell;
+
+use embassy_embedded_hal::flash::partition::BlockingPartition;
+use embassy_sync::blocking_mutex::raw::NoopRawMutex;
+use embassy_sync::blocking_mutex::Mutex;
+use embedded_storage::nor_flash::{NorFlash, NorFlashError, NorFlashErrorKind};
+
+use crate::{State, BOOT_MAGIC, STATE_ERASE_VALUE, 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())
+    }
+}
+
+/// Bootloader flash configuration holding the three flashes used by the bootloader
+///
+/// If only a single flash is actually used, then that flash should be partitioned into three partitions before use.
+/// The easiest way to do this is to use [`BootLoaderConfig::from_linkerfile_blocking`] which will partition
+/// the provided flash according to symbols defined in the linkerfile.
+pub struct BootLoaderConfig<ACTIVE, DFU, STATE> {
+    /// Flash type used for the active partition - the partition which will be booted from.
+    pub active: ACTIVE,
+    /// Flash type used for the dfu partition - the partition which will be swapped in when requested.
+    pub dfu: DFU,
+    /// Flash type used for the state partition.
+    pub state: STATE,
+}
+
+impl<'a, FLASH: NorFlash>
+    BootLoaderConfig<
+        BlockingPartition<'a, NoopRawMutex, FLASH>,
+        BlockingPartition<'a, NoopRawMutex, FLASH>,
+        BlockingPartition<'a, NoopRawMutex, FLASH>,
+    >
+{
+    /// Create a bootloader config from the flash and address symbols defined in the linkerfile
+    // #[cfg(target_os = "none")]
+    pub fn from_linkerfile_blocking(flash: &'a Mutex<NoopRawMutex, RefCell<FLASH>>) -> Self {
+        extern "C" {
+            static __bootloader_state_start: u32;
+            static __bootloader_state_end: u32;
+            static __bootloader_active_start: u32;
+            static __bootloader_active_end: u32;
+            static __bootloader_dfu_start: u32;
+            static __bootloader_dfu_end: u32;
+        }
+
+        let active = unsafe {
+            let start = &__bootloader_active_start as *const u32 as u32;
+            let end = &__bootloader_active_end as *const u32 as u32;
+            trace!("ACTIVE: 0x{:x} - 0x{:x}", start, end);
+
+            BlockingPartition::new(flash, start, end - start)
+        };
+        let dfu = unsafe {
+            let start = &__bootloader_dfu_start as *const u32 as u32;
+            let end = &__bootloader_dfu_end as *const u32 as u32;
+            trace!("DFU: 0x{:x} - 0x{:x}", start, end);
+
+            BlockingPartition::new(flash, start, end - start)
+        };
+        let state = unsafe {
+            let start = &__bootloader_state_start as *const u32 as u32;
+            let end = &__bootloader_state_end as *const u32 as u32;
+            trace!("STATE: 0x{:x} - 0x{:x}", start, end);
+
+            BlockingPartition::new(flash, start, end - start)
+        };
+
+        Self { active, dfu, state }
+    }
+}
+
+/// BootLoader works with any flash implementing embedded_storage.
+pub struct BootLoader<ACTIVE: NorFlash, DFU: NorFlash, STATE: NorFlash> {
+    active: ACTIVE,
+    dfu: DFU,
+    /// 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: STATE,
+}
+
+impl<ACTIVE: NorFlash, DFU: NorFlash, STATE: NorFlash> BootLoader<ACTIVE, DFU, STATE> {
+    /// Get the page size which is the "unit of operation" within the bootloader.
+    const PAGE_SIZE: u32 = if ACTIVE::ERASE_SIZE > DFU::ERASE_SIZE {
+        ACTIVE::ERASE_SIZE as u32
+    } else {
+        DFU::ERASE_SIZE as u32
+    };
+
+    /// Create a new instance of a bootloader with the flash 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(config: BootLoaderConfig<ACTIVE, DFU, STATE>) -> Self {
+        Self {
+            active: config.active,
+            dfu: config.dfu,
+            state: config.state,
+        }
+    }
+
+    /// 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(&mut self, aligned_buf: &mut [u8]) -> Result<State, BootError> {
+        // Ensure we have enough progress pages to store copy progress
+        assert_eq!(0, Self::PAGE_SIZE % aligned_buf.len() as u32);
+        assert_eq!(0, Self::PAGE_SIZE % ACTIVE::WRITE_SIZE as u32);
+        assert_eq!(0, Self::PAGE_SIZE % ACTIVE::ERASE_SIZE as u32);
+        assert_eq!(0, Self::PAGE_SIZE % DFU::WRITE_SIZE as u32);
+        assert_eq!(0, Self::PAGE_SIZE % DFU::ERASE_SIZE as u32);
+        assert!(aligned_buf.len() >= STATE::WRITE_SIZE);
+        assert_eq!(0, aligned_buf.len() % ACTIVE::WRITE_SIZE);
+        assert_eq!(0, aligned_buf.len() % DFU::WRITE_SIZE);
+
+        assert_partitions(&self.active, &self.dfu, &self.state, Self::PAGE_SIZE);
+
+        // Copy contents from partition N to active
+        let state = self.read_state(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(aligned_buf)? {
+                trace!("Swapping");
+                self.swap(aligned_buf)?;
+                trace!("Swapping done");
+            } else {
+                trace!("Reverting");
+                self.revert(aligned_buf)?;
+
+                let state_word = &mut aligned_buf[..STATE::WRITE_SIZE];
+
+                // Invalidate progress
+                state_word.fill(!STATE_ERASE_VALUE);
+                self.state.write(STATE::WRITE_SIZE as u32, state_word)?;
+
+                // Clear magic and progress
+                self.state.erase(0, self.state.capacity() as u32)?;
+
+                // Set magic
+                state_word.fill(BOOT_MAGIC);
+                self.state.write(0, state_word)?;
+            }
+        }
+        Ok(state)
+    }
+
+    fn is_swapped(&mut self, aligned_buf: &mut [u8]) -> Result<bool, BootError> {
+        let page_count = self.active.capacity() / Self::PAGE_SIZE as usize;
+        let progress = self.current_progress(aligned_buf)?;
+
+        Ok(progress >= page_count * 2)
+    }
+
+    fn current_progress(&mut self, aligned_buf: &mut [u8]) -> Result<usize, BootError> {
+        let write_size = STATE::WRITE_SIZE as u32;
+        let max_index = ((self.state.capacity() - STATE::WRITE_SIZE) / STATE::WRITE_SIZE) - 2;
+        let state_word = &mut aligned_buf[..write_size as usize];
+
+        self.state.read(write_size, state_word)?;
+        if state_word.iter().any(|&b| b != STATE_ERASE_VALUE) {
+            // Progress is invalid
+            return Ok(max_index);
+        }
+
+        for index in 0..max_index {
+            self.state.read((2 + index) as u32 * write_size, state_word)?;
+
+            if state_word.iter().any(|&b| b == STATE_ERASE_VALUE) {
+                return Ok(index);
+            }
+        }
+        Ok(max_index)
+    }
+
+    fn update_progress(&mut self, progress_index: usize, aligned_buf: &mut [u8]) -> Result<(), BootError> {
+        let state_word = &mut aligned_buf[..STATE::WRITE_SIZE];
+        state_word.fill(!STATE_ERASE_VALUE);
+        self.state
+            .write((2 + progress_index) as u32 * STATE::WRITE_SIZE as u32, state_word)?;
+        Ok(())
+    }
+
+    fn copy_page_once_to_active(
+        &mut self,
+        progress_index: usize,
+        from_offset: u32,
+        to_offset: u32,
+        aligned_buf: &mut [u8],
+    ) -> Result<(), BootError> {
+        if self.current_progress(aligned_buf)? <= progress_index {
+            let page_size = Self::PAGE_SIZE as u32;
+
+            self.active.erase(to_offset, to_offset + page_size)?;
+
+            for offset_in_page in (0..page_size).step_by(aligned_buf.len()) {
+                self.dfu.read(from_offset + offset_in_page as u32, aligned_buf)?;
+                self.active.write(to_offset + offset_in_page as u32, aligned_buf)?;
+            }
+
+            self.update_progress(progress_index, aligned_buf)?;
+        }
+        Ok(())
+    }
+
+    fn copy_page_once_to_dfu(
+        &mut self,
+        progress_index: usize,
+        from_offset: u32,
+        to_offset: u32,
+        aligned_buf: &mut [u8],
+    ) -> Result<(), BootError> {
+        if self.current_progress(aligned_buf)? <= progress_index {
+            let page_size = Self::PAGE_SIZE as u32;
+
+            self.dfu.erase(to_offset as u32, to_offset + page_size)?;
+
+            for offset_in_page in (0..page_size).step_by(aligned_buf.len()) {
+                self.active.read(from_offset + offset_in_page as u32, aligned_buf)?;
+                self.dfu.write(to_offset + offset_in_page as u32, aligned_buf)?;
+            }
+
+            self.update_progress(progress_index, aligned_buf)?;
+        }
+        Ok(())
+    }
+
+    fn swap(&mut self, aligned_buf: &mut [u8]) -> Result<(), BootError> {
+        let page_count = self.active.capacity() as u32 / Self::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) * Self::PAGE_SIZE;
+            let dfu_to_offset = (page_count - page_num) * Self::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, aligned_buf)?;
+
+            // Copy DFU page to the active page
+            let active_to_offset = (page_count - 1 - page_num) * Self::PAGE_SIZE;
+            let dfu_from_offset = (page_count - 1 - page_num) * Self::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, aligned_buf)?;
+        }
+
+        Ok(())
+    }
+
+    fn revert(&mut self, aligned_buf: &mut [u8]) -> Result<(), BootError> {
+        let page_count = self.active.capacity() as u32 / Self::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 * Self::PAGE_SIZE;
+            let dfu_to_offset = page_num * Self::PAGE_SIZE;
+            self.copy_page_once_to_dfu(progress_index, active_from_offset, dfu_to_offset, aligned_buf)?;
+
+            // Copy the DFU page back to the active page
+            let active_to_offset = page_num * Self::PAGE_SIZE;
+            let dfu_from_offset = (page_num + 1) * Self::PAGE_SIZE;
+            self.copy_page_once_to_active(progress_index + 1, dfu_from_offset, active_to_offset, aligned_buf)?;
+        }
+
+        Ok(())
+    }
+
+    fn read_state(&mut self, aligned_buf: &mut [u8]) -> Result<State, BootError> {
+        let state_word = &mut aligned_buf[..STATE::WRITE_SIZE];
+        self.state.read(0, state_word)?;
+
+        if !state_word.iter().any(|&b| b != SWAP_MAGIC) {
+            Ok(State::Swap)
+        } else {
+            Ok(State::Boot)
+        }
+    }
+}
+
+fn assert_partitions<ACTIVE: NorFlash, DFU: NorFlash, STATE: NorFlash>(
+    active: &ACTIVE,
+    dfu: &DFU,
+    state: &STATE,
+    page_size: u32,
+) {
+    assert_eq!(active.capacity() as u32 % page_size, 0);
+    assert_eq!(dfu.capacity() as u32 % page_size, 0);
+    assert!(dfu.capacity() as u32 - active.capacity() as u32 >= page_size);
+    assert!(2 + 2 * (active.capacity() as u32 / page_size) <= state.capacity() as u32 / STATE::WRITE_SIZE as u32);
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::mem_flash::MemFlash;
+
+    #[test]
+    #[should_panic]
+    fn test_range_asserts() {
+        const ACTIVE_SIZE: usize = 4194304 - 4096;
+        const DFU_SIZE: usize = 4194304;
+        const STATE_SIZE: usize = 4096;
+        static ACTIVE: MemFlash<ACTIVE_SIZE, 4, 4> = MemFlash::new(0xFF);
+        static DFU: MemFlash<DFU_SIZE, 4, 4> = MemFlash::new(0xFF);
+        static STATE: MemFlash<STATE_SIZE, 4, 4> = MemFlash::new(0xFF);
+        assert_partitions(&ACTIVE, &DFU, &STATE, 4096);
+    }
+}
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/asynch.rs b/embassy-boot/boot/src/firmware_updater/asynch.rs
new file mode 100644
index 000000000..20731ee0a
--- /dev/null
+++ b/embassy-boot/boot/src/firmware_updater/asynch.rs
@@ -0,0 +1,299 @@
+use digest::Digest;
+#[cfg(target_os = "none")]
+use embassy_embedded_hal::flash::partition::Partition;
+#[cfg(target_os = "none")]
+use embassy_sync::blocking_mutex::raw::NoopRawMutex;
+use embedded_storage_async::nor_flash::NorFlash;
+
+use super::FirmwareUpdaterConfig;
+use crate::{FirmwareUpdaterError, State, BOOT_MAGIC, STATE_ERASE_VALUE, SWAP_MAGIC};
+
+/// FirmwareUpdater is an application API for interacting with the BootLoader without the ability to
+/// 'mess up' the internal bootloader state
+pub struct FirmwareUpdater<DFU: NorFlash, STATE: NorFlash> {
+    dfu: DFU,
+    state: STATE,
+}
+
+#[cfg(target_os = "none")]
+impl<'a, FLASH: NorFlash>
+    FirmwareUpdaterConfig<Partition<'a, NoopRawMutex, FLASH>, Partition<'a, NoopRawMutex, FLASH>>
+{
+    /// Create a firmware updater config from the flash and address symbols defined in the linkerfile
+    pub fn from_linkerfile(flash: &'a embassy_sync::mutex::Mutex<NoopRawMutex, FLASH>) -> 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 {
+            let start = &__bootloader_dfu_start as *const u32 as u32;
+            let end = &__bootloader_dfu_end as *const u32 as u32;
+            trace!("DFU: 0x{:x} - 0x{:x}", start, end);
+
+            Partition::new(flash, start, end - start)
+        };
+        let state = unsafe {
+            let start = &__bootloader_state_start as *const u32 as u32;
+            let end = &__bootloader_state_end as *const u32 as u32;
+            trace!("STATE: 0x{:x} - 0x{:x}", start, end);
+
+            Partition::new(flash, start, end - start)
+        };
+
+        Self { dfu, state }
+    }
+}
+
+impl<DFU: NorFlash, STATE: NorFlash> FirmwareUpdater<DFU, STATE> {
+    /// Create a firmware updater instance with partition ranges for the update and state partitions.
+    pub fn new(config: FirmwareUpdaterConfig<DFU, STATE>) -> Self {
+        Self {
+            dfu: config.dfu,
+            state: config.state,
+        }
+    }
+
+    // Make sure we are running a booted firmware to avoid reverting to a bad state.
+    async fn verify_booted(&mut self, aligned: &mut [u8]) -> Result<(), FirmwareUpdaterError> {
+        assert_eq!(aligned.len(), STATE::WRITE_SIZE);
+        if self.get_state(aligned).await? == State::Boot {
+            Ok(())
+        } else {
+            Err(FirmwareUpdaterError::BadState)
+        }
+    }
+
+    /// 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(&mut self, aligned: &mut [u8]) -> Result<State, FirmwareUpdaterError> {
+        self.state.read(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 STATE::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(
+        &mut self,
+        _public_key: &[u8],
+        _signature: &[u8],
+        _update_len: u32,
+        _aligned: &mut [u8],
+    ) -> Result<(), FirmwareUpdaterError> {
+        assert_eq!(_aligned.len(), STATE::WRITE_SIZE);
+        assert!(_update_len <= self.dfu.capacity() as u32);
+
+        self.verify_booted(_aligned).await?;
+
+        #[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>(_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>(_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).await
+    }
+
+    /// Verify the update in DFU with any digest.
+    pub async fn hash<D: Digest>(
+        &mut self,
+        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(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 STATE::WRITE_SIZE, and follow the alignment rules for the flash being written to.
+    #[cfg(not(feature = "_verify"))]
+    pub async fn mark_updated(&mut self, aligned: &mut [u8]) -> Result<(), FirmwareUpdaterError> {
+        assert_eq!(aligned.len(), STATE::WRITE_SIZE);
+        self.set_magic(aligned, SWAP_MAGIC).await
+    }
+
+    /// Mark firmware boot successful and stop rollback on reset.
+    ///
+    /// # Safety
+    ///
+    /// The `aligned` buffer must have a size of STATE::WRITE_SIZE, and follow the alignment rules for the flash being written to.
+    pub async fn mark_booted(&mut self, aligned: &mut [u8]) -> Result<(), FirmwareUpdaterError> {
+        assert_eq!(aligned.len(), STATE::WRITE_SIZE);
+        self.set_magic(aligned, BOOT_MAGIC).await
+    }
+
+    async fn set_magic(&mut self, aligned: &mut [u8], magic: u8) -> Result<(), FirmwareUpdaterError> {
+        self.state.read(0, aligned).await?;
+
+        if aligned.iter().any(|&b| b != magic) {
+            // Read progress validity
+            self.state.read(STATE::WRITE_SIZE as u32, aligned).await?;
+
+            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::WRITE_SIZE as u32, aligned).await?;
+            }
+
+            // Clear magic and progress
+            self.state.erase(0, self.state.capacity() as u32).await?;
+
+            // Set magic
+            aligned.fill(magic);
+            self.state.write(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(
+        &mut self,
+        aligned: &mut [u8],
+        offset: usize,
+        data: &[u8],
+    ) -> Result<(), FirmwareUpdaterError> {
+        assert!(data.len() >= DFU::ERASE_SIZE);
+        assert_eq!(aligned.len(), STATE::WRITE_SIZE);
+
+        self.verify_booted(aligned).await?;
+
+        self.dfu.erase(offset as u32, (offset + data.len()) as u32).await?;
+
+        self.dfu.write(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.
+    ///
+    /// # Safety
+    ///
+    /// The `aligned` buffer must have a size of STATE::WRITE_SIZE, and follow the alignment rules for the flash being written to.
+    pub async fn prepare_update(&mut self, aligned: &mut [u8]) -> Result<&mut DFU, FirmwareUpdaterError> {
+        assert_eq!(aligned.len(), STATE::WRITE_SIZE);
+        self.verify_booted(aligned).await?;
+
+        self.dfu.erase(0, self.dfu.capacity() as u32).await?;
+
+        Ok(&mut self.dfu)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use embassy_embedded_hal::flash::partition::Partition;
+    use embassy_sync::blocking_mutex::raw::NoopRawMutex;
+    use embassy_sync::mutex::Mutex;
+    use futures::executor::block_on;
+    use sha1::{Digest, Sha1};
+
+    use super::*;
+    use crate::mem_flash::MemFlash;
+
+    #[test]
+    fn can_verify_sha1() {
+        let flash = Mutex::<NoopRawMutex, _>::new(MemFlash::<131072, 4096, 8>::default());
+        let state = Partition::new(&flash, 0, 4096);
+        let dfu = Partition::new(&flash, 65536, 65536);
+        let mut aligned = [0; 8];
+
+        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(FirmwareUpdaterConfig { dfu, state });
+        block_on(updater.write_firmware(&mut aligned, 0, to_write.as_slice())).unwrap();
+        let mut chunk_buf = [0; 2];
+        let mut hash = [0; 20];
+        block_on(updater.hash::<Sha1>(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/firmware_updater/blocking.rs b/embassy-boot/boot/src/firmware_updater/blocking.rs
new file mode 100644
index 000000000..f03f53e4d
--- /dev/null
+++ b/embassy-boot/boot/src/firmware_updater/blocking.rs
@@ -0,0 +1,302 @@
+use digest::Digest;
+#[cfg(target_os = "none")]
+use embassy_embedded_hal::flash::partition::BlockingPartition;
+#[cfg(target_os = "none")]
+use embassy_sync::blocking_mutex::raw::NoopRawMutex;
+use embedded_storage::nor_flash::NorFlash;
+
+use super::FirmwareUpdaterConfig;
+use crate::{FirmwareUpdaterError, State, BOOT_MAGIC, STATE_ERASE_VALUE, SWAP_MAGIC};
+
+/// Blocking FirmwareUpdater is an application API for interacting with the BootLoader without the ability to
+/// 'mess up' the internal bootloader state
+pub struct BlockingFirmwareUpdater<DFU: NorFlash, STATE: NorFlash> {
+    dfu: DFU,
+    state: STATE,
+}
+
+#[cfg(target_os = "none")]
+impl<'a, FLASH: NorFlash>
+    FirmwareUpdaterConfig<BlockingPartition<'a, NoopRawMutex, FLASH>, BlockingPartition<'a, NoopRawMutex, FLASH>>
+{
+    /// Create a firmware updater config from the flash and address symbols defined in the linkerfile
+    pub fn from_linkerfile_blocking(
+        flash: &'a embassy_sync::blocking_mutex::Mutex<NoopRawMutex, core::cell::RefCell<FLASH>>,
+    ) -> 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 {
+            let start = &__bootloader_dfu_start as *const u32 as u32;
+            let end = &__bootloader_dfu_end as *const u32 as u32;
+            trace!("DFU: 0x{:x} - 0x{:x}", start, end);
+
+            BlockingPartition::new(flash, start, end - start)
+        };
+        let state = unsafe {
+            let start = &__bootloader_state_start as *const u32 as u32;
+            let end = &__bootloader_state_end as *const u32 as u32;
+            trace!("STATE: 0x{:x} - 0x{:x}", start, end);
+
+            BlockingPartition::new(flash, start, end - start)
+        };
+
+        Self { dfu, state }
+    }
+}
+
+impl<DFU: NorFlash, STATE: NorFlash> BlockingFirmwareUpdater<DFU, STATE> {
+    /// Create a firmware updater instance with partition ranges for the update and state partitions.
+    pub fn new(config: FirmwareUpdaterConfig<DFU, STATE>) -> Self {
+        Self {
+            dfu: config.dfu,
+            state: config.state,
+        }
+    }
+
+    // Make sure we are running a booted firmware to avoid reverting to a bad state.
+    fn verify_booted(&mut self, aligned: &mut [u8]) -> Result<(), FirmwareUpdaterError> {
+        assert_eq!(aligned.len(), STATE::WRITE_SIZE);
+        if self.get_state(aligned)? == State::Boot {
+            Ok(())
+        } else {
+            Err(FirmwareUpdaterError::BadState)
+        }
+    }
+
+    /// 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(&mut self, aligned: &mut [u8]) -> Result<State, FirmwareUpdaterError> {
+        self.state.read(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 STATE::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(
+        &mut self,
+        _public_key: &[u8],
+        _signature: &[u8],
+        _update_len: u32,
+        _aligned: &mut [u8],
+    ) -> Result<(), FirmwareUpdaterError> {
+        assert_eq!(_aligned.len(), STATE::WRITE_SIZE);
+        assert!(_update_len <= self.dfu.capacity() as u32);
+
+        self.verify_booted(_aligned)?;
+
+        #[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>(_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::<Sha512>(_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(_aligned, SWAP_MAGIC)
+    }
+
+    /// Verify the update in DFU with any digest.
+    pub fn hash<D: Digest>(
+        &mut self,
+        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(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 STATE::WRITE_SIZE, and follow the alignment rules for the flash being written to.
+    #[cfg(not(feature = "_verify"))]
+    pub fn mark_updated(&mut self, aligned: &mut [u8]) -> Result<(), FirmwareUpdaterError> {
+        assert_eq!(aligned.len(), STATE::WRITE_SIZE);
+        self.set_magic(aligned, SWAP_MAGIC)
+    }
+
+    /// Mark firmware boot successful and stop rollback on reset.
+    ///
+    /// # Safety
+    ///
+    /// The `aligned` buffer must have a size of STATE::WRITE_SIZE, and follow the alignment rules for the flash being written to.
+    pub fn mark_booted(&mut self, aligned: &mut [u8]) -> Result<(), FirmwareUpdaterError> {
+        assert_eq!(aligned.len(), STATE::WRITE_SIZE);
+        self.set_magic(aligned, BOOT_MAGIC)
+    }
+
+    fn set_magic(&mut self, aligned: &mut [u8], magic: u8) -> Result<(), FirmwareUpdaterError> {
+        self.state.read(0, aligned)?;
+
+        if aligned.iter().any(|&b| b != magic) {
+            // Read progress validity
+            self.state.read(STATE::WRITE_SIZE as u32, aligned)?;
+
+            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::WRITE_SIZE as u32, aligned)?;
+            }
+
+            // Clear magic and progress
+            self.state.erase(0, self.state.capacity() as u32)?;
+
+            // Set magic
+            aligned.fill(magic);
+            self.state.write(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(
+        &mut self,
+        aligned: &mut [u8],
+        offset: usize,
+        data: &[u8],
+    ) -> Result<(), FirmwareUpdaterError> {
+        assert!(data.len() >= DFU::ERASE_SIZE);
+        assert_eq!(aligned.len(), STATE::WRITE_SIZE);
+        self.verify_booted(aligned)?;
+
+        self.dfu.erase(offset as u32, (offset + data.len()) as u32)?;
+
+        self.dfu.write(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` allows for an optimized API in
+    /// exchange for added complexity.
+    ///
+    /// # Safety
+    ///
+    /// The `aligned` buffer must have a size of STATE::WRITE_SIZE, and follow the alignment rules for the flash being written to.
+    pub fn prepare_update(&mut self, aligned: &mut [u8]) -> Result<&mut DFU, FirmwareUpdaterError> {
+        assert_eq!(aligned.len(), STATE::WRITE_SIZE);
+        self.verify_booted(aligned)?;
+        self.dfu.erase(0, self.dfu.capacity() as u32)?;
+
+        Ok(&mut self.dfu)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use core::cell::RefCell;
+
+    use embassy_embedded_hal::flash::partition::BlockingPartition;
+    use embassy_sync::blocking_mutex::raw::NoopRawMutex;
+    use embassy_sync::blocking_mutex::Mutex;
+    use sha1::{Digest, Sha1};
+
+    use super::*;
+    use crate::mem_flash::MemFlash;
+
+    #[test]
+    fn can_verify_sha1() {
+        let flash = Mutex::<NoopRawMutex, _>::new(RefCell::new(MemFlash::<131072, 4096, 8>::default()));
+        let state = BlockingPartition::new(&flash, 0, 4096);
+        let dfu = BlockingPartition::new(&flash, 65536, 65536);
+
+        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 = BlockingFirmwareUpdater::new(FirmwareUpdaterConfig { dfu, state });
+        let mut aligned = [0; 8];
+        updater.write_firmware(&mut aligned, 0, to_write.as_slice()).unwrap();
+        let mut chunk_buf = [0; 2];
+        let mut hash = [0; 20];
+        updater
+            .hash::<Sha1>(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/firmware_updater/mod.rs b/embassy-boot/boot/src/firmware_updater/mod.rs
new file mode 100644
index 000000000..55ce8f363
--- /dev/null
+++ b/embassy-boot/boot/src/firmware_updater/mod.rs
@@ -0,0 +1,51 @@
+#[cfg(feature = "nightly")]
+mod asynch;
+mod blocking;
+
+#[cfg(feature = "nightly")]
+pub use asynch::FirmwareUpdater;
+pub use blocking::BlockingFirmwareUpdater;
+use embedded_storage::nor_flash::{NorFlashError, NorFlashErrorKind};
+
+/// Firmware updater flash configuration holding the two flashes used by the updater
+///
+/// If only a single flash is actually used, then that flash should be partitioned into two partitions before use.
+/// The easiest way to do this is to use [`FirmwareUpdaterConfig::from_linkerfile`] or [`FirmwareUpdaterConfig::from_linkerfile_blocking`] which will partition
+/// the provided flash according to symbols defined in the linkerfile.
+pub struct FirmwareUpdaterConfig<DFU, STATE> {
+    /// The dfu flash partition
+    pub dfu: DFU,
+    /// The state flash partition
+    pub state: STATE,
+}
+
+/// Errors returned by FirmwareUpdater
+#[derive(Debug)]
+pub enum FirmwareUpdaterError {
+    /// Error from flash.
+    Flash(NorFlashErrorKind),
+    /// Signature errors.
+    Signature(signature::Error),
+    /// Bad state.
+    BadState,
+}
+
+#[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(_)"),
+            FirmwareUpdaterError::BadState => defmt::write!(fmt, "FirmwareUpdaterError::BadState"),
+        }
+    }
+}
+
+impl<E> From<E> for FirmwareUpdaterError
+where
+    E: NorFlashError,
+{
+    fn from(error: E) -> Self {
+        FirmwareUpdaterError::Flash(error.kind())
+    }
+}
diff --git a/embassy-boot/boot/src/lib.rs b/embassy-boot/boot/src/lib.rs
index 51e1056cf..016362b86 100644
--- a/embassy-boot/boot/src/lib.rs
+++ b/embassy-boot/boot/src/lib.rs
@@ -1,674 +1,76 @@
-#![feature(type_alias_impl_trait)]
-#![feature(generic_associated_types)]
-#![feature(generic_const_exprs)]
-#![allow(incomplete_features)]
+#![cfg_attr(feature = "nightly", feature(async_fn_in_trait))]
 #![no_std]
-///! embassy-boot is a bootloader and firmware updater for embedded devices with flash
-///! storage implemented using embedded-storage
-///!
-///! The bootloader works in conjunction with the firmware application, and only has the
-///! ability to manage two flash banks with an active and a updatable part. It implements
-///! a swap algorithm that is power-failure safe, and allows reverting to the previous
-///! version of the firmware, should the application crash and fail to mark itself as booted.
-///!
-///! This library is intended to be used by platform-specific bootloaders, such as embassy-boot-nrf,
-///! which defines the limits and flash type for that particular platform.
-///!
+#![warn(missing_docs)]
+#![doc = include_str!("../README.md")]
 mod fmt;
 
-use embedded_storage::nor_flash::{NorFlash, NorFlashError, NorFlashErrorKind, ReadNorFlash};
-use embedded_storage_async::nor_flash::AsyncNorFlash;
+mod boot_loader;
+mod digest_adapters;
+mod firmware_updater;
+#[cfg(test)]
+mod mem_flash;
+#[cfg(test)]
+mod test_flash;
 
-const BOOT_MAGIC: u8 = 0xD0;
-const SWAP_MAGIC: u8 = 0xF0;
+// The expected value of the flash after an erase
+// TODO: Use the value provided by NorFlash when available
+pub(crate) const STATE_ERASE_VALUE: u8 = 0xFF;
+pub use boot_loader::{BootError, BootLoader, BootLoaderConfig};
+#[cfg(feature = "nightly")]
+pub use firmware_updater::FirmwareUpdater;
+pub use firmware_updater::{BlockingFirmwareUpdater, FirmwareUpdaterConfig, FirmwareUpdaterError};
 
-#[derive(Copy, Clone, Debug)]
-#[cfg_attr(feature = "defmt", derive(defmt::Format))]
-pub struct Partition {
-    pub from: usize,
-    pub to: usize,
-}
+pub(crate) const BOOT_MAGIC: u8 = 0xD0;
+pub(crate) const SWAP_MAGIC: u8 = 0xF0;
 
-impl Partition {
-    pub const fn new(from: usize, to: usize) -> Self {
-        Self { from, to }
-    }
-    pub const fn len(&self) -> usize {
-        self.to - self.from
-    }
-}
-
-#[derive(PartialEq, Debug)]
+/// The state of the bootloader after running prepare.
+#[derive(PartialEq, Eq, Debug)]
 #[cfg_attr(feature = "defmt", derive(defmt::Format))]
 pub enum State {
+    /// Bootloader is ready to boot the active partition.
     Boot,
+    /// Bootloader has swapped the active partition with the dfu partition and will attempt boot.
     Swap,
 }
 
-#[derive(PartialEq, Debug)]
-pub enum BootError {
-    Flash(NorFlashErrorKind),
-    BadMagic,
-}
-
-impl<E> From<E> for BootError
-where
-    E: NorFlashError,
-{
-    fn from(error: E) -> Self {
-        BootError::Flash(error.kind())
-    }
-}
-
-pub trait FlashConfig {
-    const BLOCK_SIZE: usize;
-    const ERASE_VALUE: u8;
-    type FLASH: NorFlash + ReadNorFlash;
-
-    fn flash(&mut self) -> &mut Self::FLASH;
-}
-
-/// Trait defining the flash handles used for active and DFU partition
-pub trait FlashProvider {
-    type STATE: FlashConfig;
-    type ACTIVE: FlashConfig;
-    type DFU: FlashConfig;
-
-    /// 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.
-///
-/// The PAGE_SIZE const parameter must be a multiple of the ACTIVE and DFU page sizes.
-pub struct BootLoader<const PAGE_SIZE: usize> {
-    // 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<const PAGE_SIZE: usize> BootLoader<PAGE_SIZE> {
-    pub fn new(active: Partition, dfu: Partition, state: Partition) -> Self {
-        assert_eq!(active.len() % PAGE_SIZE, 0);
-        assert_eq!(dfu.len() % PAGE_SIZE, 0);
-        // DFU partition must have an extra page
-        assert!(dfu.len() - active.len() >= PAGE_SIZE);
-        Self { active, dfu, state }
-    }
-
-    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: FlashProvider>(&mut self, p: &mut P) -> Result<State, BootError>
-    where
-        [(); <<P as FlashProvider>::STATE as FlashConfig>::FLASH::WRITE_SIZE]:,
-        [(); <<P as FlashProvider>::ACTIVE as FlashConfig>::FLASH::ERASE_SIZE]:,
-    {
-        // Ensure we have enough progress pages to store copy progress
-        assert!(
-            self.active.len() / PAGE_SIZE
-                <= (self.state.len() - <<P as FlashProvider>::STATE as FlashConfig>::FLASH::WRITE_SIZE)
-                    / <<P as FlashProvider>::STATE as FlashConfig>::FLASH::WRITE_SIZE
-        );
-
-        // Copy contents from partition N to active
-        let state = self.read_state(p.state())?;
-        match 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.state())? {
-                    trace!("Swapping");
-                    self.swap(p)?;
-                    trace!("Swapping done");
-                } else {
-                    trace!("Reverting");
-                    self.revert(p)?;
-
-                    // Overwrite magic and reset progress
-                    let fstate = p.state().flash();
-                    let aligned = Aligned(
-                        [!P::STATE::ERASE_VALUE; <<P as FlashProvider>::STATE as FlashConfig>::FLASH::WRITE_SIZE],
-                    );
-                    fstate.write(self.state.from as u32, &aligned.0)?;
-                    fstate.erase(self.state.from as u32, self.state.to as u32)?;
-                    let aligned =
-                        Aligned([BOOT_MAGIC; <<P as FlashProvider>::STATE as FlashConfig>::FLASH::WRITE_SIZE]);
-                    fstate.write(self.state.from as u32, &aligned.0)?;
-                }
-            }
-            _ => {}
-        }
-        Ok(state)
-    }
-
-    fn is_swapped<P: FlashConfig>(&mut self, p: &mut P) -> Result<bool, BootError>
-    where
-        [(); P::FLASH::WRITE_SIZE]:,
-    {
-        let page_count = self.active.len() / P::FLASH::ERASE_SIZE;
-        let progress = self.current_progress(p)?;
-
-        Ok(progress >= page_count * 2)
-    }
-
-    fn current_progress<P: FlashConfig>(&mut self, p: &mut P) -> Result<usize, BootError>
-    where
-        [(); P::FLASH::WRITE_SIZE]:,
-    {
-        let write_size = P::FLASH::WRITE_SIZE;
-        let max_index = ((self.state.len() - write_size) / write_size) - 1;
-        let flash = p.flash();
-        let mut aligned = Aligned([!P::ERASE_VALUE; P::FLASH::WRITE_SIZE]);
-        for i in 0..max_index {
-            flash.read((self.state.from + write_size + i * write_size) as u32, &mut aligned.0)?;
-            if aligned.0 == [P::ERASE_VALUE; P::FLASH::WRITE_SIZE] {
-                return Ok(i);
-            }
-        }
-        Ok(max_index)
-    }
-
-    fn update_progress<P: FlashConfig>(&mut self, idx: usize, p: &mut P) -> Result<(), BootError>
-    where
-        [(); P::FLASH::WRITE_SIZE]:,
-    {
-        let flash = p.flash();
-        let write_size = P::FLASH::WRITE_SIZE;
-        let w = self.state.from + write_size + idx * write_size;
-        let aligned = Aligned([!P::ERASE_VALUE; P::FLASH::WRITE_SIZE]);
-        flash.write(w as u32, &aligned.0)?;
-        Ok(())
-    }
-
-    fn active_addr(&self, n: usize) -> usize {
-        self.active.from + n * PAGE_SIZE
-    }
-
-    fn dfu_addr(&self, n: usize) -> usize {
-        self.dfu.from + n * PAGE_SIZE
-    }
-
-    fn copy_page_once_to_active<P: FlashProvider>(
-        &mut self,
-        idx: usize,
-        from_page: usize,
-        to_page: usize,
-        p: &mut P,
-    ) -> Result<(), BootError>
-    where
-        [(); <<P as FlashProvider>::STATE as FlashConfig>::FLASH::WRITE_SIZE]:,
-    {
-        let mut buf: [u8; PAGE_SIZE] = [0; PAGE_SIZE];
-        if self.current_progress(p.state())? <= idx {
-            let mut offset = from_page;
-            for chunk in buf.chunks_mut(P::DFU::BLOCK_SIZE) {
-                p.dfu().flash().read(offset as u32, chunk)?;
-                offset += chunk.len();
-            }
-
-            p.active().flash().erase(to_page as u32, (to_page + PAGE_SIZE) as u32)?;
-
-            let mut offset = to_page;
-            for chunk in buf.chunks(P::ACTIVE::BLOCK_SIZE) {
-                p.active().flash().write(offset as u32, &chunk)?;
-                offset += chunk.len();
-            }
-            self.update_progress(idx, p.state())?;
-        }
-        Ok(())
-    }
-
-    fn copy_page_once_to_dfu<P: FlashProvider>(
-        &mut self,
-        idx: usize,
-        from_page: usize,
-        to_page: usize,
-        p: &mut P,
-    ) -> Result<(), BootError>
-    where
-        [(); <<P as FlashProvider>::STATE as FlashConfig>::FLASH::WRITE_SIZE]:,
-    {
-        let mut buf: [u8; PAGE_SIZE] = [0; PAGE_SIZE];
-        if self.current_progress(p.state())? <= idx {
-            let mut offset = from_page;
-            for chunk in buf.chunks_mut(P::ACTIVE::BLOCK_SIZE) {
-                p.active().flash().read(offset as u32, chunk)?;
-                offset += chunk.len();
-            }
-
-            p.dfu().flash().erase(to_page as u32, (to_page + PAGE_SIZE) as u32)?;
-
-            let mut offset = to_page;
-            for chunk in buf.chunks(P::DFU::BLOCK_SIZE) {
-                p.dfu().flash().write(offset as u32, chunk)?;
-                offset += chunk.len();
-            }
-            self.update_progress(idx, p.state())?;
-        }
-        Ok(())
-    }
-
-    fn swap<P: FlashProvider>(&mut self, p: &mut P) -> Result<(), BootError>
-    where
-        [(); <<P as FlashProvider>::STATE as FlashConfig>::FLASH::WRITE_SIZE]:,
-    {
-        let page_count = self.active.len() / PAGE_SIZE;
-        trace!("Page count: {}", page_count);
-        for page in 0..page_count {
-            trace!("COPY PAGE {}", page);
-            // Copy active page to the 'next' DFU page.
-            let active_page = self.active_addr(page_count - 1 - page);
-            let dfu_page = self.dfu_addr(page_count - page);
-            //trace!("Copy active {} to dfu {}", active_page, dfu_page);
-            self.copy_page_once_to_dfu(page * 2, active_page, dfu_page, p)?;
-
-            // Copy DFU page to the active page
-            let active_page = self.active_addr(page_count - 1 - page);
-            let dfu_page = self.dfu_addr(page_count - 1 - page);
-            //trace!("Copy dfy {} to active {}", dfu_page, active_page);
-            self.copy_page_once_to_active(page * 2 + 1, dfu_page, active_page, p)?;
-        }
-
-        Ok(())
-    }
-
-    fn revert<P: FlashProvider>(&mut self, p: &mut P) -> Result<(), BootError>
-    where
-        [(); <<P as FlashProvider>::STATE as FlashConfig>::FLASH::WRITE_SIZE]:,
-    {
-        let page_count = self.active.len() / PAGE_SIZE;
-        for page in 0..page_count {
-            // Copy the bad active page to the DFU page
-            let active_page = self.active_addr(page);
-            let dfu_page = self.dfu_addr(page);
-            self.copy_page_once_to_dfu(page_count * 2 + page * 2, active_page, dfu_page, p)?;
-
-            // Copy the DFU page back to the active page
-            let active_page = self.active_addr(page);
-            let dfu_page = self.dfu_addr(page + 1);
-            self.copy_page_once_to_active(page_count * 2 + page * 2 + 1, dfu_page, active_page, p)?;
-        }
-
-        Ok(())
-    }
-
-    fn read_state<P: FlashConfig>(&mut self, p: &mut P) -> Result<State, BootError>
-    where
-        [(); P::FLASH::WRITE_SIZE]:,
-    {
-        let mut magic: [u8; P::FLASH::WRITE_SIZE] = [0; P::FLASH::WRITE_SIZE];
-        let flash = p.flash();
-        flash.read(self.state.from as u32, &mut magic)?;
-
-        if magic == [SWAP_MAGIC; P::FLASH::WRITE_SIZE] {
-            Ok(State::Swap)
-        } else {
-            Ok(State::Boot)
-        }
-    }
-}
-
-/// Convenience provider that uses a single flash for everything
-pub struct SingleFlashProvider<'a, F, const ERASE_VALUE: u8 = 0xFF>
-where
-    F: NorFlash + ReadNorFlash,
-{
-    config: SingleFlashConfig<'a, F, ERASE_VALUE>,
-}
-
-impl<'a, F, const ERASE_VALUE: u8> SingleFlashProvider<'a, F, ERASE_VALUE>
-where
-    F: NorFlash + ReadNorFlash,
-{
-    pub fn new(flash: &'a mut F) -> Self {
-        Self {
-            config: SingleFlashConfig { flash },
-        }
-    }
-}
-
-pub struct SingleFlashConfig<'a, F, const ERASE_VALUE: u8 = 0xFF>
-where
-    F: NorFlash + ReadNorFlash,
-{
-    flash: &'a mut F,
-}
-
-impl<'a, F> FlashProvider for SingleFlashProvider<'a, F>
-where
-    F: NorFlash + ReadNorFlash,
-{
-    type STATE = SingleFlashConfig<'a, F>;
-    type ACTIVE = SingleFlashConfig<'a, F>;
-    type DFU = SingleFlashConfig<'a, F>;
-
-    fn active(&mut self) -> &mut Self::STATE {
-        &mut self.config
-    }
-    fn dfu(&mut self) -> &mut Self::ACTIVE {
-        &mut self.config
-    }
-    fn state(&mut self) -> &mut Self::DFU {
-        &mut self.config
-    }
-}
-
-impl<'a, F, const ERASE_VALUE: u8> FlashConfig for SingleFlashConfig<'a, F, ERASE_VALUE>
-where
-    F: NorFlash + ReadNorFlash,
-{
-    const BLOCK_SIZE: usize = F::ERASE_SIZE;
-    const ERASE_VALUE: u8 = ERASE_VALUE;
-    type FLASH = F;
-    fn flash(&mut self) -> &mut F {
-        self.flash
-    }
-}
-
-/// Convenience provider that uses a single flash for everything
-pub struct MultiFlashProvider<'a, ACTIVE, STATE, DFU>
-where
-    ACTIVE: NorFlash + ReadNorFlash,
-    STATE: NorFlash + ReadNorFlash,
-    DFU: NorFlash + ReadNorFlash,
-{
-    active: SingleFlashConfig<'a, ACTIVE>,
-    state: SingleFlashConfig<'a, STATE>,
-    dfu: SingleFlashConfig<'a, DFU>,
-}
-
-impl<'a, ACTIVE, STATE, DFU> MultiFlashProvider<'a, ACTIVE, STATE, DFU>
-where
-    ACTIVE: NorFlash + ReadNorFlash,
-    STATE: NorFlash + ReadNorFlash,
-    DFU: NorFlash + ReadNorFlash,
-{
-    pub fn new(active: &'a mut ACTIVE, state: &'a mut STATE, dfu: &'a mut DFU) -> Self {
-        Self {
-            active: SingleFlashConfig { flash: active },
-            state: SingleFlashConfig { flash: state },
-            dfu: SingleFlashConfig { flash: dfu },
-        }
-    }
-}
-
-impl<'a, ACTIVE, STATE, DFU> FlashProvider for MultiFlashProvider<'a, ACTIVE, STATE, DFU>
-where
-    ACTIVE: NorFlash + ReadNorFlash,
-    STATE: NorFlash + ReadNorFlash,
-    DFU: NorFlash + ReadNorFlash,
-{
-    type STATE = SingleFlashConfig<'a, STATE>;
-    type ACTIVE = SingleFlashConfig<'a, ACTIVE>;
-    type DFU = SingleFlashConfig<'a, DFU>;
-
-    fn active(&mut self) -> &mut Self::ACTIVE {
-        &mut self.active
-    }
-    fn dfu(&mut self) -> &mut Self::DFU {
-        &mut self.dfu
-    }
-    fn state(&mut self) -> &mut Self::STATE {
-        &mut self.state
-    }
-}
-
-/// 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,
-}
-
-// NOTE: Aligned to the largest write size supported by flash
+/// Buffer aligned to 32 byte boundary, largest known alignment requirement for embassy-boot.
 #[repr(align(32))]
-pub struct Aligned<const N: usize>([u8; N]);
-
-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 {
-    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()
-    }
-
-    /// Instruct bootloader that DFU should commence at next boot.
-    /// Must be provided with an aligned buffer to use for reading and writing magic;
-    pub async fn update<F: AsyncNorFlash>(&mut self, flash: &mut F) -> Result<(), F::Error>
-    where
-        [(); F::WRITE_SIZE]:,
-    {
-        let mut aligned = Aligned([0; { F::WRITE_SIZE }]);
-        self.set_magic(&mut aligned.0, SWAP_MAGIC, flash).await
-    }
+pub struct AlignedBuffer<const N: usize>(pub [u8; N]);
 
-    /// Mark firmware boot successfully
-    pub async fn mark_booted<F: AsyncNorFlash>(&mut self, flash: &mut F) -> Result<(), F::Error>
-    where
-        [(); F::WRITE_SIZE]:,
-    {
-        let mut aligned = Aligned([0; { F::WRITE_SIZE }]);
-        self.set_magic(&mut aligned.0, BOOT_MAGIC, flash).await
+impl<const N: usize> AsRef<[u8]> for AlignedBuffer<N> {
+    fn as_ref(&self) -> &[u8] {
+        &self.0
     }
+}
 
-    async fn set_magic<F: AsyncNorFlash>(
-        &mut self,
-        aligned: &mut [u8],
-        magic: u8,
-        flash: &mut F,
-    ) -> Result<(), F::Error> {
-        flash.read(self.state.from as u32, aligned).await?;
-
-        if aligned.iter().find(|&&b| b != magic).is_some() {
-            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 to a region of the DFU page
-    pub async fn write_firmware<F: AsyncNorFlash>(
-        &mut self,
-        offset: usize,
-        data: &[u8],
-        flash: &mut F,
-        block_size: usize,
-    ) -> Result<(), F::Error> {
-        assert!(data.len() >= F::ERASE_SIZE);
-
-        trace!(
-            "Writing firmware at offset 0x{:x} len {}",
-            self.dfu.from + offset,
-            data.len()
-        );
-
-        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()
-        );
-
-        let mut write_offset = self.dfu.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(())
+impl<const N: usize> AsMut<[u8]> for AlignedBuffer<N> {
+    fn as_mut(&mut self) -> &mut [u8] {
+        &mut self.0
     }
 }
 
 #[cfg(test)]
 mod tests {
-    use core::convert::Infallible;
-    use core::future::Future;
+    #![allow(unused_imports)]
 
-    use embedded_storage::nor_flash::ErrorType;
-    use embedded_storage_async::nor_flash::AsyncReadNorFlash;
+    use embedded_storage::nor_flash::{NorFlash, ReadNorFlash};
+    #[cfg(feature = "nightly")]
+    use embedded_storage_async::nor_flash::NorFlash as AsyncNorFlash;
     use futures::executor::block_on;
 
     use super::*;
+    use crate::boot_loader::BootLoaderConfig;
+    use crate::firmware_updater::FirmwareUpdaterConfig;
+    use crate::mem_flash::MemFlash;
+    #[cfg(feature = "nightly")]
+    use crate::test_flash::AsyncTestFlash;
+    use crate::test_flash::BlockingTestFlash;
 
     /*
     #[test]
     fn test_bad_magic() {
         let mut flash = MemFlash([0xff; 131072]);
-        let mut flash = SingleFlashProvider::new(&mut flash);
+        let mut flash = SingleFlashConfig::new(&mut flash);
 
         let mut bootloader = BootLoader::<4096>::new(ACTIVE, DFU, STATE);
 
@@ -681,281 +83,229 @@ mod tests {
 
     #[test]
     fn test_boot_state() {
-        const STATE: Partition = Partition::new(0, 4096);
-        const ACTIVE: Partition = Partition::new(4096, 61440);
-        const DFU: Partition = Partition::new(61440, 122880);
+        let flash = BlockingTestFlash::new(BootLoaderConfig {
+            active: MemFlash::<57344, 4096, 4>::default(),
+            dfu: MemFlash::<61440, 4096, 4>::default(),
+            state: MemFlash::<4096, 4096, 4>::default(),
+        });
 
-        let mut flash = MemFlash::<131072, 4096, 4>([0xff; 131072]);
-        flash.0[0..4].copy_from_slice(&[BOOT_MAGIC; 4]);
-        let mut flash = SingleFlashProvider::new(&mut flash);
+        flash.state().write(0, &[BOOT_MAGIC; 4]).unwrap();
 
-        let mut bootloader: BootLoader<4096> = BootLoader::new(ACTIVE, DFU, STATE);
+        let mut bootloader = BootLoader::new(BootLoaderConfig {
+            active: flash.active(),
+            dfu: flash.dfu(),
+            state: flash.state(),
+        });
 
-        assert_eq!(State::Boot, bootloader.prepare_boot(&mut flash).unwrap());
+        let mut page = [0; 4096];
+        assert_eq!(State::Boot, bootloader.prepare_boot(&mut page).unwrap());
     }
 
     #[test]
+    #[cfg(all(feature = "nightly", not(feature = "_verify")))]
     fn test_swap_state() {
-        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 original: [u8; ACTIVE.len()] = [rand::random::<u8>(); ACTIVE.len()];
-        let update: [u8; DFU.len()] = [rand::random::<u8>(); DFU.len()];
-
-        for i in ACTIVE.from..ACTIVE.to {
-            flash.0[i] = original[i - ACTIVE.from];
-        }
-
-        let mut bootloader: BootLoader<4096> = 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.update(&mut flash)).unwrap();
-
-        assert_eq!(
-            State::Swap,
-            bootloader
-                .prepare_boot(&mut SingleFlashProvider::new(&mut flash))
-                .unwrap()
-        );
-
-        for i in ACTIVE.from..ACTIVE.to {
-            assert_eq!(flash.0[i], update[i - ACTIVE.from], "Index {}", i);
-        }
-
+        const FIRMWARE_SIZE: usize = 57344;
+        let flash = AsyncTestFlash::new(BootLoaderConfig {
+            active: MemFlash::<FIRMWARE_SIZE, 4096, 4>::default(),
+            dfu: MemFlash::<61440, 4096, 4>::default(),
+            state: MemFlash::<4096, 4096, 4>::default(),
+        });
+
+        const ORIGINAL: [u8; FIRMWARE_SIZE] = [0x55; FIRMWARE_SIZE];
+        const UPDATE: [u8; FIRMWARE_SIZE] = [0xAA; FIRMWARE_SIZE];
+        let mut aligned = [0; 4];
+
+        block_on(flash.active().erase(0, ORIGINAL.len() as u32)).unwrap();
+        block_on(flash.active().write(0, &ORIGINAL)).unwrap();
+
+        let mut updater = FirmwareUpdater::new(FirmwareUpdaterConfig {
+            dfu: flash.dfu(),
+            state: flash.state(),
+        });
+        block_on(updater.write_firmware(&mut aligned, 0, &UPDATE)).unwrap();
+        block_on(updater.mark_updated(&mut aligned)).unwrap();
+
+        // Writing after marking updated is not allowed until marked as booted.
+        let res: Result<(), FirmwareUpdaterError> = block_on(updater.write_firmware(&mut aligned, 0, &UPDATE));
+        assert!(matches!(res, Err::<(), _>(FirmwareUpdaterError::BadState)));
+
+        let flash = flash.into_blocking();
+        let mut bootloader = BootLoader::new(BootLoaderConfig {
+            active: flash.active(),
+            dfu: flash.dfu(),
+            state: flash.state(),
+        });
+
+        let mut page = [0; 1024];
+        assert_eq!(State::Swap, bootloader.prepare_boot(&mut page).unwrap());
+
+        let mut read_buf = [0; FIRMWARE_SIZE];
+        flash.active().read(0, &mut read_buf).unwrap();
+        assert_eq!(UPDATE, read_buf);
         // 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.dfu().read(4096, &mut read_buf).unwrap();
+        assert_eq!(ORIGINAL, read_buf);
 
         // Running again should cause a revert
-        assert_eq!(
-            State::Swap,
-            bootloader
-                .prepare_boot(&mut SingleFlashProvider::new(&mut flash))
-                .unwrap()
-        );
+        assert_eq!(State::Swap, bootloader.prepare_boot(&mut page).unwrap());
 
-        for i in ACTIVE.from..ACTIVE.to {
-            assert_eq!(flash.0[i], original[i - ACTIVE.from], "Index {}", i);
-        }
-
-        // Last page is untouched
-        for i in DFU.from..DFU.to - 4096 {
-            assert_eq!(flash.0[i], update[i - DFU.from], "Index {}", i);
-        }
+        let mut read_buf = [0; FIRMWARE_SIZE];
+        flash.active().read(0, &mut read_buf).unwrap();
+        assert_eq!(ORIGINAL, read_buf);
+        // Last DFU page is untouched
+        flash.dfu().read(0, &mut read_buf).unwrap();
+        assert_eq!(UPDATE, read_buf);
 
         // Mark as booted
-        block_on(updater.mark_booted(&mut flash)).unwrap();
-        assert_eq!(
-            State::Boot,
-            bootloader
-                .prepare_boot(&mut SingleFlashProvider::new(&mut flash))
-                .unwrap()
-        );
+        let flash = flash.into_async();
+        let mut updater = FirmwareUpdater::new(FirmwareUpdaterConfig {
+            dfu: flash.dfu(),
+            state: flash.state(),
+        });
+        block_on(updater.mark_booted(&mut aligned)).unwrap();
+
+        let flash = flash.into_blocking();
+        let mut bootloader = BootLoader::new(BootLoaderConfig {
+            active: flash.active(),
+            dfu: flash.dfu(),
+            state: flash.state(),
+        });
+        assert_eq!(State::Boot, bootloader.prepare_boot(&mut page).unwrap());
     }
 
     #[test]
-    fn test_separate_flash_active_page_biggest() {
-        const STATE: Partition = Partition::new(2048, 4096);
-        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 original: [u8; ACTIVE.len()] = [rand::random::<u8>(); ACTIVE.len()];
-        let update: [u8; DFU.len()] = [rand::random::<u8>(); DFU.len()];
-
-        for i in ACTIVE.from..ACTIVE.to {
-            active.0[i] = original[i - ACTIVE.from];
-        }
-
-        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.update(&mut state)).unwrap();
-
-        let mut bootloader: BootLoader<4096> = BootLoader::new(ACTIVE, DFU, STATE);
-        assert_eq!(
-            State::Swap,
-            bootloader
-                .prepare_boot(&mut MultiFlashProvider::new(&mut active, &mut state, &mut dfu,))
-                .unwrap()
-        );
-
-        for i in ACTIVE.from..ACTIVE.to {
-            assert_eq!(active.0[i], update[i - ACTIVE.from], "Index {}", i);
-        }
-
+    #[cfg(all(feature = "nightly", not(feature = "_verify")))]
+    fn test_swap_state_active_page_biggest() {
+        const FIRMWARE_SIZE: usize = 12288;
+        let flash = AsyncTestFlash::new(BootLoaderConfig {
+            active: MemFlash::<12288, 4096, 8>::random(),
+            dfu: MemFlash::<16384, 2048, 8>::random(),
+            state: MemFlash::<2048, 128, 4>::random(),
+        });
+
+        const ORIGINAL: [u8; FIRMWARE_SIZE] = [0x55; FIRMWARE_SIZE];
+        const UPDATE: [u8; FIRMWARE_SIZE] = [0xAA; FIRMWARE_SIZE];
+        let mut aligned = [0; 4];
+
+        block_on(flash.active().erase(0, ORIGINAL.len() as u32)).unwrap();
+        block_on(flash.active().write(0, &ORIGINAL)).unwrap();
+
+        let mut updater = FirmwareUpdater::new(FirmwareUpdaterConfig {
+            dfu: flash.dfu(),
+            state: flash.state(),
+        });
+        block_on(updater.write_firmware(&mut aligned, 0, &UPDATE)).unwrap();
+        block_on(updater.mark_updated(&mut aligned)).unwrap();
+
+        let flash = flash.into_blocking();
+        let mut bootloader = BootLoader::new(BootLoaderConfig {
+            active: flash.active(),
+            dfu: flash.dfu(),
+            state: flash.state(),
+        });
+
+        let mut page = [0; 4096];
+        assert_eq!(State::Swap, bootloader.prepare_boot(&mut page).unwrap());
+
+        let mut read_buf = [0; FIRMWARE_SIZE];
+        flash.active().read(0, &mut read_buf).unwrap();
+        assert_eq!(UPDATE, read_buf);
         // 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);
-        }
+        flash.dfu().read(4096, &mut read_buf).unwrap();
+        assert_eq!(ORIGINAL, read_buf);
     }
 
     #[test]
-    fn test_separate_flash_dfu_page_biggest() {
-        const STATE: Partition = Partition::new(2048, 4096);
-        const ACTIVE: Partition = Partition::new(4096, 16384);
-        const DFU: Partition = Partition::new(0, 16384);
-
-        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 original: [u8; ACTIVE.len()] = [rand::random::<u8>(); ACTIVE.len()];
-        let update: [u8; DFU.len()] = [rand::random::<u8>(); DFU.len()];
-
-        for i in ACTIVE.from..ACTIVE.to {
-            active.0[i] = original[i - ACTIVE.from];
-        }
-
-        let mut updater = FirmwareUpdater::new(DFU, STATE);
-
-        let mut offset = 0;
-        for chunk in update.chunks(4096) {
-            block_on(updater.write_firmware(offset, &chunk, &mut dfu, chunk.len())).unwrap();
-            offset += chunk.len();
-        }
-        block_on(updater.update(&mut state)).unwrap();
-
-        let mut bootloader: BootLoader<4096> = BootLoader::new(ACTIVE, DFU, STATE);
-        assert_eq!(
-            State::Swap,
-            bootloader
-                .prepare_boot(&mut MultiFlashProvider::new(&mut active, &mut state, &mut dfu,))
-                .unwrap()
-        );
-
-        for i in ACTIVE.from..ACTIVE.to {
-            assert_eq!(active.0[i], update[i - ACTIVE.from], "Index {}", i);
-        }
-
+    #[cfg(all(feature = "nightly", not(feature = "_verify")))]
+    fn test_swap_state_dfu_page_biggest() {
+        const FIRMWARE_SIZE: usize = 12288;
+        let flash = AsyncTestFlash::new(BootLoaderConfig {
+            active: MemFlash::<FIRMWARE_SIZE, 2048, 4>::random(),
+            dfu: MemFlash::<16384, 4096, 8>::random(),
+            state: MemFlash::<2048, 128, 4>::random(),
+        });
+
+        const ORIGINAL: [u8; FIRMWARE_SIZE] = [0x55; FIRMWARE_SIZE];
+        const UPDATE: [u8; FIRMWARE_SIZE] = [0xAA; FIRMWARE_SIZE];
+        let mut aligned = [0; 4];
+
+        block_on(flash.active().erase(0, ORIGINAL.len() as u32)).unwrap();
+        block_on(flash.active().write(0, &ORIGINAL)).unwrap();
+
+        let mut updater = FirmwareUpdater::new(FirmwareUpdaterConfig {
+            dfu: flash.dfu(),
+            state: flash.state(),
+        });
+        block_on(updater.write_firmware(&mut aligned, 0, &UPDATE)).unwrap();
+        block_on(updater.mark_updated(&mut aligned)).unwrap();
+
+        let flash = flash.into_blocking();
+        let mut bootloader = BootLoader::new(BootLoaderConfig {
+            active: flash.active(),
+            dfu: flash.dfu(),
+            state: flash.state(),
+        });
+        let mut page = [0; 4096];
+        assert_eq!(State::Swap, bootloader.prepare_boot(&mut page).unwrap());
+
+        let mut read_buf = [0; FIRMWARE_SIZE];
+        flash.active().read(0, &mut read_buf).unwrap();
+        assert_eq!(UPDATE, read_buf);
         // 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);
-        }
-    }
-
-    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(())
-        }
+        flash.dfu().read(4096, &mut read_buf).unwrap();
+        assert_eq!(ORIGINAL, read_buf);
     }
 
-    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 = 4;
-
-        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> AsyncReadNorFlash
-        for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
-    {
-        const READ_SIZE: usize = 4;
-
-        type ReadFuture<'a> = impl Future<Output = Result<(), Self::Error>> + 'a;
-        fn read<'a>(&'a mut self, offset: u32, buf: &'a mut [u8]) -> Self::ReadFuture<'a> {
-            async move {
-                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;
-
-        type EraseFuture<'a> = impl Future<Output = Result<(), Self::Error>> + 'a;
-        fn erase<'a>(&'a mut self, from: u32, to: u32) -> Self::EraseFuture<'a> {
-            async move {
-                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(())
-            }
-        }
-
-        type WriteFuture<'a> = impl Future<Output = Result<(), Self::Error>> + 'a;
-        fn write<'a>(&'a mut self, offset: u32, data: &'a [u8]) -> Self::WriteFuture<'a> {
-            info!("Writing {} bytes to 0x{:x}", data.len(), offset);
-            async move {
-                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(())
-            }
-        }
+    #[test]
+    #[cfg(all(feature = "nightly", feature = "_verify"))]
+    fn test_verify() {
+        // The following key setup is based on:
+        // https://docs.rs/ed25519-dalek/latest/ed25519_dalek/#example
+
+        use ed25519_dalek::Keypair;
+        use rand::rngs::OsRng;
+
+        let mut csprng = OsRng {};
+        let keypair: Keypair = Keypair::generate(&mut csprng);
+
+        use ed25519_dalek::{Digest, Sha512, Signature, Signer};
+        let firmware: &[u8] = b"This are bytes that would otherwise be firmware bytes for DFU.";
+        let mut digest = Sha512::new();
+        digest.update(&firmware);
+        let message = digest.finalize();
+        let signature: Signature = keypair.sign(&message);
+
+        use ed25519_dalek::PublicKey;
+        let public_key: PublicKey = keypair.public;
+
+        // Setup flash
+        let flash = BlockingTestFlash::new(BootLoaderConfig {
+            active: MemFlash::<0, 0, 0>::default(),
+            dfu: MemFlash::<4096, 4096, 4>::default(),
+            state: MemFlash::<4096, 4096, 4>::default(),
+        });
+
+        let firmware_len = firmware.len();
+
+        let mut write_buf = [0; 4096];
+        write_buf[0..firmware_len].copy_from_slice(firmware);
+        flash.dfu().write(0, &write_buf).unwrap();
+
+        // On with the test
+        let flash = flash.into_async();
+        let mut updater = FirmwareUpdater::new(FirmwareUpdaterConfig {
+            dfu: flash.dfu(),
+            state: flash.state(),
+        });
+
+        let mut aligned = [0; 4];
+
+        assert!(block_on(updater.verify_and_mark_updated(
+            &public_key.to_bytes(),
+            &signature.to_bytes(),
+            firmware_len as u32,
+            &mut aligned,
+        ))
+        .is_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..2728e9720
--- /dev/null
+++ b/embassy-boot/boot/src/mem_flash.rs
@@ -0,0 +1,173 @@
+#![allow(unused)]
+
+use core::ops::{Bound, Range, RangeBounds};
+
+use embedded_storage::nor_flash::{ErrorType, NorFlash, NorFlashError, NorFlashErrorKind, ReadNorFlash};
+#[cfg(feature = "nightly")]
+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,
+        }
+    }
+
+    fn read(&mut self, offset: u32, bytes: &mut [u8]) -> Result<(), MemFlashError> {
+        let len = bytes.len();
+        bytes.copy_from_slice(&self.mem[offset as usize..offset as usize + len]);
+        Ok(())
+    }
+
+    fn write(&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);
+
+        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(())
+    }
+
+    fn erase(&mut self, from: u32, to: u32) -> Result<(), MemFlashError> {
+        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(())
+    }
+
+    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(())
+    }
+}
+
+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> {
+        self.read(offset, bytes)
+    }
+
+    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 write(&mut self, offset: u32, bytes: &[u8]) -> Result<(), Self::Error> {
+        self.write(offset, bytes)
+    }
+
+    fn erase(&mut self, from: u32, to: u32) -> Result<(), Self::Error> {
+        self.erase(from, to)
+    }
+}
+
+#[cfg(feature = "nightly")]
+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.read(offset, bytes)
+    }
+
+    fn capacity(&self) -> usize {
+        SIZE
+    }
+}
+
+#[cfg(feature = "nightly")]
+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 write(&mut self, offset: u32, bytes: &[u8]) -> Result<(), Self::Error> {
+        self.write(offset, bytes)
+    }
+
+    async fn erase(&mut self, from: u32, to: u32) -> Result<(), Self::Error> {
+        self.erase(from, to)
+    }
+}
diff --git a/embassy-boot/boot/src/test_flash/asynch.rs b/embassy-boot/boot/src/test_flash/asynch.rs
new file mode 100644
index 000000000..3ac9e71ab
--- /dev/null
+++ b/embassy-boot/boot/src/test_flash/asynch.rs
@@ -0,0 +1,64 @@
+use embassy_embedded_hal::flash::partition::Partition;
+use embassy_sync::blocking_mutex::raw::NoopRawMutex;
+use embassy_sync::mutex::Mutex;
+use embedded_storage_async::nor_flash::NorFlash;
+
+use crate::BootLoaderConfig;
+
+pub struct AsyncTestFlash<ACTIVE, DFU, STATE>
+where
+    ACTIVE: NorFlash,
+    DFU: NorFlash,
+    STATE: NorFlash,
+{
+    active: Mutex<NoopRawMutex, ACTIVE>,
+    dfu: Mutex<NoopRawMutex, DFU>,
+    state: Mutex<NoopRawMutex, STATE>,
+}
+
+impl<ACTIVE, DFU, STATE> AsyncTestFlash<ACTIVE, DFU, STATE>
+where
+    ACTIVE: NorFlash,
+    DFU: NorFlash,
+    STATE: NorFlash,
+{
+    pub fn new(config: BootLoaderConfig<ACTIVE, DFU, STATE>) -> Self {
+        Self {
+            active: Mutex::new(config.active),
+            dfu: Mutex::new(config.dfu),
+            state: Mutex::new(config.state),
+        }
+    }
+
+    pub fn active(&self) -> Partition<NoopRawMutex, ACTIVE> {
+        Self::create_partition(&self.active)
+    }
+
+    pub fn dfu(&self) -> Partition<NoopRawMutex, DFU> {
+        Self::create_partition(&self.dfu)
+    }
+
+    pub fn state(&self) -> Partition<NoopRawMutex, STATE> {
+        Self::create_partition(&self.state)
+    }
+
+    fn create_partition<T: NorFlash>(mutex: &Mutex<NoopRawMutex, T>) -> Partition<NoopRawMutex, T> {
+        Partition::new(mutex, 0, mutex.try_lock().unwrap().capacity() as u32)
+    }
+}
+
+impl<ACTIVE, DFU, STATE> AsyncTestFlash<ACTIVE, DFU, STATE>
+where
+    ACTIVE: NorFlash + embedded_storage::nor_flash::NorFlash,
+    DFU: NorFlash + embedded_storage::nor_flash::NorFlash,
+    STATE: NorFlash + embedded_storage::nor_flash::NorFlash,
+{
+    pub fn into_blocking(self) -> super::BlockingTestFlash<ACTIVE, DFU, STATE> {
+        let config = BootLoaderConfig {
+            active: self.active.into_inner(),
+            dfu: self.dfu.into_inner(),
+            state: self.state.into_inner(),
+        };
+        super::BlockingTestFlash::new(config)
+    }
+}
diff --git a/embassy-boot/boot/src/test_flash/blocking.rs b/embassy-boot/boot/src/test_flash/blocking.rs
new file mode 100644
index 000000000..ba33c9208
--- /dev/null
+++ b/embassy-boot/boot/src/test_flash/blocking.rs
@@ -0,0 +1,69 @@
+use core::cell::RefCell;
+
+use embassy_embedded_hal::flash::partition::BlockingPartition;
+use embassy_sync::blocking_mutex::raw::NoopRawMutex;
+use embassy_sync::blocking_mutex::Mutex;
+use embedded_storage::nor_flash::NorFlash;
+
+use crate::BootLoaderConfig;
+
+pub struct BlockingTestFlash<ACTIVE, DFU, STATE>
+where
+    ACTIVE: NorFlash,
+    DFU: NorFlash,
+    STATE: NorFlash,
+{
+    active: Mutex<NoopRawMutex, RefCell<ACTIVE>>,
+    dfu: Mutex<NoopRawMutex, RefCell<DFU>>,
+    state: Mutex<NoopRawMutex, RefCell<STATE>>,
+}
+
+impl<ACTIVE, DFU, STATE> BlockingTestFlash<ACTIVE, DFU, STATE>
+where
+    ACTIVE: NorFlash,
+    DFU: NorFlash,
+    STATE: NorFlash,
+{
+    pub fn new(config: BootLoaderConfig<ACTIVE, DFU, STATE>) -> Self {
+        Self {
+            active: Mutex::new(RefCell::new(config.active)),
+            dfu: Mutex::new(RefCell::new(config.dfu)),
+            state: Mutex::new(RefCell::new(config.state)),
+        }
+    }
+
+    pub fn active(&self) -> BlockingPartition<NoopRawMutex, ACTIVE> {
+        Self::create_partition(&self.active)
+    }
+
+    pub fn dfu(&self) -> BlockingPartition<NoopRawMutex, DFU> {
+        Self::create_partition(&self.dfu)
+    }
+
+    pub fn state(&self) -> BlockingPartition<NoopRawMutex, STATE> {
+        Self::create_partition(&self.state)
+    }
+
+    pub fn create_partition<T: NorFlash>(
+        mutex: &Mutex<NoopRawMutex, RefCell<T>>,
+    ) -> BlockingPartition<NoopRawMutex, T> {
+        BlockingPartition::new(mutex, 0, mutex.lock(|f| f.borrow().capacity()) as u32)
+    }
+}
+
+#[cfg(feature = "nightly")]
+impl<ACTIVE, DFU, STATE> BlockingTestFlash<ACTIVE, DFU, STATE>
+where
+    ACTIVE: NorFlash + embedded_storage_async::nor_flash::NorFlash,
+    DFU: NorFlash + embedded_storage_async::nor_flash::NorFlash,
+    STATE: NorFlash + embedded_storage_async::nor_flash::NorFlash,
+{
+    pub fn into_async(self) -> super::AsyncTestFlash<ACTIVE, DFU, STATE> {
+        let config = BootLoaderConfig {
+            active: self.active.into_inner().into_inner(),
+            dfu: self.dfu.into_inner().into_inner(),
+            state: self.state.into_inner().into_inner(),
+        };
+        super::AsyncTestFlash::new(config)
+    }
+}
diff --git a/embassy-boot/boot/src/test_flash/mod.rs b/embassy-boot/boot/src/test_flash/mod.rs
new file mode 100644
index 000000000..a0672322e
--- /dev/null
+++ b/embassy-boot/boot/src/test_flash/mod.rs
@@ -0,0 +1,7 @@
+#[cfg(feature = "nightly")]
+mod asynch;
+mod blocking;
+
+#[cfg(feature = "nightly")]
+pub(crate) use asynch::AsyncTestFlash;
+pub(crate) use blocking::BlockingTestFlash;
diff --git a/embassy-boot/nrf/Cargo.toml b/embassy-boot/nrf/Cargo.toml
index 234393e7c..8186a9958 100644
--- a/embassy-boot/nrf/Cargo.toml
+++ b/embassy-boot/nrf/Cargo.toml
@@ -3,6 +3,7 @@ edition = "2021"
 name = "embassy-boot-nrf"
 version = "0.1.0"
 description = "Bootloader lib for nRF chips"
+license = "MIT OR Apache-2.0"
 
 [package.metadata.embassy_docs]
 src_base = "https://github.com/embassy-rs/embassy/blob/embassy-boot-nrf-v$VERSION/embassy-boot/nrf/src/"
@@ -16,12 +17,12 @@ target = "thumbv7em-none-eabi"
 defmt = { version = "0.3", optional = true }
 
 embassy-sync = { path = "../../embassy-sync" }
-embassy-nrf = { path = "../../embassy-nrf", default-features = false, features = ["nightly"] }
+embassy-nrf = { path = "../../embassy-nrf" }
 embassy-boot = { path = "../boot", default-features = false }
 cortex-m = { version = "0.7.6" }
 cortex-m-rt = { version = "0.7" }
 embedded-storage = "0.3.0"
-embedded-storage-async = "0.3.0"
+embedded-storage-async = { version = "0.4.0", optional = true }
 cfg-if = "1.0.0"
 
 nrf-softdevice-mbr = { version = "0.1.0", git = "https://github.com/embassy-rs/nrf-softdevice.git", branch = "master", optional = true }
@@ -35,3 +36,8 @@ defmt = [
 softdevice = [
     "nrf-softdevice-mbr",
 ]
+nightly = [
+    "dep:embedded-storage-async",
+    "embassy-boot/nightly",
+    "embassy-nrf/nightly"
+]
diff --git a/embassy-boot/nrf/README.md b/embassy-boot/nrf/README.md
new file mode 100644
index 000000000..fe581823d
--- /dev/null
+++ b/embassy-boot/nrf/README.md
@@ -0,0 +1,26 @@
+# embassy-boot-nrf
+
+An [Embassy](https://embassy.dev) project.
+
+An adaptation of `embassy-boot` for nRF. 
+
+## Features
+
+* Load applications with or without the softdevice.
+* Configure bootloader partitions based on linker script.
+* Using watchdog timer to detect application failure.
+
+
+## Minimum supported Rust version (MSRV)
+
+`embassy-boot-nrf` is guaranteed to compile on the latest stable Rust version at the time of release. It might compile with older versions but that may change in any new patch release.
+
+## License
+
+This work is licensed under either of
+
+- Apache License, Version 2.0 ([LICENSE-APACHE](LICENSE-APACHE) or
+  <http://www.apache.org/licenses/LICENSE-2.0>)
+- MIT license ([LICENSE-MIT](LICENSE-MIT) or <http://opensource.org/licenses/MIT>)
+
+at your option.
diff --git a/embassy-boot/nrf/src/lib.rs b/embassy-boot/nrf/src/lib.rs
index 2d6b837cb..bb702073c 100644
--- a/embassy-boot/nrf/src/lib.rs
+++ b/embassy-boot/nrf/src/lib.rs
@@ -1,88 +1,63 @@
 #![no_std]
-#![feature(generic_associated_types)]
-#![feature(type_alias_impl_trait)]
-#![allow(incomplete_features)]
-#![feature(generic_const_exprs)]
-
+#![warn(missing_docs)]
+#![doc = include_str!("../README.md")]
 mod fmt;
 
-pub use embassy_boot::{FirmwareUpdater, FlashConfig, FlashProvider, Partition, SingleFlashProvider};
+#[cfg(feature = "nightly")]
+pub use embassy_boot::FirmwareUpdater;
+pub use embassy_boot::{AlignedBuffer, BlockingFirmwareUpdater, BootLoaderConfig, FirmwareUpdaterConfig};
 use embassy_nrf::nvmc::{Nvmc, PAGE_SIZE};
 use embassy_nrf::peripherals::WDT;
 use embassy_nrf::wdt;
 use embedded_storage::nor_flash::{ErrorType, NorFlash, ReadNorFlash};
 
-pub struct BootLoader {
-    boot: embassy_boot::BootLoader<PAGE_SIZE>,
+/// A bootloader for nRF devices.
+pub struct BootLoader<ACTIVE: NorFlash, DFU: NorFlash, STATE: NorFlash, const BUFFER_SIZE: usize = PAGE_SIZE> {
+    boot: embassy_boot::BootLoader<ACTIVE, DFU, STATE>,
+    aligned_buf: AlignedBuffer<BUFFER_SIZE>,
 }
 
-impl BootLoader {
-    /// Create a new bootloader instance using parameters from linker script
-    pub fn default() -> Self {
-        extern "C" {
-            static __bootloader_state_start: u32;
-            static __bootloader_state_end: u32;
-            static __bootloader_active_start: u32;
-            static __bootloader_active_end: u32;
-            static __bootloader_dfu_start: u32;
-            static __bootloader_dfu_end: u32;
-        }
-
-        let active = unsafe {
-            Partition::new(
-                &__bootloader_active_start as *const u32 as usize,
-                &__bootloader_active_end as *const u32 as usize,
-            )
-        };
-        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!("ACTIVE: 0x{:x} - 0x{:x}", active.from, active.to);
-        trace!("DFU: 0x{:x} - 0x{:x}", dfu.from, dfu.to);
-        trace!("STATE: 0x{:x} - 0x{:x}", state.from, state.to);
-
-        Self::new(active, dfu, state)
-    }
-
+impl<ACTIVE: NorFlash, DFU: NorFlash, STATE: NorFlash, const BUFFER_SIZE: usize>
+    BootLoader<ACTIVE, DFU, STATE, 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 {
+    pub fn new(config: BootLoaderConfig<ACTIVE, DFU, STATE>) -> Self {
         Self {
-            boot: embassy_boot::BootLoader::new(active, dfu, state),
+            boot: embassy_boot::BootLoader::new(config),
+            aligned_buf: AlignedBuffer([0; BUFFER_SIZE]),
         }
     }
 
-    /// Boots the application without softdevice mechanisms
-    pub fn prepare<F: FlashProvider>(&mut self, flash: &mut F) -> usize
-    where
-        [(); <<F as FlashProvider>::STATE as FlashConfig>::FLASH::WRITE_SIZE]:,
-        [(); <<F as FlashProvider>::ACTIVE as FlashConfig>::FLASH::ERASE_SIZE]:,
-    {
-        match self.boot.prepare_boot(flash) {
-            Ok(_) => self.boot.boot_address(),
-            Err(_) => panic!("boot prepare error!"),
-        }
+    /// Inspect the bootloader state and perform actions required before booting, such as swapping
+    /// firmware.
+    pub fn prepare(&mut self) {
+        self.boot
+            .prepare_boot(&mut self.aligned_buf.0)
+            .expect("Boot prepare error");
     }
 
+    /// Boots the application without softdevice mechanisms.
+    ///
+    /// # Safety
+    ///
+    /// This modifies the stack pointer and reset vector and will run code placed in the active partition.
     #[cfg(not(feature = "softdevice"))]
-    pub unsafe fn load(&mut self, start: usize) -> ! {
+    pub unsafe fn load(self, start: u32) -> ! {
+        core::mem::drop(self.boot);
+
         let mut p = cortex_m::Peripherals::steal();
         p.SCB.invalidate_icache();
-        p.SCB.vtor.write(start as u32);
+        p.SCB.vtor.write(start);
         cortex_m::asm::bootload(start as *const u32)
     }
 
+    /// Boots the application assuming softdevice is present.
+    ///
+    /// # Safety
+    ///
+    /// This modifies the stack pointer and reset vector and will run code placed in the active partition.
     #[cfg(feature = "softdevice")]
-    pub unsafe fn load(&mut self, _app: usize) -> ! {
+    pub unsafe fn load(&mut self, _app: u32) -> ! {
         use nrf_softdevice_mbr as mbr;
         const NRF_SUCCESS: u32 = 0;
 
@@ -137,11 +112,7 @@ pub struct WatchdogFlash<'d> {
 
 impl<'d> WatchdogFlash<'d> {
     /// Start a new watchdog with a given flash and WDT peripheral and a timeout
-    pub fn start(flash: Nvmc<'d>, wdt: WDT, timeout: u32) -> Self {
-        let mut config = wdt::Config::default();
-        config.timeout_ticks = 32768 * timeout; // timeout seconds
-        config.run_during_sleep = true;
-        config.run_during_debug_halt = false;
+    pub fn start(flash: Nvmc<'d>, wdt: WDT, config: wdt::Config) -> Self {
         let (_wdt, [wdt]) = match wdt::Watchdog::try_new(wdt, config) {
             Ok(x) => x,
             Err(_) => {
diff --git a/embassy-boot/rp/Cargo.toml b/embassy-boot/rp/Cargo.toml
new file mode 100644
index 000000000..5147392ce
--- /dev/null
+++ b/embassy-boot/rp/Cargo.toml
@@ -0,0 +1,79 @@
+[package]
+edition = "2021"
+name = "embassy-boot-rp"
+version = "0.1.0"
+description = "Bootloader lib for RP2040 chips"
+license = "MIT OR Apache-2.0"
+
+[package.metadata.embassy_docs]
+src_base = "https://github.com/embassy-rs/embassy/blob/embassy-boot-rp-v$VERSION/src/"
+src_base_git = "https://github.com/embassy-rs/embassy/blob/$COMMIT/embassy-boot/rp/src/"
+target = "thumbv6m-none-eabi"
+
+[lib]
+
+[dependencies]
+defmt = { version = "0.3", optional = true }
+defmt-rtt = { version = "0.4", optional = true }
+log = { version = "0.4", optional = true }
+
+embassy-sync = { path = "../../embassy-sync" }
+embassy-rp = { path = "../../embassy-rp", default-features = false }
+embassy-boot = { path = "../boot", default-features = false }
+embassy-time = { path = "../../embassy-time" }
+
+cortex-m = { version = "0.7.6" }
+cortex-m-rt = { version = "0.7" }
+embedded-storage = "0.3.0"
+embedded-storage-async = { version = "0.4.0", optional = true }
+cfg-if = "1.0.0"
+
+[features]
+defmt = [
+    "dep:defmt",
+    "embassy-boot/defmt",
+    "embassy-rp/defmt",
+]
+log = [
+    "dep:log",
+    "embassy-boot/log",
+    "embassy-rp/log",
+]
+debug = ["defmt-rtt"]
+nightly = [
+    "dep:embedded-storage-async",
+    "embassy-boot/nightly",
+    "embassy-rp/nightly",
+    "embassy-time/nightly"
+]
+
+[profile.dev]
+debug = 2
+debug-assertions = true
+incremental = false
+opt-level = 'z'
+overflow-checks = true
+
+[profile.release]
+codegen-units = 1
+debug = 2
+debug-assertions = false
+incremental = false
+lto = 'fat'
+opt-level = 'z'
+overflow-checks = false
+
+# do not optimize proc-macro crates = faster builds from scratch
+[profile.dev.build-override]
+codegen-units = 8
+debug = false
+debug-assertions = false
+opt-level = 0
+overflow-checks = false
+
+[profile.release.build-override]
+codegen-units = 8
+debug = false
+debug-assertions = false
+opt-level = 0
+overflow-checks = false
diff --git a/embassy-boot/rp/README.md b/embassy-boot/rp/README.md
new file mode 100644
index 000000000..315d655e3
--- /dev/null
+++ b/embassy-boot/rp/README.md
@@ -0,0 +1,26 @@
+# embassy-boot-rp
+
+An [Embassy](https://embassy.dev) project.
+
+An adaptation of `embassy-boot` for RP2040.
+
+NOTE: The applications using this bootloader should not link with the `link-rp.x` linker script.
+
+## Features
+
+* Configure bootloader partitions based on linker script.
+* Load applications from active partition.
+
+## Minimum supported Rust version (MSRV)
+
+`embassy-boot-rp` is guaranteed to compile on the latest stable Rust version at the time of release. It might compile with older versions but that may change in any new patch release.
+
+## License
+
+This work is licensed under either of
+
+- Apache License, Version 2.0 ([LICENSE-APACHE](LICENSE-APACHE) or
+  <http://www.apache.org/licenses/LICENSE-2.0>)
+- MIT license ([LICENSE-MIT](LICENSE-MIT) or <http://opensource.org/licenses/MIT>)
+
+at your option.
diff --git a/embassy-boot/rp/build.rs b/embassy-boot/rp/build.rs
new file mode 100644
index 000000000..2cbc7ef5e
--- /dev/null
+++ b/embassy-boot/rp/build.rs
@@ -0,0 +1,8 @@
+use std::env;
+
+fn main() {
+    let target = env::var("TARGET").unwrap();
+    if target.starts_with("thumbv6m-") {
+        println!("cargo:rustc-cfg=armv6m");
+    }
+}
diff --git a/embassy-boot/rp/src/fmt.rs b/embassy-boot/rp/src/fmt.rs
new file mode 100644
index 000000000..066970813
--- /dev/null
+++ b/embassy-boot/rp/src/fmt.rs
@@ -0,0 +1,225 @@
+#![macro_use]
+#![allow(unused_macros)]
+
+#[cfg(all(feature = "defmt", feature = "log"))]
+compile_error!("You may not enable both `defmt` and `log` features.");
+
+macro_rules! assert {
+    ($($x:tt)*) => {
+        {
+            #[cfg(not(feature = "defmt"))]
+            ::core::assert!($($x)*);
+            #[cfg(feature = "defmt")]
+            ::defmt::assert!($($x)*);
+        }
+    };
+}
+
+macro_rules! assert_eq {
+    ($($x:tt)*) => {
+        {
+            #[cfg(not(feature = "defmt"))]
+            ::core::assert_eq!($($x)*);
+            #[cfg(feature = "defmt")]
+            ::defmt::assert_eq!($($x)*);
+        }
+    };
+}
+
+macro_rules! assert_ne {
+    ($($x:tt)*) => {
+        {
+            #[cfg(not(feature = "defmt"))]
+            ::core::assert_ne!($($x)*);
+            #[cfg(feature = "defmt")]
+            ::defmt::assert_ne!($($x)*);
+        }
+    };
+}
+
+macro_rules! debug_assert {
+    ($($x:tt)*) => {
+        {
+            #[cfg(not(feature = "defmt"))]
+            ::core::debug_assert!($($x)*);
+            #[cfg(feature = "defmt")]
+            ::defmt::debug_assert!($($x)*);
+        }
+    };
+}
+
+macro_rules! debug_assert_eq {
+    ($($x:tt)*) => {
+        {
+            #[cfg(not(feature = "defmt"))]
+            ::core::debug_assert_eq!($($x)*);
+            #[cfg(feature = "defmt")]
+            ::defmt::debug_assert_eq!($($x)*);
+        }
+    };
+}
+
+macro_rules! debug_assert_ne {
+    ($($x:tt)*) => {
+        {
+            #[cfg(not(feature = "defmt"))]
+            ::core::debug_assert_ne!($($x)*);
+            #[cfg(feature = "defmt")]
+            ::defmt::debug_assert_ne!($($x)*);
+        }
+    };
+}
+
+macro_rules! todo {
+    ($($x:tt)*) => {
+        {
+            #[cfg(not(feature = "defmt"))]
+            ::core::todo!($($x)*);
+            #[cfg(feature = "defmt")]
+            ::defmt::todo!($($x)*);
+        }
+    };
+}
+
+macro_rules! unreachable {
+    ($($x:tt)*) => {
+        {
+            #[cfg(not(feature = "defmt"))]
+            ::core::unreachable!($($x)*);
+            #[cfg(feature = "defmt")]
+            ::defmt::unreachable!($($x)*);
+        }
+    };
+}
+
+macro_rules! panic {
+    ($($x:tt)*) => {
+        {
+            #[cfg(not(feature = "defmt"))]
+            ::core::panic!($($x)*);
+            #[cfg(feature = "defmt")]
+            ::defmt::panic!($($x)*);
+        }
+    };
+}
+
+macro_rules! trace {
+    ($s:literal $(, $x:expr)* $(,)?) => {
+        {
+            #[cfg(feature = "log")]
+            ::log::trace!($s $(, $x)*);
+            #[cfg(feature = "defmt")]
+            ::defmt::trace!($s $(, $x)*);
+            #[cfg(not(any(feature = "log", feature="defmt")))]
+            let _ = ($( & $x ),*);
+        }
+    };
+}
+
+macro_rules! debug {
+    ($s:literal $(, $x:expr)* $(,)?) => {
+        {
+            #[cfg(feature = "log")]
+            ::log::debug!($s $(, $x)*);
+            #[cfg(feature = "defmt")]
+            ::defmt::debug!($s $(, $x)*);
+            #[cfg(not(any(feature = "log", feature="defmt")))]
+            let _ = ($( & $x ),*);
+        }
+    };
+}
+
+macro_rules! info {
+    ($s:literal $(, $x:expr)* $(,)?) => {
+        {
+            #[cfg(feature = "log")]
+            ::log::info!($s $(, $x)*);
+            #[cfg(feature = "defmt")]
+            ::defmt::info!($s $(, $x)*);
+            #[cfg(not(any(feature = "log", feature="defmt")))]
+            let _ = ($( & $x ),*);
+        }
+    };
+}
+
+macro_rules! warn {
+    ($s:literal $(, $x:expr)* $(,)?) => {
+        {
+            #[cfg(feature = "log")]
+            ::log::warn!($s $(, $x)*);
+            #[cfg(feature = "defmt")]
+            ::defmt::warn!($s $(, $x)*);
+            #[cfg(not(any(feature = "log", feature="defmt")))]
+            let _ = ($( & $x ),*);
+        }
+    };
+}
+
+macro_rules! error {
+    ($s:literal $(, $x:expr)* $(,)?) => {
+        {
+            #[cfg(feature = "log")]
+            ::log::error!($s $(, $x)*);
+            #[cfg(feature = "defmt")]
+            ::defmt::error!($s $(, $x)*);
+            #[cfg(not(any(feature = "log", feature="defmt")))]
+            let _ = ($( & $x ),*);
+        }
+    };
+}
+
+#[cfg(feature = "defmt")]
+macro_rules! unwrap {
+    ($($x:tt)*) => {
+        ::defmt::unwrap!($($x)*)
+    };
+}
+
+#[cfg(not(feature = "defmt"))]
+macro_rules! unwrap {
+    ($arg:expr) => {
+        match $crate::fmt::Try::into_result($arg) {
+            ::core::result::Result::Ok(t) => t,
+            ::core::result::Result::Err(e) => {
+                ::core::panic!("unwrap of `{}` failed: {:?}", ::core::stringify!($arg), e);
+            }
+        }
+    };
+    ($arg:expr, $($msg:expr),+ $(,)? ) => {
+        match $crate::fmt::Try::into_result($arg) {
+            ::core::result::Result::Ok(t) => t,
+            ::core::result::Result::Err(e) => {
+                ::core::panic!("unwrap of `{}` failed: {}: {:?}", ::core::stringify!($arg), ::core::format_args!($($msg,)*), e);
+            }
+        }
+    }
+}
+
+#[derive(Debug, Copy, Clone, Eq, PartialEq)]
+pub struct NoneError;
+
+pub trait Try {
+    type Ok;
+    type Error;
+    fn into_result(self) -> Result<Self::Ok, Self::Error>;
+}
+
+impl<T> Try for Option<T> {
+    type Ok = T;
+    type Error = NoneError;
+
+    #[inline]
+    fn into_result(self) -> Result<T, NoneError> {
+        self.ok_or(NoneError)
+    }
+}
+
+impl<T, E> Try for Result<T, E> {
+    type Ok = T;
+    type Error = E;
+
+    #[inline]
+    fn into_result(self) -> Self {
+        self
+    }
+}
diff --git a/embassy-boot/rp/src/lib.rs b/embassy-boot/rp/src/lib.rs
new file mode 100644
index 000000000..25329f9e9
--- /dev/null
+++ b/embassy-boot/rp/src/lib.rs
@@ -0,0 +1,102 @@
+#![no_std]
+#![warn(missing_docs)]
+#![doc = include_str!("../README.md")]
+mod fmt;
+
+#[cfg(feature = "nightly")]
+pub use embassy_boot::FirmwareUpdater;
+pub use embassy_boot::{AlignedBuffer, BlockingFirmwareUpdater, BootLoaderConfig, FirmwareUpdaterConfig, State};
+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<ACTIVE: NorFlash, DFU: NorFlash, STATE: NorFlash, const BUFFER_SIZE: usize = ERASE_SIZE> {
+    boot: embassy_boot::BootLoader<ACTIVE, DFU, STATE>,
+    aligned_buf: AlignedBuffer<BUFFER_SIZE>,
+}
+
+impl<ACTIVE: NorFlash, DFU: NorFlash, STATE: NorFlash, const BUFFER_SIZE: usize>
+    BootLoader<ACTIVE, DFU, STATE, BUFFER_SIZE>
+{
+    /// Create a new bootloader instance using the supplied partitions for active, dfu and state.
+    pub fn new(config: BootLoaderConfig<ACTIVE, DFU, STATE>) -> Self {
+        Self {
+            boot: embassy_boot::BootLoader::new(config),
+            aligned_buf: AlignedBuffer([0; BUFFER_SIZE]),
+        }
+    }
+
+    /// Inspect the bootloader state and perform actions required before booting, such as swapping
+    /// firmware.
+    pub fn prepare(&mut self) {
+        self.boot
+            .prepare_boot(self.aligned_buf.as_mut())
+            .expect("Boot prepare error");
+    }
+
+    /// Boots the application.
+    ///
+    /// # Safety
+    ///
+    /// This modifies the stack pointer and reset vector and will run code placed in the active partition.
+    pub unsafe fn load(self, start: u32) -> ! {
+        core::mem::drop(self.boot);
+
+        trace!("Loading app at 0x{:x}", start);
+        #[allow(unused_mut)]
+        let mut p = cortex_m::Peripherals::steal();
+        #[cfg(not(armv6m))]
+        p.SCB.invalidate_icache();
+        p.SCB.vtor.write(start);
+
+        cortex_m::asm::bootload(start as *const u32)
+    }
+}
+
+/// A flash implementation that will feed a watchdog when touching flash.
+pub struct WatchdogFlash<'d, const SIZE: usize> {
+    flash: Flash<'d, FLASH, SIZE>,
+    watchdog: Watchdog,
+}
+
+impl<'d, const SIZE: usize> WatchdogFlash<'d, SIZE> {
+    /// Start a new watchdog with a given flash and watchdog peripheral and a timeout
+    pub fn start(flash: FLASH, watchdog: WATCHDOG, timeout: Duration) -> Self {
+        let flash: Flash<'_, FLASH, SIZE> = Flash::new(flash);
+        let mut watchdog = Watchdog::new(watchdog);
+        watchdog.start(timeout);
+        Self { flash, watchdog }
+    }
+}
+
+impl<'d, const SIZE: usize> ErrorType for WatchdogFlash<'d, SIZE> {
+    type Error = <Flash<'d, FLASH, SIZE> as ErrorType>::Error;
+}
+
+impl<'d, const SIZE: usize> NorFlash for WatchdogFlash<'d, SIZE> {
+    const WRITE_SIZE: usize = <Flash<'d, FLASH, SIZE> as NorFlash>::WRITE_SIZE;
+    const ERASE_SIZE: usize = <Flash<'d, FLASH, SIZE> as NorFlash>::ERASE_SIZE;
+
+    fn erase(&mut self, from: u32, to: u32) -> Result<(), Self::Error> {
+        self.watchdog.feed();
+        self.flash.erase(from, to)
+    }
+    fn write(&mut self, offset: u32, data: &[u8]) -> Result<(), Self::Error> {
+        self.watchdog.feed();
+        self.flash.write(offset, data)
+    }
+}
+
+impl<'d, const SIZE: usize> ReadNorFlash for WatchdogFlash<'d, SIZE> {
+    const READ_SIZE: usize = <Flash<'d, FLASH, SIZE> as ReadNorFlash>::READ_SIZE;
+    fn read(&mut self, offset: u32, data: &mut [u8]) -> Result<(), Self::Error> {
+        self.watchdog.feed();
+        self.flash.read(offset, data)
+    }
+    fn capacity(&self) -> usize {
+        self.flash.capacity()
+    }
+}
diff --git a/embassy-boot/stm32/Cargo.toml b/embassy-boot/stm32/Cargo.toml
index ad4657e0d..99a6b8e0e 100644
--- a/embassy-boot/stm32/Cargo.toml
+++ b/embassy-boot/stm32/Cargo.toml
@@ -3,6 +3,7 @@ edition = "2021"
 name = "embassy-boot-stm32"
 version = "0.1.0"
 description = "Bootloader lib for STM32 chips"
+license = "MIT OR Apache-2.0"
 
 [package.metadata.embassy_docs]
 src_base = "https://github.com/embassy-rs/embassy/blob/embassy-boot-nrf-v$VERSION/embassy-boot/stm32/src/"
@@ -14,16 +15,16 @@ target = "thumbv7em-none-eabi"
 
 [dependencies]
 defmt = { version = "0.3", optional = true }
-defmt-rtt = { version = "0.3", optional = true }
+defmt-rtt = { version = "0.4", optional = true }
 log = { version = "0.4", optional = true }
 
 embassy-sync = { path = "../../embassy-sync" }
-embassy-stm32 = { path = "../../embassy-stm32", default-features = false, features = ["nightly"] }
+embassy-stm32 = { path = "../../embassy-stm32", default-features = false }
 embassy-boot = { path = "../boot", default-features = false }
 cortex-m = { version = "0.7.6" }
 cortex-m-rt = { version = "0.7" }
 embedded-storage = "0.3.0"
-embedded-storage-async = "0.3.0"
+embedded-storage-async = { version = "0.4.0", optional = true }
 cfg-if = "1.0.0"
 
 [features]
@@ -38,6 +39,11 @@ log = [
     "embassy-stm32/log",
 ]
 debug = ["defmt-rtt"]
+nightly = [
+    "dep:embedded-storage-async",
+    "embassy-boot/nightly",
+    "embassy-stm32/nightly"
+]
 
 [profile.dev]
 debug = 2
diff --git a/embassy-boot/stm32/README.md b/embassy-boot/stm32/README.md
index a82b730b9..b4d7ba5a4 100644
--- a/embassy-boot/stm32/README.md
+++ b/embassy-boot/stm32/README.md
@@ -1,11 +1,24 @@
-# Bootloader for STM32
+# embassy-boot-stm32
 
-The bootloader uses `embassy-boot` to interact with the flash.
+An [Embassy](https://embassy.dev) project.
 
-# Usage
+An adaptation of `embassy-boot` for STM32.
 
-Flash the bootloader
+## Features
 
-```
-cargo flash --features embassy-stm32/stm32wl55jc-cm4 --release --chip STM32WLE5JCIx
-```
+* Configure bootloader partitions based on linker script.
+* Load applications from active partition.
+
+## Minimum supported Rust version (MSRV)
+
+`embassy-boot-stm32` is guaranteed to compile on the latest stable Rust version at the time of release. It might compile with older versions but that may change in any new patch release.
+
+## License
+
+This work is licensed under either of
+
+- Apache License, Version 2.0 ([LICENSE-APACHE](LICENSE-APACHE) or
+  <http://www.apache.org/licenses/LICENSE-2.0>)
+- MIT license ([LICENSE-MIT](LICENSE-MIT) or <http://opensource.org/licenses/MIT>)
+
+at your option.
diff --git a/embassy-boot/stm32/src/lib.rs b/embassy-boot/stm32/src/lib.rs
index 5a4f2d058..069de0d1a 100644
--- a/embassy-boot/stm32/src/lib.rs
+++ b/embassy-boot/stm32/src/lib.rs
@@ -1,82 +1,52 @@
 #![no_std]
-#![feature(generic_associated_types)]
-#![feature(type_alias_impl_trait)]
-#![allow(incomplete_features)]
-#![feature(generic_const_exprs)]
-
+#![warn(missing_docs)]
+#![doc = include_str!("../README.md")]
 mod fmt;
 
-pub use embassy_boot::{FirmwareUpdater, FlashConfig, FlashProvider, Partition, SingleFlashProvider, State};
+#[cfg(feature = "nightly")]
+pub use embassy_boot::FirmwareUpdater;
+pub use embassy_boot::{AlignedBuffer, BlockingFirmwareUpdater, BootLoaderConfig, FirmwareUpdaterConfig, State};
 use embedded_storage::nor_flash::NorFlash;
 
-pub struct BootLoader<const PAGE_SIZE: usize> {
-    boot: embassy_boot::BootLoader<PAGE_SIZE>,
+/// A bootloader for STM32 devices.
+pub struct BootLoader<ACTIVE: NorFlash, DFU: NorFlash, STATE: NorFlash, const BUFFER_SIZE: usize> {
+    boot: embassy_boot::BootLoader<ACTIVE, DFU, STATE>,
+    aligned_buf: AlignedBuffer<BUFFER_SIZE>,
 }
 
-impl<const PAGE_SIZE: usize> BootLoader<PAGE_SIZE> {
-    /// Create a new bootloader instance using parameters from linker script
-    pub fn default() -> Self {
-        extern "C" {
-            static __bootloader_state_start: u32;
-            static __bootloader_state_end: u32;
-            static __bootloader_active_start: u32;
-            static __bootloader_active_end: u32;
-            static __bootloader_dfu_start: u32;
-            static __bootloader_dfu_end: u32;
-        }
-
-        let active = unsafe {
-            Partition::new(
-                &__bootloader_active_start as *const u32 as usize,
-                &__bootloader_active_end as *const u32 as usize,
-            )
-        };
-        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!("ACTIVE: 0x{:x} - 0x{:x}", active.from, active.to);
-        trace!("DFU: 0x{:x} - 0x{:x}", dfu.from, dfu.to);
-        trace!("STATE: 0x{:x} - 0x{:x}", state.from, state.to);
-
-        Self::new(active, dfu, state)
-    }
-
+impl<ACTIVE: NorFlash, DFU: NorFlash, STATE: NorFlash, const BUFFER_SIZE: usize>
+    BootLoader<ACTIVE, DFU, STATE, 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 {
+    pub fn new(config: BootLoaderConfig<ACTIVE, DFU, STATE>) -> Self {
         Self {
-            boot: embassy_boot::BootLoader::new(active, dfu, state),
+            boot: embassy_boot::BootLoader::new(config),
+            aligned_buf: AlignedBuffer([0; BUFFER_SIZE]),
         }
     }
 
-    /// Boots the application
-    pub fn prepare<F: FlashProvider>(&mut self, flash: &mut F) -> usize
-    where
-        [(); <<F as FlashProvider>::STATE as FlashConfig>::FLASH::WRITE_SIZE]:,
-        [(); <<F as FlashProvider>::ACTIVE as FlashConfig>::FLASH::ERASE_SIZE]:,
-    {
-        match self.boot.prepare_boot(flash) {
-            Ok(_) => embassy_stm32::flash::FLASH_BASE + self.boot.boot_address(),
-            Err(_) => panic!("boot prepare error!"),
-        }
+    /// Inspect the bootloader state and perform actions required before booting, such as swapping
+    /// firmware.
+    pub fn prepare(&mut self) {
+        self.boot
+            .prepare_boot(self.aligned_buf.as_mut())
+            .expect("Boot prepare error");
     }
 
-    pub unsafe fn load(&mut self, start: usize) -> ! {
+    /// Boots the application.
+    ///
+    /// # Safety
+    ///
+    /// This modifies the stack pointer and reset vector and will run code placed in the active partition.
+    pub unsafe fn load(self, start: u32) -> ! {
+        core::mem::drop(self.boot);
+
         trace!("Loading app at 0x{:x}", start);
         #[allow(unused_mut)]
         let mut p = cortex_m::Peripherals::steal();
         #[cfg(not(armv6m))]
         p.SCB.invalidate_icache();
-        p.SCB.vtor.write(start as u32);
+        p.SCB.vtor.write(start);
 
         cortex_m::asm::bootload(start as *const u32)
     }