Remove the usage of the local Partition type in BootLoader

2 files changed, 112 insertions, 229 deletions

diff --git a/embassy-boot/boot/Cargo.toml b/embassy-boot/boot/Cargo.toml
index f641d5e1c..3fdf69c5e 100644
--- a/embassy-boot/boot/Cargo.toml
+++ b/embassy-boot/boot/Cargo.toml
@@ -39,6 +39,7 @@ env_logger = "0.9"
 rand = "0.7" # ed25519-dalek v1.0.1 depends on this exact version
 futures = { version = "0.3", features = ["executor"] }
 sha1 = "0.10.5"
+embassy-embedded-hal = { version = "0.1.0", path = "../../embassy-embedded-hal" }
 
 [dev-dependencies.ed25519-dalek]
 default_features = false
diff --git a/embassy-boot/boot/src/boot_loader.rs b/embassy-boot/boot/src/boot_loader.rs
index b959de2c4..50eb5e668 100644
--- a/embassy-boot/boot/src/boot_loader.rs
+++ b/embassy-boot/boot/src/boot_loader.rs
@@ -1,6 +1,8 @@
 use embedded_storage::nor_flash::{ErrorType, NorFlash, NorFlashError, NorFlashErrorKind, ReadNorFlash};
 
-use crate::{Partition, State, BOOT_MAGIC, SWAP_MAGIC};
+use crate::{State, BOOT_MAGIC, SWAP_MAGIC};
+
+const STATE_ERASE_VALUE: u8 = 0xFF;
 
 /// Errors returned by bootloader
 #[derive(PartialEq, Eq, Debug)]
@@ -30,65 +32,39 @@ where
     }
 }
 
