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    /// Erasing the given target partition
227    pub async fn erase(&mut self, target: &str) -> Result<(), NusbFastBootError> {
228        let cmd = FastBootCommand::Erase(target);
229        self.execute(cmd).await.map(|v| {
230            trace!("Erase ok: {v}");
231        })
232    }
233
234    /// Reboot the device
235    pub async fn reboot(&mut self) -> Result<(), NusbFastBootError> {
236        let cmd = FastBootCommand::<&str>::Reboot;
237        self.execute(cmd).await.map(|v| {
238            trace!("Reboot ok: {v}");
239        })
240    }
241
242    /// Reboot the device to the bootloader
243    pub async fn reboot_bootloader(&mut self) -> Result<(), NusbFastBootError> {
244        let cmd = FastBootCommand::<&str>::RebootBootloader;
245        self.execute(cmd).await.map(|v| {
246            trace!("Reboot ok: {v}");
247        })
248    }
249
250    /// Retrieve all variables
251    pub async fn get_all_vars(&mut self) -> Result<HashMap<String, String>, NusbFastBootError> {
252        let cmd = FastBootCommand::GetVar("all");
253        self.send_command(cmd).await?;
254        let mut vars = HashMap::new();
255        loop {
256            let resp = self.read_response().await?;
257            trace!("Response: {:?}", resp);
258            match resp {
259                FastBootResponse::Info(i) => {
260                    let Some((key, value)) = i.rsplit_once(':') else {
261                        warn!("Failed to parse variable: {i}");
262                        continue;
263                    };
264                    vars.insert(key.trim().to_string(), value.trim().to_string());
265                }
266                FastBootResponse::Text(t) => info!("Text: {}", t),
267                FastBootResponse::Data(_) => {
268                    return Err(NusbFastBootError::FastbootUnexpectedReply)
269                }
270                FastBootResponse::Okay(_) => {
271                    return Ok(vars);
272                }
273                FastBootResponse::Fail(fail) => {
274                    return Err(NusbFastBootError::FastbootFailed(fail))
275                }
276            }
277        }
278    }
279}
280
281/// Error during data download
282#[derive(Debug, Error)]
283pub enum DownloadError {
284    #[error("Trying to complete while nothing was Queued")]
285    NothingQueued,
286    #[error("Incorrect data length: expected {expected}, got {actual}")]
287    IncorrectDataLength { actual: u32, expected: u32 },
288    #[error(transparent)]
289    Nusb(#[from] NusbFastBootError),
290}
291
292/// Data download helper
293///
294/// To success stream data over usb it needs to be sent in blocks that are multiple of the max
295/// endpoint size, otherwise the receiver may complain. It also should only send as much data as
296/// was indicate in the DATA command.
297///
298/// This helper ensures both invariants are met. To do this data needs to be sent by using
299/// [DataDownload::extend_from_slice] or [DataDownload::get_mut_data], after sending the data [DataDownload::finish] should be called to
300/// validate and finalize.
301pub struct DataDownload<'s> {
302    fastboot: &'s mut NusbFastBoot,
303    queue: nusb::transfer::Queue<Vec<u8>>,
304    size: u32,
305    left: u32,
306    current: Vec<u8>,
307}
308
309impl<'s> DataDownload<'s> {
310    fn new(fastboot: &'s mut NusbFastBoot, size: u32) -> DataDownload<'s> {
311        let queue = fastboot.interface.bulk_out_queue(fastboot.ep_out);
312        let current = Self::allocate_buffer(fastboot.max_out);
313        Self {
314            fastboot,
315            queue,
316            size,
317            left: size,
318            current,
319        }
320    }
321}
322
323impl DataDownload<'_> {
324    /// Total size of the data transfer
325    pub fn size(&self) -> u32 {
326        self.size
327    }
328
329    /// Data left to be sent/queued
330    pub fn left(&self) -> u32 {
331        self.left
332    }
333
334    /// Extend the streaming from a slice
335    ///
336    /// This will copy all provided data and send it out if enough is collected. The total amount
337    /// of data being sent should not exceed the download size
338    pub async fn extend_from_slice(&mut self, mut data: &[u8]) -> Result<(), DownloadError> {
339        self.update_size(data.len() as u32)?;
340        loop {
341            let left = self.current.capacity() - self.current.len();
342            if left >= data.len() {
343                self.current.extend_from_slice(data);
344                break;
345            } else {
346                self.current.extend_from_slice(&data[0..left]);
347                self.next_buffer().await?;
348                data = &data[left..];
349            }
350        }
351        Ok(())
352    }
353
354    /// This will provide a mutable reference to a [u8] of at most `max` size. The returned slice
355    /// should be completely filled with data to be downloaded to the device
356    ///
357    /// The total amount of data should not exceed the download size
358    pub async fn get_mut_data(&mut self, max: usize) -> Result<&mut [u8], DownloadError> {
359        if self.current.capacity() == self.current.len() {
360            self.next_buffer().await?;
361        }
362
363        let left = self.current.capacity() - self.current.len();
364        let size = left.min(max);
365        self.update_size(size as u32)?;
366
367        let len = self.current.len();
368        self.current.resize(len + size, 0);
369        Ok(&mut self.current[len..])
370    }
371
372    fn update_size(&mut self, size: u32) -> Result<(), DownloadError> {
373        if size > self.left {
374            return Err(DownloadError::IncorrectDataLength {
375                expected: self.size,
376                actual: size - self.left + self.size,
377            });
378        }
379        self.left -= size;
380        Ok(())
381    }
382
383    fn allocate_buffer(max_out: usize) -> Vec<u8> {
384        // Allocate about 1Mb of buffer ensuring it's always a multiple of the maximum out packet
385        // size
386        let size = (1024usize * 1024).next_multiple_of(max_out);
387        Vec::with_capacity(size)
388    }
389
390    async fn next_buffer(&mut self) -> Result<(), DownloadError> {
391        let mut next = if self.queue.pending() < 3 {
392            Self::allocate_buffer(self.fastboot.max_out)
393        } else {
394            let r = self.queue.next_complete().await;
395            r.status.map_err(NusbFastBootError::from)?;
396            let mut data = r.data.reuse();
397            data.truncate(0);
398            data
399        };
400        std::mem::swap(&mut next, &mut self.current);
401        self.queue.submit(next);
402        Ok(())
403    }
404
405    /// Finish all pending transfer
406    ///
407    /// This should only be called if all data has been queued up (matching the total size)
408    #[instrument(skip_all, err)]
409    pub async fn finish(mut self) -> Result<(), DownloadError> {
410        if self.left != 0 {
411            return Err(DownloadError::IncorrectDataLength {
412                expected: self.size,
413                actual: self.size - self.left,
414            });
415        }
416
417        if !self.current.is_empty() {
418            let current = std::mem::take(&mut self.current);
419            self.queue.submit(current);
420        }
421
422        while self.queue.pending() > 0 {
423            self.queue
424                .next_complete()
425                .await
426                .status
427                .map_err(NusbFastBootError::from)?;
428        }
429
430        self.fastboot.handle_responses().await?;
431        Ok(())
432    }
433}
fastboot_protocol/nusb.rs

fastboot_protocol/
nusb.rs
