6 files changed, 468 insertions, 439 deletions

diff --git a/embassy-boot/boot/src/boot_loader.rs b/embassy-boot/boot/src/boot_loader.rs
index ad6735112..9d047f778 100644
--- a/embassy-boot/boot/src/boot_loader.rs
+++ b/embassy-boot/boot/src/boot_loader.rs
@@ -31,7 +31,7 @@ where
 }
 
 /// Extension of the embedded-storage flash type information with block size and erase value.
-pub trait Flash: NorFlash + ReadNorFlash {
+pub trait Flash: NorFlash {
     /// The block size that should be used when writing to flash. For most builtin flashes, this is the same as the erase
     /// size of the flash, but for external QSPI flash modules, this can be lower.
     const BLOCK_SIZE: usize;
@@ -60,9 +60,11 @@ pub trait FlashConfig {
 /// different page sizes and flash write sizes.
 pub struct BootLoader {
     // Page with current state of bootloader. The state partition has the following format:
-    // | Range          | Description                                                                      |
-    // | 0 - WRITE_SIZE | Magic indicating bootloader state. BOOT_MAGIC means boot, SWAP_MAGIC means swap. |
-    // | WRITE_SIZE - N | Progress index used while swapping or reverting                                  |
+    // All ranges are in multiples of WRITE_SIZE bytes.
+    // | Range    | Description                                                                      |
+    // | 0..1     | Magic indicating bootloader state. BOOT_MAGIC means boot, SWAP_MAGIC means swap. |
+    // | 1..2     | Progress validity. ERASE_VALUE means valid, !ERASE_VALUE means invalid.          |
+    // | 2..2 + N | Progress index used while swapping or reverting                                  |
     state: Partition,
     // Location of the partition which will be booted from
     active: Partition,
@@ -79,7 +81,7 @@ impl BootLoader {
         Self { active, dfu, state }
     }
 
-    /// Return the boot address for the active partition.
+    /// Return the offset of the active partition into the active flash.
     pub fn boot_address(&self) -> usize {
         self.active.from
     }
@@ -192,14 +194,19 @@ impl BootLoader {
                 trace!("Reverting");
                 self.revert(p, magic, page)?;
 
-                // Overwrite magic and reset progress
-                let fstate = p.state();
+                let state_flash = p.state();
+
+                // Invalidate progress
                 magic.fill(!P::STATE::ERASE_VALUE);
-                fstate.write(self.state.from as u32, magic)?;
-                fstate.erase(self.state.from as u32, self.state.to as u32)?;
+                self.state
+                    .write_blocking(state_flash, P::STATE::WRITE_SIZE as u32, magic)?;
+
+                // Clear magic and progress
+                self.state.wipe_blocking(state_flash)?;
 
+                // Set magic
                 magic.fill(BOOT_MAGIC);
-                fstate.write(self.state.from as u32, magic)?;
+                self.state.write_blocking(state_flash, 0, magic)?;
             }
         }
         Ok(state)
@@ -215,62 +222,61 @@ impl BootLoader {
 
     fn current_progress<P: FlashConfig>(&mut self, config: &mut P, aligned: &mut [u8]) -> Result<usize, BootError> {
         let write_size = aligned.len();
-        let max_index = ((self.state.len() - write_size) / write_size) - 1;
+        let max_index = ((self.state.len() - write_size) / write_size) - 2;
         aligned.fill(!P::STATE::ERASE_VALUE);
 
-        let flash = config.state();
-        for i in 0..max_index {
-            flash.read((self.state.from + write_size + i * write_size) as u32, aligned)?;
+        let state_flash = config.state();
+
+        self.state
+            .read_blocking(state_flash, P::STATE::WRITE_SIZE as u32, aligned)?;
+        if aligned.iter().any(|&b| b != P::STATE::ERASE_VALUE) {
+            // Progress is invalid
+            return Ok(max_index);
+        }
+
+        for index in 0..max_index {
+            self.state
+                .read_blocking(state_flash, (2 + index) as u32 * P::STATE::WRITE_SIZE as u32, aligned)?;
 
             if aligned.iter().any(|&b| b == P::STATE::ERASE_VALUE) {
-                return Ok(i);
+                return Ok(index);
             }
         }
         Ok(max_index)
     }
 
-    fn update_progress<P: FlashConfig>(&mut self, idx: usize, p: &mut P, magic: &mut [u8]) -> Result<(), BootError> {
-        let flash = p.state();
-        let write_size = magic.len();
-        let w = self.state.from + write_size + idx * write_size;
-
+    fn update_progress<P: FlashConfig>(&mut self, index: usize, p: &mut P, magic: &mut [u8]) -> Result<(), BootError> {
         let aligned = magic;
         aligned.fill(!P::STATE::ERASE_VALUE);
-        flash.write(w as u32, aligned)?;
+        self.state
+            .write_blocking(p.state(), (2 + index) as u32 * P::STATE::WRITE_SIZE as u32, aligned)?;
         Ok(())
     }
 
-    fn active_addr(&self, n: usize, page_size: usize) -> usize {
-        self.active.from + n * page_size
-    }
-
-    fn dfu_addr(&self, n: usize, page_size: usize) -> usize {
-        self.dfu.from + n * page_size
-    }
-
     fn copy_page_once_to_active<P: FlashConfig>(
         &mut self,
         idx: usize,
-        from_page: usize,
-        to_page: usize,
+        from_offset: u32,
+        to_offset: u32,
         p: &mut P,
         magic: &mut [u8],
         page: &mut [u8],
     ) -> Result<(), BootError> {
         let buf = page;
         if self.current_progress(p, magic)? <= idx {
-            let mut offset = from_page;
+            let mut offset = from_offset;
             for chunk in buf.chunks_mut(P::DFU::BLOCK_SIZE) {
-                p.dfu().read(offset as u32, chunk)?;
-                offset += chunk.len();
+                self.dfu.read_blocking(p.dfu(), offset, chunk)?;
+                offset += chunk.len() as u32;
             }
 
-            p.active().erase(to_page as u32, (to_page + buf.len()) as u32)?;
+            self.active
+                .erase_blocking(p.active(), to_offset, to_offset + buf.len() as u32)?;
 
-            let mut offset = to_page;
+            let mut offset = to_offset;
             for chunk in buf.chunks(P::ACTIVE::BLOCK_SIZE) {
-                p.active().write(offset as u32, chunk)?;
-                offset += chunk.len();
+                self.active.write_blocking(p.active(), offset, chunk)?;
+                offset += chunk.len() as u32;
             }
             self.update_progress(idx, p, magic)?;
         }
@@ -280,26 +286,27 @@ impl BootLoader {
     fn copy_page_once_to_dfu<P: FlashConfig>(
         &mut self,
         idx: usize,
-        from_page: usize,
-        to_page: usize,
+        from_offset: u32,
+        to_offset: u32,
         p: &mut P,
         magic: &mut [u8],
         page: &mut [u8],
     ) -> Result<(), BootError> {
         let buf = page;
         if self.current_progress(p, magic)? <= idx {
-            let mut offset = from_page;
+            let mut offset = from_offset;
             for chunk in buf.chunks_mut(P::ACTIVE::BLOCK_SIZE) {
-                p.active().read(offset as u32, chunk)?;
-                offset += chunk.len();
+                self.active.read_blocking(p.active(), offset, chunk)?;
+                offset += chunk.len() as u32;
             }
 
-            p.dfu().erase(to_page as u32, (to_page + buf.len()) as u32)?;
+            self.dfu
+                .erase_blocking(p.dfu(), to_offset as u32, to_offset + buf.len() as u32)?;
 
-            let mut offset = to_page;
+            let mut offset = to_offset;
             for chunk in buf.chunks(P::DFU::BLOCK_SIZE) {
-                p.dfu().write(offset as u32, chunk)?;
-                offset += chunk.len();
+                self.dfu.write_blocking(p.dfu(), offset, chunk)?;
+                offset += chunk.len() as u32;
             }
             self.update_progress(idx, p, magic)?;
         }
@@ -312,17 +319,20 @@ impl BootLoader {
         trace!("Page count: {}", page_count);
         for page_num in 0..page_count {
             trace!("COPY PAGE {}", page_num);
+
+            let idx = page_num * 2;
+
             // Copy active page to the 'next' DFU page.
-            let active_page = self.active_addr(page_count - 1 - page_num, page_size);
-            let dfu_page = self.dfu_addr(page_count - page_num, page_size);
-            //trace!("Copy active {} to dfu {}", active_page, dfu_page);
-            self.copy_page_once_to_dfu(page_num * 2, active_page, dfu_page, p, magic, page)?;
+            let active_from_offset = ((page_count - 1 - page_num) * page_size) as u32;
+            let dfu_to_offset = ((page_count - page_num) * page_size) as u32;
+            //trace!("Copy active {} to dfu {}", active_from_offset, dfu_to_offset);
+            self.copy_page_once_to_dfu(idx, active_from_offset, dfu_to_offset, p, magic, page)?;
 
             // Copy DFU page to the active page
-            let active_page = self.active_addr(page_count - 1 - page_num, page_size);
-            let dfu_page = self.dfu_addr(page_count - 1 - page_num, page_size);
-            //trace!("Copy dfy {} to active {}", dfu_page, active_page);
-            self.copy_page_once_to_active(page_num * 2 + 1, dfu_page, active_page, p, magic, page)?;
+            let active_to_offset = ((page_count - 1 - page_num) * page_size) as u32;
+            let dfu_from_offset = ((page_count - 1 - page_num) * page_size) as u32;
+            //trace!("Copy dfy {} to active {}", dfu_from_offset, active_to_offset);
+            self.copy_page_once_to_active(idx + 1, dfu_from_offset, active_to_offset, p, magic, page)?;
         }
 
         Ok(())
@@ -332,23 +342,24 @@ impl BootLoader {
         let page_size = page.len();
         let page_count = self.active.len() / page_size;
         for page_num in 0..page_count {
+            let idx = page_count * 2 + page_num * 2;
+
             // Copy the bad active page to the DFU page
-            let active_page = self.active_addr(page_num, page_size);
-            let dfu_page = self.dfu_addr(page_num, page_size);
-            self.copy_page_once_to_dfu(page_count * 2 + page_num * 2, active_page, dfu_page, p, magic, page)?;
+            let active_from_offset = (page_num * page_size) as u32;
+            let dfu_to_offset = (page_num * page_size) as u32;
+            self.copy_page_once_to_dfu(idx, active_from_offset, dfu_to_offset, p, magic, page)?;
 
             // Copy the DFU page back to the active page
-            let active_page = self.active_addr(page_num, page_size);
-            let dfu_page = self.dfu_addr(page_num + 1, page_size);
-            self.copy_page_once_to_active(page_count * 2 + page_num * 2 + 1, dfu_page, active_page, p, magic, page)?;
+            let active_to_offset = (page_num * page_size) as u32;
+            let dfu_from_offset = ((page_num + 1) * page_size) as u32;
+            self.copy_page_once_to_active(idx + 1, dfu_from_offset, active_to_offset, p, magic, page)?;
         }
 
         Ok(())
     }
 
     fn read_state<P: FlashConfig>(&mut self, config: &mut P, magic: &mut [u8]) -> Result<State, BootError> {
-        let flash = config.state();
-        flash.read(self.state.from as u32, magic)?;
+        self.state.read_blocking(config.state(), 0, magic)?;
 
         if !magic.iter().any(|&b| b != SWAP_MAGIC) {
             Ok(State::Swap)
@@ -362,7 +373,7 @@ fn assert_partitions(active: Partition, dfu: Partition, state: Partition, page_s
     assert_eq!(active.len() % page_size, 0);
     assert_eq!(dfu.len() % page_size, 0);
     assert!(dfu.len() - active.len() >= page_size);
-    assert!(2 * (active.len() / page_size) <= (state.len() - write_size) / write_size);
+    assert!(2 + 2 * (active.len() / page_size) <= state.len() / write_size);
 }
 
 /// A flash wrapper implementing the Flash and embedded_storage traits.
diff --git a/embassy-boot/boot/src/firmware_updater.rs b/embassy-boot/boot/src/firmware_updater.rs
index 2d8712277..2b5cc72fa 100644
--- a/embassy-boot/boot/src/firmware_updater.rs
+++ b/embassy-boot/boot/src/firmware_updater.rs
@@ -1,7 +1,7 @@
 use embedded_storage::nor_flash::{NorFlash, NorFlashError, NorFlashErrorKind};
 use embedded_storage_async::nor_flash::NorFlash as AsyncNorFlash;
 
-use crate::{FirmwareWriter, Partition, State, BOOT_MAGIC, SWAP_MAGIC};
+use crate::{Partition, State, BOOT_MAGIC, SWAP_MAGIC};
 
 /// Errors returned by FirmwareUpdater
 #[derive(Debug)]
@@ -84,10 +84,10 @@ impl FirmwareUpdater {
     /// `mark_booted`.
     pub async fn get_state<F: AsyncNorFlash>(
         &mut self,
-        flash: &mut F,
+        state_flash: &mut F,
         aligned: &mut [u8],
     ) -> Result<State, FirmwareUpdaterError> {
-        flash.read(self.state.from as u32, aligned).await?;
+        self.state.read(state_flash, 0, aligned).await?;
 
         if !aligned.iter().any(|&b| b != SWAP_MAGIC) {
             Ok(State::Swap)
@@ -115,17 +115,16 @@ impl FirmwareUpdater {
     #[cfg(feature = "_verify")]
     pub async fn verify_and_mark_updated<F: AsyncNorFlash>(
         &mut self,
-        _flash: &mut F,
+        _state_and_dfu_flash: &mut F,
         _public_key: &[u8],
         _signature: &[u8],
         _update_len: usize,
         _aligned: &mut [u8],
     ) -> Result<(), FirmwareUpdaterError> {
-        let _end = self.dfu.from + _update_len;
         let _read_size = _aligned.len();
 
         assert_eq!(_aligned.len(), F::WRITE_SIZE);
-        assert!(_end <= self.dfu.to);
+        assert!(_update_len <= self.dfu.len());
 
         #[cfg(feature = "ed25519-dalek")]
         {
@@ -137,21 +136,10 @@ impl FirmwareUpdater {
             let signature = Signature::from_bytes(_signature).map_err(into_signature_error)?;
 
             let mut digest = Sha512::new();
-
-            let mut offset = self.dfu.from;
-            let last_offset = _end / _read_size * _read_size;
-
-            while offset < last_offset {
-                _flash.read(offset as u32, _aligned).await?;
-                digest.update(&_aligned);
-                offset += _read_size;
-            }
-
-            let remaining = _end % _read_size;
-
-            if remaining > 0 {
-                _flash.read(last_offset as u32, _aligned).await?;
-                digest.update(&_aligned[0..remaining]);
+            for offset in (0.._update_len).step_by(_aligned.len()) {
+                self.dfu.read(_state_and_dfu_flash, offset as u32, _aligned).await?;
+                let len = core::cmp::min(_update_len - offset, _aligned.len());
+                digest.update(&_aligned[..len]);
             }
 
             public_key
@@ -173,21 +161,10 @@ impl FirmwareUpdater {
             let signature = Signature::try_from(&signature).map_err(into_signature_error)?;
 
             let mut digest = Sha512::new();
-
-            let mut offset = self.dfu.from;
-            let last_offset = _end / _read_size * _read_size;
-
-            while offset < last_offset {
-                _flash.read(offset as u32, _aligned).await?;
-                digest.update(&_aligned);
-                offset += _read_size;
-            }
-
-            let remaining = _end % _read_size;
-
-            if remaining > 0 {
-                _flash.read(last_offset as u32, _aligned).await?;
-                digest.update(&_aligned[0..remaining]);
+            for offset in (0.._update_len).step_by(_aligned.len()) {
+                self.dfu.read(_state_and_dfu_flash, offset as u32, _aligned).await?;
+                let len = core::cmp::min(_update_len - offset, _aligned.len());
+                digest.update(&_aligned[..len]);
             }
 
             let message = digest.finalize();
@@ -202,7 +179,7 @@ impl FirmwareUpdater {
             r.map_err(into_signature_error)?
         }
 
-        self.set_magic(_aligned, SWAP_MAGIC, _flash).await
+        self.set_magic(_aligned, SWAP_MAGIC, _state_and_dfu_flash).await
     }
 
     /// Mark to trigger firmware swap on next boot.
@@ -213,11 +190,11 @@ impl FirmwareUpdater {
     #[cfg(not(feature = "_verify"))]
     pub async fn mark_updated<F: AsyncNorFlash>(
         &mut self,
-        flash: &mut F,
+        state_flash: &mut F,
         aligned: &mut [u8],
     ) -> Result<(), FirmwareUpdaterError> {
         assert_eq!(aligned.len(), F::WRITE_SIZE);
-        self.set_magic(aligned, SWAP_MAGIC, flash).await
+        self.set_magic(aligned, SWAP_MAGIC, state_flash).await
     }
 
     /// Mark firmware boot successful and stop rollback on reset.
@@ -227,29 +204,42 @@ impl FirmwareUpdater {
     /// The `aligned` buffer must have a size of F::WRITE_SIZE, and follow the alignment rules for the flash being written to.
     pub async fn mark_booted<F: AsyncNorFlash>(
         &mut self,
-        flash: &mut F,
+        state_flash: &mut F,
         aligned: &mut [u8],
     ) -> Result<(), FirmwareUpdaterError> {
         assert_eq!(aligned.len(), F::WRITE_SIZE);
-        self.set_magic(aligned, BOOT_MAGIC, flash).await
+        self.set_magic(aligned, BOOT_MAGIC, state_flash).await
     }
 
     async fn set_magic<F: AsyncNorFlash>(
         &mut self,
         aligned: &mut [u8],
         magic: u8,
-        flash: &mut F,
+        state_flash: &mut F,
     ) -> Result<(), FirmwareUpdaterError> {
-        flash.read(self.state.from as u32, aligned).await?;
+        self.state.read(state_flash, 0, aligned).await?;
 
         if aligned.iter().any(|&b| b != magic) {
-            aligned.fill(0);
+            // Read progress validity
+            self.state.read(state_flash, F::WRITE_SIZE as u32, aligned).await?;
+
+            // FIXME: Do not make this assumption.
+            const STATE_ERASE_VALUE: u8 = 0xFF;
+
+            if aligned.iter().any(|&b| b != STATE_ERASE_VALUE) {
+                // The current progress validity marker is invalid
+            } else {
+                // Invalidate progress
+                aligned.fill(!STATE_ERASE_VALUE);
+                self.state.write(state_flash, F::WRITE_SIZE as u32, aligned).await?;
+            }
 
-            flash.write(self.state.from as u32, aligned).await?;
-            flash.erase(self.state.from as u32, self.state.to as u32).await?;
+            // Clear magic and progress
+            self.state.wipe(state_flash).await?;
 
+            // Set magic
             aligned.fill(magic);
-            flash.write(self.state.from as u32, aligned).await?;
+            self.state.write(state_flash, 0, aligned).await?;
         }
         Ok(())
     }
@@ -265,45 +255,31 @@ impl FirmwareUpdater {
         &mut self,
         offset: usize,
         data: &[u8],
-        flash: &mut F,
-        block_size: usize,
+        dfu_flash: &mut F,
     ) -> Result<(), FirmwareUpdaterError> {
         assert!(data.len() >= F::ERASE_SIZE);
 
-        flash
-            .erase(
-                (self.dfu.from + offset) as u32,
-                (self.dfu.from + offset + data.len()) as u32,
-            )
+        self.dfu
+            .erase(dfu_flash, offset as u32, (offset + data.len()) as u32)
             .await?;
 
-        trace!(
-            "Erased from {} to {}",
-            self.dfu.from + offset,
-            self.dfu.from + offset + data.len()
-        );
-
-        FirmwareWriter(self.dfu)
-            .write_block(offset, data, flash, block_size)
-            .await?;
+        self.dfu.write(dfu_flash, offset as u32, data).await?;
 
         Ok(())
     }
 
     /// Prepare for an incoming DFU update by erasing the entire DFU area and
-    /// returning a `FirmwareWriter`.
+    /// returning its `Partition`.
     ///
     /// Using this instead of `write_firmware` allows for an optimized API in
     /// exchange for added complexity.
     pub async fn prepare_update<F: AsyncNorFlash>(
         &mut self,
-        flash: &mut F,
-    ) -> Result<FirmwareWriter, FirmwareUpdaterError> {
-        flash.erase((self.dfu.from) as u32, (self.dfu.to) as u32).await?;
+        dfu_flash: &mut F,
+    ) -> Result<Partition, FirmwareUpdaterError> {
+        self.dfu.wipe(dfu_flash).await?;
 
-        trace!("Erased from {} to {}", self.dfu.from, self.dfu.to);
-
-        Ok(FirmwareWriter(self.dfu))
+        Ok(self.dfu)
     }
 
     //
@@ -317,10 +293,10 @@ impl FirmwareUpdater {
     /// `mark_booted`.
     pub fn get_state_blocking<F: NorFlash>(
         &mut self,
-        flash: &mut F,
+        state_flash: &mut F,
         aligned: &mut [u8],
     ) -> Result<State, FirmwareUpdaterError> {
-        flash.read(self.state.from as u32, aligned)?;
+        self.state.read_blocking(state_flash, 0, aligned)?;
 
         if !aligned.iter().any(|&b| b != SWAP_MAGIC) {
             Ok(State::Swap)
@@ -348,7 +324,7 @@ impl FirmwareUpdater {
     #[cfg(feature = "_verify")]
     pub fn verify_and_mark_updated_blocking<F: NorFlash>(
         &mut self,
-        _flash: &mut F,
+        _state_and_dfu_flash: &mut F,
         _public_key: &[u8],
         _signature: &[u8],
         _update_len: usize,
@@ -370,21 +346,10 @@ impl FirmwareUpdater {
             let signature = Signature::from_bytes(_signature).map_err(into_signature_error)?;
 
             let mut digest = Sha512::new();
-
-            let mut offset = self.dfu.from;
-            let last_offset = _end / _read_size * _read_size;
-
-            while offset < last_offset {
-                _flash.read(offset as u32, _aligned)?;
-                digest.update(&_aligned);
-                offset += _read_size;
-            }
-
-            let remaining = _end % _read_size;
-
-            if remaining > 0 {
-                _flash.read(last_offset as u32, _aligned)?;
-                digest.update(&_aligned[0..remaining]);
+            for offset in (0.._update_len).step_by(_aligned.len()) {
+                self.dfu.read_blocking(_state_and_dfu_flash, offset as u32, _aligned)?;
+                let len = core::cmp::min(_update_len - offset, _aligned.len());
+                digest.update(&_aligned[..len]);
             }
 
             public_key
@@ -406,21 +371,10 @@ impl FirmwareUpdater {
             let signature = Signature::try_from(&signature).map_err(into_signature_error)?;
 
             let mut digest = Sha512::new();
-
-            let mut offset = self.dfu.from;
-            let last_offset = _end / _read_size * _read_size;
-
-            while offset < last_offset {
-                _flash.read(offset as u32, _aligned)?;
-                digest.update(&_aligned);
-                offset += _read_size;
-            }
-
-            let remaining = _end % _read_size;
-
-            if remaining > 0 {
-                _flash.read(last_offset as u32, _aligned)?;
-                digest.update(&_aligned[0..remaining]);
+            for offset in (0.._update_len).step_by(_aligned.len()) {
+                self.dfu.read_blocking(_state_and_dfu_flash, offset as u32, _aligned)?;
+                let len = core::cmp::min(_update_len - offset, _aligned.len());
+                digest.update(&_aligned[..len]);
             }
 
             let message = digest.finalize();
@@ -435,7 +389,7 @@ impl FirmwareUpdater {
             r.map_err(into_signature_error)?
         }
 
-        self.set_magic_blocking(_aligned, SWAP_MAGIC, _flash)
+        self.set_magic_blocking(_aligned, SWAP_MAGIC, _state_and_dfu_flash)
     }
 
     /// Mark to trigger firmware swap on next boot.
@@ -446,11 +400,11 @@ impl FirmwareUpdater {
     #[cfg(not(feature = "_verify"))]
     pub fn mark_updated_blocking<F: NorFlash>(
         &mut self,
-        flash: &mut F,
+        state_flash: &mut F,
         aligned: &mut [u8],
     ) -> Result<(), FirmwareUpdaterError> {
         assert_eq!(aligned.len(), F::WRITE_SIZE);
-        self.set_magic_blocking(aligned, SWAP_MAGIC, flash)
+        self.set_magic_blocking(aligned, SWAP_MAGIC, state_flash)
     }
 
     /// Mark firmware boot successful and stop rollback on reset.
@@ -460,29 +414,42 @@ impl FirmwareUpdater {
     /// The `aligned` buffer must have a size of F::WRITE_SIZE, and follow the alignment rules for the flash being written to.
     pub fn mark_booted_blocking<F: NorFlash>(
         &mut self,
-        flash: &mut F,
+        state_flash: &mut F,
         aligned: &mut [u8],
     ) -> Result<(), FirmwareUpdaterError> {
         assert_eq!(aligned.len(), F::WRITE_SIZE);
-        self.set_magic_blocking(aligned, BOOT_MAGIC, flash)
+        self.set_magic_blocking(aligned, BOOT_MAGIC, state_flash)
     }
 
     fn set_magic_blocking<F: NorFlash>(
         &mut self,
         aligned: &mut [u8],
         magic: u8,
-        flash: &mut F,
+        state_flash: &mut F,
     ) -> Result<(), FirmwareUpdaterError> {
-        flash.read(self.state.from as u32, aligned)?;
+        self.state.read_blocking(state_flash, 0, aligned)?;
 
         if aligned.iter().any(|&b| b != magic) {
-            aligned.fill(0);
+            // Read progress validity
+            self.state.read_blocking(state_flash, F::WRITE_SIZE as u32, aligned)?;
+
+            // FIXME: Do not make this assumption.
+            const STATE_ERASE_VALUE: u8 = 0xFF;
+
+            if aligned.iter().any(|&b| b != STATE_ERASE_VALUE) {
+                // The current progress validity marker is invalid
+            } else {
+                // Invalidate progress
+                aligned.fill(!STATE_ERASE_VALUE);
+                self.state.write_blocking(state_flash, F::WRITE_SIZE as u32, aligned)?;
+            }
 
-            flash.write(self.state.from as u32, aligned)?;
-            flash.erase(self.state.from as u32, self.state.to as u32)?;
+            // Clear magic and progress
+            self.state.wipe_blocking(state_flash)?;
 
+            // Set magic
             aligned.fill(magic);
-            flash.write(self.state.from as u32, aligned)?;
+            self.state.write_blocking(state_flash, 0, aligned)?;
         }
         Ok(())
     }
@@ -498,40 +465,26 @@ impl FirmwareUpdater {
         &mut self,
         offset: usize,
         data: &[u8],
-        flash: &mut F,
-        block_size: usize,
+        dfu_flash: &mut F,
     ) -> Result<(), FirmwareUpdaterError> {
         assert!(data.len() >= F::ERASE_SIZE);
 
-        flash.erase(
-            (self.dfu.from + offset) as u32,
-            (self.dfu.from + offset + data.len()) as u32,
-        )?;
+        self.dfu
+            .erase_blocking(dfu_flash, offset as u32, (offset + data.len()) as u32)?;
 
-        trace!(
-            "Erased from {} to {}",
-            self.dfu.from + offset,
-            self.dfu.from + offset + data.len()
-        );
-
-        FirmwareWriter(self.dfu).write_block_blocking(offset, data, flash, block_size)?;
+        self.dfu.write_blocking(dfu_flash, offset as u32, data)?;
 
         Ok(())
     }
 
     /// Prepare for an incoming DFU update by erasing the entire DFU area and
-    /// returning a `FirmwareWriter`.
+    /// returning its `Partition`.
     ///
     /// Using this instead of `write_firmware_blocking` allows for an optimized
     /// API in exchange for added complexity.
-    pub fn prepare_update_blocking<F: NorFlash>(
-        &mut self,
-        flash: &mut F,
-    ) -> Result<FirmwareWriter, FirmwareUpdaterError> {
-        flash.erase((self.dfu.from) as u32, (self.dfu.to) as u32)?;
-
-        trace!("Erased from {} to {}", self.dfu.from, self.dfu.to);
+    pub fn prepare_update_blocking<F: NorFlash>(&mut self, flash: &mut F) -> Result<Partition, FirmwareUpdaterError> {
+        self.dfu.wipe_blocking(flash)?;
 
-        Ok(FirmwareWriter(self.dfu))
+        Ok(self.dfu)
     }
 }
diff --git a/embassy-boot/boot/src/firmware_writer.rs b/embassy-boot/boot/src/firmware_writer.rs
deleted file mode 100644
index f992021bb..000000000
--- a/embassy-boot/boot/src/firmware_writer.rs
+++ /dev/null
@@ -1,97 +0,0 @@
-use embedded_storage::nor_flash::NorFlash;
-use embedded_storage_async::nor_flash::NorFlash as AsyncNorFlash;
-
-use crate::Partition;
-
-/// FirmwareWriter allows writing blocks to an already erased flash.
-pub struct FirmwareWriter(pub(crate) Partition);
-
-impl FirmwareWriter {
-    /// Write data to a flash page.
-    ///
-    /// The buffer must follow alignment requirements of the target flash and a multiple of page size big.
-    ///
-    /// # Safety
-    ///
-    /// Failing to meet alignment and size requirements may result in a panic.
-    pub async fn write_block<F: AsyncNorFlash>(
-        &mut self,
-        offset: usize,
-        data: &[u8],
-        flash: &mut F,
-        block_size: usize,
-    ) -> Result<(), F::Error> {
-        trace!(
-            "Writing firmware at offset 0x{:x} len {}",
-            self.0.from + offset,
-            data.len()
-        );
-
-        let mut write_offset = self.0.from + offset;
-        for chunk in data.chunks(block_size) {
-            trace!("Wrote chunk at {}: {:?}", write_offset, chunk);
-            flash.write(write_offset as u32, chunk).await?;
-            write_offset += chunk.len();
-        }
-        /*
-        trace!("Wrote data, reading back for verification");
-
-        let mut buf: [u8; 4096] = [0; 4096];
-        let mut data_offset = 0;
-        let mut read_offset = self.dfu.from + offset;
-        for chunk in buf.chunks_mut(block_size) {
-            flash.read(read_offset as u32, chunk).await?;
-            trace!("Read chunk at {}: {:?}", read_offset, chunk);
-            assert_eq!(&data[data_offset..data_offset + block_size], chunk);
-            read_offset += chunk.len();
-            data_offset += chunk.len();
-        }
-        */
-
-        Ok(())
-    }
-
-    /// Write data to a flash page.
-    ///
-    /// The buffer must follow alignment requirements of the target flash and a multiple of page size big.
-    ///
-    /// # Safety
-    ///
-    /// Failing to meet alignment and size requirements may result in a panic.
-    pub fn write_block_blocking<F: NorFlash>(
-        &mut self,
-        offset: usize,
-        data: &[u8],
-        flash: &mut F,
-        block_size: usize,
-    ) -> Result<(), F::Error> {
-        trace!(
-            "Writing firmware at offset 0x{:x} len {}",
-            self.0.from + offset,
-            data.len()
-        );
-
-        let mut write_offset = self.0.from + offset;
-        for chunk in data.chunks(block_size) {
-            trace!("Wrote chunk at {}: {:?}", write_offset, chunk);
-            flash.write(write_offset as u32, chunk)?;
-            write_offset += chunk.len();
-        }
-        /*
-        trace!("Wrote data, reading back for verification");
-
-        let mut buf: [u8; 4096] = [0; 4096];
-        let mut data_offset = 0;
-        let mut read_offset = self.dfu.from + offset;
-        for chunk in buf.chunks_mut(block_size) {
-            flash.read(read_offset as u32, chunk).await?;
-            trace!("Read chunk at {}: {:?}", read_offset, chunk);
-            assert_eq!(&data[data_offset..data_offset + block_size], chunk);
-            read_offset += chunk.len();
-            data_offset += chunk.len();
-        }
-        */
-
-        Ok(())
-    }
-}
diff --git a/embassy-boot/boot/src/lib.rs b/embassy-boot/boot/src/lib.rs
index a2259411f..d53c613a3 100644
--- a/embassy-boot/boot/src/lib.rs
+++ b/embassy-boot/boot/src/lib.rs
@@ -7,12 +7,11 @@ mod fmt;
 
 mod boot_loader;
 mod firmware_updater;
-mod firmware_writer;
+mod mem_flash;
 mod partition;
 
 pub use boot_loader::{BootError, BootFlash, BootLoader, Flash, FlashConfig, MultiFlashConfig, SingleFlashConfig};
 pub use firmware_updater::{FirmwareUpdater, FirmwareUpdaterError};
-pub use firmware_writer::FirmwareWriter;
 pub use partition::Partition;
 
 pub(crate) const BOOT_MAGIC: u8 = 0xD0;
@@ -46,13 +45,10 @@ impl<const N: usize> AsMut<[u8]> for AlignedBuffer<N> {
 
 #[cfg(test)]
 mod tests {
-    use core::convert::Infallible;
-
-    use embedded_storage::nor_flash::{ErrorType, NorFlash, ReadNorFlash};
-    use embedded_storage_async::nor_flash::{NorFlash as AsyncNorFlash, ReadNorFlash as AsyncReadNorFlash};
     use futures::executor::block_on;
 
     use super::*;
+    use crate::mem_flash::MemFlash;
 
     /*
     #[test]
@@ -75,8 +71,8 @@ mod tests {
         const ACTIVE: Partition = Partition::new(4096, 61440);
         const DFU: Partition = Partition::new(61440, 122880);
 
-        let mut flash = MemFlash::<131072, 4096, 4>([0xff; 131072]);
-        flash.0[0..4].copy_from_slice(&[BOOT_MAGIC; 4]);
+        let mut flash = MemFlash::<131072, 4096, 4>::default();
+        flash.mem[0..4].copy_from_slice(&[BOOT_MAGIC; 4]);
         let mut flash = SingleFlashConfig::new(&mut flash);
 
         let mut bootloader: BootLoader = BootLoader::new(ACTIVE, DFU, STATE);
@@ -95,21 +91,21 @@ mod tests {
         const STATE: Partition = Partition::new(0, 4096);
         const ACTIVE: Partition = Partition::new(4096, 61440);
         const DFU: Partition = Partition::new(61440, 122880);
-        let mut flash = MemFlash::<131072, 4096, 4>([0xff; 131072]);
+        let mut flash = MemFlash::<131072, 4096, 4>::random();
 
         let original: [u8; ACTIVE.len()] = [rand::random::<u8>(); ACTIVE.len()];
         let update: [u8; DFU.len()] = [rand::random::<u8>(); DFU.len()];
         let mut aligned = [0; 4];
 
         for i in ACTIVE.from..ACTIVE.to {
-            flash.0[i] = original[i - ACTIVE.from];
+            flash.mem[i] = original[i - ACTIVE.from];
         }
 
         let mut bootloader: BootLoader = BootLoader::new(ACTIVE, DFU, STATE);
         let mut updater = FirmwareUpdater::new(DFU, STATE);
         let mut offset = 0;
         for chunk in update.chunks(4096) {
-            block_on(updater.write_firmware(offset, chunk, &mut flash, 4096)).unwrap();
+            block_on(updater.write_firmware(offset, chunk, &mut flash)).unwrap();
             offset += chunk.len();
         }
         block_on(updater.mark_updated(&mut flash, &mut aligned)).unwrap();
@@ -124,12 +120,12 @@ mod tests {
         );
 
         for i in ACTIVE.from..ACTIVE.to {
-            assert_eq!(flash.0[i], update[i - ACTIVE.from], "Index {}", i);
+            assert_eq!(flash.mem[i], update[i - ACTIVE.from], "Index {}", i);
         }
 
         // 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);
+            assert_eq!(flash.mem[i], original[i - DFU.from - 4096], "Index {}", i);
         }
 
         // Running again should cause a revert
@@ -141,12 +137,12 @@ mod tests {
         );
 
         for i in ACTIVE.from..ACTIVE.to {
-            assert_eq!(flash.0[i], original[i - ACTIVE.from], "Index {}", i);
+            assert_eq!(flash.mem[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);
+            assert_eq!(flash.mem[i], update[i - DFU.from], "Index {}", i);
         }
 
         // Mark as booted
@@ -166,23 +162,23 @@ mod tests {
         const ACTIVE: Partition = Partition::new(4096, 16384);
         const DFU: Partition = Partition::new(0, 16384);
 
-        let mut active = MemFlash::<16384, 4096, 8>([0xff; 16384]);
-        let mut dfu = MemFlash::<16384, 2048, 8>([0xff; 16384]);
-        let mut state = MemFlash::<4096, 128, 4>([0xff; 4096]);
+        let mut active = MemFlash::<16384, 4096, 8>::random();
+        let mut dfu = MemFlash::<16384, 2048, 8>::random();
+        let mut state = MemFlash::<4096, 128, 4>::random();
         let mut aligned = [0; 4];
 
         let original: [u8; ACTIVE.len()] = [rand::random::<u8>(); ACTIVE.len()];
         let update: [u8; DFU.len()] = [rand::random::<u8>(); DFU.len()];
 
         for i in ACTIVE.from..ACTIVE.to {
-            active.0[i] = original[i - ACTIVE.from];
+            active.mem[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();
+            block_on(updater.write_firmware(offset, chunk, &mut dfu)).unwrap();
             offset += chunk.len();
         }
         block_on(updater.mark_updated(&mut state, &mut aligned)).unwrap();
@@ -203,12 +199,12 @@ mod tests {
         );
 
         for i in ACTIVE.from..ACTIVE.to {
-            assert_eq!(active.0[i], update[i - ACTIVE.from], "Index {}", i);
+            assert_eq!(active.mem[i], update[i - ACTIVE.from], "Index {}", i);
         }
 
         // 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);
+            assert_eq!(dfu.mem[i], original[i - DFU.from - 4096], "Index {}", i);
         }
     }
 
@@ -220,22 +216,22 @@ mod tests {
         const DFU: Partition = Partition::new(0, 16384);
 
         let mut aligned = [0; 4];
-        let mut active = MemFlash::<16384, 2048, 4>([0xff; 16384]);
-        let mut dfu = MemFlash::<16384, 4096, 8>([0xff; 16384]);
-        let mut state = MemFlash::<4096, 128, 4>([0xff; 4096]);
+        let mut active = MemFlash::<16384, 2048, 4>::random();
+        let mut dfu = MemFlash::<16384, 4096, 8>::random();
+        let mut state = MemFlash::<4096, 128, 4>::random();
 
         let original: [u8; ACTIVE.len()] = [rand::random::<u8>(); ACTIVE.len()];
         let update: [u8; DFU.len()] = [rand::random::<u8>(); DFU.len()];
 
         for i in ACTIVE.from..ACTIVE.to {
-            active.0[i] = original[i - ACTIVE.from];
+            active.mem[i] = original[i - ACTIVE.from];
         }
 
         let mut updater = FirmwareUpdater::new(DFU, STATE);
 
         let mut offset = 0;
         for chunk in update.chunks(4096) {
-            block_on(updater.write_firmware(offset, chunk, &mut dfu, chunk.len())).unwrap();
+            block_on(updater.write_firmware(offset, chunk, &mut dfu)).unwrap();
             offset += chunk.len();
         }
         block_on(updater.mark_updated(&mut state, &mut aligned)).unwrap();
@@ -255,12 +251,12 @@ mod tests {
         );
 
         for i in ACTIVE.from..ACTIVE.to {
-            assert_eq!(active.0[i], update[i - ACTIVE.from], "Index {}", i);
+            assert_eq!(active.mem[i], update[i - ACTIVE.from], "Index {}", i);
         }
 
         // 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);
+            assert_eq!(dfu.mem[i], original[i - DFU.from - 4096], "Index {}", i);
         }
     }
 
@@ -290,13 +286,13 @@ mod tests {
 
         const STATE: Partition = Partition::new(0, 4096);
         const DFU: Partition = Partition::new(4096, 8192);
-        let mut flash = MemFlash::<8192, 4096, 4>([0xff; 8192]);
+        let mut flash = MemFlash::<8192, 4096, 4>::default();
 
         let firmware_len = firmware.len();
 
         let mut write_buf = [0; 4096];
         write_buf[0..firmware_len].copy_from_slice(firmware);
-        NorFlash::write(&mut flash, DFU.from as u32, &write_buf).unwrap();
+        DFU.write_blocking(&mut flash, 0, &write_buf).unwrap();
 
         // On with the test
 
@@ -313,112 +309,4 @@ mod tests {
         ))
         .is_ok());
     }
-    struct MemFlash<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize>([u8; SIZE]);
-
-    impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> NorFlash
-        for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
-    {
-        const WRITE_SIZE: usize = WRITE_SIZE;
-        const ERASE_SIZE: usize = ERASE_SIZE;
-        fn erase(&mut self, from: u32, to: u32) -> Result<(), Self::Error> {
-            let from = from as usize;
-            let to = to as usize;
-            assert!(from % ERASE_SIZE == 0);
-            assert!(to % ERASE_SIZE == 0, "To: {}, erase size: {}", to, ERASE_SIZE);
-            for i in from..to {
-                self.0[i] = 0xFF;
-            }
-            Ok(())
-        }
-
-        fn write(&mut self, offset: u32, data: &[u8]) -> Result<(), Self::Error> {
-            assert!(data.len() % WRITE_SIZE == 0);
-            assert!(offset as usize % WRITE_SIZE == 0);
-            assert!(offset as usize + data.len() <= SIZE);
-
-            self.0[offset as usize..offset as usize + data.len()].copy_from_slice(data);
-
-            Ok(())
-        }
-    }
-
-    impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> ErrorType
-        for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
-    {
-        type Error = Infallible;
-    }
-
-    impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> ReadNorFlash
-        for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
-    {
-        const READ_SIZE: usize = 1;
-
-        fn read(&mut self, offset: u32, buf: &mut [u8]) -> Result<(), Self::Error> {
-            let len = buf.len();
-            buf[..].copy_from_slice(&self.0[offset as usize..offset as usize + len]);
-            Ok(())
-        }
-
-        fn capacity(&self) -> usize {
-            SIZE
-        }
-    }
-
-    impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> super::Flash
-        for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
-    {
-        const BLOCK_SIZE: usize = ERASE_SIZE;
-        const ERASE_VALUE: u8 = 0xFF;
-    }
-
-    impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> AsyncReadNorFlash
-        for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
-    {
-        const READ_SIZE: usize = 1;
-
-        async fn read(&mut self, offset: u32, buf: &mut [u8]) -> Result<(), Self::Error> {
-            let len = buf.len();
-            buf[..].copy_from_slice(&self.0[offset as usize..offset as usize + len]);
-            Ok(())
-        }
-
-        fn capacity(&self) -> usize {
-            SIZE
-        }
-    }
-
-    impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> AsyncNorFlash
-        for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
-    {
-        const WRITE_SIZE: usize = WRITE_SIZE;
-        const ERASE_SIZE: usize = ERASE_SIZE;
-
-        async fn erase(&mut self, from: u32, to: u32) -> Result<(), Self::Error> {
-            let from = from as usize;
-            let to = to as usize;
-            assert!(from % ERASE_SIZE == 0);
-            assert!(to % ERASE_SIZE == 0);
-            for i in from..to {
-                self.0[i] = 0xFF;
-            }
-            Ok(())
-        }
-
-        async fn write(&mut self, offset: u32, data: &[u8]) -> Result<(), Self::Error> {
-            info!("Writing {} bytes to 0x{:x}", data.len(), offset);
-            assert!(data.len() % WRITE_SIZE == 0);
-            assert!(offset as usize % WRITE_SIZE == 0);
-            assert!(
-                offset as usize + data.len() <= SIZE,
-                "OFFSET: {}, LEN: {}, FLASH SIZE: {}",
-                offset,
-                data.len(),
-                SIZE
-            );
-
-            self.0[offset as usize..offset as usize + data.len()].copy_from_slice(data);
-
-            Ok(())
-        }
-    }
 }
diff --git a/embassy-boot/boot/src/mem_flash.rs b/embassy-boot/boot/src/mem_flash.rs
new file mode 100644
index 000000000..828aad9d9
--- /dev/null
+++ b/embassy-boot/boot/src/mem_flash.rs
@@ -0,0 +1,156 @@
+#![allow(unused)]
+
+use core::ops::{Bound, Range, RangeBounds};
+
+use embedded_storage::nor_flash::{ErrorType, NorFlash, NorFlashError, NorFlashErrorKind, ReadNorFlash};
+use embedded_storage_async::nor_flash::{NorFlash as AsyncNorFlash, ReadNorFlash as AsyncReadNorFlash};
+
+use crate::Flash;
+
+pub struct MemFlash<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> {
+    pub mem: [u8; SIZE],
+    pub pending_write_successes: Option<usize>,
+}
+
+#[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,
+        }
+    }
+}
+
+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> Flash
+    for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
+{
+    const BLOCK_SIZE: usize = ERASE_SIZE;
+    const ERASE_VALUE: u8 = 0xFF;
+}
+
+impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> ErrorType
+    for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
+{
+    type Error = MemFlashError;
+}
+
+impl NorFlashError for MemFlashError {
+    fn kind(&self) -> NorFlashErrorKind {
+        NorFlashErrorKind::Other
+    }
+}
+
+impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> ReadNorFlash
+    for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
+{
+    const READ_SIZE: usize = 1;
+
+    fn read(&mut self, offset: u32, bytes: &mut [u8]) -> Result<(), Self::Error> {
+        let len = bytes.len();
+        bytes.copy_from_slice(&self.mem[offset as usize..offset as usize + len]);
+        Ok(())
+    }
+
+    fn capacity(&self) -> usize {
+        SIZE
+    }
+}
+
+impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> NorFlash
+    for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
+{
+    const WRITE_SIZE: usize = WRITE_SIZE;
+    const ERASE_SIZE: usize = ERASE_SIZE;
+
+    fn erase(&mut self, from: u32, to: u32) -> Result<(), Self::Error> {
+        let from = from as usize;
+        let to = to as usize;
+        assert!(from % ERASE_SIZE == 0);
+        assert!(to % ERASE_SIZE == 0, "To: {}, erase size: {}", to, ERASE_SIZE);
+        for i in from..to {
+            self.mem[i] = 0xFF;
+        }
+        Ok(())
+    }
+
+    fn write(&mut self, offset: u32, bytes: &[u8]) -> Result<(), Self::Error> {
+        let offset = offset as usize;
+        assert!(bytes.len() % WRITE_SIZE == 0);
+        assert!(offset % WRITE_SIZE == 0);
+        assert!(offset + bytes.len() <= SIZE);
+
+        if let Some(pending_successes) = self.pending_write_successes {
+            if pending_successes > 0 {
+                self.pending_write_successes = Some(pending_successes - 1);
+            } else {
+                return Err(MemFlashError);
+            }
+        }
+
+        for ((offset, mem_byte), new_byte) in self
+            .mem
+            .iter_mut()
+            .enumerate()
+            .skip(offset)
+            .take(bytes.len())
+            .zip(bytes)
+        {
+            assert_eq!(0xFF, *mem_byte, "Offset {} is not erased", offset);
+            *mem_byte = *new_byte;
+        }
+
+        Ok(())
+    }
+}
+
+impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> AsyncReadNorFlash
+    for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
+{
+    const READ_SIZE: usize = 1;
+
+    async fn read(&mut self, offset: u32, bytes: &mut [u8]) -> Result<(), Self::Error> {
+        <Self as ReadNorFlash>::read(self, offset, bytes)
+    }
+
+    fn capacity(&self) -> usize {
+        <Self as ReadNorFlash>::capacity(self)
+    }
+}
+
+impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> AsyncNorFlash
+    for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
+{
+    const WRITE_SIZE: usize = WRITE_SIZE;
+    const ERASE_SIZE: usize = ERASE_SIZE;
+
+    async fn erase(&mut self, from: u32, to: u32) -> Result<(), Self::Error> {
+        <Self as NorFlash>::erase(self, from, to)
+    }
+
+    async fn write(&mut self, offset: u32, bytes: &[u8]) -> Result<(), Self::Error> {
+        <Self as NorFlash>::write(self, offset, bytes)
+    }
+}
diff --git a/embassy-boot/boot/src/partition.rs b/embassy-boot/boot/src/partition.rs
index 46f80a23c..ac6b0ed0f 100644
--- a/embassy-boot/boot/src/partition.rs
+++ b/embassy-boot/boot/src/partition.rs
@@ -1,10 +1,13 @@
+use embedded_storage::nor_flash::{NorFlash, ReadNorFlash};
+use embedded_storage_async::nor_flash::{NorFlash as AsyncNorFlash, ReadNorFlash as AsyncReadNorFlash};
+
 /// A region in flash used by the bootloader.
 #[derive(Copy, Clone, Debug)]
 #[cfg_attr(feature = "defmt", derive(defmt::Format))]
 pub struct Partition {
-    /// Start of the flash region.
+    /// The offset into the flash where the partition starts.
     pub from: usize,
-    /// End of the flash region.
+    /// The offset into the flash where the partition ends.
     pub to: usize,
 }
 
@@ -14,9 +17,124 @@ impl Partition {
         Self { from, to }
     }
 
-    /// Return the length of the partition
+    /// Return the size of the partition
     #[allow(clippy::len_without_is_empty)]
     pub const fn len(&self) -> usize {
         self.to - self.from
     }
+
+    /// Read from the partition on the provided flash
+    pub async fn read<F: AsyncReadNorFlash>(
+        &self,
+        flash: &mut F,
+        offset: u32,
+        bytes: &mut [u8],
+    ) -> Result<(), F::Error> {
+        let offset = self.from as u32 + offset;
+        flash.read(offset, bytes).await
+    }
+
+    /// Write to the partition on the provided flash
+    pub async fn write<F: AsyncNorFlash>(&self, flash: &mut F, offset: u32, bytes: &[u8]) -> Result<(), F::Error> {
+        let offset = self.from as u32 + offset;
+        flash.write(offset, bytes).await?;
+        trace!("Wrote from 0x{:x} len {}", offset, bytes.len());
+        Ok(())
+    }
+
+    /// Erase part of the partition on the provided flash
+    pub async fn erase<F: AsyncNorFlash>(&self, flash: &mut F, from: u32, to: u32) -> Result<(), F::Error> {
+        let from = self.from as u32 + from;
+        let to = self.from as u32 + to;
+        flash.erase(from, to).await?;
+        trace!("Erased from 0x{:x} to 0x{:x}", from, to);
+        Ok(())
+    }
+
+    /// Erase the entire partition
+    pub(crate) async fn wipe<F: AsyncNorFlash>(&self, flash: &mut F) -> Result<(), F::Error> {
+        let from = self.from as u32;
+        let to = self.to as u32;
+        flash.erase(from, to).await?;
+        trace!("Wiped from 0x{:x} to 0x{:x}", from, to);
+        Ok(())
+    }
+
+    /// Read from the partition on the provided flash
+    pub fn read_blocking<F: ReadNorFlash>(&self, flash: &mut F, offset: u32, bytes: &mut [u8]) -> Result<(), F::Error> {
+        let offset = self.from as u32 + offset;
+        flash.read(offset, bytes)
+    }
+
+    /// Write to the partition on the provided flash
+    pub fn write_blocking<F: NorFlash>(&self, flash: &mut F, offset: u32, bytes: &[u8]) -> Result<(), F::Error> {
+        let offset = self.from as u32 + offset;
+        flash.write(offset, bytes)?;
+        trace!("Wrote from 0x{:x} len {}", offset, bytes.len());
+        Ok(())
+    }
+
+    /// Erase part of the partition on the provided flash
+    pub fn erase_blocking<F: NorFlash>(&self, flash: &mut F, from: u32, to: u32) -> Result<(), F::Error> {
+        let from = self.from as u32 + from;
+        let to = self.from as u32 + to;
+        flash.erase(from, to)?;
+        trace!("Erased from 0x{:x} to 0x{:x}", from, to);
+        Ok(())
+    }
+
+    /// Erase the entire partition
+    pub(crate) fn wipe_blocking<F: NorFlash>(&self, flash: &mut F) -> Result<(), F::Error> {
+        let from = self.from as u32;
+        let to = self.to as u32;
+        flash.erase(from, to)?;
+        trace!("Wiped from 0x{:x} to 0x{:x}", from, to);
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::mem_flash::MemFlash;
+    use crate::Partition;
+
+    #[test]
+    fn can_erase() {
+        let mut flash = MemFlash::<1024, 64, 4>::new(0x00);
+        let partition = Partition::new(256, 512);
+
+        partition.erase_blocking(&mut flash, 64, 192).unwrap();
+
+        for (index, byte) in flash.mem.iter().copied().enumerate().take(256 + 64) {
+            assert_eq!(0x00, byte, "Index {}", index);
+        }
+
+        for (index, byte) in flash.mem.iter().copied().enumerate().skip(256 + 64).take(128) {
+            assert_eq!(0xFF, byte, "Index {}", index);
+        }
+
+        for (index, byte) in flash.mem.iter().copied().enumerate().skip(256 + 64 + 128) {
+            assert_eq!(0x00, byte, "Index {}", index);
+        }
+    }
+
+    #[test]
+    fn can_wipe() {
+        let mut flash = MemFlash::<1024, 64, 4>::new(0x00);
+        let partition = Partition::new(256, 512);
+
+        partition.wipe_blocking(&mut flash).unwrap();
+
+        for (index, byte) in flash.mem.iter().copied().enumerate().take(256) {
+            assert_eq!(0x00, byte, "Index {}", index);
+        }
+
+        for (index, byte) in flash.mem.iter().copied().enumerate().skip(256).take(256) {
+            assert_eq!(0xFF, byte, "Index {}", index);
+        }
+
+        for (index, byte) in flash.mem.iter().copied().enumerate().skip(512) {
+            assert_eq!(0x00, byte, "Index {}", index);
+        }
+    }
 }