Skip to main content

embassy_usb/class/dfu/
dfu_mode.rs

1use embassy_usb_driver::Driver;
2
3use super::consts::{
4    APPN_SPEC_SUBCLASS_DFU, DESC_DFU_FUNCTIONAL, DFU_PROTOCOL_DFU, DfuAttributes, Request, State, Status,
5    USB_CLASS_APPN_SPEC,
6};
7use crate::control::{InResponse, OutResponse, Recipient, Request as ControlRequest, RequestType};
8use crate::{Builder, FunctionBuilder};
9
10/// Handler trait for DFU bootloader mode.
11///
12/// Implement this trait to handle firmware download operations.
13pub trait Handler {
14    /// Called when a firmware download starts.
15    ///
16    /// This is called when the first DFU_DNLOAD request is received.
17    ///
18    /// Returns `Ok(())` on success, or a `Status` error on failure.
19    fn start(&mut self) -> Result<(), Status>;
20
21    /// Called to write a chunk of firmware data.
22    ///
23    /// Returns `Ok(())` on success, or a `Status` error on failure.
24    fn write(&mut self, data: &[u8]) -> Result<(), Status>;
25
26    /// Called when the firmware download is complete.
27    ///
28    /// This is called when a zero-length DFU_DNLOAD is received, indicating
29    /// the end of the firmware data. This is where you would typically
30    /// mark the firmware as ready to boot.
31    ///
32    /// Returns `Ok(())` on success, or a `Status` error on failure.
33    fn finish(&mut self) -> Result<(), Status>;
34
35    /// Called at the end of the DFU procedure.
36    ///
37    /// This is typically where you would perform a system reset to boot
38    /// the new firmware after a successful download.
39    fn system_reset(&mut self);
40}
41
42/// Internal state for USB DFU
43pub struct DfuState<H: Handler> {
44    handler: H,
45    attrs: DfuAttributes,
46    state: State,
47    status: Status,
48    next_block_num: usize,
49}
50
51impl<'d, H: Handler> DfuState<H> {
52    /// Create a new DFU instance to handle DFU transfers.
53    pub fn new(handler: H, attrs: DfuAttributes) -> Self {
54        Self {
55            handler,
56            attrs,
57            state: State::DfuIdle,
58            status: Status::Ok,
59            next_block_num: 0,
60        }
61    }
62
63    /// Set DFU instance into firmware error
64    ///
65    /// This is typically called when the stored firmware
66    /// is corrupt and can not be booted.
67    pub fn set_to_firmware_error(&mut self) {
68        self.reset_state();
69        self.state = State::Error;
70        self.status = Status::ErrFirmware;
71    }
72
73    fn reset_state(&mut self) {
74        self.next_block_num = 0;
75        self.state = State::DfuIdle;
76        self.status = Status::Ok;
77    }
78}
79
80impl<H: Handler> crate::Handler for DfuState<H> {
81    fn reset(&mut self) {
82        if matches!(self.state, State::ManifestSync | State::Manifest) {
83            self.handler.system_reset();
84        }
85    }
86
87    fn control_out(&mut self, req: ControlRequest, data: &[u8]) -> Option<OutResponse> {
88        if (req.request_type, req.recipient) != (RequestType::Class, Recipient::Interface) {
89            return None;
90        }
91        match Request::try_from(req.request) {
92            Ok(Request::Abort) => {
93                info!("Abort requested");
94                self.reset_state();
95                Some(OutResponse::Accepted)
96            }
97            Ok(Request::Dnload) if self.attrs.contains(DfuAttributes::CAN_DOWNLOAD) => {
98                if req.value as usize != self.next_block_num {
99                    error!("expected next block num {}, got {}", self.next_block_num, req.value);
100                    self.state = State::Error;
101                    self.status = Status::ErrUnknown;
102                    return Some(OutResponse::Rejected);
103                }
104
105                if req.value == 0 {
106                    match self.handler.start() {
107                        Ok(_) => {
108                            self.state = State::Download;
109                        }
110                        Err(e) => {
111                            self.state = State::Error;
112                            self.status = e;
113                            return Some(OutResponse::Rejected);
114                        }
115                    }
116                }
117
118                if req.length == 0 {
119                    match self.handler.finish() {
120                        Ok(_) => {
121                            self.status = Status::Ok;
122                            self.state = State::ManifestSync;
123                        }
124                        Err(e) => {
125                            self.state = State::Error;
126                            self.status = e;
127                        }
128                    }
129                } else {
130                    if self.state != State::Download {
131                        error!("Unexpected DNLOAD while chip is waiting for a GETSTATUS");
132                        self.status = Status::ErrUnknown;
133                        self.state = State::Error;
134                        return Some(OutResponse::Rejected);
135                    }
136                    match self.handler.write(data) {
137                        Ok(_) => {
138                            self.status = Status::Ok;
139                            self.state = State::DlSync;
140                            self.next_block_num += 1;
141                        }
142                        Err(e) => {
143                            self.state = State::Error;
144                            self.status = e;
145                        }
146                    }
147                }
148
149                Some(OutResponse::Accepted)
150            }
151            Ok(Request::Detach) => Some(OutResponse::Accepted), // Device is already in DFU mode
152            Ok(Request::ClrStatus) => {
153                info!("Clear status requested");
154                self.reset_state();
155                Some(OutResponse::Accepted)
156            }
157            _ => {
158                debug!("Unknown OUT request {:?}", req);
159                None
160            }
161        }
162    }
163
164    fn control_in<'a>(&'a mut self, req: ControlRequest, buf: &'a mut [u8]) -> Option<InResponse<'a>> {
165        if (req.request_type, req.recipient) != (RequestType::Class, Recipient::Interface) {
166            return None;
167        }
168        match Request::try_from(req.request) {
169            Ok(Request::GetStatus) => {
170                match self.state {
171                    State::DlSync => self.state = State::Download,
172                    State::ManifestSync if self.attrs.contains(DfuAttributes::MANIFESTATION_TOLERANT) => {
173                        self.state = State::DfuIdle
174                    }
175                    State::ManifestSync => {
176                        // Technically we would be in ManifestWaitReset after responding with
177                        // Manifest, but ManifestWaitReset isn't meant to be seen by the host
178                        // anyways.
179                        self.state = State::Manifest;
180                        if self.attrs.contains(DfuAttributes::WILL_DETACH) {
181                            self.reset();
182                        }
183                    }
184                    _ => {}
185                }
186                //TODO: Configurable poll timeout, ability to add string for Vendor error
187                buf[0..6].copy_from_slice(&[self.status as u8, 0x32, 0x00, 0x00, self.state as u8, 0x00]);
188                Some(InResponse::Accepted(&buf[0..6]))
189            }
190            Ok(Request::GetState) => {
191                buf[0] = self.state as u8;
192                Some(InResponse::Accepted(&buf[0..1]))
193            }
194            Ok(Request::Upload) if self.attrs.contains(DfuAttributes::CAN_UPLOAD) => {
195                //TODO: FirmwareUpdater does not provide a way of reading the active partition, can't upload.
196                Some(InResponse::Rejected)
197            }
198            _ => {
199                debug!("Unknown IN request {:?}", req);
200                None
201            }
202        }
203    }
204}
205
206/// An implementation of the USB DFU 1.1 protocol
207///
208/// This function will add a DFU interface descriptor to the provided Builder, and register the provided Control as a handler for the USB device
209/// The handler is responsive to DFU GetState, GetStatus, Abort, and ClrStatus commands, as well as Download if configured by the user.
210///
211/// Once the host has initiated a DFU download operation, the chunks sent by the host will be written to the DFU partition.
212/// Once the final sync in the manifestation phase has been received, the handler will trigger a system reset to swap the new firmware.
213pub fn usb_dfu<'d, D: Driver<'d>, H: Handler>(
214    builder: &mut Builder<'d, D>,
215    state: &'d mut DfuState<H>,
216    max_write_size: usize,
217    func_modifier: impl Fn(&mut FunctionBuilder<'_, 'd, D>),
218) {
219    let mut func = builder.function(0x00, 0x00, 0x00);
220
221    // Here we give users the opportunity to add their own function level MSOS headers for instance.
222    // This is useful when DFU functionality is part of a composite USB device.
223    func_modifier(&mut func);
224
225    let mut iface = func.interface();
226    let mut alt = iface.alt_setting(USB_CLASS_APPN_SPEC, APPN_SPEC_SUBCLASS_DFU, DFU_PROTOCOL_DFU, None);
227    alt.descriptor(
228        DESC_DFU_FUNCTIONAL,
229        &[
230            state.attrs.bits(),
231            0xc4,
232            0x09, // 2500ms timeout, doesn't affect operation as DETACH not necessary in bootloader code
233            (max_write_size & 0xff) as u8,
234            ((max_write_size & 0xff00) >> 8) as u8,
235            0x10,
236            0x01, // DFU 1.1
237        ],
238    );
239
240    drop(func);
241    builder.handler(state);
242}