fastboot_protocol/
nusb.rs

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

fastboot_protocol/
nusb.rs
