Move MemFlash to separate module and add verify_erased_before_write verification

3 files changed, 244 insertions, 142 deletions

diff --git a/embassy-boot/boot/src/lib.rs b/embassy-boot/boot/src/lib.rs
index 4c28d7aa4..a5795781f 100644
--- a/embassy-boot/boot/src/lib.rs
+++ b/embassy-boot/boot/src/lib.rs
@@ -8,6 +8,7 @@ mod fmt;
 mod boot_loader;
 mod firmware_updater;
 mod firmware_writer;
+mod mem_flash;
 mod partition;
 
 pub use boot_loader::{BootError, BootFlash, BootLoader, Flash, FlashConfig, MultiFlashConfig, SingleFlashConfig};
@@ -46,13 +47,10 @@ impl<const N: usize> AsMut<[u8]> for AlignedBuffer<N> {
 
 #[cfg(test)]
 mod tests {
-    use core::convert::Infallible;
-
-    use embedded_storage::nor_flash::{ErrorType, NorFlash, ReadNorFlash};
-    use embedded_storage_async::nor_flash::{NorFlash as AsyncNorFlash, ReadNorFlash as AsyncReadNorFlash};
     use futures::executor::block_on;
 
     use super::*;
+    use crate::mem_flash::MemFlash;
 
     /*
     #[test]
@@ -75,8 +73,8 @@ mod tests {
         const ACTIVE: Partition = Partition::new(4096, 61440);
         const DFU: Partition = Partition::new(61440, 122880);
 
-        let mut flash = MemFlash::<131072, 4096, 4>([0xff; 131072]);
-        flash.0[0..4].copy_from_slice(&[BOOT_MAGIC; 4]);
+        let mut flash = MemFlash::<131072, 4096, 4>::default();
+        flash.mem[0..4].copy_from_slice(&[BOOT_MAGIC; 4]);
         let mut flash = SingleFlashConfig::new(&mut flash);
 
         let mut bootloader: BootLoader = BootLoader::new(ACTIVE, DFU, STATE);
@@ -95,14 +93,14 @@ 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>([0xff; 131072]);
+        let mut flash = MemFlash::<131072, 4096, 4>::random().with_limited_erase_before_write_verification(4..);
 
         let original: [u8; ACTIVE.len()] = [rand::random::<u8>(); ACTIVE.len()];
         let update: [u8; DFU.len()] = [rand::random::<u8>(); DFU.len()];
         let mut aligned = [0; 4];
 
         for i in ACTIVE.from..ACTIVE.to {
-            flash.0[i] = original[i - ACTIVE.from];
+            flash.mem[i] = original[i - ACTIVE.from];
         }
 
         let mut bootloader: BootLoader = BootLoader::new(ACTIVE, DFU, STATE);
@@ -124,12 +122,12 @@ mod tests {
         );
 
         for i in ACTIVE.from..ACTIVE.to {
-            assert_eq!(flash.0[i], update[i - ACTIVE.from], "Index {}", i);
+            assert_eq!(flash.mem[i], update[i - ACTIVE.from], "Index {}", i);
         }
 
         // First DFU page is untouched
         for i in DFU.from + 4096..DFU.to {
-            assert_eq!(flash.0[i], original[i - DFU.from - 4096], "Index {}", i);
+            assert_eq!(flash.mem[i], original[i - DFU.from - 4096], "Index {}", i);
         }
 
         // Running again should cause a revert
@@ -141,12 +139,12 @@ mod tests {
         );
 
         for i in ACTIVE.from..ACTIVE.to {
-            assert_eq!(flash.0[i], original[i - ACTIVE.from], "Index {}", i);
+            assert_eq!(flash.mem[i], original[i - ACTIVE.from], "Index {}", i);
         }
 
         // Last page is untouched
         for i in DFU.from..DFU.to - 4096 {
-            assert_eq!(flash.0[i], update[i - DFU.from], "Index {}", i);
+            assert_eq!(flash.mem[i], update[i - DFU.from], "Index {}", i);
         }
 
         // Mark as booted
@@ -166,16 +164,16 @@ mod tests {
         const ACTIVE: Partition = Partition::new(4096, 16384);
         const DFU: Partition = Partition::new(0, 16384);
 
-        let mut active = MemFlash::<16384, 4096, 8>([0xff; 16384]);
-        let mut dfu = MemFlash::<16384, 2048, 8>([0xff; 16384]);
-        let mut state = MemFlash::<4096, 128, 4>([0xff; 4096]);
+        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 aligned = [0; 4];
 
         let original: [u8; ACTIVE.len()] = [rand::random::<u8>(); ACTIVE.len()];
         let update: [u8; DFU.len()] = [rand::random::<u8>(); DFU.len()];
 
         for i in ACTIVE.from..ACTIVE.to {
-            active.0[i] = original[i - ACTIVE.from];
+            active.mem[i] = original[i - ACTIVE.from];
         }
 
         let mut updater = FirmwareUpdater::new(DFU, STATE);
@@ -203,12 +201,12 @@ mod tests {
         );
 
         for i in ACTIVE.from..ACTIVE.to {
-            assert_eq!(active.0[i], update[i - ACTIVE.from], "Index {}", i);
+            assert_eq!(active.mem[i], update[i - ACTIVE.from], "Index {}", i);
         }
 
         // First DFU page is untouched
         for i in DFU.from + 4096..DFU.to {
-            assert_eq!(dfu.0[i], original[i - DFU.from - 4096], "Index {}", i);
+            assert_eq!(dfu.mem[i], original[i - DFU.from - 4096], "Index {}", i);
         }
     }
 
@@ -220,15 +218,15 @@ mod tests {
         const DFU: Partition = Partition::new(0, 16384);
 
         let mut aligned = [0; 4];
-        let mut active = MemFlash::<16384, 2048, 4>([0xff; 16384]);
-        let mut dfu = MemFlash::<16384, 4096, 8>([0xff; 16384]);
-        let mut state = MemFlash::<4096, 128, 4>([0xff; 4096]);
+        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 original: [u8; ACTIVE.len()] = [rand::random::<u8>(); ACTIVE.len()];
         let update: [u8; DFU.len()] = [rand::random::<u8>(); DFU.len()];
 
         for i in ACTIVE.from..ACTIVE.to {
-            active.0[i] = original[i - ACTIVE.from];
+            active.mem[i] = original[i - ACTIVE.from];
         }
 
         let mut updater = FirmwareUpdater::new(DFU, STATE);
@@ -255,12 +253,12 @@ mod tests {
         );
 
         for i in ACTIVE.from..ACTIVE.to {
-            assert_eq!(active.0[i], update[i - ACTIVE.from], "Index {}", i);
+            assert_eq!(active.mem[i], update[i - ACTIVE.from], "Index {}", i);
         }
 
         // First DFU page is untouched
         for i in DFU.from + 4096..DFU.to {
-            assert_eq!(dfu.0[i], original[i - DFU.from - 4096], "Index {}", i);
+            assert_eq!(dfu.mem[i], original[i - DFU.from - 4096], "Index {}", i);
         }
     }
 
@@ -313,113 +311,4 @@ mod tests {
         ))
         .is_ok());
     }
-
-    pub struct MemFlash<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize>(pub [u8; SIZE]);
-
-    impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> NorFlash
-        for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
-    {
-        const WRITE_SIZE: usize = WRITE_SIZE;
-        const ERASE_SIZE: usize = ERASE_SIZE;
-        fn erase(&mut self, from: u32, to: u32) -> Result<(), Self::Error> {
-            let from = from as usize;
-            let to = to as usize;
-            assert!(from % ERASE_SIZE == 0);
-            assert!(to % ERASE_SIZE == 0, "To: {}, erase size: {}", to, ERASE_SIZE);
-            for i in from..to {
-                self.0[i] = 0xFF;
-            }
-            Ok(())
-        }
-
-        fn write(&mut self, offset: u32, data: &[u8]) -> Result<(), Self::Error> {
-            assert!(data.len() % WRITE_SIZE == 0);
-            assert!(offset as usize % WRITE_SIZE == 0);
-            assert!(offset as usize + data.len() <= SIZE);
-
-            self.0[offset as usize..offset as usize + data.len()].copy_from_slice(data);
-
-            Ok(())
-        }
-    }
-
-    impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> ErrorType
-        for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
-    {
-        type Error = Infallible;
-    }
-
-    impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> ReadNorFlash
-        for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
-    {
-        const READ_SIZE: usize = 1;
-
-        fn read(&mut self, offset: u32, buf: &mut [u8]) -> Result<(), Self::Error> {
-            let len = buf.len();
-            buf[..].copy_from_slice(&self.0[offset as usize..offset as usize + len]);
-            Ok(())
-        }
-
-        fn capacity(&self) -> usize {
-            SIZE
-        }
-    }
-
-    impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> super::Flash
-        for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
-    {
-        const BLOCK_SIZE: usize = ERASE_SIZE;
-        const ERASE_VALUE: u8 = 0xFF;
-    }
-
-    impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> AsyncReadNorFlash
-        for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
-    {
-        const READ_SIZE: usize = 1;
-
-        async fn read(&mut self, offset: u32, buf: &mut [u8]) -> Result<(), Self::Error> {
-            let len = buf.len();
-            buf[..].copy_from_slice(&self.0[offset as usize..offset as usize + len]);
-            Ok(())
-        }
-
-        fn capacity(&self) -> usize {
-            SIZE
-        }
-    }
-
-    impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> AsyncNorFlash
-        for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
-    {
-        const WRITE_SIZE: usize = WRITE_SIZE;
-        const ERASE_SIZE: usize = ERASE_SIZE;
-
-        async fn erase(&mut self, from: u32, to: u32) -> Result<(), Self::Error> {
-            let from = from as usize;
-            let to = to as usize;
-            assert!(from % ERASE_SIZE == 0);
-            assert!(to % ERASE_SIZE == 0);
-            for i in from..to {
-                self.0[i] = 0xFF;
-            }
-            Ok(())
-        }
-
-        async fn write(&mut self, offset: u32, data: &[u8]) -> Result<(), Self::Error> {
-            info!("Writing {} bytes to 0x{:x}", data.len(), offset);
-            assert!(data.len() % WRITE_SIZE == 0);
-            assert!(offset as usize % WRITE_SIZE == 0);
-            assert!(
-                offset as usize + data.len() <= SIZE,
-                "OFFSET: {}, LEN: {}, FLASH SIZE: {}",
-                offset,
-                data.len(),
-                SIZE
-            );
-
-            self.0[offset as usize..offset as usize + data.len()].copy_from_slice(data);
-
-            Ok(())
-        }
-    }
 }
diff --git a/embassy-boot/boot/src/mem_flash.rs b/embassy-boot/boot/src/mem_flash.rs
new file mode 100644
index 000000000..e87ccd37a
--- /dev/null
+++ b/embassy-boot/boot/src/mem_flash.rs
@@ -0,0 +1,213 @@
+#![allow(unused)]
+
+use core::ops::{Bound, Range, RangeBounds};
+
+use embedded_storage::nor_flash::{ErrorType, NorFlash, NorFlashError, NorFlashErrorKind, ReadNorFlash};
+use embedded_storage_async::nor_flash::{NorFlash as AsyncNorFlash, ReadNorFlash as AsyncReadNorFlash};
+
+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>,
+}
+
+#[derive(Debug)]
+pub struct MemFlashError;
+
+impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE> {
+    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,
+        }
+    }
+
+    #[cfg(test)]
+    pub fn random() -> Self {
+        let mut mem = [0; SIZE];
+        for byte in mem.iter_mut() {
+            *byte = rand::random::<u8>();
+        }
+        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
+    for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
+{
+    fn default() -> Self {
+        Self::new(0xFF)
+    }
+}
+
+impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> Flash
+    for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
+{
+    const BLOCK_SIZE: usize = ERASE_SIZE;
+    const ERASE_VALUE: u8 = 0xFF;
+}
+
+impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> ErrorType
+    for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
+{
+    type Error = MemFlashError;
+}
+
+impl NorFlashError for MemFlashError {
+    fn kind(&self) -> NorFlashErrorKind {
+        NorFlashErrorKind::Other
+    }
+}
+
+impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> ReadNorFlash
+    for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
+{
+    const READ_SIZE: usize = 1;
+
+    fn read(&mut self, offset: u32, bytes: &mut [u8]) -> Result<(), Self::Error> {
+        let len = bytes.len();
+        bytes.copy_from_slice(&self.mem[offset as usize..offset as usize + len]);
+        Ok(())
+    }
+
+    fn capacity(&self) -> usize {
+        SIZE
+    }
+}
+
+impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> NorFlash
+    for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
+{
+    const WRITE_SIZE: usize = WRITE_SIZE;
+    const ERASE_SIZE: usize = ERASE_SIZE;
+
+    fn erase(&mut self, from: u32, to: u32) -> Result<(), Self::Error> {
+        let from = from as usize;
+        let to = to as usize;
+        assert!(from % ERASE_SIZE == 0);
+        assert!(to % ERASE_SIZE == 0, "To: {}, erase size: {}", to, ERASE_SIZE);
+        for i in from..to {
+            self.mem[i] = 0xFF;
+        }
+        Ok(())
+    }
+
+    fn write(&mut self, offset: u32, bytes: &[u8]) -> Result<(), Self::Error> {
+        let offset = offset as usize;
+        assert!(bytes.len() % WRITE_SIZE == 0);
+        assert!(offset % WRITE_SIZE == 0);
+        assert!(offset + bytes.len() <= SIZE);
+
+        if let Some(pending_successes) = self.pending_write_successes {
+            if pending_successes > 0 {
+                self.pending_write_successes = Some(pending_successes - 1);
+            } else {
+                return Err(MemFlashError);
+            }
+        }
+
+        for ((offset, mem_byte), new_byte) in self
+            .mem
+            .iter_mut()
+            .enumerate()
+            .skip(offset)
+            .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;
+            }
+        }
+
+        Ok(())
+    }
+}
+
+impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> AsyncReadNorFlash
+    for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
+{
+    const READ_SIZE: usize = 1;
+
+    async fn read(&mut self, offset: u32, bytes: &mut [u8]) -> Result<(), Self::Error> {
+        <Self as ReadNorFlash>::read(self, offset, bytes)
+    }
+
+    fn capacity(&self) -> usize {
+        <Self as ReadNorFlash>::capacity(self)
+    }
+}
+
+impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> AsyncNorFlash
+    for MemFlash<SIZE, ERASE_SIZE, WRITE_SIZE>
+{
+    const WRITE_SIZE: usize = WRITE_SIZE;
+    const ERASE_SIZE: usize = ERASE_SIZE;
+
+    async fn erase(&mut self, from: u32, to: u32) -> Result<(), Self::Error> {
+        <Self as NorFlash>::erase(self, from, to)
+    }
+
+    async fn write(&mut self, offset: u32, bytes: &[u8]) -> Result<(), Self::Error> {
+        <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]);
+    }
+}
diff --git a/embassy-boot/boot/src/partition.rs b/embassy-boot/boot/src/partition.rs
index 3ccd4dd76..1157e8cd8 100644
--- a/embassy-boot/boot/src/partition.rs
+++ b/embassy-boot/boot/src/partition.rs
@@ -105,45 +105,45 @@ impl Partition {
 
 #[cfg(test)]
 mod tests {
-    use crate::tests::MemFlash;
+    use crate::mem_flash::MemFlash;
     use crate::Partition;
 
     #[test]
     fn can_erase() {
-        let mut flash = MemFlash::<1024, 64, 4>([0x00; 1024]);
+        let mut flash = MemFlash::<1024, 64, 4>::new(0x00);
         let partition = Partition::new(256, 512);
 
         partition.erase_blocking(&mut flash, 64, 192).unwrap();
 
-        for (index, byte) in flash.0.iter().copied().enumerate().take(256 + 64) {
+        for (index, byte) in flash.mem.iter().copied().enumerate().take(256 + 64) {
             assert_eq!(0x00, byte, "Index {}", index);
         }
 
-        for (index, byte) in flash.0.iter().copied().enumerate().skip(256 + 64).take(128) {
+        for (index, byte) in flash.mem.iter().copied().enumerate().skip(256 + 64).take(128) {
             assert_eq!(0xFF, byte, "Index {}", index);
         }
 
-        for (index, byte) in flash.0.iter().copied().enumerate().skip(256 + 64 + 128) {
+        for (index, byte) in flash.mem.iter().copied().enumerate().skip(256 + 64 + 128) {
             assert_eq!(0x00, byte, "Index {}", index);
         }
     }
 
     #[test]
     fn can_wipe() {
-        let mut flash = MemFlash::<1024, 64, 4>([0x00; 1024]);
+        let mut flash = MemFlash::<1024, 64, 4>::new(0x00);
         let partition = Partition::new(256, 512);
 
         partition.wipe_blocking(&mut flash).unwrap();
 
-        for (index, byte) in flash.0.iter().copied().enumerate().take(256) {
+        for (index, byte) in flash.mem.iter().copied().enumerate().take(256) {
             assert_eq!(0x00, byte, "Index {}", index);
         }
 
-        for (index, byte) in flash.0.iter().copied().enumerate().skip(256).take(256) {
+        for (index, byte) in flash.mem.iter().copied().enumerate().skip(256).take(256) {
             assert_eq!(0xFF, byte, "Index {}", index);
         }
 
-        for (index, byte) in flash.0.iter().copied().enumerate().skip(512) {
+        for (index, byte) in flash.mem.iter().copied().enumerate().skip(512) {
             assert_eq!(0x00, byte, "Index {}", index);
         }
     }