From cde6e2b58be4910c73b3738bb3b0e264c8ac6049 Mon Sep 17 00:00:00 2001
From: William Yager <will@yager.io>
Date: Sun, 3 Mar 2024 18:50:27 -0500
Subject: [PATCH] ok

---
 embassy-boot-nrf/src/lib.rs | 17 ++++++++++++-----
 embassy-boot-rp/src/lib.rs  | 15 +++++++++++----
 2 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/embassy-boot-nrf/src/lib.rs b/embassy-boot-nrf/src/lib.rs
index 5b20a93c6..6996a92f8 100644
--- a/embassy-boot-nrf/src/lib.rs
+++ b/embassy-boot-nrf/src/lib.rs
@@ -4,8 +4,8 @@
 mod fmt;
 
 pub use embassy_boot::{
-    AlignedBuffer, BlockingFirmwareState, BlockingFirmwareUpdater, BootLoaderConfig, FirmwareState, FirmwareUpdater,
-    FirmwareUpdaterConfig,
+    AlignedBuffer, BlockingFirmwareState, BlockingFirmwareUpdater, BootError, BootLoaderConfig, FirmwareState,
+    FirmwareUpdater, FirmwareUpdaterConfig,
 };
 use embassy_nrf::nvmc::PAGE_SIZE;
 use embassy_nrf::peripherals::WDT;
@@ -16,14 +16,21 @@ use embedded_storage::nor_flash::{ErrorType, NorFlash, ReadNorFlash};
 pub struct BootLoader<const BUFFER_SIZE: usize = PAGE_SIZE>;
 
 impl<const BUFFER_SIZE: usize> BootLoader<BUFFER_SIZE> {
-    /// Inspect the bootloader state and perform actions required before booting, such as swapping firmware.
+    /// Inspect the bootloader state and perform actions required before booting, such as swapping firmware
     pub fn prepare<ACTIVE: NorFlash, DFU: NorFlash, STATE: NorFlash>(
         config: BootLoaderConfig<ACTIVE, DFU, STATE>,
     ) -> Self {
+        Self::try_prepare::<ACTIVE, DFU, STATE>(config).expect("Boot prepare error")
+    }
+
+    /// Inspect the bootloader state and perform actions required before booting, such as swapping firmware
+    pub fn try_prepare<ACTIVE: NorFlash, DFU: NorFlash, STATE: NorFlash>(
+        config: BootLoaderConfig<ACTIVE, DFU, STATE>,
+    ) -> Result<Self, BootError> {
         let mut aligned_buf = AlignedBuffer([0; BUFFER_SIZE]);
         let mut boot = embassy_boot::BootLoader::new(config);
-        boot.prepare_boot(&mut aligned_buf.0).expect("Boot prepare error");
-        Self
+        let state = boot.prepare_boot(aligned_buf.as_mut())?;
+        Ok(Self)
     }
 
     /// Boots the application without softdevice mechanisms.
diff --git a/embassy-boot-rp/src/lib.rs b/embassy-boot-rp/src/lib.rs
index 07a5b3f4d..d88e6dfc6 100644
--- a/embassy-boot-rp/src/lib.rs
+++ b/embassy-boot-rp/src/lib.rs
@@ -4,8 +4,8 @@
 mod fmt;
 
 pub use embassy_boot::{
-    AlignedBuffer, BlockingFirmwareState, BlockingFirmwareUpdater, BootLoaderConfig, FirmwareState, FirmwareUpdater,
-    FirmwareUpdaterConfig, State,
+    AlignedBuffer, BlockingFirmwareState, BlockingFirmwareUpdater, BootError, BootLoaderConfig, FirmwareState,
+    FirmwareUpdater, FirmwareUpdaterConfig, State,
 };
 use embassy_rp::flash::{Blocking, Flash, ERASE_SIZE};
 use embassy_rp::peripherals::{FLASH, WATCHDOG};
@@ -21,10 +21,17 @@ impl<const BUFFER_SIZE: usize> BootLoader<BUFFER_SIZE> {
     pub fn prepare<ACTIVE: NorFlash, DFU: NorFlash, STATE: NorFlash>(
         config: BootLoaderConfig<ACTIVE, DFU, STATE>,
     ) -> Self {
+        Self::try_prepare::<ACTIVE, DFU, STATE>(config).expect("Boot prepare error")
+    }
+
+    /// Inspect the bootloader state and perform actions required before booting, such as swapping firmware
+    pub fn try_prepare<ACTIVE: NorFlash, DFU: NorFlash, STATE: NorFlash>(
+        config: BootLoaderConfig<ACTIVE, DFU, STATE>,
+    ) -> Result<Self, BootError> {
         let mut aligned_buf = AlignedBuffer([0; BUFFER_SIZE]);
         let mut boot = embassy_boot::BootLoader::new(config);
-        boot.prepare_boot(aligned_buf.as_mut()).expect("Boot prepare error");
-        Self
+        let state = boot.prepare_boot(aligned_buf.as_mut())?;
+        Ok(Self { state })
     }
 
     /// Boots the application.