Ensure bootloader state is parsed correctly

3 files changed, 20 insertions, 14 deletions

diff --git a/embassy-boot/src/firmware_updater/asynch.rs b/embassy-boot/src/firmware_updater/asynch.rs
index b23857e2f..d9d15b004 100644
--- a/embassy-boot/src/firmware_updater/asynch.rs
+++ b/embassy-boot/src/firmware_updater/asynch.rs
@@ -304,12 +304,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))
     }
 
     /// 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 5f64b4be9..08062b0d0 100644
--- a/embassy-boot/src/firmware_updater/blocking.rs
+++ b/embassy-boot/src/firmware_updater/blocking.rs
@@ -339,14 +339,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/lib.rs b/embassy-boot/src/lib.rs
index 7d5cc58f9..e2c4cf771 100644
--- a/embassy-boot/src/lib.rs
+++ b/embassy-boot/src/lib.rs
@@ -44,6 +44,24 @@ pub enum State {
     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]);