5 files changed, 85 insertions, 56 deletions

diff --git a/embassy-boot/boot/src/boot_loader.rs b/embassy-boot/boot/src/boot_loader.rs
index a8c19197b..e568001bc 100644
--- a/embassy-boot/boot/src/boot_loader.rs
+++ b/embassy-boot/boot/src/boot_loader.rs
@@ -5,7 +5,7 @@ use embassy_sync::blocking_mutex::raw::NoopRawMutex;
 use embassy_sync::blocking_mutex::Mutex;
 use embedded_storage::nor_flash::{NorFlash, NorFlashError, NorFlashErrorKind};
 
-use crate::{State, BOOT_MAGIC, STATE_ERASE_VALUE, SWAP_MAGIC};
+use crate::{State, BOOT_MAGIC, DFU_DETACH_MAGIC, STATE_ERASE_VALUE, SWAP_MAGIC};
 
 /// Errors returned by bootloader
 #[derive(PartialEq, Eq, Debug)]
@@ -135,51 +135,44 @@ impl<ACTIVE: NorFlash, DFU: NorFlash, STATE: NorFlash> BootLoader<ACTIVE, DFU, S
     /// The provided aligned_buf argument must satisfy any alignment requirements
     /// given by the partition flashes. All flash operations will use this buffer.
     ///
-    /// SWAPPING
+    /// ## SWAPPING
     ///
     /// Assume a flash size of 3 pages for the active partition, and 4 pages for the DFU partition.
     /// The swap index contains the copy progress, as to allow continuation of the copy process on
     /// power failure. The index counter is represented within 1 or more pages (depending on total
-    /// flash size), where a page X is considered swapped if index at location (X + WRITE_SIZE)
+    /// flash size), where a page X is considered swapped if index at location (`X + WRITE_SIZE`)
     /// contains a zero value. This ensures that index updates can be performed atomically and
     /// avoid a situation where the wrong index value is set (page write size is "atomic").
     ///
-    /// +-----------+------------+--------+--------+--------+--------+
+    ///
     /// | Partition | Swap Index | Page 0 | Page 1 | Page 3 | Page 4 |
-    /// +-----------+------------+--------+--------+--------+--------+
+    /// |-----------|------------|--------|--------|--------|--------|
     /// |    Active |          0 |      1 |      2 |      3 |      - |
     /// |       DFU |          0 |      3 |      2 |      1 |      X |
-    /// +-----------+------------+--------+--------+--------+--------+
     ///
     /// The algorithm starts by copying 'backwards', and after the first step, the layout is
     /// as follows:
     ///
-    /// +-----------+------------+--------+--------+--------+--------+
     /// | Partition | Swap Index | Page 0 | Page 1 | Page 3 | Page 4 |
-    /// +-----------+------------+--------+--------+--------+--------+
+    /// |-----------|------------|--------|--------|--------|--------|
     /// |    Active |          1 |      1 |      2 |      1 |      - |
     /// |       DFU |          1 |      3 |      2 |      1 |      3 |
-    /// +-----------+------------+--------+--------+--------+--------+
     ///
     /// The next iteration performs the same steps
     ///
-    /// +-----------+------------+--------+--------+--------+--------+
     /// | Partition | Swap Index | Page 0 | Page 1 | Page 3 | Page 4 |
-    /// +-----------+------------+--------+--------+--------+--------+
+    /// |-----------|------------|--------|--------|--------|--------|
     /// |    Active |          2 |      1 |      2 |      1 |      - |
     /// |       DFU |          2 |      3 |      2 |      2 |      3 |
-    /// +-----------+------------+--------+--------+--------+--------+
     ///
     /// And again until we're done
     ///
-    /// +-----------+------------+--------+--------+--------+--------+
     /// | Partition | Swap Index | Page 0 | Page 1 | Page 3 | Page 4 |
