3 files changed, 99 insertions, 95 deletions

diff --git a/embassy-boot/boot/src/firmware_updater.rs b/embassy-boot/boot/src/firmware_updater.rs
index f3d76a56c..1f2a6d24e 100644
--- a/embassy-boot/boot/src/firmware_updater.rs
+++ b/embassy-boot/boot/src/firmware_updater.rs
@@ -74,21 +74,18 @@ impl FirmwareUpdater {
         Self { dfu, state }
     }
 
-    //
-    // Blocking API
-    //
-
     /// Obtain the current state.
     ///
     /// This is useful to check if the bootloader has just done a swap, in order
     /// to do verifications and self-tests of the new image before calling
     /// `mark_booted`.
-    pub fn get_state_blocking<F: NorFlash>(
+    #[cfg(feature = "nightly")]
+    pub async fn get_state<F: AsyncNorFlash>(
         &mut self,
         state_flash: &mut F,
         aligned: &mut [u8],
     ) -> Result<State, FirmwareUpdaterError> {
-        self.state.read_blocking(state_flash, 0, aligned)?;
+        self.state.read(state_flash, 0, aligned).await?;
 
         if !aligned.iter().any(|&b| b != SWAP_MAGIC) {
             Ok(State::Swap)
@@ -113,8 +110,8 @@ impl FirmwareUpdater {
     ///
     /// The `_aligned` buffer must have a size of F::WRITE_SIZE, and follow the alignment rules for the flash being read from
     /// and written to.
-    #[cfg(feature = "_verify")]
-    pub fn verify_and_mark_updated_blocking<F: NorFlash>(
+    #[cfg(all(feature = "_verify", feature = "nightly"))]
+    pub async fn verify_and_mark_updated<F: AsyncNorFlash>(
         &mut self,
         _state_and_dfu_flash: &mut F,
         _public_key: &[u8],
@@ -137,7 +134,8 @@ impl FirmwareUpdater {
             let signature = Signature::from_bytes(_signature).map_err(into_signature_error)?;
 
             let mut message = [0; 64];
-            self.hash_blocking::<_, Sha512>(_state_and_dfu_flash, _update_len, _aligned, &mut message)?;
+            self.hash::<_, Sha512>(_state_and_dfu_flash, _update_len, _aligned, &mut message)
+                .await?;
 
             public_key.verify(&message, &signature).map_err(into_signature_error)?
         }
@@ -158,7 +156,8 @@ impl FirmwareUpdater {
             let signature = Signature::try_from(&signature).map_err(into_signature_error)?;
 
             let mut message = [0; 64];
-            self.hash_blocking::<_, Sha512>(_state_and_dfu_flash, _update_len, _aligned, &mut message)?;
+            self.hash::<_, Sha512>(_state_and_dfu_flash, _update_len, _aligned, &mut message)
+                .await?;
 
             let r = public_key.verify(&message, &signature);
             trace!(
@@ -171,11 +170,12 @@ impl FirmwareUpdater {
             r.map_err(into_signature_error)?
         }
 
-        self.set_magic_blocking(_aligned, SWAP_MAGIC, _state_and_dfu_flash)
+        self.set_magic(_aligned, SWAP_MAGIC, _state_and_dfu_flash).await
     }
 
     /// Verify the update in DFU with any digest.
-    pub fn hash_blocking<F: NorFlash, D: Digest>(
+    #[cfg(feature = "nightly")]
+    pub async fn hash<F: AsyncNorFlash, D: Digest>(
         &mut self,
         dfu_flash: &mut F,
         update_len: u32,
@@ -184,7 +184,7 @@ impl FirmwareUpdater {
     ) -> Result<(), FirmwareUpdaterError> {
         let mut digest = D::new();
         for offset in (0..update_len).step_by(chunk_buf.len()) {
-            self.dfu.read_blocking(dfu_flash, offset, chunk_buf)?;
+            self.dfu.read(dfu_flash, offset, chunk_buf).await?;
             let len = core::cmp::min((update_len - offset) as usize, chunk_buf.len());
             digest.update(&chunk_buf[..len]);
         }
@@ -197,14 +197,14 @@ impl FirmwareUpdater {
     /// # Safety
     ///
     /// The `aligned` buffer must have a size of F::WRITE_SIZE, and follow the alignment rules for the flash being written to.
-    #[cfg(not(feature = "_verify"))]
-    pub fn mark_updated_blocking<F: NorFlash>(
+    #[cfg(all(feature = "nightly", not(feature = "_verify")))]
+    pub async fn mark_updated<F: AsyncNorFlash>(
         &mut self,
         state_flash: &mut F,
         aligned: &mut [u8],
     ) -> Result<(), FirmwareUpdaterError> {
         assert_eq!(aligned.len(), F::WRITE_SIZE);
-        self.set_magic_blocking(aligned, SWAP_MAGIC, state_flash)
+        self.set_magic(aligned, SWAP_MAGIC, state_flash).await
     }
 
     /// Mark firmware boot successful and stop rollback on reset.
@@ -212,26 +212,28 @@ impl FirmwareUpdater {
     /// # Safety
     ///
     /// The `aligned` buffer must have a size of F::WRITE_SIZE, and follow the alignment rules for the flash being written to.
-    pub fn mark_booted_blocking<F: NorFlash>(
+    #[cfg(feature = "nightly")]
+    pub async fn mark_booted<F: AsyncNorFlash>(
         &mut self,
         state_flash: &mut F,
         aligned: &mut [u8],
     ) -> Result<(), FirmwareUpdaterError> {
         assert_eq!(aligned.len(), F::WRITE_SIZE);
-        self.set_magic_blocking(aligned, BOOT_MAGIC, state_flash)
+        self.set_magic(aligned, BOOT_MAGIC, state_flash).await
     }
 
-    fn set_magic_blocking<F: NorFlash>(
+    #[cfg(feature = "nightly")]
+    async fn set_magic<F: AsyncNorFlash>(
         &mut self,
         aligned: &mut [u8],
         magic: u8,
         state_flash: &mut F,
     ) -> Result<(), FirmwareUpdaterError> {
-        self.state.read_blocking(state_flash, 0, aligned)?;
+        self.state.read(state_flash, 0, aligned).await?;
 
         if aligned.iter().any(|&b| b != magic) {
             // Read progress validity
-            self.state.read_blocking(state_flash, F::WRITE_SIZE as u32, aligned)?;
+            self.state.read(state_flash, F::WRITE_SIZE as u32, aligned).await?;
 
             // FIXME: Do not make this assumption.
             const STATE_ERASE_VALUE: u8 = 0xFF;
@@ -241,15 +243,15 @@ impl FirmwareUpdater {
             } else {
                 // Invalidate progress
                 aligned.fill(!STATE_ERASE_VALUE);
-                self.state.write_blocking(state_flash, F::WRITE_SIZE as u32, aligned)?;
+                self.state.write(state_flash, F::WRITE_SIZE as u32, aligned).await?;
             }
 
             // Clear magic and progress
-            self.state.wipe_blocking(state_flash)?;
+            self.state.wipe(state_flash).await?;
 
             // Set magic
             aligned.fill(magic);
-            self.state.write_blocking(state_flash, 0, aligned)?;
+            self.state.write(state_flash, 0, aligned).await?;
         }
         Ok(())
     }
@@ -261,7 +263,8 @@ impl FirmwareUpdater {
     /// # Safety
     ///
     /// Failing to meet alignment and size requirements may result in a panic.
-    pub fn write_firmware_blocking<F: NorFlash>(
+    #[cfg(feature = "nightly")]
+    pub async fn write_firmware<F: AsyncNorFlash>(
         &mut self,
         offset: usize,
         data: &[u8],
@@ -270,9 +273,10 @@ impl FirmwareUpdater {
         assert!(data.len() >= F::ERASE_SIZE);
 
         self.dfu
-            .erase_blocking(dfu_flash, offset as u32, (offset + data.len()) as u32)?;
+            .erase(dfu_flash, offset as u32, (offset + data.len()) as u32)
+            .await?;
 
-        self.dfu.write_blocking(dfu_flash, offset as u32, data)?;
+        self.dfu.write(dfu_flash, offset as u32, data).await?;
 
         Ok(())
     }
@@ -280,29 +284,33 @@ impl FirmwareUpdater {
     /// Prepare for an incoming DFU update by erasing the entire DFU area and
     /// returning its `Partition`.
     ///
-    /// Using this instead of `write_firmware_blocking` allows for an optimized
-    /// API in exchange for added complexity.
-    pub fn prepare_update_blocking<F: NorFlash>(&mut self, flash: &mut F) -> Result<Partition, FirmwareUpdaterError> {
-        self.dfu.wipe_blocking(flash)?;
+    /// Using this instead of `write_firmware` allows for an optimized API in
+    /// exchange for added complexity.
+    #[cfg(feature = "nightly")]
+    pub async fn prepare_update<F: AsyncNorFlash>(
+        &mut self,
+        dfu_flash: &mut F,
+    ) -> Result<Partition, FirmwareUpdaterError> {
+        self.dfu.wipe(dfu_flash).await?;
 
         Ok(self.dfu)
     }
-}
 
-// Async API
-#[cfg(feature = "nightly")]
-impl FirmwareUpdater {
+    //
+    // Blocking API
+    //
+
     /// Obtain the current state.
     ///
     /// This is useful to check if the bootloader has just done a swap, in order
     /// to do verifications and self-tests of the new image before calling
     /// `mark_booted`.
-    pub async fn get_state<F: AsyncNorFlash>(
+    pub fn get_state_blocking<F: NorFlash>(
         &mut self,
         state_flash: &mut F,
         aligned: &mut [u8],
     ) -> Result<State, FirmwareUpdaterError> {
-        self.state.read(state_flash, 0, aligned).await?;
+        self.state.read_blocking(state_flash, 0, aligned)?;
 
         if !aligned.iter().any(|&b| b != SWAP_MAGIC) {
             Ok(State::Swap)
@@ -328,7 +336,7 @@ impl FirmwareUpdater {
     /// The `_aligned` buffer must have a size of F::WRITE_SIZE, and follow the alignment rules for the flash being read from
     /// and written to.
     #[cfg(feature = "_verify")]
-    pub async fn verify_and_mark_updated<F: AsyncNorFlash>(
+    pub fn verify_and_mark_updated_blocking<F: NorFlash>(
         &mut self,
         _state_and_dfu_flash: &mut F,
         _public_key: &[u8],
@@ -351,8 +359,7 @@ impl FirmwareUpdater {
             let signature = Signature::from_bytes(_signature).map_err(into_signature_error)?;
 
             let mut message = [0; 64];
-            self.hash::<_, Sha512>(_state_and_dfu_flash, _update_len, _aligned, &mut message)
-                .await?;
+            self.hash_blocking::<_, Sha512>(_state_and_dfu_flash, _update_len, _aligned, &mut message)?;
 
             public_key.verify(&message, &signature).map_err(into_signature_error)?
         }
@@ -373,8 +380,7 @@ impl FirmwareUpdater {
             let signature = Signature::try_from(&signature).map_err(into_signature_error)?;
 
             let mut message = [0; 64];
-            self.hash::<_, Sha512>(_state_and_dfu_flash, _update_len, _aligned, &mut message)
-                .await?;
+            self.hash_blocking::<_, Sha512>(_state_and_dfu_flash, _update_len, _aligned, &mut message)?;
 
             let r = public_key.verify(&message, &signature);
             trace!(
@@ -387,11 +393,11 @@ impl FirmwareUpdater {
             r.map_err(into_signature_error)?
         }
 
-        self.set_magic(_aligned, SWAP_MAGIC, _state_and_dfu_flash).await
+        self.set_magic_blocking(_aligned, SWAP_MAGIC, _state_and_dfu_flash)
     }
 
     /// Verify the update in DFU with any digest.
-    pub async fn hash<F: AsyncNorFlash, D: Digest>(
+    pub fn hash_blocking<F: NorFlash, D: Digest>(
         &mut self,
         dfu_flash: &mut F,
         update_len: u32,
@@ -400,7 +406,7 @@ impl FirmwareUpdater {
     ) -> Result<(), FirmwareUpdaterError> {
         let mut digest = D::new();
         for offset in (0..update_len).step_by(chunk_buf.len()) {
-            self.dfu.read(dfu_flash, offset, chunk_buf).await?;
+            self.dfu.read_blocking(dfu_flash, offset, chunk_buf)?;
             let len = core::cmp::min((update_len - offset) as usize, chunk_buf.len());
             digest.update(&chunk_buf[..len]);
         }
@@ -414,13 +420,13 @@ impl FirmwareUpdater {
     ///
     /// The `aligned` buffer must have a size of F::WRITE_SIZE, and follow the alignment rules for the flash being written to.
     #[cfg(not(feature = "_verify"))]
-    pub async fn mark_updated<F: AsyncNorFlash>(
+    pub fn mark_updated_blocking<F: NorFlash>(
         &mut self,
         state_flash: &mut F,
         aligned: &mut [u8],
     ) -> Result<(), FirmwareUpdaterError> {
         assert_eq!(aligned.len(), F::WRITE_SIZE);
-        self.set_magic(aligned, SWAP_MAGIC, state_flash).await
+        self.set_magic_blocking(aligned, SWAP_MAGIC, state_flash)
     }
 
     /// Mark firmware boot successful and stop rollback on reset.
@@ -428,26 +434,26 @@ impl FirmwareUpdater {
     /// # Safety
     ///
     /// The `aligned` buffer must have a size of F::WRITE_SIZE, and follow the alignment rules for the flash being written to.
-    pub async fn mark_booted<F: AsyncNorFlash>(
+    pub fn mark_booted_blocking<F: NorFlash>(
         &mut self,
         state_flash: &mut F,
         aligned: &mut [u8],
     ) -> Result<(), FirmwareUpdaterError> {
         assert_eq!(aligned.len(), F::WRITE_SIZE);
-        self.set_magic(aligned, BOOT_MAGIC, state_flash).await
+        self.set_magic_blocking(aligned, BOOT_MAGIC, state_flash)
     }
 
-    async fn set_magic<F: AsyncNorFlash>(
+    fn set_magic_blocking<F: NorFlash>(
         &mut self,
         aligned: &mut [u8],
         magic: u8,
         state_flash: &mut F,
     ) -> Result<(), FirmwareUpdaterError> {
-        self.state.read(state_flash, 0, aligned).await?;
+        self.state.read_blocking(state_flash, 0, aligned)?;
 
         if aligned.iter().any(|&b| b != magic) {
             // Read progress validity
-            self.state.read(state_flash, F::WRITE_SIZE as u32, aligned).await?;
+            self.state.read_blocking(state_flash, F::WRITE_SIZE as u32, aligned)?;
 
             // FIXME: Do not make this assumption.
             const STATE_ERASE_VALUE: u8 = 0xFF;
@@ -457,15 +463,15 @@ impl FirmwareUpdater {
             } else {
                 // Invalidate progress
                 aligned.fill(!STATE_ERASE_VALUE);
-                self.state.write(state_flash, F::WRITE_SIZE as u32, aligned).await?;
+                self.state.write_blocking(state_flash, F::WRITE_SIZE as u32, aligned)?;
             }
 
             // Clear magic and progress
-            self.state.wipe(state_flash).await?;
+            self.state.wipe_blocking(state_flash)?;
 
             // Set magic
             aligned.fill(magic);
-            self.state.write(state_flash, 0, aligned).await?;
+            self.state.write_blocking(state_flash, 0, aligned)?;
         }
         Ok(())
     }
@@ -477,7 +483,7 @@ impl FirmwareUpdater {
     /// # Safety
     ///
     /// Failing to meet alignment and size requirements may result in a panic.
-    pub async fn write_firmware<F: AsyncNorFlash>(
+    pub fn write_firmware_blocking<F: NorFlash>(
         &mut self,
         offset: usize,
         data: &[u8],
@@ -486,10 +492,9 @@ impl FirmwareUpdater {
         assert!(data.len() >= F::ERASE_SIZE);
 
         self.dfu
-            .erase(dfu_flash, offset as u32, (offset + data.len()) as u32)
-            .await?;
+            .erase_blocking(dfu_flash, offset as u32, (offset + data.len()) as u32)?;
 
-        self.dfu.write(dfu_flash, offset as u32, data).await?;
+        self.dfu.write_blocking(dfu_flash, offset as u32, data)?;
 
         Ok(())
     }
@@ -497,20 +502,16 @@ impl FirmwareUpdater {
     /// Prepare for an incoming DFU update by erasing the entire DFU area and
     /// returning its `Partition`.
     ///
-    /// Using this instead of `write_firmware` allows for an optimized API in
-    /// exchange for added complexity.
-    pub async fn prepare_update<F: AsyncNorFlash>(
-        &mut self,
-        dfu_flash: &mut F,
-    ) -> Result<Partition, FirmwareUpdaterError> {
-        self.dfu.wipe(dfu_flash).await?;
+    /// Using this instead of `write_firmware_blocking` allows for an optimized
+    /// API in exchange for added complexity.
+    pub fn prepare_update_blocking<F: NorFlash>(&mut self, flash: &mut F) -> Result<Partition, FirmwareUpdaterError> {
+        self.dfu.wipe_blocking(flash)?;
 
         Ok(self.dfu)
     }
 }
 
-
-    #[cfg(test)]
+#[cfg(test)]
 mod tests {
     use futures::executor::block_on;
     use sha1::{Digest, Sha1};
diff --git a/embassy-boot/boot/src/partition.rs b/embassy-boot/boot/src/partition.rs
index 67bd7abba..7b56a8240 100644
--- a/embassy-boot/boot/src/partition.rs
+++ b/embassy-boot/boot/src/partition.rs
@@ -24,75 +24,74 @@ impl Partition {
     }
 
     /// 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> {
+    #[cfg(feature = "nightly")]
+    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)
+        flash.read(offset, bytes).await
     }
 
     /// 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> {
+    #[cfg(feature = "nightly")]
+    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)?;
+        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 fn erase_blocking<F: NorFlash>(&self, flash: &mut F, from: u32, to: u32) -> Result<(), F::Error> {
+    #[cfg(feature = "nightly")]
+    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)?;
+        flash.erase(from, to).await?;
         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> {
+    #[cfg(feature = "nightly")]
+    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)?;
+        flash.erase(from, to).await?;
         trace!("Wiped from 0x{:x} to 0x{:x}", from, to);
         Ok(())
     }
-}
-
-// Async API
-#[cfg(feature = "nightly")]
-impl Partition {
 
     /// 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> {
+    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).await
+        flash.read(offset, bytes)
     }
 
     /// 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> {
+    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).await?;
+        flash.write(offset, bytes)?;
         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> {
+    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).await?;
+        flash.erase(from, to)?;
         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> {
+    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).await?;
+        flash.erase(from, to)?;
         trace!("Wiped from 0x{:x} to 0x{:x}", from, to);
         Ok(())
     }
diff --git a/embassy-boot/nrf/Cargo.toml b/embassy-boot/nrf/Cargo.toml
index 90a36d014..05ab87896 100644
--- a/embassy-boot/nrf/Cargo.toml
+++ b/embassy-boot/nrf/Cargo.toml
@@ -37,4 +37,8 @@ defmt = [
 softdevice = [
     "nrf-softdevice-mbr",
 ]
-nightly = ["dep:embedded-storage-async", "embassy-boot/nightly", "embassy-nrf/nightly"]
+nightly = [
+    "dep:embedded-storage-async",
+    "embassy-boot/nightly",
+    "embassy-nrf/nightly"
+]