1 files changed, 323 insertions, 0 deletions

diff --git a/embassy-boot/src/lib.rs b/embassy-boot/src/lib.rs
new file mode 100644
index 000000000..b4f03e01e
--- /dev/null
+++ b/embassy-boot/src/lib.rs
@@ -0,0 +1,323 @@
+#![no_std]
+#![allow(async_fn_in_trait)]
+#![warn(missing_docs)]
+#![doc = include_str!("../README.md")]
+mod fmt;
+
+mod boot_loader;
+mod digest_adapters;
+mod firmware_updater;
+#[cfg(test)]
+mod mem_flash;
+#[cfg(test)]
+mod test_flash;
+
+// The expected value of the flash after an erase
+// TODO: Use the value provided by NorFlash when available
+pub(crate) const STATE_ERASE_VALUE: u8 = 0xFF;
+pub use boot_loader::{BootError, BootLoader, BootLoaderConfig};
+pub use firmware_updater::{
+    BlockingFirmwareState, BlockingFirmwareUpdater, FirmwareState, FirmwareUpdater, FirmwareUpdaterConfig,
+    FirmwareUpdaterError,
+};
+
+pub(crate) const BOOT_MAGIC: u8 = 0xD0;
+pub(crate) const SWAP_MAGIC: u8 = 0xF0;
+pub(crate) const DFU_DETACH_MAGIC: u8 = 0xE0;
+
+/// The state of the bootloader after running prepare.
+#[derive(PartialEq, Eq, Debug)]
+#[cfg_attr(feature = "defmt", derive(defmt::Format))]
+pub enum State {
+    /// Bootloader is ready to boot the active partition.
+    Boot,
+    /// Bootloader has swapped the active partition with the dfu partition and will attempt boot.
+    Swap,
+    /// Application has received a request to reboot into DFU mode to apply an update.
+    DfuDetach,
+}
+
+/// Buffer aligned to 32 byte boundary, largest known alignment requirement for embassy-boot.
+#[repr(align(32))]
+pub struct AlignedBuffer<const N: usize>(pub [u8; N]);
+
+impl<const N: usize> AsRef<[u8]> for AlignedBuffer<N> {
+    fn as_ref(&self) -> &[u8] {
+        &self.0
+    }
+}
+
+impl<const N: usize> AsMut<[u8]> for AlignedBuffer<N> {
+    fn as_mut(&mut self) -> &mut [u8] {
+        &mut self.0
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    #![allow(unused_imports)]
+
+    use embedded_storage::nor_flash::{NorFlash, ReadNorFlash};
+    use embedded_storage_async::nor_flash::NorFlash as AsyncNorFlash;
+    use futures::executor::block_on;
+
+    use super::*;
+    use crate::boot_loader::BootLoaderConfig;
+    use crate::firmware_updater::FirmwareUpdaterConfig;
+    use crate::mem_flash::MemFlash;
+    use crate::test_flash::{AsyncTestFlash, BlockingTestFlash};
+
+    /*
+    #[test]
+    fn test_bad_magic() {
+        let mut flash = MemFlash([0xff; 131072]);
+        let mut flash = SingleFlashConfig::new(&mut flash);
+
+        let mut bootloader = BootLoader::<4096>::new(ACTIVE, DFU, STATE);
+
+        assert_eq!(
+            bootloader.prepare_boot(&mut flash),
+            Err(BootError::BadMagic)
+        );
+    }
+    */
+
+    #[test]
+    fn test_boot_state() {
+        let flash = BlockingTestFlash::new(BootLoaderConfig {
+            active: MemFlash::<57344, 4096, 4>::default(),
+            dfu: MemFlash::<61440, 4096, 4>::default(),
+            state: MemFlash::<4096, 4096, 4>::default(),
+        });
+
+        flash.state().write(0, &[BOOT_MAGIC; 4]).unwrap();
+
+        let mut bootloader = BootLoader::new(BootLoaderConfig {
+            active: flash.active(),
+            dfu: flash.dfu(),
+            state: flash.state(),
+        });
+
+        let mut page = [0; 4096];
+        assert_eq!(State::Boot, bootloader.prepare_boot(&mut page).unwrap());
+    }
+
+    #[test]
+    #[cfg(not(feature = "_verify"))]
+    fn test_swap_state() {
+        const FIRMWARE_SIZE: usize = 57344;
+        let flash = AsyncTestFlash::new(BootLoaderConfig {
+            active: MemFlash::<FIRMWARE_SIZE, 4096, 4>::default(),
+            dfu: MemFlash::<61440, 4096, 4>::default(),
+            state: MemFlash::<4096, 4096, 4>::default(),
+        });
+
+        const ORIGINAL: [u8; FIRMWARE_SIZE] = [0x55; FIRMWARE_SIZE];
+        const UPDATE: [u8; FIRMWARE_SIZE] = [0xAA; FIRMWARE_SIZE];
+        let mut aligned = [0; 4];
+
+        block_on(flash.active().erase(0, ORIGINAL.len() as u32)).unwrap();
+        block_on(flash.active().write(0, &ORIGINAL)).unwrap();
+
+        let mut updater = FirmwareUpdater::new(
+            FirmwareUpdaterConfig {
+                dfu: flash.dfu(),
+                state: flash.state(),
+            },
+            &mut aligned,
+        );
+        block_on(updater.write_firmware(0, &UPDATE)).unwrap();
+        block_on(updater.mark_updated()).unwrap();
+
+        // Writing after marking updated is not allowed until marked as booted.
+        let res: Result<(), FirmwareUpdaterError> = block_on(updater.write_firmware(0, &UPDATE));
+        assert!(matches!(res, Err::<(), _>(FirmwareUpdaterError::BadState)));
+
+        let flash = flash.into_blocking();
+        let mut bootloader = BootLoader::new(BootLoaderConfig {
+            active: flash.active(),
+            dfu: flash.dfu(),
+            state: flash.state(),
+        });
+
+        let mut page = [0; 1024];
+        assert_eq!(State::Swap, bootloader.prepare_boot(&mut page).unwrap());
+
+        let mut read_buf = [0; FIRMWARE_SIZE];
+        flash.active().read(0, &mut read_buf).unwrap();
+        assert_eq!(UPDATE, read_buf);
+        // First DFU page is untouched
+        flash.dfu().read(4096, &mut read_buf).unwrap();
+        assert_eq!(ORIGINAL, read_buf);
+
+        // Running again should cause a revert
+        assert_eq!(State::Swap, bootloader.prepare_boot(&mut page).unwrap());
+
+        let mut read_buf = [0; FIRMWARE_SIZE];
+        flash.active().read(0, &mut read_buf).unwrap();
+        assert_eq!(ORIGINAL, read_buf);
+        // Last DFU page is untouched
+        flash.dfu().read(0, &mut read_buf).unwrap();
+        assert_eq!(UPDATE, read_buf);
+
+        // Mark as booted
+        let flash = flash.into_async();
+        let mut updater = FirmwareUpdater::new(
+            FirmwareUpdaterConfig {
+                dfu: flash.dfu(),
+                state: flash.state(),
+            },
+            &mut aligned,
+        );
+        block_on(updater.mark_booted()).unwrap();
+
+        let flash = flash.into_blocking();
+        let mut bootloader = BootLoader::new(BootLoaderConfig {
+            active: flash.active(),
+            dfu: flash.dfu(),
+            state: flash.state(),
+        });
+        assert_eq!(State::Boot, bootloader.prepare_boot(&mut page).unwrap());
+    }
+
+    #[test]
+    #[cfg(not(feature = "_verify"))]
+    fn test_swap_state_active_page_biggest() {
+        const FIRMWARE_SIZE: usize = 12288;
+        let flash = AsyncTestFlash::new(BootLoaderConfig {
+            active: MemFlash::<12288, 4096, 8>::random(),
+            dfu: MemFlash::<16384, 2048, 8>::random(),
+            state: MemFlash::<2048, 128, 4>::random(),
+        });
+
+        const ORIGINAL: [u8; FIRMWARE_SIZE] = [0x55; FIRMWARE_SIZE];
+        const UPDATE: [u8; FIRMWARE_SIZE] = [0xAA; FIRMWARE_SIZE];
+        let mut aligned = [0; 4];
+
+        block_on(flash.active().erase(0, ORIGINAL.len() as u32)).unwrap();
+        block_on(flash.active().write(0, &ORIGINAL)).unwrap();
+
+        let mut updater = FirmwareUpdater::new(
+            FirmwareUpdaterConfig {
+                dfu: flash.dfu(),
+                state: flash.state(),
+            },
+            &mut aligned,
+        );
+        block_on(updater.write_firmware(0, &UPDATE)).unwrap();
+        block_on(updater.mark_updated()).unwrap();
+
+        let flash = flash.into_blocking();
+        let mut bootloader = BootLoader::new(BootLoaderConfig {
+            active: flash.active(),
+            dfu: flash.dfu(),
+            state: flash.state(),
+        });
+
+        let mut page = [0; 4096];
+        assert_eq!(State::Swap, bootloader.prepare_boot(&mut page).unwrap());
+
+        let mut read_buf = [0; FIRMWARE_SIZE];
+        flash.active().read(0, &mut read_buf).unwrap();
+        assert_eq!(UPDATE, read_buf);
+        // First DFU page is untouched
+        flash.dfu().read(4096, &mut read_buf).unwrap();
+        assert_eq!(ORIGINAL, read_buf);
+    }
+
+    #[test]
+    #[cfg(not(feature = "_verify"))]
+    fn test_swap_state_dfu_page_biggest() {
+        const FIRMWARE_SIZE: usize = 12288;
+        let flash = AsyncTestFlash::new(BootLoaderConfig {
+            active: MemFlash::<FIRMWARE_SIZE, 2048, 4>::random(),
+            dfu: MemFlash::<16384, 4096, 8>::random(),
+            state: MemFlash::<2048, 128, 4>::random(),
+        });
+
+        const ORIGINAL: [u8; FIRMWARE_SIZE] = [0x55; FIRMWARE_SIZE];
+        const UPDATE: [u8; FIRMWARE_SIZE] = [0xAA; FIRMWARE_SIZE];
+        let mut aligned = [0; 4];
+
+        block_on(flash.active().erase(0, ORIGINAL.len() as u32)).unwrap();
+        block_on(flash.active().write(0, &ORIGINAL)).unwrap();
+
+        let mut updater = FirmwareUpdater::new(
+            FirmwareUpdaterConfig {
+                dfu: flash.dfu(),
+                state: flash.state(),
+            },
+            &mut aligned,
+        );
+        block_on(updater.write_firmware(0, &UPDATE)).unwrap();
+        block_on(updater.mark_updated()).unwrap();
+
+        let flash = flash.into_blocking();
+        let mut bootloader = BootLoader::new(BootLoaderConfig {
+            active: flash.active(),
+            dfu: flash.dfu(),
+            state: flash.state(),
+        });
+        let mut page = [0; 4096];
+        assert_eq!(State::Swap, bootloader.prepare_boot(&mut page).unwrap());
+
+        let mut read_buf = [0; FIRMWARE_SIZE];
+        flash.active().read(0, &mut read_buf).unwrap();
+        assert_eq!(UPDATE, read_buf);
+        // First DFU page is untouched
+        flash.dfu().read(4096, &mut read_buf).unwrap();
+        assert_eq!(ORIGINAL, read_buf);
+    }
+
+    #[test]
+    #[cfg(feature = "_verify")]
+    fn test_verify() {
+        // The following key setup is based on:
+        // https://docs.rs/ed25519-dalek/latest/ed25519_dalek/#example
+
+        use ed25519_dalek::{Digest, Sha512, Signature, Signer, SigningKey, VerifyingKey};
+        use rand::rngs::OsRng;
+
+        let mut csprng = OsRng {};
+        let keypair = SigningKey::generate(&mut csprng);
+
+        let firmware: &[u8] = b"This are bytes that would otherwise be firmware bytes for DFU.";
+        let mut digest = Sha512::new();
+        digest.update(&firmware);
+        let message = digest.finalize();
+        let signature: Signature = keypair.sign(&message);
+
+        let public_key = keypair.verifying_key();
+
+        // Setup flash
+        let flash = BlockingTestFlash::new(BootLoaderConfig {
+            active: MemFlash::<0, 0, 0>::default(),
+            dfu: MemFlash::<4096, 4096, 4>::default(),
+            state: MemFlash::<4096, 4096, 4>::default(),
+        });
+
+        let firmware_len = firmware.len();
+
+        let mut write_buf = [0; 4096];
+        write_buf[0..firmware_len].copy_from_slice(firmware);
+        flash.dfu().write(0, &write_buf).unwrap();
+
+        // On with the test
+        let flash = flash.into_async();
+        let mut aligned = [0; 4];
+        let mut updater = FirmwareUpdater::new(
+            FirmwareUpdaterConfig {
+                dfu: flash.dfu(),
+                state: flash.state(),
+            },
+            &mut aligned,
+        );
+
+        assert!(block_on(updater.verify_and_mark_updated(
+            &public_key.to_bytes(),
+            &signature.to_bytes(),
+            firmware_len as u32,
+        ))
+        .is_ok());
+    }
+}