5 files changed, 64 insertions, 4 deletions

diff --git a/embassy-boot/rp/Cargo.toml b/embassy-boot/rp/Cargo.toml
index 93099b233..ffc36a4e0 100644
--- a/embassy-boot/rp/Cargo.toml
+++ b/embassy-boot/rp/Cargo.toml
@@ -20,6 +20,8 @@ log = { version = "0.4", optional = true }
 embassy-sync = { path = "../../embassy-sync" }
 embassy-rp = { path = "../../embassy-rp", default-features = false, features = ["nightly"] }
 embassy-boot = { path = "../boot", default-features = false }
+embassy-time = { path = "../../embassy-time", features = ["nightly"] }
+
 cortex-m = { version = "0.7.6" }
 cortex-m-rt = { version = "0.7" }
 embedded-storage = "0.3.0"
diff --git a/embassy-boot/rp/src/lib.rs b/embassy-boot/rp/src/lib.rs
index 85fc81827..6df34133e 100644
--- a/embassy-boot/rp/src/lib.rs
+++ b/embassy-boot/rp/src/lib.rs
@@ -5,7 +5,11 @@
 mod fmt;
 
 pub use embassy_boot::{AlignedBuffer, BootFlash, FirmwareUpdater, FlashConfig, Partition, SingleFlashConfig, State};
-use embassy_rp::flash::{ERASE_SIZE, WRITE_SIZE};
+use embassy_rp::flash::{Flash, ERASE_SIZE, WRITE_SIZE};
+use embassy_rp::peripherals::{FLASH, WATCHDOG};
+use embassy_rp::watchdog::Watchdog;
+use embassy_time::Duration;
+use embedded_storage::nor_flash::{ErrorType, NorFlash, ReadNorFlash};
 
 /// A bootloader for RP2040 devices.
 pub struct BootLoader {
@@ -88,3 +92,48 @@ impl Default for BootLoader {
         Self::new(active, dfu, state)
     }
 }
+
+/// A flash implementation that will feed a watchdog when touching flash.
+pub struct WatchdogFlash<'d, const SIZE: usize> {
+    flash: Flash<'d, FLASH, SIZE>,
+    watchdog: Watchdog,
+}
+
+impl<'d, const SIZE: usize> WatchdogFlash<'d, SIZE> {
+    /// Start a new watchdog with a given flash and watchdog peripheral and a timeout
+    pub fn start(flash: FLASH, watchdog: WATCHDOG, timeout: Duration) -> Self {
+        let flash: Flash<'_, FLASH, SIZE> = Flash::new(flash);
+        let mut watchdog = Watchdog::new(watchdog);
+        watchdog.start(timeout);
+        Self { flash, watchdog }
+    }
+}
+
+impl<'d, const SIZE: usize> ErrorType for WatchdogFlash<'d, SIZE> {
+    type Error = <Flash<'d, FLASH, SIZE> as ErrorType>::Error;
+}
+
+impl<'d, const SIZE: usize> NorFlash for WatchdogFlash<'d, SIZE> {
+    const WRITE_SIZE: usize = <Flash<'d, FLASH, SIZE> as NorFlash>::WRITE_SIZE;
+    const ERASE_SIZE: usize = <Flash<'d, FLASH, SIZE> as NorFlash>::ERASE_SIZE;
+
+    fn erase(&mut self, from: u32, to: u32) -> Result<(), Self::Error> {
+        self.watchdog.feed();
+        self.flash.erase(from, to)
+    }
+    fn write(&mut self, offset: u32, data: &[u8]) -> Result<(), Self::Error> {
+        self.watchdog.feed();
+        self.flash.write(offset, data)
+    }
+}
+
+impl<'d, const SIZE: usize> ReadNorFlash for WatchdogFlash<'d, SIZE> {
+    const READ_SIZE: usize = <Flash<'d, FLASH, SIZE> as ReadNorFlash>::READ_SIZE;
+    fn read(&mut self, offset: u32, data: &mut [u8]) -> Result<(), Self::Error> {
+        self.watchdog.feed();
+        self.flash.read(offset, data)
+    }
+    fn capacity(&self) -> usize {
+        self.flash.capacity()
+    }
+}
diff --git a/examples/boot/application/rp/src/bin/a.rs b/examples/boot/application/rp/src/bin/a.rs
index 3736c9141..e3ac634c2 100644
--- a/examples/boot/application/rp/src/bin/a.rs
+++ b/examples/boot/application/rp/src/bin/a.rs
@@ -7,6 +7,7 @@ use embassy_boot_rp::*;
 use embassy_executor::Spawner;
 use embassy_rp::flash::Flash;
 use embassy_rp::gpio::{Level, Output};
