10 files changed, 618 insertions, 165 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 dd22b2fa7..2b5cc72fa 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?;
         }
@@ -417,11 +430,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 428e7ca2b..d53c613a3 100644
--- a/embassy-boot/boot/src/lib.rs
+++ b/embassy-boot/boot/src/lib.rs
@@ -7,6 +7,7 @@ mod fmt;
 
 mod boot_loader;
 mod firmware_updater;
+mod mem_flash;
 mod partition;
 
 pub use boot_loader::{BootError, BootFlash, BootLoader, Flash, FlashConfig, MultiFlashConfig, SingleFlashConfig};
@@ -44,13 +45,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]
@@ -73,8 +71,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);
@@ -93,14 +91,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();
 
         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);
@@ -122,12 +120,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
@@ -139,12 +137,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
@@ -164,16 +162,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();
         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);
@@ -201,12 +199,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);
         }
     }
 
@@ -218,15 +216,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();
 
         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);
@@ -253,12 +251,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);
         }
     }
 
@@ -288,13 +286,13 @@ mod tests {
 
         const STATE: Partition = Partition::new(0, 4096);
         const DFU: Partition = Partition::new(4096, 8192);
-        let mut flash = MemFlash::<8192, 4096, 4>([0xff; 8192]);
+        let mut flash = MemFlash::<8192, 4096, 4>::default();
 
         let firmware_len = firmware.len();
 
         let mut write_buf = [0; 4096];
         write_buf[0..firmware_len].copy_from_slice(firmware);
-        NorFlash::write(&mut flash, DFU.from as u32, &write_buf).unwrap();
+        DFU.write_blocking(&mut flash, 0, &write_buf).unwrap();
 
         // On with the test
 
@@ -311,113 +309,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..828aad9d9
--- /dev/null
+++ b/embassy-boot/boot/src/mem_flash.rs
@@ -0,0 +1,156 @@
+#![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 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],
+            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,
+            pending_write_successes: None,
+        }
+    }
+}
+
+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)
+        {
+            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)
+    }
+}
diff --git a/embassy-boot/boot/src/partition.rs b/embassy-boot/boot/src/partition.rs
index 217c457fc..ac6b0ed0f 100644
--- a/embassy-boot/boot/src/partition.rs
+++ b/embassy-boot/boot/src/partition.rs
@@ -95,45 +95,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);
         }
     }
