5 files changed, 93 insertions, 111 deletions

diff --git a/embassy-boot/boot/src/boot_loader.rs b/embassy-boot/boot/src/boot_loader.rs
index 2412427c0..b959de2c4 100644
--- a/embassy-boot/boot/src/boot_loader.rs
+++ b/embassy-boot/boot/src/boot_loader.rs
@@ -50,13 +50,13 @@ pub trait FlashConfig {
 }
 
 trait FlashConfigEx {
-    fn page_size() -> usize;
+    fn page_size() -> u32;
 }
 
 impl<T: FlashConfig> FlashConfigEx for T {
     /// Get the page size which is the "unit of operation" within the bootloader.
-    fn page_size() -> usize {
-        core::cmp::max(T::ACTIVE::ERASE_SIZE, T::DFU::ERASE_SIZE)
+    fn page_size() -> u32 {
+        core::cmp::max(T::ACTIVE::ERASE_SIZE, T::DFU::ERASE_SIZE) as u32
     }
 }
 
@@ -86,7 +86,7 @@ impl BootLoader {
 
     /// Return the offset of the active partition into the active flash.
     pub fn boot_address(&self) -> usize {
-        self.active.from
+        self.active.from as usize
     }
 
     /// Perform necessary boot preparations like swapping images.
@@ -177,11 +177,11 @@ impl BootLoader {
     ///
     pub fn prepare_boot<P: FlashConfig>(&mut self, p: &mut P, aligned_buf: &mut [u8]) -> Result<State, BootError> {
         // Ensure we have enough progress pages to store copy progress
-        assert_eq!(0, P::page_size() % aligned_buf.len());
-        assert_eq!(0, P::page_size() % P::ACTIVE::WRITE_SIZE);
-        assert_eq!(0, P::page_size() % P::ACTIVE::ERASE_SIZE);
-        assert_eq!(0, P::page_size() % P::DFU::WRITE_SIZE);
-        assert_eq!(0, P::page_size() % P::DFU::ERASE_SIZE);
+        assert_eq!(0, P::page_size() % aligned_buf.len() as u32);
+        assert_eq!(0, P::page_size() % P::ACTIVE::WRITE_SIZE as u32);
+        assert_eq!(0, P::page_size() % P::ACTIVE::ERASE_SIZE as u32);
+        assert_eq!(0, P::page_size() % P::DFU::WRITE_SIZE as u32);
+        assert_eq!(0, P::page_size() % P::DFU::ERASE_SIZE as u32);
         assert!(aligned_buf.len() >= P::STATE::WRITE_SIZE);
         assert_eq!(0, aligned_buf.len() % P::ACTIVE::WRITE_SIZE);
         assert_eq!(0, aligned_buf.len() % P::DFU::WRITE_SIZE);
@@ -222,30 +222,27 @@ impl BootLoader {
     }
 
     fn is_swapped<P: FlashConfig>(&mut self, p: &mut P, aligned_buf: &mut [u8]) -> Result<bool, BootError> {
-        let page_count = self.active.len() / P::page_size();
+        let page_count = (self.active.size() / P::page_size()) as usize;
         let progress = self.current_progress(p, aligned_buf)?;
 
         Ok(progress >= page_count * 2)
     }
 
     fn current_progress<P: FlashConfig>(&mut self, config: &mut P, aligned_buf: &mut [u8]) -> Result<usize, BootError> {
-        let max_index = ((self.state.len() - P::STATE::WRITE_SIZE) / P::STATE::WRITE_SIZE) - 2;
+        let write_size = P::STATE::WRITE_SIZE as u32;
+        let max_index = (((self.state.size() - write_size) / write_size) - 2) as usize;
         let state_flash = config.state();
-        let state_word = &mut aligned_buf[..P::STATE::WRITE_SIZE];
+        let state_word = &mut aligned_buf[..write_size as usize];
 
-        self.state
-            .read_blocking(state_flash, P::STATE::WRITE_SIZE as u32, state_word)?;
+        self.state.read_blocking(state_flash, write_size, state_word)?;
         if state_word.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,
-                (2 + index) as u32 * P::STATE::WRITE_SIZE as u32,
-                state_word,
-            )?;
+            self.state
+                .read_blocking(state_flash, (2 + index) as u32 * write_size, state_word)?;
 
             if state_word.iter().any(|&b| b == P::STATE_ERASE_VALUE) {
                 return Ok(index);
@@ -256,26 +253,29 @@ impl BootLoader {
 
     fn update_progress<P: FlashConfig>(
         &mut self,
-        index: usize,
+        progress_index: usize,
         p: &mut P,
         aligned_buf: &mut [u8],
     ) -> Result<(), BootError> {
         let state_word = &mut aligned_buf[..P::STATE::WRITE_SIZE];
         state_word.fill(!P::STATE_ERASE_VALUE);
-        self.state
-            .write_blocking(p.state(), (2 + index) as u32 * P::STATE::WRITE_SIZE as u32, state_word)?;
+        self.state.write_blocking(
+            p.state(),
+            (2 + progress_index) as u32 * P::STATE::WRITE_SIZE as u32,
+            state_word,
+        )?;
         Ok(())
     }
 
     fn copy_page_once_to_active<P: FlashConfig>(
         &mut self,
-        idx: usize,
+        progress_index: usize,
         from_offset: u32,
         to_offset: u32,
         p: &mut P,
         aligned_buf: &mut [u8],
     ) -> Result<(), BootError> {
-        if self.current_progress(p, aligned_buf)? <= idx {
+        if self.current_progress(p, aligned_buf)? <= progress_index {
             let page_size = P::page_size() as u32;
 
             self.active
@@ -288,20 +288,20 @@ impl BootLoader {
                     .write_blocking(p.active(), to_offset + offset_in_page as u32, aligned_buf)?;
             }
 
-            self.update_progress(idx, p, aligned_buf)?;
+            self.update_progress(progress_index, p, aligned_buf)?;
         }
         Ok(())
     }
 
     fn copy_page_once_to_dfu<P: FlashConfig>(
         &mut self,
-        idx: usize,
+        progress_index: usize,
         from_offset: u32,
         to_offset: u32,
         p: &mut P,
         aligned_buf: &mut [u8],
     ) -> Result<(), BootError> {
-        if self.current_progress(p, aligned_buf)? <= idx {
+        if self.current_progress(p, aligned_buf)? <= progress_index {
             let page_size = P::page_size() as u32;
 
             self.dfu
@@ -314,31 +314,28 @@ impl BootLoader {
                     .write_blocking(p.dfu(), to_offset + offset_in_page as u32, aligned_buf)?;
             }
 
-            self.update_progress(idx, p, aligned_buf)?;
+            self.update_progress(progress_index, p, aligned_buf)?;
         }
         Ok(())
     }
 
     fn swap<P: FlashConfig>(&mut self, p: &mut P, aligned_buf: &mut [u8]) -> Result<(), BootError> {
         let page_size = P::page_size();
-        let page_count = self.active.len() / page_size;
-        trace!("Page count: {}", page_count);
+        let page_count = self.active.size() / page_size;
         for page_num in 0..page_count {
-            trace!("COPY PAGE {}", page_num);
-
-            let idx = page_num * 2;
+            let progress_index = (page_num * 2) as usize;
 
             // Copy active page to the 'next' DFU page.
-            let active_from_offset = ((page_count - 1 - page_num) * page_size) as u32;
-            let dfu_to_offset = ((page_count - page_num) * page_size) as u32;
+            let active_from_offset = (page_count - 1 - page_num) * page_size;
+            let dfu_to_offset = (page_count - page_num) * page_size;
             //trace!("Copy active {} to dfu {}", active_from_offset, dfu_to_offset);
-            self.copy_page_once_to_dfu(idx, active_from_offset, dfu_to_offset, p, aligned_buf)?;
+            self.copy_page_once_to_dfu(progress_index, active_from_offset, dfu_to_offset, p, aligned_buf)?;
 
             // Copy DFU page to the active page
-            let active_to_offset = ((page_count - 1 - page_num) * page_size) as u32;
-            let dfu_from_offset = ((page_count - 1 - page_num) * page_size) as u32;
+            let active_to_offset = (page_count - 1 - page_num) * page_size;
+            let dfu_from_offset = (page_count - 1 - page_num) * page_size;
             //trace!("Copy dfy {} to active {}", dfu_from_offset, active_to_offset);
-            self.copy_page_once_to_active(idx + 1, dfu_from_offset, active_to_offset, p, aligned_buf)?;
+            self.copy_page_once_to_active(progress_index + 1, dfu_from_offset, active_to_offset, p, aligned_buf)?;
         }
 
         Ok(())
@@ -346,19 +343,19 @@ impl BootLoader {
 
     fn revert<P: FlashConfig>(&mut self, p: &mut P, aligned_buf: &mut [u8]) -> Result<(), BootError> {
         let page_size = P::page_size();
-        let page_count = self.active.len() / page_size;
+        let page_count = self.active.size() / page_size;
         for page_num in 0..page_count {
-            let idx = page_count * 2 + page_num * 2;
+            let progress_index = (page_count * 2 + page_num * 2) as usize;
 
             // Copy the bad active page to the DFU page
-            let active_from_offset = (page_num * page_size) as u32;
-            let dfu_to_offset = (page_num * page_size) as u32;
-            self.copy_page_once_to_dfu(idx, active_from_offset, dfu_to_offset, p, aligned_buf)?;
+            let active_from_offset = page_num * page_size;
+            let dfu_to_offset = page_num * page_size;
+            self.copy_page_once_to_dfu(progress_index, active_from_offset, dfu_to_offset, p, aligned_buf)?;
 
             // Copy the DFU page back to the active page
-            let active_to_offset = (page_num * page_size) as u32;
-            let dfu_from_offset = ((page_num + 1) * page_size) as u32;
-            self.copy_page_once_to_active(idx + 1, dfu_from_offset, active_to_offset, p, aligned_buf)?;
+            let active_to_offset = page_num * page_size;
+            let dfu_from_offset = (page_num + 1) * page_size;
+            self.copy_page_once_to_active(progress_index + 1, dfu_from_offset, active_to_offset, p, aligned_buf)?;
         }
 
         Ok(())
@@ -376,11 +373,11 @@ impl BootLoader {
     }
 }
 
-fn assert_partitions(active: Partition, dfu: Partition, state: Partition, page_size: usize, write_size: usize) {
-    assert_eq!(active.len() % page_size, 0);
-    assert_eq!(dfu.len() % page_size, 0);
-    assert!(dfu.len() - active.len() >= page_size);
-    assert!(2 + 2 * (active.len() / page_size) <= state.len() / write_size);
+fn assert_partitions(active: Partition, dfu: Partition, state: Partition, page_size: u32, state_write_size: usize) {
+    assert_eq!(active.size() % page_size, 0);
+    assert_eq!(dfu.size() % page_size, 0);
+    assert!(dfu.size() - active.size() >= page_size);
+    assert!(2 + 2 * (active.size() / page_size) <= state.size() / state_write_size as u32);
 }
 
 /// 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 2b5cc72fa..93d4a4c12 100644
--- a/embassy-boot/boot/src/firmware_updater.rs
+++ b/embassy-boot/boot/src/firmware_updater.rs
@@ -49,14 +49,14 @@ impl Default for FirmwareUpdater {
 
         let dfu = unsafe {
             Partition::new(
-                &__bootloader_dfu_start as *const u32 as usize,
-                &__bootloader_dfu_end as *const u32 as usize,
+                &__bootloader_dfu_start as *const u32 as u32,
+                &__bootloader_dfu_end as *const u32 as u32,
             )
         };
         let state = unsafe {
             Partition::new(
-                &__bootloader_state_start as *const u32 as usize,
-                &__bootloader_state_end as *const u32 as usize,
+                &__bootloader_state_start as *const u32 as u32,
+                &__bootloader_state_end as *const u32 as u32,
             )
         };
 
@@ -121,10 +121,8 @@ impl FirmwareUpdater {
         _update_len: usize,
         _aligned: &mut [u8],
     ) -> Result<(), FirmwareUpdaterError> {
-        let _read_size = _aligned.len();
-
         assert_eq!(_aligned.len(), F::WRITE_SIZE);
-        assert!(_update_len <= self.dfu.len());
+        assert!(_update_len as u32 <= self.dfu.size());
 
         #[cfg(feature = "ed25519-dalek")]
         {
@@ -330,11 +328,8 @@ impl FirmwareUpdater {
         _update_len: usize,
         _aligned: &mut [u8],
     ) -> Result<(), FirmwareUpdaterError> {
-        let _end = self.dfu.from + _update_len;
-        let _read_size = _aligned.len();
-
         assert_eq!(_aligned.len(), F::WRITE_SIZE);
-        assert!(_end <= self.dfu.to);
+        assert!(_update_len as u32 <= self.dfu.size());
 
         #[cfg(feature = "ed25519-dalek")]
         {
diff --git a/embassy-boot/boot/src/lib.rs b/embassy-boot/boot/src/lib.rs
index 3109f2b47..8b94b6bdc 100644
--- a/embassy-boot/boot/src/lib.rs
+++ b/embassy-boot/boot/src/lib.rs
@@ -89,13 +89,11 @@ mod tests {
         const DFU: Partition = Partition::new(61440, 122880);
         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 original = [rand::random::<u8>(); ACTIVE.size() as usize];
+        let update = [rand::random::<u8>(); ACTIVE.size() as usize];
         let mut aligned = [0; 4];
 
-        for i in ACTIVE.from..ACTIVE.to {
-            flash.mem[i] = original[i - ACTIVE.from];
-        }
+        flash.program(ACTIVE.from, &original).unwrap();
 
         let mut bootloader: BootLoader = BootLoader::new(ACTIVE, DFU, STATE);
         let mut updater = FirmwareUpdater::new(DFU, STATE);
@@ -110,14 +108,9 @@ mod tests {
                 .unwrap()
         );
 
-        for i in ACTIVE.from..ACTIVE.to {
-            assert_eq!(flash.mem[i], update[i - ACTIVE.from], "Index {}", i);
-        }
-
+        flash.assert_eq(ACTIVE.from, &update);
         // First DFU page is untouched
-        for i in DFU.from + 4096..DFU.to {
-            assert_eq!(flash.mem[i], original[i - DFU.from - 4096], "Index {}", i);
-        }
+        flash.assert_eq(DFU.from + 4096, &original);
 
         // Running again should cause a revert
         assert_eq!(
@@ -127,14 +120,9 @@ mod tests {
                 .unwrap()
         );
 
-        for i in ACTIVE.from..ACTIVE.to {
-            assert_eq!(flash.mem[i], original[i - ACTIVE.from], "Index {}", i);
-        }
-
+        flash.assert_eq(ACTIVE.from, &original);
         // Last page is untouched
-        for i in DFU.from..DFU.to - 4096 {
-            assert_eq!(flash.mem[i], update[i - DFU.from], "Index {}", i);
-        }
+        flash.assert_eq(DFU.from, &update);
 
         // Mark as booted
         block_on(updater.mark_booted(&mut flash, &mut aligned)).unwrap();
@@ -158,12 +146,10 @@ mod tests {
         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()];
+        let original = [rand::random::<u8>(); ACTIVE.size() as usize];
+        let update = [rand::random::<u8>(); ACTIVE.size() as usize];
 
-        for i in ACTIVE.from..ACTIVE.to {
-            active.mem[i] = original[i - ACTIVE.from];
-        }
+        active.program(ACTIVE.from, &original).unwrap();
 
         let mut updater = FirmwareUpdater::new(DFU, STATE);
 
@@ -180,14 +166,9 @@ mod tests {
                 .unwrap()
         );
 
-        for i in ACTIVE.from..ACTIVE.to {
-            assert_eq!(active.mem[i], update[i - ACTIVE.from], "Index {}", i);
-        }
-
+        active.assert_eq(ACTIVE.from, &update);
         // First DFU page is untouched
-        for i in DFU.from + 4096..DFU.to {
-            assert_eq!(dfu.mem[i], original[i - DFU.from - 4096], "Index {}", i);
-        }
+        dfu.assert_eq(DFU.from + 4096, &original);
     }
 
     #[test]
@@ -202,12 +183,10 @@ mod tests {
         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()];
+        let original = [rand::random::<u8>(); ACTIVE.size() as usize];
+        let update = [rand::random::<u8>(); ACTIVE.size() as usize];
 
-        for i in ACTIVE.from..ACTIVE.to {
-            active.mem[i] = original[i - ACTIVE.from];
-        }
+        active.program(ACTIVE.from, &original).unwrap();
 
         let mut updater = FirmwareUpdater::new(DFU, STATE);
 
@@ -226,14 +205,9 @@ mod tests {
                 .unwrap()
         );
 
-        for i in ACTIVE.from..ACTIVE.to {
-            assert_eq!(active.mem[i], update[i - ACTIVE.from], "Index {}", i);
-        }
-
+        active.assert_eq(ACTIVE.from, &update);
         // First DFU page is untouched
-        for i in DFU.from + 4096..DFU.to {
-            assert_eq!(dfu.mem[i], original[i - DFU.from - 4096], "Index {}", i);
-        }
+        dfu.assert_eq(DFU.from + 4096, &original);
     }
 
     #[test]
diff --git a/embassy-boot/boot/src/mem_flash.rs b/embassy-boot/boot/src/mem_flash.rs
index dd85405c8..c62379b24 100644
--- a/embassy-boot/boot/src/mem_flash.rs
+++ b/embassy-boot/boot/src/mem_flash.rs
@@ -32,6 +32,23 @@ impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> MemFla
             pending_write_successes: None,
         }
     }
+
+    pub fn program(&mut self, offset: u32, bytes: &[u8]) -> Result<(), MemFlashError> {
+        let offset = offset as usize;
+        assert!(bytes.len() % WRITE_SIZE == 0);
+        assert!(offset % WRITE_SIZE == 0);
+        assert!(offset + bytes.len() <= SIZE);
+
+        self.mem[offset..offset + bytes.len()].copy_from_slice(bytes);
+
+        Ok(())
+    }
+
+    pub fn assert_eq(&self, offset: u32, expectation: &[u8]) {
+        for i in 0..expectation.len() {
+            assert_eq!(self.mem[offset as usize + i], expectation[i], "Index {}", i);
+        }
+    }
 }
 
 impl<const SIZE: usize, const ERASE_SIZE: usize, const WRITE_SIZE: usize> Default
diff --git a/embassy-boot/boot/src/partition.rs b/embassy-boot/boot/src/partition.rs
index ac6b0ed0f..7529059b6 100644
--- a/embassy-boot/boot/src/partition.rs
+++ b/embassy-boot/boot/src/partition.rs
@@ -6,20 +6,19 @@ use embedded_storage_async::nor_flash::{NorFlash as AsyncNorFlash, ReadNorFlash
 #[cfg_attr(feature = "defmt", derive(defmt::Format))]
 pub struct Partition {
     /// The offset into the flash where the partition starts.
-    pub from: usize,
+    pub from: u32,
     /// The offset into the flash where the partition ends.
-    pub to: usize,
+    pub to: u32,
 }
 
 impl Partition {
     /// Create a new partition with the provided range
-    pub const fn new(from: usize, to: usize) -> Self {
+    pub const fn new(from: u32, to: u32) -> Self {
         Self { from, to }
     }
 
     /// Return the size of the partition
-    #[allow(clippy::len_without_is_empty)]
-    pub const fn len(&self) -> usize {
+    pub const fn size(&self) -> u32 {
         self.to - self.from
     }