-/// Trait defining the flash handles used for active and DFU partition.
-pub trait FlashConfig {
-    /// The erase value of the state flash. Typically the default of 0xFF is used, but some flashes use a different value.
-    const STATE_ERASE_VALUE: u8 = 0xFF;
+/// BootLoader works with any flash implementing embedded_storage.
+pub struct BootLoader<ACTIVE: NorFlash, DFU: NorFlash, STATE: NorFlash> {
+    /// Flash type used for the active partition - the partition which will be booted from.
+    active: ACTIVE,
+    /// Flash type used for the dfu partition - he partition which will be swapped in when requested.
+    dfu: DFU,
     /// Flash type used for the state partition.
-    type STATE: NorFlash;
-    /// Flash type used for the active partition.
-    type ACTIVE: NorFlash;
-    /// Flash type used for the dfu partition.
-    type DFU: NorFlash;
-
-    /// Return flash instance used to write/read to/from active partition.
-    fn active(&mut self) -> &mut Self::ACTIVE;
-    /// Return flash instance used to write/read to/from dfu partition.
-    fn dfu(&mut self) -> &mut Self::DFU;
-    /// Return flash instance used to write/read to/from bootloader state.
-    fn state(&mut self) -> &mut Self::STATE;
-}
-
-trait FlashConfigEx {
-    fn page_size() -> u32;
+    ///
+    /// The state partition has the following format:
+    /// 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: STATE,
 }
 
-impl<T: FlashConfig> FlashConfigEx for T {
+impl<ACTIVE: NorFlash, DFU: NorFlash, STATE: NorFlash> BootLoader<ACTIVE, DFU, STATE> {
     /// Get the page size which is the "unit of operation" within the bootloader.
-    fn page_size() -> u32 {
-        core::cmp::max(T::ACTIVE::ERASE_SIZE, T::DFU::ERASE_SIZE) as u32
-    }
-}
+    const PAGE_SIZE: u32 = if ACTIVE::ERASE_SIZE > DFU::ERASE_SIZE {
+        ACTIVE::ERASE_SIZE as u32
+    } else {
+        DFU::ERASE_SIZE as u32
+    };
 
-/// BootLoader works with any flash implementing embedded_storage.
-pub struct BootLoader {
-    // Page with current state of bootloader. The state partition has the following format:
-    // 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,
-    // Location of the partition which will be swapped in when requested
-    dfu: Partition,
-}
-
-impl BootLoader {
-    /// Create a new instance of a bootloader with the given partitions.
+    /// Create a new instance of a bootloader with the flash partitions.
     ///
     /// - All partitions must be aligned with the PAGE_SIZE const generic parameter.
     /// - The dfu partition must be at least PAGE_SIZE bigger than the active partition.
-    pub fn new(active: Partition, dfu: Partition, state: Partition) -> Self {
+    pub fn new(active: ACTIVE, dfu: DFU, state: STATE) -> Self {
         Self { active, dfu, state }
     }
 
-    /// Return the offset of the active partition into the active flash.
-    pub fn boot_address(&self) -> usize {
-        self.active.from as usize
-    }
-
     /// Perform necessary boot preparations like swapping images.
     ///
     /// The DFU partition is assumed to be 1 page bigger than the active partition for the swap
@@ -175,195 +151,174 @@ impl BootLoader {
     /// |       DFU |            3 |      3 |      2 |      1 |      3 |
     /// +-----------+--------------+--------+--------+--------+--------+
     ///
-    pub fn prepare_boot<P: FlashConfig>(&mut self, p: &mut P, aligned_buf: &mut [u8]) -> Result<State, BootError> {
+    pub fn prepare_boot(&mut self, 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() 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);
-        assert_partitions(self.active, self.dfu, self.state, P::page_size(), P::STATE::WRITE_SIZE);
+        assert_eq!(0, Self::PAGE_SIZE % aligned_buf.len() as u32);
+        assert_eq!(0, Self::PAGE_SIZE % ACTIVE::WRITE_SIZE as u32);
+        assert_eq!(0, Self::PAGE_SIZE % ACTIVE::ERASE_SIZE as u32);
+        assert_eq!(0, Self::PAGE_SIZE % DFU::WRITE_SIZE as u32);
+        assert_eq!(0, Self::PAGE_SIZE % DFU::ERASE_SIZE as u32);
+        assert!(aligned_buf.len() >= STATE::WRITE_SIZE);
+        assert_eq!(0, aligned_buf.len() % ACTIVE::WRITE_SIZE);
+        assert_eq!(0, aligned_buf.len() % DFU::WRITE_SIZE);
+
+        assert_partitions(&self.active, &self.dfu, &self.state, Self::PAGE_SIZE);
 
         // Copy contents from partition N to active
-        let state = self.read_state(p, aligned_buf)?;
+        let state = self.read_state(aligned_buf)?;
         if state == State::Swap {
             //
             // Check if we already swapped. If we're in the swap state, this means we should revert
             // since the app has failed to mark boot as successful
             //
-            if !self.is_swapped(p, aligned_buf)? {
+            if !self.is_swapped(aligned_buf)? {
                 trace!("Swapping");
-                self.swap(p, aligned_buf)?;
+                self.swap(aligned_buf)?;
                 trace!("Swapping done");
             } else {
                 trace!("Reverting");
-                self.revert(p, aligned_buf)?;
+                self.revert(aligned_buf)?;
 
-                let state_flash = p.state();
-                let state_word = &mut aligned_buf[..P::STATE::WRITE_SIZE];
+                let state_word = &mut aligned_buf[..STATE::WRITE_SIZE];
 
                 // Invalidate progress
-                state_word.fill(!P::STATE_ERASE_VALUE);
-                self.state
-                    .write_blocking(state_flash, P::STATE::WRITE_SIZE as u32, state_word)?;
+                state_word.fill(!STATE_ERASE_VALUE);
+                self.state.write(STATE::WRITE_SIZE as u32, state_word)?;
 
                 // Clear magic and progress
-                self.state.wipe_blocking(state_flash)?;
+                self.state.erase(0, self.state.capacity() as u32)?;
 
                 // Set magic
                 state_word.fill(BOOT_MAGIC);
-                self.state.write_blocking(state_flash, 0, state_word)?;
+                self.state.write(0, state_word)?;
             }
         }
         Ok(state)
     }
 
-    fn is_swapped<P: FlashConfig>(&mut self, p: &mut P, aligned_buf: &mut [u8]) -> Result<bool, BootError> {
-        let page_count = (self.active.size() / P::page_size()) as usize;
-        let progress = self.current_progress(p, aligned_buf)?;
+    fn is_swapped(&mut self, aligned_buf: &mut [u8]) -> Result<bool, BootError> {
+        let page_count = self.active.capacity() / Self::PAGE_SIZE as usize;
+        let progress = self.current_progress(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 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();
+    fn current_progress(&mut self, aligned_buf: &mut [u8]) -> Result<usize, BootError> {
+        let write_size = STATE::WRITE_SIZE as u32;
+        let max_index = ((self.state.capacity() - STATE::WRITE_SIZE) / STATE::WRITE_SIZE) - 2;
         let state_word = &mut aligned_buf[..write_size as usize];
 
-        self.state.read_blocking(state_flash, write_size, state_word)?;
-        if state_word.iter().any(|&b| b != P::STATE_ERASE_VALUE) {
+        self.state.read(write_size, state_word)?;
+        if state_word.iter().any(|&b| b != 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 * write_size, state_word)?;
+            self.state.read((2 + index) as u32 * write_size, state_word)?;
 
-            if state_word.iter().any(|&b| b == P::STATE_ERASE_VALUE) {
+            if state_word.iter().any(|&b| b == STATE_ERASE_VALUE) {
                 return Ok(index);
             }
         }
         Ok(max_index)
     }
 
-    fn update_progress<P: FlashConfig>(
-        &mut self,
-        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 + progress_index) as u32 * P::STATE::WRITE_SIZE as u32,
-            state_word,
-        )?;
+    fn update_progress(&mut self, progress_index: usize, aligned_buf: &mut [u8]) -> Result<(), BootError> {
+        let state_word = &mut aligned_buf[..STATE::WRITE_SIZE];
+        state_word.fill(!STATE_ERASE_VALUE);
+        self.state
+            .write((2 + progress_index) as u32 * STATE::WRITE_SIZE as u32, state_word)?;
         Ok(())
     }
 
-    fn copy_page_once_to_active<P: FlashConfig>(
+    fn copy_page_once_to_active(
         &mut self,
         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)? <= progress_index {
-            let page_size = P::page_size() as u32;
+        if self.current_progress(aligned_buf)? <= progress_index {
+            let page_size = Self::PAGE_SIZE as u32;
 
-            self.active
-                .erase_blocking(p.active(), to_offset, to_offset + page_size)?;
+            self.active.erase(to_offset, to_offset + page_size)?;
 
             for offset_in_page in (0..page_size).step_by(aligned_buf.len()) {
-                self.dfu
-                    .read_blocking(p.dfu(), from_offset + offset_in_page as u32, aligned_buf)?;
-                self.active
-                    .write_blocking(p.active(), to_offset + offset_in_page as u32, aligned_buf)?;
+                self.dfu.read(from_offset + offset_in_page as u32, aligned_buf)?;
+                self.active.write(to_offset + offset_in_page as u32, aligned_buf)?;
             }
 
-            self.update_progress(progress_index, p, aligned_buf)?;
+            self.update_progress(progress_index, aligned_buf)?;
         }
         Ok(())
     }
 
-    fn copy_page_once_to_dfu<P: FlashConfig>(
+    fn copy_page_once_to_dfu(
         &mut self,
         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)? <= progress_index {
-            let page_size = P::page_size() as u32;
+        if self.current_progress(aligned_buf)? <= progress_index {
+            let page_size = Self::PAGE_SIZE as u32;
 
-            self.dfu
-                .erase_blocking(p.dfu(), to_offset as u32, to_offset + page_size)?;
+            self.dfu.erase(to_offset as u32, to_offset + page_size)?;
 
             for offset_in_page in (0..page_size).step_by(aligned_buf.len()) {
-                self.active
-                    .read_blocking(p.active(), from_offset + offset_in_page as u32, aligned_buf)?;
-                self.dfu
-                    .write_blocking(p.dfu(), to_offset + offset_in_page as u32, aligned_buf)?;
+                self.active.read(from_offset + offset_in_page as u32, aligned_buf)?;
+                self.dfu.write(to_offset + offset_in_page as u32, aligned_buf)?;
             }
 
-            self.update_progress(progress_index, p, aligned_buf)?;
+            self.update_progress(progress_index, 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.size() / page_size;
+    fn swap(&mut self, aligned_buf: &mut [u8]) -> Result<(), BootError> {
+        let page_count = self.active.capacity() as u32 / Self::PAGE_SIZE;
         for page_num in 0..page_count {
             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;
-            let dfu_to_offset = (page_count - page_num) * page_size;
+            let active_from_offset = (page_count - 1 - page_num) * Self::PAGE_SIZE;
+            let dfu_to_offset = (page_count - page_num) * Self::PAGE_SIZE;
             //trace!("Copy active {} to dfu {}", active_from_offset, dfu_to_offset);
-            self.copy_page_once_to_dfu(progress_index, active_from_offset, dfu_to_offset, p, aligned_buf)?;
+            self.copy_page_once_to_dfu(progress_index, active_from_offset, dfu_to_offset, aligned_buf)?;
 
             // Copy DFU page to the active page
-            let active_to_offset = (page_count - 1 - page_num) * page_size;
-            let dfu_from_offset = (page_count - 1 - page_num) * page_size;
+            let active_to_offset = (page_count - 1 - page_num) * Self::PAGE_SIZE;
+            let dfu_from_offset = (page_count - 1 - page_num) * Self::PAGE_SIZE;
             //trace!("Copy dfy {} to active {}", dfu_from_offset, active_to_offset);
-            self.copy_page_once_to_active(progress_index + 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, aligned_buf)?;
         }
 
         Ok(())
     }
 
-    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.size() / page_size;
+    fn revert(&mut self, aligned_buf: &mut [u8]) -> Result<(), BootError> {
+        let page_count = self.active.capacity() as u32 / Self::PAGE_SIZE;
         for page_num in 0..page_count {
             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;
-            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)?;
+            let active_from_offset = page_num * Self::PAGE_SIZE;
+            let dfu_to_offset = page_num * Self::PAGE_SIZE;
+            self.copy_page_once_to_dfu(progress_index, active_from_offset, dfu_to_offset, aligned_buf)?;
 
             // Copy the DFU page back to the active page
-            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)?;
+            let active_to_offset = page_num * Self::PAGE_SIZE;
+            let dfu_from_offset = (page_num + 1) * Self::PAGE_SIZE;
+            self.copy_page_once_to_active(progress_index + 1, dfu_from_offset, active_to_offset, aligned_buf)?;
         }
 
         Ok(())
     }
 
-    fn read_state<P: FlashConfig>(&mut self, config: &mut P, aligned_buf: &mut [u8]) -> Result<State, BootError> {
-        let state_word = &mut aligned_buf[..P::STATE::WRITE_SIZE];
-        self.state.read_blocking(config.state(), 0, state_word)?;
+    fn read_state(&mut self, aligned_buf: &mut [u8]) -> Result<State, BootError> {
+        let state_word = &mut aligned_buf[..STATE::WRITE_SIZE];
+        self.state.read(0, state_word)?;
 
         if !state_word.iter().any(|&b| b != SWAP_MAGIC) {
             Ok(State::Swap)
@@ -373,11 +328,16 @@ impl BootLoader {
     }
 }
 
-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);
+fn assert_partitions<ACTIVE: NorFlash, DFU: NorFlash, STATE: NorFlash>(
+    active: &ACTIVE,
+    dfu: &DFU,
+    state: &STATE,
+    page_size: u32,
+) {
+    assert_eq!(active.capacity() as u32 % page_size, 0);
+    assert_eq!(dfu.capacity() as u32 % page_size, 0);
+    assert!(dfu.capacity() as u32 - active.capacity() as u32 >= page_size);
+    assert!(2 + 2 * (active.capacity() as u32 / page_size) <= state.capacity() as u32 / STATE::WRITE_SIZE as u32);
 }
 
 /// A flash wrapper implementing the Flash and embedded_storage traits.
@@ -436,98 +396,20 @@ where
     }
 }
 
-/// Convenience provider that uses a single flash for all partitions.
-pub struct SingleFlashConfig<'a, F>
-where
-    F: NorFlash,
-{
-    flash: &'a mut F,
-}
-
-impl<'a, F> SingleFlashConfig<'a, F>
-where
-    F: NorFlash,
-{
-    /// Create a provider for a single flash.
-    pub fn new(flash: &'a mut F) -> Self {
-        Self { flash }
-    }
-}
-
-impl<'a, F> FlashConfig for SingleFlashConfig<'a, F>
-where
-    F: NorFlash,
-{
-    type STATE = F;
-    type ACTIVE = F;
-    type DFU = F;
-
-    fn active(&mut self) -> &mut Self::STATE {
-        self.flash
-    }
-    fn dfu(&mut self) -> &mut Self::ACTIVE {
-        self.flash
-    }
-    fn state(&mut self) -> &mut Self::DFU {
-        self.flash
-    }
-}
-
-/// Convenience flash provider that uses separate flash instances for each partition.
-pub struct MultiFlashConfig<'a, ACTIVE, STATE, DFU>
-where
-    ACTIVE: NorFlash,
-    STATE: NorFlash,
-    DFU: NorFlash,
-{
-    active: &'a mut ACTIVE,
-    state: &'a mut STATE,
-    dfu: &'a mut DFU,
-}
-
-impl<'a, ACTIVE, STATE, DFU> MultiFlashConfig<'a, ACTIVE, STATE, DFU>
-where
-    ACTIVE: NorFlash,
-    STATE: NorFlash,
-    DFU: NorFlash,
-{
-    /// Create a new flash provider with separate configuration for all three partitions.
-    pub fn new(active: &'a mut ACTIVE, state: &'a mut STATE, dfu: &'a mut DFU) -> Self {
-        Self { active, state, dfu }
-    }
-}
-
-impl<'a, ACTIVE, STATE, DFU> FlashConfig for MultiFlashConfig<'a, ACTIVE, STATE, DFU>
-where
-    ACTIVE: NorFlash,
-    STATE: NorFlash,
-    DFU: NorFlash,
-{
-    type STATE = STATE;
-    type ACTIVE = ACTIVE;
-    type DFU = DFU;
-
-    fn active(&mut self) -> &mut Self::ACTIVE {
-        self.active
-    }
-    fn dfu(&mut self) -> &mut Self::DFU {
-        self.dfu
-    }
-    fn state(&mut self) -> &mut Self::STATE {
-        self.state
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::mem_flash::MemFlash;
 
     #[test]
     #[should_panic]
     fn test_range_asserts() {
-        const ACTIVE: Partition = Partition::new(4096, 4194304);
-        const DFU: Partition = Partition::new(4194304, 2 * 4194304);
-        const STATE: Partition = Partition::new(0, 4096);
-        assert_partitions(ACTIVE, DFU, STATE, 4096, 4);
+        const ACTIVE_SIZE: usize = 4194304 - 4096;
+        const DFU_SIZE: usize = 4194304;
+        const STATE_SIZE: usize = 4096;
+        static ACTIVE: MemFlash<ACTIVE_SIZE, 4, 4> = MemFlash::new(0xFF);
+        static DFU: MemFlash<DFU_SIZE, 4, 4> = MemFlash::new(0xFF);
+        static STATE: MemFlash<STATE_SIZE, 4, 4> = MemFlash::new(0xFF);
+        assert_partitions(&ACTIVE, &DFU, &STATE, 4096);
     }
 }