Prevent accidental revert when using firmware updater

This change prevents accidentally overwriting the previous firmware before
the new one has been marked as booted.

1 files changed, 9 insertions, 3 deletions

diff --git a/embassy-boot/boot/src/lib.rs b/embassy-boot/boot/src/lib.rs
index 45a87bd0e..016362b86 100644
--- a/embassy-boot/boot/src/lib.rs
+++ b/embassy-boot/boot/src/lib.rs
@@ -51,6 +51,8 @@ impl<const N: usize> AsMut<[u8]> for AlignedBuffer<N> {
 
 #[cfg(test)]
 mod tests {
+    #![allow(unused_imports)]
+
     use embedded_storage::nor_flash::{NorFlash, ReadNorFlash};
     #[cfg(feature = "nightly")]
     use embedded_storage_async::nor_flash::NorFlash as AsyncNorFlash;
@@ -120,9 +122,13 @@ mod tests {
             dfu: flash.dfu(),
             state: flash.state(),
         });
-        block_on(updater.write_firmware(0, &UPDATE)).unwrap();
+        block_on(updater.write_firmware(&mut aligned, 0, &UPDATE)).unwrap();
         block_on(updater.mark_updated(&mut aligned)).unwrap();
 
+        // Writing after marking updated is not allowed until marked as booted.
+        let res: Result<(), FirmwareUpdaterError> = block_on(updater.write_firmware(&mut aligned, 0, &UPDATE));
+        assert!(matches!(res, Err::<(), _>(FirmwareUpdaterError::BadState)));
+
         let flash = flash.into_blocking();
         let mut bootloader = BootLoader::new(BootLoaderConfig {
             active: flash.active(),
@@ -188,7 +194,7 @@ mod tests {
             dfu: flash.dfu(),
             state: flash.state(),
         });
-        block_on(updater.write_firmware(0, &UPDATE)).unwrap();
+        block_on(updater.write_firmware(&mut aligned, 0, &UPDATE)).unwrap();
         block_on(updater.mark_updated(&mut aligned)).unwrap();
 
         let flash = flash.into_blocking();
@@ -230,7 +236,7 @@ mod tests {
             dfu: flash.dfu(),
             state: flash.state(),
         });
-        block_on(updater.write_firmware(0, &UPDATE)).unwrap();
+        block_on(updater.write_firmware(&mut aligned, 0, &UPDATE)).unwrap();
         block_on(updater.mark_updated(&mut aligned)).unwrap();
 
         let flash = flash.into_blocking();