-    /// +-----------+------------+--------+--------+--------+--------+
+    /// |-----------|------------|--------|--------|--------|--------|
     /// |    Active |          3 |      3 |      2 |      1 |      - |
     /// |       DFU |          3 |      3 |      1 |      2 |      3 |
-    /// +-----------+------------+--------+--------+--------+--------+
     ///
-    /// REVERTING
+    /// ## REVERTING
     ///
     /// The reverting algorithm uses the swap index to discover that images were swapped, but that
     /// the application failed to mark the boot successful. In this case, the revert algorithm will
@@ -190,28 +183,21 @@ impl<ACTIVE: NorFlash, DFU: NorFlash, STATE: NorFlash> BootLoader<ACTIVE, DFU, S
     ///
     /// The revert algorithm works forwards, by starting copying into the 'unused' DFU page at the start.
     ///
-    /// +-----------+--------------+--------+--------+--------+--------+
     /// | Partition | Revert Index | Page 0 | Page 1 | Page 3 | Page 4 |
-    //*/
-    /// +-----------+--------------+--------+--------+--------+--------+
+    /// |-----------|--------------|--------|--------|--------|--------|
     /// |    Active |            3 |      1 |      2 |      1 |      - |
     /// |       DFU |            3 |      3 |      1 |      2 |      3 |
-    /// +-----------+--------------+--------+--------+--------+--------+
     ///
     ///
-    /// +-----------+--------------+--------+--------+--------+--------+
     /// | Partition | Revert Index | Page 0 | Page 1 | Page 3 | Page 4 |
-    /// +-----------+--------------+--------+--------+--------+--------+
+    /// |-----------|--------------|--------|--------|--------|--------|
     /// |    Active |            3 |      1 |      2 |      1 |      - |
     /// |       DFU |            3 |      3 |      2 |      2 |      3 |
-    /// +-----------+--------------+--------+--------+--------+--------+
     ///
-    /// +-----------+--------------+--------+--------+--------+--------+
     /// | Partition | Revert Index | Page 0 | Page 1 | Page 3 | Page 4 |
-    /// +-----------+--------------+--------+--------+--------+--------+
+    /// |-----------|--------------|--------|--------|--------|--------|
     /// |    Active |            3 |      1 |      2 |      3 |      - |
     /// |       DFU |            3 |      3 |      2 |      1 |      3 |
