fastboot_protocol/
nusb.rs

1use std::{collections::HashMap, fmt::Display, io::Write};
2
3use nusb::transfer::RequestBuffer;
4pub use nusb::{transfer::TransferError, Device, DeviceInfo, Interface};
5use thiserror::Error;
6use tracing::{info, warn};
7use tracing::{instrument, trace};
8
9use crate::protocol::FastBootResponse;
10use crate::protocol::{FastBootCommand, FastBootResponseParseError};
11
12/// List fastboot devices
13pub fn devices() -> Result<impl Iterator<Item = DeviceInfo>, nusb::Error> {
14    Ok(nusb::list_devices()?.filter(|d| NusbFastBoot::find_fastboot_interface(d).is_some()))
15}
16
17/// Fastboot communication errors
18#[derive(Debug, Error)]
19pub enum NusbFastBootError {
20    #[error("Transfer error: {0}")]
21    Transfer(#[from] nusb::transfer::TransferError),
22    #[error("Fastboot client failure: {0}")]
23    FastbootFailed(String),
24    #[error("Unexpected fastboot response")]
25    FastbootUnexpectedReply,
26    #[error("Unknown fastboot response: {0}")]
27    FastbootParseError(#[from] FastBootResponseParseError),
28}
29
30/// Errors when opening the fastboot device
31#[derive(Debug, Error)]
32pub enum NusbFastBootOpenError {
33    #[error("Failed to open device: {0}")]
34    Device(std::io::Error),
35    #[error("Failed to claim interface: {0}")]
36    Interface(std::io::Error),
37    #[error("Failed to find interface for fastboot")]
38    MissingInterface,
39    #[error("Failed to find required endpoints for fastboot")]
40    MissingEndpoints,
41    #[error("Unknown fastboot response: {0}")]
42    FastbootParseError(#[from] FastBootResponseParseError),
43}
44
45/// Nusb fastboot client
46pub struct NusbFastBoot {
47    interface: nusb::Interface,
48    ep_out: u8,
49    max_out: usize,
50    ep_in: u8,
51    max_in: usize,
52}
53
54impl NusbFastBoot {
55    /// Find fastboot interface within a USB device
56    pub fn find_fastboot_interface(info: &DeviceInfo) -> Option<u8> {
57        info.interfaces().find_map(|i| {
58            if i.class() == 0xff && i.subclass() == 0x42 && i.protocol() == 0x3 {
59                Some(i.interface_number())
60            } else {
61                None
62            }
63        })
64    }
65
66    /// Create a fastboot client based on a USB interface. Interface is assumed to be a fastboot
67    /// interface
68    #[tracing::instrument(skip_all, err)]
69    pub fn from_interface(interface: nusb::Interface) -> Result<Self, NusbFastBootOpenError> {
70        let (ep_out, max_out, ep_in, max_in) = interface
71            .descriptors()
72            .find_map(|alt| {
73                // Requires one bulk IN and one bulk OUT
74                let (ep_out, max_out) = alt.endpoints().find_map(|end| {
75                    if end.transfer_type() == nusb::transfer::EndpointType::Bulk
76                        && end.direction() == nusb::transfer::Direction::Out
77                    {
78                        Some((end.address(), end.max_packet_size()))
79                    } else {
80                        None
81                    }
82                })?;
83                let (ep_in, max_in) = alt.endpoints().find_map(|end| {
84                    if end.transfer_type() == nusb::transfer::EndpointType::Bulk
85                        && end.direction() == nusb::transfer::Direction::In
86                    {
87                        Some((end.address(), end.max_packet_size()))
88                    } else {
89                        None
90                    }
91                })?;
92                Some((ep_out, max_out, ep_in, max_in))
93            })
94            .ok_or(NusbFastBootOpenError::MissingEndpoints)?;
95        trace!(
96            "Fastboot endpoints: OUT: {} (max: {}), IN: {} (max: {})",
97            ep_out,
98            max_out,
99            ep_in,
100            max_in
101        );
102        Ok(Self {
103            interface,
104            ep_out,
105            max_out,
106            ep_in,
107            max_in,
108        })
109    }
110
111    /// Create a fastboot client based on a USB device. Interface number must be the fastboot
112    /// interface
113    #[tracing::instrument(skip_all, err)]
114    pub fn from_device(device: Device, interface: u8) -> Result<Self, NusbFastBootOpenError> {
115        let interface = device
116            .claim_interface(interface)
117            .map_err(NusbFastBootOpenError::Interface)?;
118        Self::from_interface(interface)
119    }
120
121    /// Create a fastboot client based on device info. The correct interface will automatically be
122    /// determined
123    #[tracing::instrument(skip_all, err)]
124    pub fn from_info(info: &DeviceInfo) -> Result<Self, NusbFastBootOpenError> {
125        let interface =
126            Self::find_fastboot_interface(info).ok_or(NusbFastBootOpenError::MissingInterface)?;
127        let device = info.open().map_err(NusbFastBootOpenError::Device)?;
128        Self::from_device(device, interface)
129    }
130
131    #[tracing::instrument(skip_all, err)]
132    async fn send_data(&mut self, data: Vec<u8>) -> Result<(), NusbFastBootError> {
133        self.interface.bulk_out(self.ep_out, data).await.status?;
134        Ok(())
135    }
136
137    async fn send_command<S: Display>(
138        &mut self,
139        cmd: FastBootCommand<S>,
140    ) -> Result<(), NusbFastBootError> {
141        let mut out = vec![];
142        // Only fails if memory allocation fails
143        out.write_fmt(format_args!("{}", cmd)).unwrap();
144        trace!(
145            "Sending command: {}",
146            std::str::from_utf8(&out).unwrap_or("Invalid utf-8")
147        );
148        self.send_data(out).await
149    }
150
151    #[tracing::instrument(skip_all, err)]
152    async fn read_response(&mut self) -> Result<FastBootResponse, FastBootResponseParseError> {
153        let req = RequestBuffer::new(self.max_in);
154        let resp = self.interface.bulk_in(self.ep_in, req).await;
155        FastBootResponse::from_bytes(&resp.data)
156    }
157
158    #[tracing::instrument(skip_all, err)]
159    async fn handle_responses(&mut self) -> Result<String, NusbFastBootError> {
160        loop {
161            let resp = self.read_response().await?;
162            trace!("Response: {:?}", resp);
163            match resp {
164                FastBootResponse::Info(_) => (),
165                FastBootResponse::Text(_) => (),
166                FastBootResponse::Data(_) => {
167                    return Err(NusbFastBootError::FastbootUnexpectedReply)
168                }
169                FastBootResponse::Okay(value) => return Ok(value),
170                FastBootResponse::Fail(fail) => {
171                    return Err(NusbFastBootError::FastbootFailed(fail))
172                }
173            }
174        }
175    }
176
177    #[tracing::instrument(skip_all, err)]
178    async fn execute<S: Display>(
179        &mut self,
180        cmd: FastBootCommand<S>,
181    ) -> Result<String, NusbFastBootError> {
182        self.send_command(cmd).await?;
183        self.handle_responses().await
184    }
185
186    /// Get the named variable
187    ///
188    /// The "all" variable is special; For that [Self::get_all_vars] should be used instead
189    pub async fn get_var(&mut self, var: &str) -> Result<String, NusbFastBootError> {
190        let cmd = FastBootCommand::GetVar(var);
191        self.execute(cmd).await
192    }
193
194    /// Prepare a download of a given size
195    ///
196    /// When successfull the [DataDownload] helper should be used to actually send the data
197    pub async fn download(&mut self, size: u32) -> Result<DataDownload, NusbFastBootError> {
198        let cmd = FastBootCommand::<&str>::Download(size);
199        self.send_command(cmd).await?;
200        loop {
201            let resp = self.read_response().await?;
202            match resp {
203                FastBootResponse::Info(i) => info!("info: {i}"),
204                FastBootResponse::Text(t) => info!("Text: {}", t),
205                FastBootResponse::Data(size) => {
206                    return Ok(DataDownload::new(self, size));
207                }
208                FastBootResponse::Okay(_) => {
209                    return Err(NusbFastBootError::FastbootUnexpectedReply)
210                }
211                FastBootResponse::Fail(fail) => {
212                    return Err(NusbFastBootError::FastbootFailed(fail))
213                }
214            }
215        }
216    }
217
218    /// Flash downloaded data to a given target partition
219    pub async fn flash(&mut self, target: &str) -> Result<(), NusbFastBootError> {
220        let cmd = FastBootCommand::Flash(target);
221        self.execute(cmd).await.map(|v| {
222            trace!("Flash ok: {v}");
223        })
224    }
225
226    /// Continue booting
227    pub async fn continue_boot(&mut self) -> Result<(), NusbFastBootError> {
228        let cmd = FastBootCommand::<&str>::Continue;
229        self.execute(cmd).await.map(|v| {
230            trace!("Continue ok: {v}");
231        })
232    }
233
234    /// Erasing the given target partition
235    pub async fn erase(&mut self, target: &str) -> Result<(), NusbFastBootError> {
236        let cmd = FastBootCommand::Erase(target);
237        self.execute(cmd).await.map(|v| {
238            trace!("Erase ok: {v}");
239        })
240    }
241
242    /// Reboot the device
243    pub async fn reboot(&mut self) -> Result<(), NusbFastBootError> {
244        let cmd = FastBootCommand::<&str>::Reboot;
245        self.execute(cmd).await.map(|v| {
246            trace!("Reboot ok: {v}");
247        })
248    }
249
250    /// Reboot the device to the bootloader
251    pub async fn reboot_bootloader(&mut self) -> Result<(), NusbFastBootError> {
252        let cmd = FastBootCommand::<&str>::RebootBootloader;
253        self.execute(cmd).await.map(|v| {
254            trace!("Reboot ok: {v}");
255        })
256    }
257
258    /// Retrieve all variables
259    pub async fn get_all_vars(&mut self) -> Result<HashMap<String, String>, NusbFastBootError> {
260        let cmd = FastBootCommand::GetVar("all");
261        self.send_command(cmd).await?;
262        let mut vars = HashMap::new();
263        loop {
264            let resp = self.read_response().await?;
265            trace!("Response: {:?}", resp);
266            match resp {
267                FastBootResponse::Info(i) => {
268                    let Some((key, value)) = i.rsplit_once(':') else {
269                        warn!("Failed to parse variable: {i}");
270                        continue;
271                    };
272                    vars.insert(key.trim().to_string(), value.trim().to_string());
273                }
274                FastBootResponse::Text(t) => info!("Text: {}", t),
275                FastBootResponse::Data(_) => {
276                    return Err(NusbFastBootError::FastbootUnexpectedReply)
277                }
278                FastBootResponse::Okay(_) => {
279                    return Ok(vars);
280                }
281                FastBootResponse::Fail(fail) => {
282                    return Err(NusbFastBootError::FastbootFailed(fail))
283                }
284            }
285        }
286    }
287}
288
289/// Error during data download
290#[derive(Debug, Error)]
291pub enum DownloadError {
292    #[error("Trying to complete while nothing was Queued")]
293    NothingQueued,
294    #[error("Incorrect data length: expected {expected}, got {actual}")]
295    IncorrectDataLength { actual: u32, expected: u32 },
296    #[error(transparent)]
297    Nusb(#[from] NusbFastBootError),
298}
299
300/// Data download helper
301///
302/// To success stream data over usb it needs to be sent in blocks that are multiple of the max
303/// endpoint size, otherwise the receiver may complain. It also should only send as much data as
304/// was indicate in the DATA command.
305///
306/// This helper ensures both invariants are met. To do this data needs to be sent by using
307/// [DataDownload::extend_from_slice] or [DataDownload::get_mut_data], after sending the data [DataDownload::finish] should be called to
308/// validate and finalize.
309pub struct DataDownload<'s> {
310    fastboot: &'s mut NusbFastBoot,
311    queue: nusb::transfer::Queue<Vec<u8>>,
312    size: u32,
313    left: u32,
314    current: Vec<u8>,
315}
316
317impl<'s> DataDownload<'s> {
318    fn new(fastboot: &'s mut NusbFastBoot, size: u32) -> DataDownload<'s> {
319        let queue = fastboot.interface.bulk_out_queue(fastboot.ep_out);
320        let current = Self::allocate_buffer(fastboot.max_out);
321        Self {
322            fastboot,
323            queue,
324            size,
325            left: size,
326            current,
327        }
328    }
329}
330
331impl DataDownload<'_> {
332    /// Total size of the data transfer
333    pub fn size(&self) -> u32 {
334        self.size
335    }
336
337    /// Data left to be sent/queued
338    pub fn left(&self) -> u32 {
339        self.left
340    }
341
342    /// Extend the streaming from a slice
343    ///
344    /// This will copy all provided data and send it out if enough is collected. The total amount
345    /// of data being sent should not exceed the download size
346    pub async fn extend_from_slice(&mut self, mut data: &[u8]) -> Result<(), DownloadError> {
347        self.update_size(data.len() as u32)?;
348        loop {
349            let left = self.current.capacity() - self.current.len();
350            if left >= data.len() {
351                self.current.extend_from_slice(data);
352                break;
353            } else {
354                self.current.extend_from_slice(&data[0..left]);
355                self.next_buffer().await?;
356                data = &data[left..];
357            }
358        }
359        Ok(())
360    }
361
362    /// This will provide a mutable reference to a [u8] of at most `max` size. The returned slice
363    /// should be completely filled with data to be downloaded to the device
364    ///
365    /// The total amount of data should not exceed the download size
366    pub async fn get_mut_data(&mut self, max: usize) -> Result<&mut [u8], DownloadError> {
367        if self.current.capacity() == self.current.len() {
368            self.next_buffer().await?;
369        }
370
371        let left = self.current.capacity() - self.current.len();
372        let size = left.min(max);
373        self.update_size(size as u32)?;
374
375        let len = self.current.len();
376        self.current.resize(len + size, 0);
377        Ok(&mut self.current[len..])
378    }
379
380    fn update_size(&mut self, size: u32) -> Result<(), DownloadError> {
381        if size > self.left {
382            return Err(DownloadError::IncorrectDataLength {
383                expected: self.size,
384                actual: size - self.left + self.size,
385            });
386        }
387        self.left -= size;
388        Ok(())
389    }
390
391    fn allocate_buffer(max_out: usize) -> Vec<u8> {
392        // Allocate about 1Mb of buffer ensuring it's always a multiple of the maximum out packet
393        // size
394        let size = (1024usize * 1024).next_multiple_of(max_out);
395        Vec::with_capacity(size)
396    }
397
398    async fn next_buffer(&mut self) -> Result<(), DownloadError> {
399        let mut next = if self.queue.pending() < 3 {
400            Self::allocate_buffer(self.fastboot.max_out)
401        } else {
402            let r = self.queue.next_complete().await;
403            r.status.map_err(NusbFastBootError::from)?;
404            let mut data = r.data.reuse();
405            data.truncate(0);
406            data
407        };
408        std::mem::swap(&mut next, &mut self.current);
409        self.queue.submit(next);
410        Ok(())
411    }
412
413    /// Finish all pending transfer
414    ///
415    /// This should only be called if all data has been queued up (matching the total size)
416    #[instrument(skip_all, err)]
417    pub async fn finish(mut self) -> Result<(), DownloadError> {
418        if self.left != 0 {
419            return Err(DownloadError::IncorrectDataLength {
420                expected: self.size,
421                actual: self.size - self.left,
422            });
423        }
424
425        if !self.current.is_empty() {
426            let current = std::mem::take(&mut self.current);
427            self.queue.submit(current);
428        }
429
430        while self.queue.pending() > 0 {
431            self.queue
432                .next_complete()
433                .await
434                .status
435                .map_err(NusbFastBootError::from)?;
436        }
437
438        self.fastboot.handle_responses().await?;
439        Ok(())
440    }
441}
fastboot_protocol/nusb.rs

fastboot_protocol/
nusb.rs