diff --git a/embassy-stm32/build.rs b/embassy-stm32/build.rs
index 3780c5a40..481fec677 100644
--- a/embassy-stm32/build.rs
+++ b/embassy-stm32/build.rs
@@ -50,10 +50,13 @@ fn main() {
                 // We *shouldn't* have singletons for these, but the HAL currently requires
                 // singletons, for using with RccPeripheral to enable/disable clocks to them.
                 "rcc" => {
-                    if r.version.starts_with("h7") {
+                    if r.version.starts_with("h7") || r.version.starts_with("f4") {
                         singletons.push("MCO1".to_string());
                         singletons.push("MCO2".to_string());
                     }
+                    if r.version.starts_with("l4") {
+                        singletons.push("MCO".to_string());
+                    }
                     singletons.push(p.name.to_string());
                 }
                 //"dbgmcu" => {}
@@ -258,6 +261,7 @@ fn main() {
         (("i2c", "SCL"), quote!(crate::i2c::SclPin)),
         (("rcc", "MCO_1"), quote!(crate::rcc::McoPin)),
         (("rcc", "MCO_2"), quote!(crate::rcc::McoPin)),
+        (("rcc", "MCO"), quote!(crate::rcc::McoPin)),
         (("dcmi", "D0"), quote!(crate::dcmi::D0Pin)),
         (("dcmi", "D1"), quote!(crate::dcmi::D1Pin)),
         (("dcmi", "D2"), quote!(crate::dcmi::D2Pin)),
@@ -447,13 +451,22 @@ fn main() {
                     // MCO is special
                     if pin.signal.starts_with("MCO_") {
                         // Supported in H7 only for now
-                        if regs.version.starts_with("h7") {
+                        if regs.version.starts_with("h7") || regs.version.starts_with("f4") {
                             peri = format_ident!("{}", pin.signal.replace("_", ""));
                         } else {
                             continue;
                         }
                     }
 
+                    if pin.signal == "MCO" {
+                        // Supported in H7 only for now
+                        if regs.version.starts_with("l4") {
+                            peri = format_ident!("MCO");
+                        } else {
+                            continue;
+                        }
+                    }
+
                     g.extend(quote! {
                         pin_trait_impl!(#tr, #peri, #pin_name, #af);
                     })
diff --git a/embassy-stm32/src/rcc/f4.rs b/embassy-stm32/src/rcc/f4.rs
index 200bcce9c..2a17eb9b0 100644
--- a/embassy-stm32/src/rcc/f4.rs
+++ b/embassy-stm32/src/rcc/f4.rs
@@ -1,8 +1,16 @@
+use core::marker::PhantomData;
+
+use embassy_hal_common::into_ref;
+use stm32_metapac::rcc::vals::{Mco1, Mco2, Mcopre};
+
 use super::sealed::RccPeripheral;
+use crate::gpio::sealed::AFType;
+use crate::gpio::Speed;
 use crate::pac::rcc::vals::{Hpre, Ppre, Sw};
 use crate::pac::{FLASH, PWR, RCC};
 use crate::rcc::{set_freqs, Clocks};
 use crate::time::Hertz;
+use crate::{peripherals, Peripheral};
 
 /// HSI speed
 pub const HSI_FREQ: Hertz = Hertz(16_000_000);
@@ -96,6 +104,164 @@ unsafe fn setup_pll(pllsrcclk: u32, use_hse: bool, pllsysclk: Option<u32>, pll48
     }
 }
 
+pub enum McoClock {
+    DIV1,
+    DIV2,
+    DIV3,
+    DIV4,
+    DIV5,
+}
+
+impl McoClock {
+    fn into_raw(&self) -> Mcopre {
+        match self {
+            McoClock::DIV1 => Mcopre::DIV1,
+            McoClock::DIV2 => Mcopre::DIV2,
+            McoClock::DIV3 => Mcopre::DIV3,
+            McoClock::DIV4 => Mcopre::DIV4,
+            McoClock::DIV5 => Mcopre::DIV5,
+        }
+    }
+}
+
+#[derive(Copy, Clone)]
+pub enum Mco1Source {
+    Hsi,
+    Lse,
+    Hse,
+    Pll,
+}
+
+impl Default for Mco1Source {
+    fn default() -> Self {
+        Self::Hsi
+    }
+}
+
+pub trait McoSource {
+    type Raw;
+
+    fn into_raw(&self) -> Self::Raw;
+}
+
+impl McoSource for Mco1Source {
+    type Raw = Mco1;
+    fn into_raw(&self) -> Self::Raw {
+        match self {
+            Mco1Source::Hsi => Mco1::HSI,
+            Mco1Source::Lse => Mco1::LSE,
+            Mco1Source::Hse => Mco1::HSE,
+            Mco1Source::Pll => Mco1::PLL,
+        }
+    }
+}
+
+#[derive(Copy, Clone)]
+pub enum Mco2Source {
+    SysClk,
+    Plli2s,
+    Hse,
+    Pll,
+}
+
+impl Default for Mco2Source {
+    fn default() -> Self {
+        Self::SysClk
+    }
+}
+
+impl McoSource for Mco2Source {
+    type Raw = Mco2;
+    fn into_raw(&self) -> Self::Raw {
+        match self {
+            Mco2Source::SysClk => Mco2::SYSCLK,
+            Mco2Source::Plli2s => Mco2::PLLI2S,
+            Mco2Source::Hse => Mco2::HSE,
+            Mco2Source::Pll => Mco2::PLL,
+        }
+    }
+}
+
+pub(crate) mod sealed {
+    use stm32_metapac::rcc::vals::Mcopre;
+    pub trait McoInstance {
+        type Source;
+        unsafe fn apply_clock_settings(source: Self::Source, prescaler: Mcopre);
+    }
+}
+
+pub trait McoInstance: sealed::McoInstance + 'static {}
+
+pin_trait!(McoPin, McoInstance);
+
+impl sealed::McoInstance for peripherals::MCO1 {
+    type Source = Mco1;
+    unsafe fn apply_clock_settings(source: Self::Source, prescaler: Mcopre) {
+        RCC.cfgr().modify(|w| {
+            w.set_mco1(source);
+            w.set_mco1pre(prescaler);
+        });
+        match source {
+            Mco1::PLL => {
+                RCC.cr().modify(|w| w.set_pllon(true));
+                while !RCC.cr().read().pllrdy() {}
+            }
+            Mco1::HSI => {
+                RCC.cr().modify(|w| w.set_hsion(true));
+                while !RCC.cr().read().hsirdy() {}
+            }
+            _ => {}
+        }
+    }
+}
+impl McoInstance for peripherals::MCO1 {}
+
+impl sealed::McoInstance for peripherals::MCO2 {
+    type Source = Mco2;
+    unsafe fn apply_clock_settings(source: Self::Source, prescaler: Mcopre) {
+        RCC.cfgr().modify(|w| {
+            w.set_mco2(source);
+            w.set_mco2pre(prescaler);
+        });
+        match source {
+            Mco2::PLL => {
+                RCC.cr().modify(|w| w.set_pllon(true));
+                while !RCC.cr().read().pllrdy() {}
+            }
+            #[cfg(not(stm32f410))]
+            Mco2::PLLI2S => {
+                RCC.cr().modify(|w| w.set_plli2son(true));
+                while !RCC.cr().read().plli2srdy() {}
+            }
+            _ => {}
+        }
+    }
+}
+impl McoInstance for peripherals::MCO2 {}
+
+pub struct Mco<'d, T: McoInstance> {
+    phantom: PhantomData<&'d mut T>,
+}
+
+impl<'d, T: McoInstance> Mco<'d, T> {
+    pub fn new(
+        _peri: impl Peripheral<P = T> + 'd,
+        pin: impl Peripheral<P = impl McoPin<T>> + 'd,
+        source: impl McoSource<Raw = T::Source>,
+        prescaler: McoClock,
+    ) -> Self {
+        into_ref!(pin);
+
+        critical_section::with(|_| unsafe {
+            T::apply_clock_settings(source.into_raw(), prescaler.into_raw());
+            pin.set_as_af(pin.af_num(), AFType::OutputPushPull);
+            pin.set_speed(Speed::VeryHigh);
+        });
+
+        Self { phantom: PhantomData }
+    }
+}
+
 unsafe fn flash_setup(sysclk: u32) {
     use crate::pac::flash::vals::Latency;
 
diff --git a/embassy-stm32/src/rcc/l4.rs b/embassy-stm32/src/rcc/l4.rs
index e650490fe..c1bf7d0cd 100644
--- a/embassy-stm32/src/rcc/l4.rs
+++ b/embassy-stm32/src/rcc/l4.rs
@@ -1,7 +1,15 @@
+use core::marker::PhantomData;
+
+use embassy_hal_common::into_ref;
+use stm32_metapac::rcc::vals::{Mcopre, Mcosel};
+
+use crate::gpio::sealed::AFType;
+use crate::gpio::Speed;
 use crate::pac::rcc::vals::{Hpre, Msirange, Pllsrc, Ppre, Sw};
 use crate::pac::{FLASH, RCC};
 use crate::rcc::{set_freqs, Clocks};
 use crate::time::Hertz;
+use crate::{peripherals, Peripheral};
 
 /// HSI speed
 pub const HSI_FREQ: Hertz = Hertz(16_000_000);
@@ -298,6 +306,131 @@ impl Default for Config {
     }
 }
 
+pub enum McoClock {
+    DIV1,
+    DIV2,
+    DIV4,
+    DIV8,
+    DIV16,
+}
+
+impl McoClock {
+    fn into_raw(&self) -> Mcopre {
+        match self {
+            McoClock::DIV1 => Mcopre::DIV1,
+            McoClock::DIV2 => Mcopre::DIV2,
+            McoClock::DIV4 => Mcopre::DIV4,
+            McoClock::DIV8 => Mcopre::DIV8,
+            McoClock::DIV16 => Mcopre::DIV16,
+        }
+    }
+}
+
+#[derive(Copy, Clone)]
+pub enum Mco1Source {
+    Disabled,
+    Lse,
+    Lsi,
+    Hse,
+    Hsi16,
+    PllClk,
+    SysClk,
+    Msi,
+    #[cfg(not(any(stm32l471, stm32l475, stm32l476, stm32l486)))]
+    Hsi48,
+}
+
+impl Default for Mco1Source {
+    fn default() -> Self {
+        Self::Hsi16
+    }
+}
+
+pub trait McoSource {
+    type Raw;
+
+    fn into_raw(&self) -> Self::Raw;
+}
+
+impl McoSource for Mco1Source {
+    type Raw = Mcosel;
+    fn into_raw(&self) -> Self::Raw {
+        match self {
+            Mco1Source::Disabled => Mcosel::NOCLOCK,
+            Mco1Source::Lse => Mcosel::LSE,
+            Mco1Source::Lsi => Mcosel::LSI,
+            Mco1Source::Hse => Mcosel::HSE,
+            Mco1Source::Hsi16 => Mcosel::HSI16,
+            Mco1Source::PllClk => Mcosel::PLL,
+            Mco1Source::SysClk => Mcosel::SYSCLK,
+            Mco1Source::Msi => Mcosel::MSI,
+            #[cfg(not(any(stm32l471, stm32l475, stm32l476, stm32l486)))]
+            Mco1Source::Hsi48 => Mcosel::HSI48,
+        }
+    }
+}
+
+pub(crate) mod sealed {
+    use stm32_metapac::rcc::vals::Mcopre;
+    pub trait McoInstance {
+        type Source;
+        unsafe fn apply_clock_settings(source: Self::Source, prescaler: Mcopre);
+    }
+}
+
+pub trait McoInstance: sealed::McoInstance + 'static {}
+
+pin_trait!(McoPin, McoInstance);
+
+impl sealed::McoInstance for peripherals::MCO {
+    type Source = Mcosel;
+
+    unsafe fn apply_clock_settings(source: Self::Source, prescaler: Mcopre) {
+        RCC.cfgr().modify(|w| {
+            w.set_mcosel(source);
+            w.set_mcopre(prescaler);
+        });
+
+        match source {
+            Mcosel::HSI16 => {
+                RCC.cr().modify(|w| w.set_hsion(true));
+                while !RCC.cr().read().hsirdy() {}
+            }
+            #[cfg(not(any(stm32l471, stm32l475, stm32l476, stm32l486)))]
+            Mcosel::HSI48 => {
+                RCC.crrcr().modify(|w| w.set_hsi48on(true));
+                while !RCC.crrcr().read().hsi48rdy() {}
+            }
+            _ => {}
+        }
+    }
+}
+
+impl McoInstance for peripherals::MCO {}
+
+pub struct Mco<'d, T: McoInstance> {
+    phantom: PhantomData<&'d mut T>,
+}
+
+impl<'d, T: McoInstance> Mco<'d, T> {
+    pub fn new(
+        _peri: impl Peripheral<P = T> + 'd,
+        pin: impl Peripheral<P = impl McoPin<T>> + 'd,
+        source: impl McoSource<Raw = T::Source>,
+        prescaler: McoClock,
+    ) -> Self {
+        into_ref!(pin);
+
+        critical_section::with(|_| unsafe {
+            T::apply_clock_settings(source.into_raw(), prescaler.into_raw());
+            pin.set_as_af(pin.af_num(), AFType::OutputPushPull);
+            pin.set_speed(Speed::VeryHigh);
+        });
+
+        Self { phantom: PhantomData }
+    }
+}
+
 pub(crate) unsafe fn init(config: Config) {
     let (sys_clk, sw) = match config.mux {
         ClockSrc::MSI(range) => {
diff --git a/examples/stm32f4/src/bin/mco.rs b/examples/stm32f4/src/bin/mco.rs
new file mode 100644
index 000000000..2b9ceebc3
--- /dev/null
+++ b/examples/stm32f4/src/bin/mco.rs
@@ -0,0 +1,30 @@
+#![no_std]
+#![no_main]
+#![feature(type_alias_impl_trait)]
+
+use defmt::*;
+use embassy_executor::Spawner;
+use embassy_stm32::gpio::{Level, Output, Speed};
+use embassy_stm32::rcc::{Mco, Mco1Source, Mco2Source, McoClock};
+use embassy_time::{Duration, Timer};
+use {defmt_rtt as _, panic_probe as _};
+
+#[embassy_executor::main]
+async fn main(_spawner: Spawner) {
+    let p = embassy_stm32::init(Default::default());
+    info!("Hello World!");
+
+    let _mco1 = Mco::new(p.MCO1, p.PA8, Mco1Source::Hsi, McoClock::DIV1);
+    let _mco2 = Mco::new(p.MCO2, p.PC9, Mco2Source::Pll, McoClock::DIV4);
+    let mut led = Output::new(p.PB7, Level::High, Speed::Low);
+
+    loop {
+        info!("high");
+        led.set_high();
+        Timer::after(Duration::from_millis(300)).await;
+
+        info!("low");
+        led.set_low();
+        Timer::after(Duration::from_millis(300)).await;
+    }
+}
diff --git a/examples/stm32l4/src/bin/mco.rs b/examples/stm32l4/src/bin/mco.rs
new file mode 100644
index 000000000..dea0c66e0
--- /dev/null
+++ b/examples/stm32l4/src/bin/mco.rs
@@ -0,0 +1,27 @@
+#![no_std]
+#![no_main]
+#![feature(type_alias_impl_trait)]
+
+use defmt::*;
+use embassy_executor::Spawner;
+use embassy_stm32::gpio::{Level, Output, Speed};
+use embassy_stm32::rcc::{Mco, Mco1Source, McoClock};
+use embassy_time::{Duration, Timer};
+use {defmt_rtt as _, panic_probe as _};
+
+#[embassy_executor::main]
+async fn main(_spawner: Spawner) {
+    let p = embassy_stm32::init(Default::default());
+    info!("Hello World!");
+
+    let _mco = Mco::new(p.MCO, p.PA8, Mco1Source::Hsi16, McoClock::DIV1);
+
+    let mut led = Output::new(p.PB14, Level::High, Speed::Low);
+
+    loop {
+        led.set_high();
+        Timer::after(Duration::from_millis(300)).await;
+        led.set_low();
+        Timer::after(Duration::from_millis(300)).await;
+    }
+}