1 files changed, 114 insertions, 0 deletions

diff --git a/embassy-usb-dfu/src/application.rs b/embassy-usb-dfu/src/application.rs
new file mode 100644
index 000000000..5a52a9fed
--- /dev/null
+++ b/embassy-usb-dfu/src/application.rs
@@ -0,0 +1,114 @@
+
+use embassy_boot::BlockingFirmwareUpdater;
+use embassy_time::{Instant, Duration};
+use embassy_usb::{Handler, control::{RequestType, Recipient, OutResponse, InResponse}, Builder, driver::Driver};
+use embedded_storage::nor_flash::NorFlash;
+
+use crate::consts::{DfuAttributes, Request, Status, State, USB_CLASS_APPN_SPEC, APPN_SPEC_SUBCLASS_DFU, DESC_DFU_FUNCTIONAL, DFU_PROTOCOL_RT};
+
+/// Internal state for the DFU class
+pub struct Control<'d, DFU: NorFlash, STATE: NorFlash> {
+    updater: BlockingFirmwareUpdater<'d, DFU, STATE>,
+    attrs: DfuAttributes,
+    state: State,
+    timeout: Option<Duration>,
+    detach_start: Option<Instant>,
+}
+
+impl<'d, DFU: NorFlash, STATE: NorFlash> Control<'d, DFU, STATE> {
+    pub fn new(updater: BlockingFirmwareUpdater<'d, DFU, STATE>, attrs: DfuAttributes) -> Self {
+        Control { updater, attrs, state: State::AppIdle, detach_start: None, timeout: None }
+    }
+}
+
+impl<'d, DFU: NorFlash, STATE: NorFlash> Handler for Control<'d, DFU, STATE> {
+    fn reset(&mut self) {
+        if let Some(start) = self.detach_start {
+            let delta = Instant::now() - start;
+            let timeout = self.timeout.unwrap();
+            #[cfg(feature = "defmt")]
+            defmt::info!("Received RESET with delta = {}, timeout = {}", delta.as_millis(), timeout.as_millis());
+            if delta < timeout {
+                self.updater.mark_dfu().expect("Failed to mark DFU mode in bootloader");
+                cortex_m::asm::dsb();
+                cortex_m::peripheral::SCB::sys_reset();
+            }
+        }
+    }
+
+    fn control_out(&mut self, req: embassy_usb::control::Request, _: &[u8]) -> Option<embassy_usb::control::OutResponse> {
+        if (req.request_type, req.recipient) != (RequestType::Class, Recipient::Interface) {
+            return None;
+        }
+
+        #[cfg(feature = "defmt")]
+        defmt::info!("Received request {}", req);
+
+        match Request::try_from(req.request) {
+            Ok(Request::Detach) => {
+                #[cfg(feature = "defmt")]
+                defmt::info!("Received DETACH, awaiting USB reset");
+                self.detach_start = Some(Instant::now());
+                self.timeout = Some(Duration::from_millis(req.value as u64));
+                self.state = State::AppDetach;
+                Some(OutResponse::Accepted)
+            }
+            _ => {
+                None
+            }
+        }
+    }
+
+    fn control_in<'a>(&'a mut self, req: embassy_usb::control::Request, buf: &'a mut [u8]) -> Option<embassy_usb::control::InResponse<'a>> {
+        if (req.request_type, req.recipient) != (RequestType::Class, Recipient::Interface) {
+            return None;
+        }
+
+        #[cfg(feature = "defmt")]
+        defmt::info!("Received request {}", req);
+
+        match Request::try_from(req.request) {
+            Ok(Request::GetStatus) => {
+                buf[0..6].copy_from_slice(&[Status::Ok as u8, 0x32, 0x00, 0x00, self.state as u8, 0x00]);
+                Some(InResponse::Accepted(buf))
+            }
+            _ => None
+        }
+    }
+}
+
+/// An implementation of the USB DFU 1.1 runtime protocol
+/// 
+/// This function will add a DFU interface descriptor to the provided Builder, and register the provided Control as a handler for the USB device. The USB builder can be used as normal once this is complete.
+/// The handler is responsive to DFU GetStatus and Detach commands. 
+/// 
+/// Once a detach command, followed by a USB reset is received by the host, a magic number will be written into the bootloader state partition to indicate that
+/// it should expose a DFU device, and a software reset will be issued.
+/// 
+/// To apply USB DFU updates, the bootloader must be capable of recognizing the DFU magic and exposing a device to handle the full DFU transaction with the host.
+pub fn usb_dfu<'d, D: Driver<'d>, DFU: NorFlash, STATE: NorFlash>(builder: &mut Builder<'d, D>, handler: &'d mut Control<'d, DFU, STATE>, timeout: Duration) {
+    #[cfg(feature = "defmt")]
+    defmt::info!("Application USB DFU initializing");
+    let mut func = builder.function(0x00, 0x00, 0x00);
+    let mut iface = func.interface();
+    let mut alt = iface.alt_setting(
+        USB_CLASS_APPN_SPEC,
+        APPN_SPEC_SUBCLASS_DFU,
+        DFU_PROTOCOL_RT,
+        None,
+    );
+    let timeout = timeout.as_millis() as u16;
+    alt.descriptor(
+        DESC_DFU_FUNCTIONAL,
+        &[
+            handler.attrs.bits(),
+            (timeout & 0xff) as u8,
+            ((timeout >> 8) & 0xff) as u8,
+            0x40, 0x00, // 64B control buffer size for application side
+            0x10, 0x01, // DFU 1.1
+        ],
+    );
+
+    drop(func);
+    builder.handler(handler); 
+}
\ No newline at end of file