7 files changed, 268 insertions, 217 deletions

diff --git a/embassy-boot/boot/src/firmware_updater/asynch.rs b/embassy-boot/boot/src/firmware_updater/asynch.rs
index 20731ee0a..ae713bb6f 100644
--- a/embassy-boot/boot/src/firmware_updater/asynch.rs
+++ b/embassy-boot/boot/src/firmware_updater/asynch.rs
@@ -10,9 +10,9 @@ use crate::{FirmwareUpdaterError, State, BOOT_MAGIC, STATE_ERASE_VALUE, SWAP_MAG
 
 /// 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> {
+pub struct FirmwareUpdater<'d, DFU: NorFlash, STATE: NorFlash> {
     dfu: DFU,
-    state: STATE,
+    state: FirmwareState<'d, STATE>,
 }
 
 #[cfg(target_os = "none")]
@@ -47,22 +47,12 @@ impl<'a, FLASH: NorFlash>
     }
 }
 
-impl<DFU: NorFlash, STATE: NorFlash> FirmwareUpdater<DFU, STATE> {
+impl<'d, DFU: NorFlash, STATE: NorFlash> FirmwareUpdater<'d, DFU, STATE> {
     /// Create a firmware updater instance with partition ranges for the update and state partitions.
-    pub fn new(config: FirmwareUpdaterConfig<DFU, STATE>) -> Self {
+    pub fn new(config: FirmwareUpdaterConfig<DFU, STATE>, aligned: &'d mut [u8]) -> 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)
+            state: FirmwareState::new(config.state, aligned),
         }
     }
 
@@ -71,14 +61,8 @@ impl<DFU: NorFlash, STATE: NorFlash> FirmwareUpdater<DFU, 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)
-        }
+    pub async fn get_state(&mut self) -> Result<State, FirmwareUpdaterError> {
+        self.state.get_state().await
     }
 
     /// Verify the DFU given a public key. If there is an error then DO NOT