-    /// +-----------+--------------+--------+--------+--------+--------+
     ///
     pub fn prepare_boot(&mut self, aligned_buf: &mut [u8]) -> Result<State, BootError> {
         // Ensure we have enough progress pages to store copy progress
@@ -224,6 +210,7 @@ impl<ACTIVE: NorFlash, DFU: NorFlash, STATE: NorFlash> BootLoader<ACTIVE, DFU, S
         assert_eq!(0, aligned_buf.len() % ACTIVE::WRITE_SIZE);
         assert_eq!(0, aligned_buf.len() % DFU::WRITE_SIZE);
 
+        // Ensure our partitions are able to handle boot operations
         assert_partitions(&self.active, &self.dfu, &self.state, Self::PAGE_SIZE);
 
         // Copy contents from partition N to active
@@ -384,6 +371,8 @@ impl<ACTIVE: NorFlash, DFU: NorFlash, STATE: NorFlash> BootLoader<ACTIVE, DFU, S
 
         if !state_word.iter().any(|&b| b != SWAP_MAGIC) {
             Ok(State::Swap)
+        } else if !state_word.iter().any(|&b| b != DFU_DETACH_MAGIC) {
+            Ok(State::DfuDetach)
         } else {
             Ok(State::Boot)
         }
@@ -398,6 +387,7 @@ fn assert_partitions<ACTIVE: NorFlash, DFU: NorFlash, STATE: NorFlash>(
 ) {
     assert_eq!(active.capacity() as u32 % page_size, 0);
     assert_eq!(dfu.capacity() as u32 % page_size, 0);
+    // DFU partition has to be bigger than ACTIVE partition to handle swap algorithm
     assert!(dfu.capacity() as u32 - active.capacity() as u32 >= page_size);
     assert!(2 + 2 * (active.capacity() as u32 / page_size) <= state.capacity() as u32 / STATE::WRITE_SIZE as u32);
 }
diff --git a/embassy-boot/boot/src/digest_adapters/ed25519_dalek.rs b/embassy-boot/boot/src/digest_adapters/ed25519_dalek.rs
index a184d1c51..2e4e03da3 100644
--- a/embassy-boot/boot/src/digest_adapters/ed25519_dalek.rs
+++ b/embassy-boot/boot/src/digest_adapters/ed25519_dalek.rs
@@ -1,6 +1,6 @@
 use digest::typenum::U64;
 use digest::{FixedOutput, HashMarker, OutputSizeUser, Update};
-use ed25519_dalek::Digest as _;
+use ed25519_dalek::Digest;
 
 pub struct Sha512(ed25519_dalek::Sha512);
 
@@ -12,7 +12,7 @@ impl Default for Sha512 {
 
 impl Update for Sha512 {
     fn update(&mut self, data: &[u8]) {
-        self.0.update(data)
+        Digest::update(&mut self.0, data)
     }
 }
 
diff --git a/embassy-boot/boot/src/firmware_updater/asynch.rs b/embassy-boot/boot/src/firmware_updater/asynch.rs
index ae713bb6f..64a4b32ec 100644
--- a/embassy-boot/boot/src/firmware_updater/asynch.rs
+++ b/embassy-boot/boot/src/firmware_updater/asynch.rs
@@ -6,7 +6,7 @@ use embassy_sync::blocking_mutex::raw::NoopRawMutex;
 use embedded_storage_async::nor_flash::NorFlash;
 
 use super::FirmwareUpdaterConfig;
-use crate::{FirmwareUpdaterError, State, BOOT_MAGIC, STATE_ERASE_VALUE, SWAP_MAGIC};
+use crate::{FirmwareUpdaterError, State, BOOT_MAGIC, DFU_DETACH_MAGIC, STATE_ERASE_VALUE, SWAP_MAGIC};
 
 /// FirmwareUpdater is an application API for interacting with the BootLoader without the ability to
 /// 'mess up' the internal bootloader state
@@ -79,8 +79,8 @@ impl<'d, DFU: NorFlash, STATE: NorFlash> FirmwareUpdater<'d, DFU, STATE> {
     #[cfg(feature = "_verify")]
     pub async fn verify_and_mark_updated(
         &mut self,
-        _public_key: &[u8],
-        _signature: &[u8],
+        _public_key: &[u8; 32],
+        _signature: &[u8; 64],
         _update_len: u32,
     ) -> Result<(), FirmwareUpdaterError> {
         assert!(_update_len <= self.dfu.capacity() as u32);
@@ -89,14 +89,14 @@ impl<'d, DFU: NorFlash, STATE: NorFlash> FirmwareUpdater<'d, DFU, STATE> {
 
         #[cfg(feature = "ed25519-dalek")]
         {
-            use ed25519_dalek::{PublicKey, Signature, SignatureError, Verifier};
+            use ed25519_dalek::{Signature, SignatureError, Verifier, VerifyingKey};
 
             use crate::digest_adapters::ed25519_dalek::Sha512;
 
             let into_signature_error = |e: SignatureError| FirmwareUpdaterError::Signature(e.into());
 
-            let public_key = PublicKey::from_bytes(_public_key).map_err(into_signature_error)?;
-            let signature = Signature::from_bytes(_signature).map_err(into_signature_error)?;
+            let public_key = VerifyingKey::from_bytes(_public_key).map_err(into_signature_error)?;
+            let signature = Signature::from_bytes(_signature);
 
             let mut chunk_buf = [0; 2];
             let mut message = [0; 64];
@@ -106,7 +106,6 @@ impl<'d, DFU: NorFlash, STATE: NorFlash> FirmwareUpdater<'d, DFU, STATE> {
         }
         #[cfg(feature = "ed25519-salty")]
         {
-            use salty::constants::{PUBLICKEY_SERIALIZED_LENGTH, SIGNATURE_SERIALIZED_LENGTH};
             use salty::{PublicKey, Signature};
 
             use crate::digest_adapters::salty::Sha512;
@@ -115,10 +114,8 @@ impl<'d, DFU: NorFlash, STATE: NorFlash> FirmwareUpdater<'d, DFU, STATE> {
                 FirmwareUpdaterError::Signature(signature::Error::default())
             }
 
-            let public_key: [u8; PUBLICKEY_SERIALIZED_LENGTH] = _public_key.try_into().map_err(into_signature_error)?;
-            let public_key = PublicKey::try_from(&public_key).map_err(into_signature_error)?;
-            let signature: [u8; SIGNATURE_SERIALIZED_LENGTH] = _signature.try_into().map_err(into_signature_error)?;
-            let signature = Signature::try_from(&signature).map_err(into_signature_error)?;
+            let public_key = PublicKey::try_from(_public_key).map_err(into_signature_error)?;
+            let signature = Signature::try_from(_signature).map_err(into_signature_error)?;
 
             let mut message = [0; 64];
             let mut chunk_buf = [0; 2];
@@ -161,6 +158,12 @@ impl<'d, DFU: NorFlash, STATE: NorFlash> FirmwareUpdater<'d, DFU, STATE> {
         self.state.mark_updated().await
     }
 
+    /// Mark to trigger USB DFU on next boot.
+    pub async fn mark_dfu(&mut self) -> Result<(), FirmwareUpdaterError> {
+        self.state.verify_booted().await?;
+        self.state.mark_dfu().await
+    }
+
     /// Mark firmware boot successful and stop rollback on reset.
     pub async fn mark_booted(&mut self) -> Result<(), FirmwareUpdaterError> {
         self.state.mark_booted().await
@@ -207,6 +210,16 @@ pub struct FirmwareState<'d, STATE> {
 }
 
 impl<'d, STATE: NorFlash> FirmwareState<'d, STATE> {
+    /// Create a firmware state instance from a FirmwareUpdaterConfig 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 read from
+    /// and written to.
+    pub fn from_config<DFU: NorFlash>(config: FirmwareUpdaterConfig<DFU, STATE>, aligned: &'d mut [u8]) -> Self {
+        Self::new(config.state, aligned)
+    }
+
     /// Create a firmware state instance with a buffer for magic content and state partition.
     ///
     /// # Safety
@@ -247,6 +260,11 @@ impl<'d, STATE: NorFlash> FirmwareState<'d, STATE> {
         self.set_magic(SWAP_MAGIC).await
     }
 
+    /// Mark to trigger USB DFU on next boot.
+    pub async fn mark_dfu(&mut self) -> Result<(), FirmwareUpdaterError> {
+        self.set_magic(DFU_DETACH_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
diff --git a/embassy-boot/boot/src/firmware_updater/blocking.rs b/embassy-boot/boot/src/firmware_updater/blocking.rs
index 76e4264a0..f1368540d 100644
--- a/embassy-boot/boot/src/firmware_updater/blocking.rs
+++ b/embassy-boot/boot/src/firmware_updater/blocking.rs
@@ -6,7 +6,7 @@ use embassy_sync::blocking_mutex::raw::NoopRawMutex;
 use embedded_storage::nor_flash::NorFlash;
 
 use super::FirmwareUpdaterConfig;
-use crate::{FirmwareUpdaterError, State, BOOT_MAGIC, STATE_ERASE_VALUE, SWAP_MAGIC};
+use crate::{FirmwareUpdaterError, State, BOOT_MAGIC, DFU_DETACH_MAGIC, STATE_ERASE_VALUE, SWAP_MAGIC};
 
 /// Blocking FirmwareUpdater is an application API for interacting with the BootLoader without the ability to
 /// 'mess up' the internal bootloader state
@@ -86,8 +86,8 @@ impl<'d, DFU: NorFlash, STATE: NorFlash> BlockingFirmwareUpdater<'d, DFU, STATE>
     #[cfg(feature = "_verify")]
     pub fn verify_and_mark_updated(
         &mut self,
-        _public_key: &[u8],
-        _signature: &[u8],
+        _public_key: &[u8; 32],
+        _signature: &[u8; 64],
         _update_len: u32,
     ) -> Result<(), FirmwareUpdaterError> {
         assert!(_update_len <= self.dfu.capacity() as u32);
@@ -96,14 +96,14 @@ impl<'d, DFU: NorFlash, STATE: NorFlash> BlockingFirmwareUpdater<'d, DFU, STATE>
 
         #[cfg(feature = "ed25519-dalek")]
         {
-            use ed25519_dalek::{PublicKey, Signature, SignatureError, Verifier};
+            use ed25519_dalek::{Signature, SignatureError, Verifier, VerifyingKey};
 
             use crate::digest_adapters::ed25519_dalek::Sha512;
 
             let into_signature_error = |e: SignatureError| FirmwareUpdaterError::Signature(e.into());
 
-            let public_key = PublicKey::from_bytes(_public_key).map_err(into_signature_error)?;
-            let signature = Signature::from_bytes(_signature).map_err(into_signature_error)?;
+            let public_key = VerifyingKey::from_bytes(_public_key).map_err(into_signature_error)?;
+            let signature = Signature::from_bytes(_signature);
 
             let mut message = [0; 64];
             let mut chunk_buf = [0; 2];
@@ -113,7 +113,6 @@ impl<'d, DFU: NorFlash, STATE: NorFlash> BlockingFirmwareUpdater<'d, DFU, STATE>
         }
         #[cfg(feature = "ed25519-salty")]
         {
-            use salty::constants::{PUBLICKEY_SERIALIZED_LENGTH, SIGNATURE_SERIALIZED_LENGTH};
             use salty::{PublicKey, Signature};
 
             use crate::digest_adapters::salty::Sha512;
@@ -122,10 +121,8 @@ impl<'d, DFU: NorFlash, STATE: NorFlash> BlockingFirmwareUpdater<'d, DFU, STATE>
                 FirmwareUpdaterError::Signature(signature::Error::default())
             }
 
-            let public_key: [u8; PUBLICKEY_SERIALIZED_LENGTH] = _public_key.try_into().map_err(into_signature_error)?;
-            let public_key = PublicKey::try_from(&public_key).map_err(into_signature_error)?;
-            let signature: [u8; SIGNATURE_SERIALIZED_LENGTH] = _signature.try_into().map_err(into_signature_error)?;
-            let signature = Signature::try_from(&signature).map_err(into_signature_error)?;
+            let public_key = PublicKey::try_from(_public_key).map_err(into_signature_error)?;
+            let signature = Signature::try_from(_signature).map_err(into_signature_error)?;
 
             let mut message = [0; 64];
             let mut chunk_buf = [0; 2];
@@ -168,6 +165,12 @@ impl<'d, DFU: NorFlash, STATE: NorFlash> BlockingFirmwareUpdater<'d, DFU, STATE>
         self.state.mark_updated()
     }
 
+    /// Mark to trigger USB DFU device on next boot.
+    pub fn mark_dfu(&mut self) -> Result<(), FirmwareUpdaterError> {
+        self.state.verify_booted()?;
+        self.state.mark_dfu()
+    }
+
     /// Mark firmware boot successful and stop rollback on reset.
     pub fn mark_booted(&mut self) -> Result<(), FirmwareUpdaterError> {
         self.state.mark_booted()
@@ -213,6 +216,16 @@ pub struct BlockingFirmwareState<'d, STATE> {
 }
 
 impl<'d, STATE: NorFlash> BlockingFirmwareState<'d, STATE> {
+    /// Creates a firmware state instance from a FirmwareUpdaterConfig, 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 read from
+    /// and written to.
+    pub fn from_config<DFU: NorFlash>(config: FirmwareUpdaterConfig<DFU, STATE>, aligned: &'d mut [u8]) -> Self {
+        Self::new(config.state, aligned)
+    }
+
     /// Create a firmware state instance with a buffer for magic content and state partition.
     ///
     /// # Safety
@@ -226,7 +239,7 @@ impl<'d, STATE: NorFlash> BlockingFirmwareState<'d, STATE> {
 
     // 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 {
+        if self.get_state()? == State::Boot || self.get_state()? == State::DfuDetach {
             Ok(())
         } else {
             Err(FirmwareUpdaterError::BadState)
@@ -243,6 +256,8 @@ impl<'d, STATE: NorFlash> BlockingFirmwareState<'d, STATE> {
 
         if !self.aligned.iter().any(|&b| b != SWAP_MAGIC) {
             Ok(State::Swap)
+        } else if !self.aligned.iter().any(|&b| b != DFU_DETACH_MAGIC) {
+            Ok(State::DfuDetach)
         } else {
             Ok(State::Boot)
         }
@@ -253,6 +268,11 @@ impl<'d, STATE: NorFlash> BlockingFirmwareState<'d, STATE> {
         self.set_magic(SWAP_MAGIC)
     }
 
+    /// Mark to trigger USB DFU on next boot.
+    pub fn mark_dfu(&mut self) -> Result<(), FirmwareUpdaterError> {
+        self.set_magic(DFU_DETACH_MAGIC)
+    }
+
     /// Mark firmware boot successful and stop rollback on reset.
     pub fn mark_booted(&mut self) -> Result<(), FirmwareUpdaterError> {
         self.set_magic(BOOT_MAGIC)
diff --git a/embassy-boot/boot/src/lib.rs b/embassy-boot/boot/src/lib.rs
index 9e70a4dca..b4f03e01e 100644
--- a/embassy-boot/boot/src/lib.rs
+++ b/embassy-boot/boot/src/lib.rs
@@ -23,6 +23,7 @@ pub use firmware_updater::{
 
 pub(crate) const BOOT_MAGIC: u8 = 0xD0;
 pub(crate) const SWAP_MAGIC: u8 = 0xF0;
+pub(crate) const DFU_DETACH_MAGIC: u8 = 0xE0;
 
 /// The state of the bootloader after running prepare.
 #[derive(PartialEq, Eq, Debug)]
@@ -32,6 +33,8 @@ pub enum State {
     Boot,
     /// Bootloader has swapped the active partition with the dfu partition and will attempt boot.
     Swap,
+    /// Application has received a request to reboot into DFU mode to apply an update.
+    DfuDetach,
 }
 
 /// Buffer aligned to 32 byte boundary, largest known alignment requirement for embassy-boot.
@@ -272,21 +275,19 @@ mod tests {
         // The following key setup is based on:
         // https://docs.rs/ed25519-dalek/latest/ed25519_dalek/#example
 
-        use ed25519_dalek::Keypair;
+        use ed25519_dalek::{Digest, Sha512, Signature, Signer, SigningKey, VerifyingKey};
         use rand::rngs::OsRng;
 
         let mut csprng = OsRng {};
-        let keypair: Keypair = Keypair::generate(&mut csprng);
+        let keypair = SigningKey::generate(&mut csprng);
 
-        use ed25519_dalek::{Digest, Sha512, Signature, Signer};
         let firmware: &[u8] = b"This are bytes that would otherwise be firmware bytes for DFU.";
         let mut digest = Sha512::new();
         digest.update(&firmware);
         let message = digest.finalize();
         let signature: Signature = keypair.sign(&message);
 
-        use ed25519_dalek::PublicKey;
-        let public_key: PublicKey = keypair.public;
+        let public_key = keypair.verifying_key();
 
         // Setup flash
         let flash = BlockingTestFlash::new(BootLoaderConfig {