+use embassy_rp::watchdog::Watchdog;
 use embassy_time::{Duration, Timer};
 #[cfg(feature = "panic-probe")]
 use panic_probe as _;
@@ -21,11 +22,16 @@ async fn main(_s: Spawner) {
     let p = embassy_rp::init(Default::default());
     let mut led = Output::new(p.PIN_25, Level::Low);
 
+    // Override bootloader watchdog
+    let mut watchdog = Watchdog::new(p.WATCHDOG);
+    watchdog.start(Duration::from_secs(8));
+
     let mut flash: Flash<_, FLASH_SIZE> = Flash::new(p.FLASH);
 
     let mut updater = FirmwareUpdater::default();
 
     Timer::after(Duration::from_secs(5)).await;
+    watchdog.feed();
     led.set_high();
     let mut offset = 0;
     let mut buf: AlignedBuffer<4096> = AlignedBuffer([0; 4096]);
@@ -43,6 +49,7 @@ async fn main(_s: Spawner) {
             .unwrap();
         offset += chunk.len();
     }
+    watchdog.feed();
     defmt::info!("firmware written, marking update");
     updater.mark_updated_blocking(&mut flash, &mut buf.0[..1]).unwrap();
     Timer::after(Duration::from_secs(2)).await;
diff --git a/examples/boot/bootloader/rp/Cargo.toml b/examples/boot/bootloader/rp/Cargo.toml
index 580ced22e..c0b576cff 100644
--- a/examples/boot/bootloader/rp/Cargo.toml
+++ b/examples/boot/bootloader/rp/Cargo.toml
@@ -11,6 +11,8 @@ defmt-rtt = { version = "0.4", optional = true }
 
 embassy-rp = { path = "../../../../embassy-rp", default-features = false, features = ["nightly"] }
 embassy-boot-rp = { path = "../../../../embassy-boot/rp", default-features = false }
+embassy-time = { path = "../../../../embassy-time", features = ["nightly"] }
+
 cortex-m = { version = "0.7.6", features = ["critical-section-single-core"] }
 cortex-m-rt = { version = "0.7" }
 embedded-storage = "0.3.0"
diff --git a/examples/boot/bootloader/rp/src/main.rs b/examples/boot/bootloader/rp/src/main.rs
index 5028ec688..fb7f0522b 100644
--- a/examples/boot/bootloader/rp/src/main.rs
+++ b/examples/boot/bootloader/rp/src/main.rs
@@ -5,8 +5,8 @@ use cortex_m_rt::{entry, exception};
 #[cfg(feature = "defmt")]
 use defmt_rtt as _;
 use embassy_boot_rp::*;
-use embassy_rp::flash::{Flash, ERASE_SIZE};
-use embassy_rp::peripherals::FLASH;
+use embassy_rp::flash::ERASE_SIZE;
+use embassy_time::Duration;
 
 const FLASH_SIZE: usize = 2 * 1024 * 1024;
 
@@ -23,7 +23,7 @@ fn main() -> ! {
     */
 
     let mut bl: BootLoader = BootLoader::default();
-    let flash: Flash<'_, FLASH, FLASH_SIZE> = Flash::new(p.FLASH);
+    let flash = WatchdogFlash::<FLASH_SIZE>::start(p.FLASH, p.WATCHDOG, Duration::from_secs(8));
     let mut flash = BootFlash::<_, ERASE_SIZE>::new(flash);
     let start = bl.prepare(&mut SingleFlashConfig::new(&mut flash));
     core::mem::drop(flash);