@@ -92,23 +76,16 @@ impl<DFU: NorFlash, STATE: NorFlash> FirmwareUpdater<DFU, STATE> {
     ///
     /// 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?;
+        self.state.verify_booted().await?;
 
         #[cfg(feature = "ed25519-dalek")]
         {
@@ -121,8 +98,9 @@ impl<DFU: NorFlash, STATE: NorFlash> FirmwareUpdater<DFU, STATE> {
             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 chunk_buf = [0; 2];
             let mut message = [0; 64];
-            self.hash::<Sha512>(_update_len, _aligned, &mut message).await?;
+            self.hash::<Sha512>(_update_len, &mut chunk_buf, &mut message).await?;
 
             public_key.verify(&message, &signature).map_err(into_signature_error)?
         }
@@ -143,7 +121,8 @@ impl<DFU: NorFlash, STATE: NorFlash> FirmwareUpdater<DFU, STATE> {
             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 mut chunk_buf = [0; 2];
+            self.hash::<Sha512>(_update_len, &mut chunk_buf, &mut message).await?;
 
             let r = public_key.verify(&message, &signature);
             trace!(
@@ -156,7 +135,7 @@ impl<DFU: NorFlash, STATE: NorFlash> FirmwareUpdater<DFU, STATE> {
             r.map_err(into_signature_error)?
         }
 
-        self.set_magic(_aligned, SWAP_MAGIC).await
+        self.state.mark_updated().await
     }
 
     /// Verify the update in DFU with any digest.
@@ -177,49 +156,14 @@ impl<DFU: NorFlash, STATE: NorFlash> FirmwareUpdater<DFU, STATE> {
     }
 
     /// 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
+    pub async fn mark_updated(&mut self) -> Result<(), FirmwareUpdaterError> {
+        self.state.mark_updated().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(())
+    pub async fn mark_booted(&mut self) -> Result<(), FirmwareUpdaterError> {
+        self.state.mark_booted().await
     }
 
     /// Write data to a flash page.
@@ -229,16 +173,10 @@ impl<DFU: NorFlash, STATE: NorFlash> FirmwareUpdater<DFU, STATE> {
     /// # 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> {
+    pub async fn write_firmware(&mut self, 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.state.verify_booted().await?;
 
         self.dfu.erase(offset as u32, (offset + data.len()) as u32).await?;
 
@@ -252,17 +190,91 @@ impl<DFU: NorFlash, STATE: NorFlash> FirmwareUpdater<DFU, STATE> {
     ///
     /// Using this instead of `write_firmware` allows for an optimized API in
     /// exchange for added complexity.
+    pub async fn prepare_update(&mut self) -> Result<&mut DFU, FirmwareUpdaterError> {
+        self.state.verify_booted().await?;
+        self.dfu.erase(0, self.dfu.capacity() as u32).await?;
+
+        Ok(&mut self.dfu)
+    }
+}
+
+/// Manages the state partition of the firmware update.
+///
+/// Can be used standalone for more fine grained control, or as part of the updater.
+pub struct FirmwareState<'d, STATE> {
+    state: STATE,
+    aligned: &'d mut [u8],
+}
+
+impl<'d, STATE: NorFlash> FirmwareState<'d, STATE> {
+    /// Create a firmware state instance with a buffer for magic content and state partition.
     ///
     /// # 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> {
+    /// 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.
+    pub fn new(state: STATE, aligned: &'d mut [u8]) -> Self {
         assert_eq!(aligned.len(), STATE::WRITE_SIZE);
-        self.verify_booted(aligned).await?;
+        Self { state, aligned }
+    }
 
-        self.dfu.erase(0, self.dfu.capacity() as u32).await?;
+    // Make sure we are running a booted firmware to avoid reverting to a bad state.
+    async fn verify_booted(&mut self) -> Result<(), FirmwareUpdaterError> {
+        if self.get_state().await? == State::Boot {
+            Ok(())
+        } else {
+            Err(FirmwareUpdaterError::BadState)
+        }
+    }
 
-        Ok(&mut self.dfu)
+    /// 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) -> Result<State, FirmwareUpdaterError> {
+        self.state.read(0, &mut self.aligned).await?;
+
+        if !self.aligned.iter().any(|&b| b != SWAP_MAGIC) {
+            Ok(State::Swap)
+        } else {
+            Ok(State::Boot)
+        }
+    }
+
+    /// Mark to trigger firmware swap on next boot.
+    pub async fn mark_updated(&mut self) -> Result<(), FirmwareUpdaterError> {
+        self.set_magic(SWAP_MAGIC).await
+    }
+
+    /// Mark firmware boot successful and stop rollback on reset.
+    pub async fn mark_booted(&mut self) -> Result<(), FirmwareUpdaterError> {
+        self.set_magic(BOOT_MAGIC).await
+    }
+
+    async fn set_magic(&mut self, magic: u8) -> Result<(), FirmwareUpdaterError> {
+        self.state.read(0, &mut self.aligned).await?;
+
+        if self.aligned.iter().any(|&b| b != magic) {
+            // Read progress validity
+            self.state.read(STATE::WRITE_SIZE as u32, &mut self.aligned).await?;
+
+            if self.aligned.iter().any(|&b| b != STATE_ERASE_VALUE) {
+                // The current progress validity marker is invalid
+            } else {
+                // Invalidate progress
+                self.aligned.fill(!STATE_ERASE_VALUE);
+                self.state.write(STATE::WRITE_SIZE as u32, &self.aligned).await?;
+            }
+
+            // Clear magic and progress
+            self.state.erase(0, self.state.capacity() as u32).await?;
+
+            // Set magic
+            self.aligned.fill(magic);
+            self.state.write(0, &self.aligned).await?;
+        }
+        Ok(())
     }
 }
 
@@ -288,8 +300,8 @@ mod tests {
         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 updater = FirmwareUpdater::new(FirmwareUpdaterConfig { dfu, state }, &mut aligned);
+        block_on(updater.write_firmware(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();
diff --git a/embassy-boot/boot/src/firmware_updater/blocking.rs b/embassy-boot/boot/src/firmware_updater/blocking.rs
index f03f53e4d..76e4264a0 100644
--- a/embassy-boot/boot/src/firmware_updater/blocking.rs
+++ b/embassy-boot/boot/src/firmware_updater/blocking.rs
@@ -10,9 +10,9 @@ use crate::{FirmwareUpdaterError, State, BOOT_MAGIC, STATE_ERASE_VALUE, SWAP_MAG
 
 /// 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> {
+pub struct BlockingFirmwareUpdater<'d, DFU: NorFlash, STATE: NorFlash> {
     dfu: DFU,
-    state: STATE,
+    state: BlockingFirmwareState<'d, STATE>,
 }
 
 #[cfg(target_os = "none")]
@@ -49,22 +49,17 @@ impl<'a, FLASH: NorFlash>
     }
 }
 
-impl<DFU: NorFlash, STATE: NorFlash> BlockingFirmwareUpdater<DFU, STATE> {
+impl<'d, DFU: NorFlash, STATE: NorFlash> BlockingFirmwareUpdater<'d, DFU, STATE> {
     /// Create a firmware updater instance with partition ranges for the update and state partitions.
-    pub fn new(config: FirmwareUpdaterConfig<DFU, STATE>) -> Self {
+    ///
+    /// # 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.
+    pub fn new(config: FirmwareUpdaterConfig<DFU, STATE>, aligned: &'d mut [u8]) -> 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)
+            state: BlockingFirmwareState::new(config.state, aligned),
         }
     }
 
@@ -73,14 +68,8 @@ impl<DFU: NorFlash, STATE: NorFlash> BlockingFirmwareUpdater<DFU, 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)
-        }
+    pub fn get_state(&mut self) -> Result<State, FirmwareUpdaterError> {
+        self.state.get_state()
     }
 
     /// Verify the DFU given a public key. If there is an error then DO NOT
@@ -94,23 +83,16 @@ impl<DFU: NorFlash, STATE: NorFlash> BlockingFirmwareUpdater<DFU, STATE> {
     ///
     /// 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)?;
+        self.state.verify_booted()?;
 
         #[cfg(feature = "ed25519-dalek")]
         {
@@ -124,7 +106,8 @@ impl<DFU: NorFlash, STATE: NorFlash> BlockingFirmwareUpdater<DFU, STATE> {
             let signature = Signature::from_bytes(_signature).map_err(into_signature_error)?;
 
             let mut message = [0; 64];
-            self.hash::<Sha512>(_update_len, _aligned, &mut message)?;
+            let mut chunk_buf = [0; 2];
+            self.hash::<Sha512>(_update_len, &mut chunk_buf, &mut message)?;
 
             public_key.verify(&message, &signature).map_err(into_signature_error)?
         }
@@ -145,7 +128,8 @@ impl<DFU: NorFlash, STATE: NorFlash> BlockingFirmwareUpdater<DFU, STATE> {
             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 mut chunk_buf = [0; 2];
+            self.hash::<Sha512>(_update_len, &mut chunk_buf, &mut message)?;
 
             let r = public_key.verify(&message, &signature);
             trace!(
@@ -158,7 +142,7 @@ impl<DFU: NorFlash, STATE: NorFlash> BlockingFirmwareUpdater<DFU, STATE> {
             r.map_err(into_signature_error)?
         }
 
-        self.set_magic(_aligned, SWAP_MAGIC)
+        self.state.mark_updated()
     }
 
     /// Verify the update in DFU with any digest.
@@ -179,49 +163,14 @@ impl<DFU: NorFlash, STATE: NorFlash> BlockingFirmwareUpdater<DFU, STATE> {
     }
 
     /// 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)
+    pub fn mark_updated(&mut self) -> Result<(), FirmwareUpdaterError> {
+        self.state.mark_updated()
     }
 
     /// 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(())
+    pub fn mark_booted(&mut self) -> Result<(), FirmwareUpdaterError> {
+        self.state.mark_booted()
     }
 
     /// Write data to a flash page.
@@ -231,15 +180,9 @@ impl<DFU: NorFlash, STATE: NorFlash> BlockingFirmwareUpdater<DFU, STATE> {
     /// # 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> {
+    pub fn write_firmware(&mut self, offset: usize, data: &[u8]) -> Result<(), FirmwareUpdaterError> {
         assert!(data.len() >= DFU::ERASE_SIZE);
-        assert_eq!(aligned.len(), STATE::WRITE_SIZE);
-        self.verify_booted(aligned)?;
+        self.state.verify_booted()?;
 
         self.dfu.erase(offset as u32, (offset + data.len()) as u32)?;
 
@@ -253,16 +196,91 @@ impl<DFU: NorFlash, STATE: NorFlash> BlockingFirmwareUpdater<DFU, STATE> {
     ///
     /// Using this instead of `write_firmware` allows for an optimized API in
     /// exchange for added complexity.
+    pub fn prepare_update(&mut self) -> Result<&mut DFU, FirmwareUpdaterError> {
+        self.state.verify_booted()?;
+        self.dfu.erase(0, self.dfu.capacity() as u32)?;
+
+        Ok(&mut self.dfu)
+    }
+}
+
+/// Manages the state partition of the firmware update.
+///
+/// Can be used standalone for more fine grained control, or as part of the updater.
+pub struct BlockingFirmwareState<'d, STATE> {
+    state: STATE,
+    aligned: &'d mut [u8],
+}
+
+impl<'d, STATE: NorFlash> BlockingFirmwareState<'d, STATE> {
+    /// Create a firmware state instance with a buffer for magic content and state partition.
     ///
     /// # 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> {
+    /// 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.
+    pub fn new(state: STATE, aligned: &'d mut [u8]) -> Self {
         assert_eq!(aligned.len(), STATE::WRITE_SIZE);
-        self.verify_booted(aligned)?;
-        self.dfu.erase(0, self.dfu.capacity() as u32)?;
+        Self { state, aligned }
+    }
 
-        Ok(&mut self.dfu)
+    // Make sure we are running a booted firmware to avoid reverting to a bad state.
+    fn verify_booted(&mut self) -> Result<(), FirmwareUpdaterError> {
+        if self.get_state()? == 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) -> Result<State, FirmwareUpdaterError> {
+        self.state.read(0, &mut self.aligned)?;
+
+        if !self.aligned.iter().any(|&b| b != SWAP_MAGIC) {
+            Ok(State::Swap)
+        } else {
+            Ok(State::Boot)
+        }
+    }
+
+    /// Mark to trigger firmware swap on next boot.
+    pub fn mark_updated(&mut self) -> Result<(), FirmwareUpdaterError> {
+        self.set_magic(SWAP_MAGIC)
+    }
+
+    /// Mark firmware boot successful and stop rollback on reset.
+    pub fn mark_booted(&mut self) -> Result<(), FirmwareUpdaterError> {
+        self.set_magic(BOOT_MAGIC)
+    }
+
+    fn set_magic(&mut self, magic: u8) -> Result<(), FirmwareUpdaterError> {
+        self.state.read(0, &mut self.aligned)?;
+
+        if self.aligned.iter().any(|&b| b != magic) {
+            // Read progress validity
+            self.state.read(STATE::WRITE_SIZE as u32, &mut self.aligned)?;
+
+            if self.aligned.iter().any(|&b| b != STATE_ERASE_VALUE) {
+                // The current progress validity marker is invalid
+            } else {
+                // Invalidate progress
+                self.aligned.fill(!STATE_ERASE_VALUE);
+                self.state.write(STATE::WRITE_SIZE as u32, &self.aligned)?;
+            }
+
+            // Clear magic and progress
+            self.state.erase(0, self.state.capacity() as u32)?;
+
+            // Set magic
+            self.aligned.fill(magic);
+            self.state.write(0, &self.aligned)?;
+        }
+        Ok(())
     }
 }
 
@@ -283,14 +301,14 @@ mod tests {
         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 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 = BlockingFirmwareUpdater::new(FirmwareUpdaterConfig { dfu, state });
-        let mut aligned = [0; 8];
-        updater.write_firmware(&mut aligned, 0, to_write.as_slice()).unwrap();
+        let mut updater = BlockingFirmwareUpdater::new(FirmwareUpdaterConfig { dfu, state }, &mut aligned);
+        updater.write_firmware(0, to_write.as_slice()).unwrap();
         let mut chunk_buf = [0; 2];
         let mut hash = [0; 20];
         updater
diff --git a/embassy-boot/boot/src/firmware_updater/mod.rs b/embassy-boot/boot/src/firmware_updater/mod.rs
index 55ce8f363..937ddcc69 100644
--- a/embassy-boot/boot/src/firmware_updater/mod.rs
+++ b/embassy-boot/boot/src/firmware_updater/mod.rs
@@ -3,8 +3,8 @@ mod asynch;
 mod blocking;
 
 #[cfg(feature = "nightly")]
-pub use asynch::FirmwareUpdater;
-pub use blocking::BlockingFirmwareUpdater;
+pub use asynch::{FirmwareState, FirmwareUpdater};
+pub use blocking::{BlockingFirmwareState, BlockingFirmwareUpdater};
 use embedded_storage::nor_flash::{NorFlashError, NorFlashErrorKind};
 
 /// Firmware updater flash configuration holding the two flashes used by the updater
diff --git a/embassy-boot/boot/src/lib.rs b/embassy-boot/boot/src/lib.rs
index 016362b86..47f7c1797 100644
--- a/embassy-boot/boot/src/lib.rs
+++ b/embassy-boot/boot/src/lib.rs
@@ -16,9 +16,11 @@ mod test_flash;
 // TODO: Use the value provided by NorFlash when available
 pub(crate) const STATE_ERASE_VALUE: u8 = 0xFF;
 pub use boot_loader::{BootError, BootLoader, BootLoaderConfig};
+pub use firmware_updater::{
+    BlockingFirmwareState, BlockingFirmwareUpdater, FirmwareUpdaterConfig, FirmwareUpdaterError,
+};
 #[cfg(feature = "nightly")]
-pub use firmware_updater::FirmwareUpdater;
-pub use firmware_updater::{BlockingFirmwareUpdater, FirmwareUpdaterConfig, FirmwareUpdaterError};
+pub use firmware_updater::{FirmwareState, FirmwareUpdater};
 
 pub(crate) const BOOT_MAGIC: u8 = 0xD0;
 pub(crate) const SWAP_MAGIC: u8 = 0xF0;
@@ -118,15 +120,18 @@ mod tests {
         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 mut updater = FirmwareUpdater::new(
+            FirmwareUpdaterConfig {
+                dfu: flash.dfu(),
+                state: flash.state(),
+            },
+            &mut aligned,
+        );
+        block_on(updater.write_firmware(0, &UPDATE)).unwrap();
+        block_on(updater.mark_updated()).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));
+        let res: Result<(), FirmwareUpdaterError> = block_on(updater.write_firmware(0, &UPDATE));
         assert!(matches!(res, Err::<(), _>(FirmwareUpdaterError::BadState)));
 
         let flash = flash.into_blocking();
@@ -158,11 +163,14 @@ mod tests {
 
         // Mark as booted
         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 mut updater = FirmwareUpdater::new(
+            FirmwareUpdaterConfig {
+                dfu: flash.dfu(),
+                state: flash.state(),
+            },
+            &mut aligned,
+        );
+        block_on(updater.mark_booted()).unwrap();
 
         let flash = flash.into_blocking();
         let mut bootloader = BootLoader::new(BootLoaderConfig {
@@ -190,12 +198,15 @@ mod tests {
         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 mut updater = FirmwareUpdater::new(
+            FirmwareUpdaterConfig {
+                dfu: flash.dfu(),
+                state: flash.state(),
+            },
+            &mut aligned,
+        );
+        block_on(updater.write_firmware(0, &UPDATE)).unwrap();
+        block_on(updater.mark_updated()).unwrap();
 
         let flash = flash.into_blocking();
         let mut bootloader = BootLoader::new(BootLoaderConfig {
@@ -232,12 +243,15 @@ mod tests {
         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 mut updater = FirmwareUpdater::new(
+            FirmwareUpdaterConfig {
+                dfu: flash.dfu(),
+                state: flash.state(),
+            },
+            &mut aligned,
+        );
+        block_on(updater.write_firmware(0, &UPDATE)).unwrap();
+        block_on(updater.mark_updated()).unwrap();
 
         let flash = flash.into_blocking();
         let mut bootloader = BootLoader::new(BootLoaderConfig {
@@ -293,18 +307,19 @@ mod tests {
 
         // 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];
+        let mut updater = FirmwareUpdater::new(
+            FirmwareUpdaterConfig {
+                dfu: flash.dfu(),
+                state: flash.state(),
+            },
+            &mut aligned,
+        );
 
         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/nrf/src/lib.rs b/embassy-boot/nrf/src/lib.rs
index 65f57fcd1..df94819fc 100644
--- a/embassy-boot/nrf/src/lib.rs
+++ b/embassy-boot/nrf/src/lib.rs
@@ -3,9 +3,11 @@
 #![doc = include_str!("../README.md")]
 mod fmt;
 
+pub use embassy_boot::{
+    AlignedBuffer, BlockingFirmwareState, BlockingFirmwareUpdater, BootLoaderConfig, FirmwareUpdaterConfig,
+};
 #[cfg(feature = "nightly")]
-pub use embassy_boot::FirmwareUpdater;
-pub use embassy_boot::{AlignedBuffer, BlockingFirmwareUpdater, BootLoaderConfig, FirmwareUpdaterConfig};
+pub use embassy_boot::{FirmwareState, FirmwareUpdater};
 use embassy_nrf::nvmc::PAGE_SIZE;
 use embassy_nrf::peripherals::WDT;
 use embassy_nrf::wdt;
diff --git a/embassy-boot/rp/src/lib.rs b/embassy-boot/rp/src/lib.rs
index 35fc104ec..f5aefa416 100644
--- a/embassy-boot/rp/src/lib.rs
+++ b/embassy-boot/rp/src/lib.rs
@@ -3,9 +3,11 @@
 #![doc = include_str!("../README.md")]
 mod fmt;
 
+pub use embassy_boot::{
+    AlignedBuffer, BlockingFirmwareState, BlockingFirmwareUpdater, BootLoaderConfig, FirmwareUpdaterConfig, State,
+};
 #[cfg(feature = "nightly")]
-pub use embassy_boot::FirmwareUpdater;
-pub use embassy_boot::{AlignedBuffer, BlockingFirmwareUpdater, BootLoaderConfig, FirmwareUpdaterConfig, State};
+pub use embassy_boot::{FirmwareState, FirmwareUpdater};
 use embassy_rp::flash::{Blocking, Flash, ERASE_SIZE};
 use embassy_rp::peripherals::{FLASH, WATCHDOG};
 use embassy_rp::watchdog::Watchdog;
diff --git a/embassy-boot/stm32/src/lib.rs b/embassy-boot/stm32/src/lib.rs
index 069de0d1a..25f029423 100644
--- a/embassy-boot/stm32/src/lib.rs
+++ b/embassy-boot/stm32/src/lib.rs
@@ -3,9 +3,11 @@
 #![doc = include_str!("../README.md")]
 mod fmt;
 
+pub use embassy_boot::{
+    AlignedBuffer, BlockingFirmwareState, BlockingFirmwareUpdater, BootLoaderConfig, FirmwareUpdaterConfig, State,
+};
 #[cfg(feature = "nightly")]
-pub use embassy_boot::FirmwareUpdater;
-pub use embassy_boot::{AlignedBuffer, BlockingFirmwareUpdater, BootLoaderConfig, FirmwareUpdaterConfig, State};
+pub use embassy_boot::{FirmwareState, FirmwareUpdater};
 use embedded_storage::nor_flash::NorFlash;
 
 /// A bootloader for STM32 devices.