Avoid write to not-erased magic

This introduces an additional marker to the state partition right after the magic which indicates whether the current progress is valid or not. Validation in tests that we never write without an erase is added.

There is currently a FIXME in the FirmwareUpdater. Let me know if we should take the erase value as a parameter. I opened a feature request in embedded-storage to get this value in the trait. Before this, the assumption about ERASE_VALUE=0xFF was the same.

4 files changed, 63 insertions, 81 deletions

diff --git a/embassy-boot/boot/src/boot_loader.rs b/embassy-boot/boot/src/boot_loader.rs
index e2e361e3c..9d047f778 100644
--- a/embassy-boot/boot/src/boot_loader.rs
+++ b/embassy-boot/boot/src/boot_loader.rs
@@ -31,7 +31,7 @@ where
 }
 
 /// Extension of the embedded-storage flash type information with block size and erase value.
-pub trait Flash: NorFlash + ReadNorFlash {
+pub trait Flash: NorFlash {
     /// The block size that should be used when writing to flash. For most builtin flashes, this is the same as the erase
     /// size of the flash, but for external QSPI flash modules, this can be lower.
     const BLOCK_SIZE: usize;
@@ -60,9 +60,11 @@ pub trait FlashConfig {
 /// different page sizes and flash write sizes.
 pub struct BootLoader {
     // Page with current state of bootloader. The state partition has the following format:
-    // | Range          | Description                                                                      |
-    // | 0 - WRITE_SIZE | Magic indicating bootloader state. BOOT_MAGIC means boot, SWAP_MAGIC means swap. |
-    // | WRITE_SIZE - N | Progress index used while swapping or reverting                                  |
+    // All ranges are in multiples of WRITE_SIZE bytes.
+    // | Range    | Description                                                                      |
+    // | 0..1     | Magic indicating bootloader state. BOOT_MAGIC means boot, SWAP_MAGIC means swap. |
+    // | 1..2     | Progress validity. ERASE_VALUE means valid, !ERASE_VALUE means invalid.          |
+    // | 2..2 + N | Progress index used while swapping or reverting                                  |
     state: Partition,
     // Location of the partition which will be booted from
     active: Partition,
@@ -192,12 +194,17 @@ impl BootLoader {
                 trace!("Reverting");
                 self.revert(p, magic, page)?;
 
-                // Overwrite magic and reset progress
                 let state_flash = p.state();
+
+                // Invalidate progress
                 magic.fill(!P::STATE::ERASE_VALUE);
-                self.state.write_blocking(state_flash, 0, magic)?;
+                self.state
+                    .write_blocking(state_flash, P::STATE::WRITE_SIZE as u32, magic)?;
+
+                // Clear magic and progress
                 self.state.wipe_blocking(state_flash)?;
 
+                // Set magic
                 magic.fill(BOOT_MAGIC);
                 self.state.write_blocking(state_flash, 0, magic)?;
             }
@@ -215,28 +222,34 @@ impl BootLoader {
 
     fn current_progress<P: FlashConfig>(&mut self, config: &mut P, aligned: &mut [u8]) -> Result<usize, BootError> {
         let write_size = aligned.len();
-        let max_index = ((self.state.len() - write_size) / write_size) - 1;
+        let max_index = ((self.state.len() - write_size) / write_size) - 2;
         aligned.fill(!P::STATE::ERASE_VALUE);
 
         let state_flash = config.state();
-        for i in 0..max_index {
+
+        self.state
+            .read_blocking(state_flash, P::STATE::WRITE_SIZE as u32, aligned)?;
+        if aligned.iter().any(|&b| b != P::STATE::ERASE_VALUE) {
+            // Progress is invalid
+            return Ok(max_index);
+        }
+
+        for index in 0..max_index {
             self.state
-                .read_blocking(state_flash, (write_size + i * write_size) as u32, aligned)?;
+                .read_blocking(state_flash, (2 + index) as u32 * P::STATE::WRITE_SIZE as u32, aligned)?;
 
             if aligned.iter().any(|&b| b == P::STATE::ERASE_VALUE) {
-                return Ok(i);
+                return Ok(index);
             }
         }
         Ok(max_index)
     }
 
-    fn update_progress<P: FlashConfig>(&mut self, idx: usize, p: &mut P, magic: &mut [u8]) -> Result<(), BootError> {
-        let write_size = magic.len();
-
+    fn update_progress<P: FlashConfig>(&mut self, index: usize, p: &mut P, magic: &mut [u8]) -> Result<(), BootError> {
         let aligned = magic;
         aligned.fill(!P::STATE::ERASE_VALUE);
         self.state
-            .write_blocking(p.state(), (write_size + idx * write_size) as u32, aligned)?;
+            .write_blocking(p.state(), (2 + index) as u32 * P::STATE::WRITE_SIZE as u32, aligned)?;
         Ok(())
     }
 
@@ -360,7 +373,7 @@ fn assert_partitions(active: Partition, dfu: Partition, state: Partition, page_s
     assert_eq!(active.len() % page_size, 0);
     assert_eq!(dfu.len() % page_size, 0);
     assert!(dfu.len() - active.len() >= page_size);
-    assert!(2 * (active.len() / page_size) <= (state.len() - write_size) / write_size);
+    assert!(2 + 2 * (active.len() / page_size) <= state.len() / write_size);
 }
 
 /// A flash wrapper implementing the Flash and embedded_storage traits.
diff --git a/embassy-boot/boot/src/firmware_updater.rs b/embassy-boot/boot/src/firmware_updater.rs
index af1ba114a..fe3c0452f 100644
--- a/embassy-boot/boot/src/firmware_updater.rs
+++ b/embassy-boot/boot/src/firmware_updater.rs
@@ -220,11 +220,24 @@ impl FirmwareUpdater {
         self.state.read(state_flash, 0, aligned).await?;
 
         if aligned.iter().any(|&b| b != magic) {
-            aligned.fill(0);
+            // Read progress validity
+            self.state.read(state_flash, F::WRITE_SIZE as u32, aligned).await?;
+
+            // FIXME: Do not make this assumption.
+            const STATE_ERASE_VALUE: u8 = 0xFF;
+
+            if aligned.iter().any(|&b| b != STATE_ERASE_VALUE) {
+                // The current progress validity marker is invalid
+            } else {
+                // Invalidate progress
+                aligned.fill(!STATE_ERASE_VALUE);
+                self.state.write(state_flash, F::WRITE_SIZE as u32, aligned).await?;
+            }
 
-            self.state.write(state_flash, 0, aligned).await?;
+            // Clear magic and progress
             self.state.wipe(state_flash).await?;
 
+            // Set magic
             aligned.fill(magic);
             self.state.write(state_flash, 0, aligned).await?;
         }
@@ -420,11 +433,24 @@ impl FirmwareUpdater {
         self.state.read_blocking(state_flash, 0, aligned)?;
 
         if aligned.iter().any(|&b| b != magic) {
-            aligned.fill(0);
+            // Read progress validity
+            self.state.read_blocking(state_flash, F::WRITE_SIZE as u32, aligned)?;
+
+            // FIXME: Do not make this assumption.
+            const STATE_ERASE_VALUE: u8 = 0xFF;
+
+            if aligned.iter().any(|&b| b != STATE_ERASE_VALUE) {
+                // The current progress validity marker is invalid
+            } else {
+                // Invalidate progress
+                aligned.fill(!STATE_ERASE_VALUE);
+                self.state.write_blocking(state_flash, F::WRITE_SIZE as u32, aligned)?;
+            }
 
-            self.state.write_blocking(state_flash, 0, aligned)?;
+            // Clear magic and progress
             self.state.wipe_blocking(state_flash)?;
 
+            // Set magic
             aligned.fill(magic);
             self.state.write_blocking(state_flash, 0, aligned)?;
         }
diff --git a/embassy-boot/boot/src/lib.rs b/embassy-boot/boot/src/lib.rs
index a5795781f..597ce3fa3 100644
--- a/embassy-boot/boot/src/lib.rs
+++ b/embassy-boot/boot/src/lib.rs
@@ -93,7 +93,7 @@ mod tests {
         const STATE: Partition = Partition::new(0, 4096);
         const ACTIVE: Partition = Partition::new(4096, 61440);
         const DFU: Partition = Partition::new(61440, 122880);
-        let mut flash = MemFlash::<131072, 4096, 4>::random().with_limited_erase_before_write_verification(4..);
+        let mut flash = MemFlash::<131072, 4096, 4>::random();
 
         let original: [u8; ACTIVE.len()] = [rand::random::<u8>(); ACTIVE.len()];
         let update: [u8; DFU.len()] = [rand::random::<u8>(); DFU.len()];
@@ -166,7 +166,7 @@ mod tests {
 
         let mut active = MemFlash::<16384, 4096, 8>::random();
         let mut dfu = MemFlash::<16384, 2048, 8>::random();
-        let mut state = MemFlash::<4096, 128, 4>::random().with_limited_erase_before_write_verification(2048 + 4..);
+        let mut state = MemFlash::<4096, 128, 4>::random();
         let mut aligned = [0; 4];
 
         let original: [u8; ACTIVE.len()] = [rand::random::<u8>(); ACTIVE.len()];
@@ -220,7 +220,7 @@ mod tests {
         let mut aligned = [0; 4];
         let mut active = MemFlash::<16384, 2048, 4>::random();
         let mut dfu = MemFlash::<16384, 4096, 8>::random();
-        let mut state = MemFlash::<4096, 128, 4>::random().with_limited_erase_before_write_verification(2048 + 4..);
+        let mut state = MemFlash::<4096, 128, 4>::random();
 
         let original: [u8; ACTIVE.len()] = [rand::random::<u8>(); ACTIVE.len()];
         let update: [u8; DFU.len()] = [rand::random::<u8>(); DFU.len()];
diff --git a/embassy-boot/boot/src/mem_flash.rs b/embassy-boot/boot/src/mem_flash.rs
index e87ccd37a..828aad9d9 100644
--- a/embassy-boot/boot/src/mem_flash.rs
+++ b/embassy-boot/boot/src/mem_flash.rs
@@ -9,8 +9,6 @@ use crate::Flash;
 
 pub struct MemFlash<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> {
     pub mem: [u8; SIZE],
-    pub allow_same_write: bool,
-    pub verify_erased_before_write: Range<usize>,
     pub pending_write_successes: Option<usize>,
 }
 
@@ -21,8 +19,6 @@ impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> MemFla
     pub const fn new(fill: u8) -> Self {
         Self {
             mem: [fill; SIZE],
-            allow_same_write: false,
-            verify_erased_before_write: 0..SIZE,
             pending_write_successes: None,
         }
     }
@@ -35,37 +31,9 @@ impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> MemFla
         }
         Self {
             mem,
-            allow_same_write: false,
-            verify_erased_before_write: 0..SIZE,
             pending_write_successes: None,
         }
     }
-
-    #[must_use]
-    pub fn allow_same_write(self, allow: bool) -> Self {
-        Self {
-            allow_same_write: allow,
-            ..self
-        }
-    }
-
-    #[must_use]
-    pub fn with_limited_erase_before_write_verification<R: RangeBounds<usize>>(self, verified_range: R) -> Self {
-        let start = match verified_range.start_bound() {
-            Bound::Included(start) => *start,
-            Bound::Excluded(start) => *start + 1,
-            Bound::Unbounded => 0,
-        };
-        let end = match verified_range.end_bound() {
-            Bound::Included(end) => *end - 1,
-            Bound::Excluded(end) => *end,
-            Bound::Unbounded => self.mem.len(),
-        };
-        Self {
-            verify_erased_before_write: start..end,
-            ..self
-        }
-    }
 }
 
 impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> Default
@@ -150,14 +118,8 @@ impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> NorFla
             .take(bytes.len())
             .zip(bytes)
         {
-            if self.allow_same_write && mem_byte == new_byte {
-                // Write does not change the flash memory which is allowed
-            } else {
-                if self.verify_erased_before_write.contains(&offset) {
-                    assert_eq!(0xFF, *mem_byte, "Offset {} is not erased", offset);
-                }
-                *mem_byte &= *new_byte;
-            }
+            assert_eq!(0xFF, *mem_byte, "Offset {} is not erased", offset);
+            *mem_byte = *new_byte;
         }
 
         Ok(())
@@ -192,22 +154,3 @@ impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> AsyncN
         <Self as NorFlash>::write(self, offset, bytes)
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use core::ops::Range;
-
-    use embedded_storage::nor_flash::NorFlash;
-
-    use super::MemFlash;
-
-    #[test]
-    fn writes_only_flip_bits_from_1_to_0() {
-        let mut flash = MemFlash::<16, 16, 1>::default().with_limited_erase_before_write_verification(0..0);
-
-        flash.write(0, &[0x55]).unwrap();
-        flash.write(0, &[0xAA]).unwrap();
-
-        assert_eq!(0x00, flash.mem[0]);
-    }
-}