5 files changed, 77 insertions, 27 deletions

diff --git a/embassy-boot/src/boot_loader.rs b/embassy-boot/src/boot_loader.rs
index 61d61b96e..5bffdc5ea 100644
--- a/embassy-boot/src/boot_loader.rs
+++ b/embassy-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, DFU_DETACH_MAGIC, STATE_ERASE_VALUE, SWAP_MAGIC};
+use crate::{State, DFU_DETACH_MAGIC, REVERT_MAGIC, STATE_ERASE_VALUE, SWAP_MAGIC};
 
 /// Errors returned by bootloader
 #[derive(PartialEq, Eq, Debug)]
@@ -276,7 +276,7 @@ impl<ACTIVE: NorFlash, DFU: NorFlash, STATE: NorFlash> BootLoader<ACTIVE, DFU, S
                 self.state.erase(0, self.state.capacity() as u32)?;
 
                 // Set magic
-                state_word.fill(BOOT_MAGIC);
+                state_word.fill(REVERT_MAGIC);
                 self.state.write(0, state_word)?;
             }
         }
@@ -411,6 +411,8 @@ impl<ACTIVE: NorFlash, DFU: NorFlash, STATE: NorFlash> BootLoader<ACTIVE, DFU, S
             Ok(State::Swap)
         } else if !state_word.iter().any(|&b| b != DFU_DETACH_MAGIC) {
             Ok(State::DfuDetach)
+        } else if !state_word.iter().any(|&b| b != REVERT_MAGIC) {
+            Ok(State::Revert)
         } else {
             Ok(State::Boot)
         }
diff --git a/embassy-boot/src/firmware_updater/asynch.rs b/embassy-boot/src/firmware_updater/asynch.rs
index 26f65f295..66e311e38 100644
--- a/embassy-boot/src/firmware_updater/asynch.rs
+++ b/embassy-boot/src/firmware_updater/asynch.rs
@@ -107,7 +107,8 @@ impl<'d, DFU: NorFlash, STATE: NorFlash> FirmwareUpdater<'d, DFU, STATE> {
             let mut message = [0; 64];
             self.hash::<Sha512>(_update_len, &mut chunk_buf, &mut message).await?;
 
-            public_key.verify(&message, &signature).map_err(into_signature_error)?
+            public_key.verify(&message, &signature).map_err(into_signature_error)?;
+            return self.state.mark_updated().await;
         }
         #[cfg(feature = "ed25519-salty")]
         {
@@ -134,10 +135,13 @@ impl<'d, DFU: NorFlash, STATE: NorFlash> FirmwareUpdater<'d, DFU, STATE> {
                 message,
                 r.is_ok()
             );
-            r.map_err(into_signature_error)?
+            r.map_err(into_signature_error)?;
+            return self.state.mark_updated().await;
+        }
+        #[cfg(not(any(feature = "ed25519-dalek", feature = "ed25519-salty")))]
+        {
+            Err(FirmwareUpdaterError::Signature(signature::Error::new()))
         }
-
-        self.state.mark_updated().await
     }
 
     /// Verify the update in DFU with any digest.
@@ -157,6 +161,17 @@ impl<'d, DFU: NorFlash, STATE: NorFlash> FirmwareUpdater<'d, DFU, STATE> {
         Ok(())
     }
 
+    /// Read a slice of data from the DFU storage peripheral, starting the read
+    /// operation at the given address offset, and reading `buf.len()` bytes.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the arguments are not aligned or out of bounds.
+    pub async fn read_dfu(&mut self, offset: u32, buf: &mut [u8]) -> Result<(), FirmwareUpdaterError> {
+        self.dfu.read(offset, buf).await?;
+        Ok(())
+    }
+
     /// Mark to trigger firmware swap on next boot.
     #[cfg(not(feature = "_verify"))]
     pub async fn mark_updated(&mut self) -> Result<(), FirmwareUpdaterError> {
@@ -285,7 +300,8 @@ impl<'d, STATE: NorFlash> FirmwareState<'d, STATE> {
 
     // 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 {
+        let state = self.get_state().await?;
+        if state == State::Boot || state == State::DfuDetach || state == State::Revert {
             Ok(())
         } else {
             Err(FirmwareUpdaterError::BadState)
@@ -299,12 +315,7 @@ impl<'d, STATE: NorFlash> FirmwareState<'d, STATE> {
     /// `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)
-        }
+        Ok(State::from(&self.aligned[..STATE::WRITE_SIZE]))
     }
 
     /// Mark to trigger firmware swap on next boot.
diff --git a/embassy-boot/src/firmware_updater/blocking.rs b/embassy-boot/src/firmware_updater/blocking.rs
index 35772a856..0fedac1ea 100644
--- a/embassy-boot/src/firmware_updater/blocking.rs
+++ b/embassy-boot/src/firmware_updater/blocking.rs
@@ -142,7 +142,8 @@ impl<'d, DFU: NorFlash, STATE: NorFlash> BlockingFirmwareUpdater<'d, DFU, STATE>
             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)?
+            public_key.verify(&message, &signature).map_err(into_signature_error)?;
+            return self.state.mark_updated();
         }
         #[cfg(feature = "ed25519-salty")]
         {
@@ -169,10 +170,13 @@ impl<'d, DFU: NorFlash, STATE: NorFlash> BlockingFirmwareUpdater<'d, DFU, STATE>
                 message,
                 r.is_ok()
             );
-            r.map_err(into_signature_error)?
+            r.map_err(into_signature_error)?;
+            return self.state.mark_updated();
+        }
+        #[cfg(not(any(feature = "ed25519-dalek", feature = "ed25519-salty")))]
+        {
+            Err(FirmwareUpdaterError::Signature(signature::Error::new()))
         }
