usb_disk_probe/
stream.rs

1use super::DISK_DIR;
2use futures::stream::Stream;
3use std::{
4    io,
5    path::Path,
6    pin::Pin,
7    task::{Context, Poll},
8};
9use tokio::fs;
10
11#[derive(Debug, Error)]
12pub enum Error {
13    #[error("failed to open /dev/disk/by-path/ directory for reading")]
14    Open(#[source] io::Error),
15    #[error("directory iteration error")]
16    Iteration(#[source] io::Error),
17    #[error("device path lacks a file name")]
18    DeviceWithoutFileName,
19    #[error("device path is not UTF-8 valid")]
20    DevicePathNotUtf8,
21}
22
23/// A stream which iterates on USB device paths to find USB storage disks in the system.
24///
25/// # Example
26///
27/// ```no_run
28/// use usb_disk_probe::stream::UsbDiskProbe;
29///
30/// use futures::stream::StreamExt;
31///
32/// fn main() {
33///     futures::executor::block_on(async move {
34///         let mut stream = UsbDiskProbe::new().await.unwrap();
35///         while let Some(device_result) = stream.next().await {
36///             let device = device_result.unwrap();
37///             println!("{}", device.display());
38///         }
39///     });
40/// }
41/// ```
42pub struct UsbDiskProbe(fs::ReadDir);
43
44impl UsbDiskProbe {
45    pub async fn new() -> Result<Self, Error> {
46        fs::read_dir(DISK_DIR)
47            .await
48            .map(|dir| Self(dir))
49            .map_err(Error::Open)
50    }
51}
52
53impl Stream for UsbDiskProbe {
54    type Item = Result<Box<Path>, Error>;
55
56    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
57        loop {
58            match self.0.poll_next_entry(cx) {
59                Poll::Pending => return Poll::Pending,
60                Poll::Ready(Ok(Some(res))) => match filter_device(res) {
61                    value @ Some(_) => return Poll::Ready(value),
62                    None => continue,
63                },
64                Poll::Ready(Ok(None)) => return Poll::Ready(None),
65                Poll::Ready(Err(why)) => return Poll::Ready(Some(Err(Error::Iteration(why)))),
66            }
67        }
68    }
69}
70
71/// Filter USB devices which are not USB devices
72fn filter_device(entry: fs::DirEntry) -> Option<Result<Box<Path>, Error>> {
73    let path = entry.path();
74
75    match path.file_name() {
76        Some(filename) => match filename.to_str().ok_or(Error::DevicePathNotUtf8) {
77            Ok(filename) => {
78                if is_usb(filename) {
79                    Some(Ok(path.into_boxed_path()))
80                } else {
81                    None
82                }
83            }
84            Err(why) => return Some(Err(why)),
85        },
86        None => return Some(Err(Error::DeviceWithoutFileName)),
87    }
88}
89
90/// Checks if a device is a USB device
91fn is_usb(filename: &str) -> bool {
92    filename.starts_with("pci-") && filename.contains("-usb-") && filename.ends_with("-0:0:0:0")
93}