From ab0a227e4c02137bc3a621907d17ede0ace4cb1d Mon Sep 17 00:00:00 2001
From: Ulf Lilleengen <ulf@digili.no>
Date: Thu, 19 Sep 2024 09:15:08 +0200
Subject: [PATCH] Ensure bootloader state is parsed correctly

---
 embassy-boot/src/firmware_updater/asynch.rs   |  7 +------
 embassy-boot/src/firmware_updater/blocking.rs |  9 +--------
 embassy-boot/src/lib.rs                       | 18 ++++++++++++++++++
 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]);