-
-        self.state.mark_updated()
     }
 
     /// Verify the update in DFU with any digest.
@@ -192,6 +196,17 @@ impl<'d, DFU: NorFlash, STATE: NorFlash> BlockingFirmwareUpdater<'d, DFU, STATE>
         Ok(())
     }
 
+    /// Read a slice of data from the DFU storage peripheral, starting the read
+    /// operation at the given address offset, and reading `buf.len()` bytes.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the arguments are not aligned or out of bounds.
+    pub fn read_dfu(&mut self, offset: u32, buf: &mut [u8]) -> Result<(), FirmwareUpdaterError> {
+        self.dfu.read(offset, buf)?;
+        Ok(())
+    }
+
     /// Mark to trigger firmware swap on next boot.
     #[cfg(not(feature = "_verify"))]
     pub fn mark_updated(&mut self) -> Result<(), FirmwareUpdaterError> {
@@ -320,7 +335,8 @@ 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 || self.get_state()? == State::DfuDetach {
+        let state = self.get_state()?;
+        if state == State::Boot || state == State::DfuDetach || state == State::Revert {
             Ok(())
         } else {
             Err(FirmwareUpdaterError::BadState)
@@ -334,14 +350,7 @@ impl<'d, STATE: NorFlash> BlockingFirmwareState<'d, STATE> {
     /// `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 if !self.aligned.iter().any(|&b| b != DFU_DETACH_MAGIC) {
-            Ok(State::DfuDetach)
-        } else {
-            Ok(State::Boot)
-        }
+        Ok(State::from(&self.aligned))
     }
 
     /// Mark to trigger firmware swap on next boot.
diff --git a/embassy-boot/src/firmware_updater/mod.rs b/embassy-boot/src/firmware_updater/mod.rs
index 4c4f4f10b..4814786bf 100644
--- a/embassy-boot/src/firmware_updater/mod.rs
+++ b/embassy-boot/src/firmware_updater/mod.rs
@@ -8,7 +8,7 @@ use embedded_storage::nor_flash::{NorFlashError, NorFlashErrorKind};
 /// Firmware updater flash configuration holding the two flashes used by the updater
 ///
 /// If only a single flash is actually used, then that flash should be partitioned into two partitions before use.
-/// The easiest way to do this is to use [`FirmwareUpdaterConfig::from_linkerfile_blocking`] or [`FirmwareUpdaterConfig::from_linkerfile_blocking`] which will partition
+/// The easiest way to do this is to use [`FirmwareUpdaterConfig::from_linkerfile`] or [`FirmwareUpdaterConfig::from_linkerfile_blocking`] which will partition
 /// the provided flash according to symbols defined in the linkerfile.
 pub struct FirmwareUpdaterConfig<DFU, STATE> {
     /// The dfu flash partition
diff --git a/embassy-boot/src/lib.rs b/embassy-boot/src/lib.rs
index b4f03e01e..e2c4cf771 100644
--- a/embassy-boot/src/lib.rs
+++ b/embassy-boot/src/lib.rs
@@ -14,13 +14,18 @@ mod test_flash;
 
 // The expected value of the flash after an erase
 // TODO: Use the value provided by NorFlash when available
+#[cfg(not(feature = "flash-erase-zero"))]
 pub(crate) const STATE_ERASE_VALUE: u8 = 0xFF;
+#[cfg(feature = "flash-erase-zero")]
+pub(crate) const STATE_ERASE_VALUE: u8 = 0x00;
+
 pub use boot_loader::{BootError, BootLoader, BootLoaderConfig};
 pub use firmware_updater::{
     BlockingFirmwareState, BlockingFirmwareUpdater, FirmwareState, FirmwareUpdater, FirmwareUpdaterConfig,
     FirmwareUpdaterError,
 };
 
+pub(crate) const REVERT_MAGIC: u8 = 0xC0;
 pub(crate) const BOOT_MAGIC: u8 = 0xD0;
 pub(crate) const SWAP_MAGIC: u8 = 0xF0;
 pub(crate) const DFU_DETACH_MAGIC: u8 = 0xE0;
@@ -33,10 +38,30 @@ pub enum State {
     Boot,
     /// Bootloader has swapped the active partition with the dfu partition and will attempt boot.
     Swap,
+    /// Bootloader has reverted the active partition with the dfu partition and will attempt boot.
+    Revert,
     /// Application has received a request to reboot into DFU mode to apply an update.
     DfuDetach,
 }
 
+impl<T> From<T> for State
+where
+    T: AsRef<[u8]>,
+{
+    fn from(magic: T) -> State {
+        let magic = magic.as_ref();
+        if !magic.iter().any(|&b| b != SWAP_MAGIC) {
+            State::Swap
+        } else if !magic.iter().any(|&b| b != REVERT_MAGIC) {
+            State::Revert
+        } else if !magic.iter().any(|&b| b != DFU_DETACH_MAGIC) {
+            State::DfuDetach
+        } else {
+            State::Boot
+        }
+    }
+}
+
 /// Buffer aligned to 32 byte boundary, largest known alignment requirement for embassy-boot.
 #[repr(align(32))]
 pub struct AlignedBuffer<const N: usize>(pub [u8; N]);
@@ -153,6 +178,9 @@ mod tests {
         // Running again should cause a revert
         assert_eq!(State::Swap, bootloader.prepare_boot(&mut page).unwrap());
 
+        // Next time we know it was reverted
+        assert_eq!(State::Revert, bootloader.prepare_boot(&mut page).unwrap());
+
         let mut read_buf = [0; FIRMWARE_SIZE];
         flash.active().read(0, &mut read_buf).unwrap();
         assert_eq!(ORIGINAL, read_buf);