Add embassy-usb-dfu

5 files changed, 453 insertions, 0 deletions

diff --git a/embassy-usb-dfu/Cargo.toml b/embassy-usb-dfu/Cargo.toml
new file mode 100644
index 000000000..62398afbc
--- /dev/null
+++ b/embassy-usb-dfu/Cargo.toml
@@ -0,0 +1,31 @@
+[package]
+edition = "2021"
+name = "embassy-usb-dfu"
+version = "0.1.0"
+description = "An implementation of the USB DFU 1.1 protocol, using embassy-boot"
+license = "MIT OR Apache-2.0"
+repository = "https://github.com/embassy-rs/embassy"
+categories = [
+    "embedded",
+    "no-std",
+    "asynchronous"
+]
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+bitflags = "2.4.1"
+cortex-m = { version = "0.7.7", features = ["inline-asm"] }
+defmt = { version = "0.3.5", optional = true }
+embassy-boot = { version = "0.1.1", path = "../embassy-boot/boot" }
+embassy-embedded-hal = { version = "0.1.0", path = "../embassy-embedded-hal" }
+embassy-futures = { version = "0.1.1", path = "../embassy-futures" }
+embassy-sync = { version = "0.5.0", path = "../embassy-sync" }
+embassy-time = { version = "0.2.0", path = "../embassy-time" }
+embassy-usb = { version = "0.1.0", path = "../embassy-usb", default-features = false }
+embedded-storage = { version = "0.3.1" }
+
+[features]
+bootloader = []
+application = []
+defmt = ["dep:defmt"]
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
diff --git a/embassy-usb-dfu/src/bootloader.rs b/embassy-usb-dfu/src/bootloader.rs
new file mode 100644
index 000000000..7bcb0b258
--- /dev/null
+++ b/embassy-usb-dfu/src/bootloader.rs
@@ -0,0 +1,196 @@
+use embassy_boot::BlockingFirmwareUpdater;
+use embassy_usb::{
+    control::{InResponse, OutResponse, Recipient, RequestType},
+    driver::Driver,
+    Builder, Handler,
+};
+use embedded_storage::nor_flash::{NorFlashErrorKind, NorFlash};
+
+use crate::consts::{DfuAttributes, Request, State, Status, USB_CLASS_APPN_SPEC, APPN_SPEC_SUBCLASS_DFU, DFU_PROTOCOL_DFU, DESC_DFU_FUNCTIONAL};
+
+/// Internal state for USB DFU
+pub struct Control<'d, DFU: NorFlash, STATE: NorFlash, const BLOCK_SIZE: usize> {
+    updater: BlockingFirmwareUpdater<'d, DFU, STATE>,
+    attrs: DfuAttributes,
+    state: State,
+    status: Status,
+    offset: usize,
+}
+
+impl<'d, DFU: NorFlash, STATE: NorFlash, const BLOCK_SIZE: usize> Control<'d, DFU, STATE, BLOCK_SIZE> {
+    pub fn new(updater: BlockingFirmwareUpdater<'d, DFU, STATE>, attrs: DfuAttributes) -> Self {
+        Self {
+            updater,
+            attrs,
+            state: State::DfuIdle,
+            status: Status::Ok,
+            offset: 0,
+        }
+    }
+
+    fn reset_state(&mut self) {
+        self.offset = 0;
+        self.state = State::DfuIdle;
+        self.status = Status::Ok;
+    }
+}
+
+impl<'d, DFU: NorFlash, STATE: NorFlash, const BLOCK_SIZE: usize> Handler for Control<'d, DFU, STATE, BLOCK_SIZE> {
+    fn control_out(
+        &mut self,
+        req: embassy_usb::control::Request,
+        data: &[u8],
+    ) -> Option<embassy_usb::control::OutResponse> {
+        if (req.request_type, req.recipient) != (RequestType::Class, Recipient::Interface) {
+            return None;
+        }
+        match Request::try_from(req.request) {
+            Ok(Request::Abort) => {
+                self.reset_state();
+                Some(OutResponse::Accepted)
+            }
+            Ok(Request::Dnload) if self.attrs.contains(DfuAttributes::CAN_DOWNLOAD) => {
+                if req.value == 0 {
+                    self.state = State::Download;
+                    self.offset = 0;
+                }
+
+                let mut buf = [0; BLOCK_SIZE];
+                buf[..data.len()].copy_from_slice(data);
+
+                if req.length == 0 {
+                    match self.updater.mark_updated() {
+                        Ok(_) => {
+                            self.status = Status::Ok;
+                            self.state = State::ManifestSync;
+                        }
+                        Err(e) => {
+                            self.state = State::Error;
+                            match e {
+                                embassy_boot::FirmwareUpdaterError::Flash(e) => match e {
+                                    NorFlashErrorKind::NotAligned => self.status = Status::ErrWrite,
+                                    NorFlashErrorKind::OutOfBounds => {
+                                        self.status = Status::ErrAddress
+                                    }
+                                    _ => self.status = Status::ErrUnknown,
+                                },
+                                embassy_boot::FirmwareUpdaterError::Signature(_) => {
+                                    self.status = Status::ErrVerify
+                                }
+                                embassy_boot::FirmwareUpdaterError::BadState => {
+                                    self.status = Status::ErrUnknown
+                                }
+                            }
+                        }
+                    }
+                } else {
+                    if self.state != State::Download {
+                        // Unexpected DNLOAD while chip is waiting for a GETSTATUS
+                        self.status = Status::ErrUnknown;
+                        self.state = State::Error;
+                        return Some(OutResponse::Rejected);
+                    }
+                    match self.updater.write_firmware(self.offset, &buf[..]) {
+                        Ok(_) => {
+                            self.status = Status::Ok;
+                            self.state = State::DlSync;
+                            self.offset += data.len();
+                        }
+                        Err(e) => {
+                            self.state = State::Error;
+                            match e {
+                                embassy_boot::FirmwareUpdaterError::Flash(e) => match e {
+                                    NorFlashErrorKind::NotAligned => self.status = Status::ErrWrite,
+                                    NorFlashErrorKind::OutOfBounds => {
+                                        self.status = Status::ErrAddress
+                                    }
+                                    _ => self.status = Status::ErrUnknown,
+                                },
+                                embassy_boot::FirmwareUpdaterError::Signature(_) => {
+                                    self.status = Status::ErrVerify
+                                }
+                                embassy_boot::FirmwareUpdaterError::BadState => {
+                                    self.status = Status::ErrUnknown
+                                }
+                            }
+                        }
+                    }
+                }
+
+                Some(OutResponse::Accepted)
+            }
+            Ok(Request::Detach) => Some(OutResponse::Accepted), // Device is already in DFU mode
+            Ok(Request::ClrStatus) => {
+                self.reset_state();
+                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;
+        }
+        match Request::try_from(req.request) {
+            Ok(Request::GetStatus) => {
+                //TODO: Configurable poll timeout, ability to add string for Vendor error
+                buf[0..6].copy_from_slice(&[self.status as u8, 0x32, 0x00, 0x00, self.state as u8, 0x00]);
+                match self.state {
+                    State::DlSync => self.state = State::Download,
+                    State::ManifestSync => cortex_m::peripheral::SCB::sys_reset(),
+                    _ => {}
+                }
+
+                Some(InResponse::Accepted(&buf[0..6]))
+            }
+            Ok(Request::GetState) => {
+                buf[0] = self.state as u8;
+                Some(InResponse::Accepted(&buf[0..1]))
+            }
+            Ok(Request::Upload) if self.attrs.contains(DfuAttributes::CAN_UPLOAD) => {
+                //TODO: FirmwareUpdater does not provide a way of reading the active partition, can't upload.
+                Some(InResponse::Rejected)
+            }
+            _ => None,
+        }
+    }
+}
+
+/// An implementation of the USB DFU 1.1 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 handler is responsive to DFU GetState, GetStatus, Abort, and ClrStatus commands, as well as Download if configured by the user.
+/// 
+/// Once the host has initiated a DFU download operation, the chunks sent by the host will be written to the DFU partition.
+/// Once the final sync in the manifestation phase has been received, the handler will trigger a system reset to swap the new firmware.
+pub fn usb_dfu<'d, D: Driver<'d>, DFU: NorFlash, STATE: NorFlash, const BLOCK_SIZE: usize>(
+    builder: &mut Builder<'d, D>,
+    handler: &'d mut Control<'d, DFU, STATE, BLOCK_SIZE>,
+) {
+    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_DFU,
+        None,
+    );
+    alt.descriptor(
+        DESC_DFU_FUNCTIONAL,
+        &[
+            handler.attrs.bits(),
+            0xc4, 0x09, // 2500ms timeout, doesn't affect operation as DETACH not necessary in bootloader code
+            (BLOCK_SIZE & 0xff) as u8,
+            ((BLOCK_SIZE & 0xff00) >> 8) as u8,
+            0x10, 0x01, // DFU 1.1
+        ],
+    );
+
+    drop(func);
+    builder.handler(handler);
+}
diff --git a/embassy-usb-dfu/src/consts.rs b/embassy-usb-dfu/src/consts.rs
new file mode 100644
index 000000000..b083af9de
--- /dev/null
+++ b/embassy-usb-dfu/src/consts.rs
@@ -0,0 +1,96 @@
+
+pub(crate) const USB_CLASS_APPN_SPEC: u8 = 0xFE;
+pub(crate) const APPN_SPEC_SUBCLASS_DFU: u8 = 0x01;
+#[allow(unused)]
+pub(crate) const DFU_PROTOCOL_DFU: u8 = 0x02;
+#[allow(unused)]
+pub(crate) const DFU_PROTOCOL_RT: u8 = 0x01;
+pub(crate) const DESC_DFU_FUNCTIONAL: u8 = 0x21;
+
+#[cfg(feature = "defmt")]
+defmt::bitflags! {
+    pub struct DfuAttributes: u8 {
+        const WILL_DETACH = 0b0000_1000;
+        const MANIFESTATION_TOLERANT = 0b0000_0100;
+        const CAN_UPLOAD = 0b0000_0010;
+        const CAN_DOWNLOAD = 0b0000_0001;
+    }
+}
+
+#[cfg(not(feature = "defmt"))]
+bitflags::bitflags! {
+    pub struct DfuAttributes: u8 {
+        const WILL_DETACH = 0b0000_1000;
+        const MANIFESTATION_TOLERANT = 0b0000_0100;
+        const CAN_UPLOAD = 0b0000_0010;
+        const CAN_DOWNLOAD = 0b0000_0001;
+    }
+}
+
+#[derive(Copy, Clone, PartialEq, Eq)]
+#[repr(u8)]
+#[allow(unused)]
+pub enum State {
+    AppIdle = 0,
+    AppDetach = 1,
+    DfuIdle = 2,
+    DlSync = 3,
+    DlBusy = 4,
+    Download = 5,
+    ManifestSync = 6,
+    Manifest = 7,
+    ManifestWaitReset = 8,
+    UploadIdle = 9,
+    Error = 10,
+}
+
+#[derive(Copy, Clone, PartialEq, Eq)]
+#[repr(u8)]
+#[allow(unused)]
+pub enum Status {
+    Ok = 0x00,
+    ErrTarget = 0x01,
+    ErrFile = 0x02,
+    ErrWrite = 0x03,
+    ErrErase = 0x04,
+    ErrCheckErased = 0x05,
+    ErrProg = 0x06,
+    ErrVerify = 0x07,
+    ErrAddress = 0x08,
+    ErrNotDone = 0x09,
+    ErrFirmware = 0x0A,
+    ErrVendor = 0x0B,
+    ErrUsbr = 0x0C,
+    ErrPor = 0x0D,
+    ErrUnknown = 0x0E,
+    ErrStalledPkt = 0x0F,
+}
+
+#[derive(Copy, Clone, PartialEq, Eq)]
+#[repr(u8)]
+pub enum Request {
+    Detach = 0,
+    Dnload = 1,
+    Upload = 2,
+    GetStatus = 3,
+    ClrStatus = 4,
+    GetState = 5,
+    Abort = 6,
+}
+
+impl TryFrom<u8> for Request {
+    type Error = ();
+
+    fn try_from(value: u8) -> Result<Self, Self::Error> {
+        match value {
+            0 => Ok(Request::Detach),
+            1 => Ok(Request::Dnload),
+            2 => Ok(Request::Upload),
+            3 => Ok(Request::GetStatus),
+            4 => Ok(Request::ClrStatus),
+            5 => Ok(Request::GetState),
+            6 => Ok(Request::Abort),
+            _ => Err(()),
+        }
+    }
+}
diff --git a/embassy-usb-dfu/src/lib.rs b/embassy-usb-dfu/src/lib.rs
new file mode 100644
index 000000000..dcdb11b2a
--- /dev/null
+++ b/embassy-usb-dfu/src/lib.rs
@@ -0,0 +1,16 @@
+#![no_std]
+
+pub mod consts;
+
+#[cfg(feature = "bootloader")]
+mod bootloader;
+#[cfg(feature = "bootloader")]
+pub use self::bootloader::*;
+
+#[cfg(feature = "application")]
+mod application;
+#[cfg(feature = "application")]
+pub use self::application::*;
+
+#[cfg(any(all(feature = "bootloader", feature = "application"), not(any(feature = "bootloader", feature = "application"))))]
+compile_error!("usb-dfu must be compiled with exactly one of `bootloader`, or `application` features");