Skip to main content

nuclei/fs/
read_dir.rs

1//! [`read_dir`]: a directory's entries, streamed. The blocking pool reads
2//! them a page of [`PAGE`] at a time, so a directory of any size costs a
3//! page of memory, never its whole listing.
4
5use core::future::Future;
6use core::pin::Pin;
7use core::task::{Context, Poll, ready};
8use std::collections::VecDeque;
9use std::ffi::OsString;
10use std::fs::{FileType, Metadata};
11use std::io;
12use std::path::{Path, PathBuf};
13
14use futures_core::Stream;
15
16use super::ops::blocking;
17use crate::task::Join;
18
19/// Entries read from the directory in one call of the blocking pool.
20const PAGE: usize = 32;
21
22/// One page and what reads the next. The std reader is boxed: it is large
23/// on Windows, and goes to the pool and back with each page.
24type Page = (
25    Option<Box<std::fs::ReadDir>>,
26    VecDeque<io::Result<DirEntry>>,
27);
28
29/// The entries of the directory at `path`, through
30/// [`ReadDir::next_entry`] or as a [`Stream`].
31pub async fn read_dir(path: impl AsRef<Path>) -> io::Result<ReadDir> {
32    let path = path.as_ref().to_path_buf();
33    let dir = blocking(move || std::fs::read_dir(path)).await?;
34    Ok(ReadDir {
35        state: State::Idle(Some(Box::new(dir)), VecDeque::new()),
36    })
37}
38
39/// A directory's entries, read a page at a time.
40#[derive(Debug)]
41pub struct ReadDir {
42    state: State,
43}
44
45#[derive(Debug)]
46enum State {
47    /// The entries read and not handed out; `None` once the directory ended.
48    Idle(
49        Option<Box<std::fs::ReadDir>>,
50        VecDeque<io::Result<DirEntry>>,
51    ),
52    /// The next page on the pool.
53    Busy(Join<Page>),
54}
55
56impl ReadDir {
57    /// The next entry, `None` at the end.
58    pub async fn next_entry(&mut self) -> io::Result<Option<DirEntry>> {
59        core::future::poll_fn(|cx| self.poll_next_entry(cx)).await
60    }
61
62    /// Polls for the next entry.
63    pub fn poll_next_entry(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<Option<DirEntry>>> {
64        loop {
65            match &mut self.state {
66                State::Idle(dir, page) => {
67                    if let Some(entry) = page.pop_front() {
68                        return Poll::Ready(entry.map(Some));
69                    }
70                    let Some(mut dir) = dir.take() else {
71                        return Poll::Ready(Ok(None));
72                    };
73                    self.state = State::Busy(
74                        crate::runtime::spawn_blocking(move || {
75                            let mut page = VecDeque::with_capacity(PAGE);
76                            for _ in 0..PAGE {
77                                match dir.next() {
78                                    Some(entry) => page.push_back(entry.map(DirEntry::new)),
79                                    None => return (None, page),
80                                }
81                            }
82                            (Some(dir), page)
83                        })
84                        .join(),
85                    );
86                }
87                State::Busy(job) => match ready!(Pin::new(job).poll(cx)) {
88                    Ok((dir, page)) => self.state = State::Idle(dir, page),
89                    Err(lost) => {
90                        self.state = State::Idle(None, VecDeque::new());
91                        return Poll::Ready(Err(lost.into()));
92                    }
93                },
94            }
95        }
96    }
97}
98
99impl Stream for ReadDir {
100    type Item = io::Result<DirEntry>;
101
102    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
103        self.get_mut()
104            .poll_next_entry(cx)
105            .map(|entry| entry.transpose())
106    }
107}
108
109/// An entry of a directory.
110#[derive(Debug)]
111pub struct DirEntry {
112    inner: std::sync::Arc<std::fs::DirEntry>,
113    /// Read with the page: on most systems the directory read gives it at
114    /// no cost.
115    file_type: Result<FileType, (io::ErrorKind, String)>,
116}
117
118impl DirEntry {
119    fn new(entry: std::fs::DirEntry) -> Self {
120        let file_type = entry
121            .file_type()
122            .map_err(|err| (err.kind(), err.to_string()));
123        Self {
124            inner: std::sync::Arc::new(entry),
125            file_type,
126        }
127    }
128
129    /// The entry's full path: the directory's path joined with its name.
130    pub fn path(&self) -> PathBuf {
131        self.inner.path()
132    }
133
134    /// The entry's name in its directory.
135    pub fn file_name(&self) -> OsString {
136        self.inner.file_name()
137    }
138
139    /// The entry's inode number.
140    #[cfg(unix)]
141    pub fn ino(&self) -> u64 {
142        std::os::unix::fs::DirEntryExt::ino(&*self.inner)
143    }
144
145    /// What kind of file the entry is, without following a link.
146    pub async fn file_type(&self) -> io::Result<FileType> {
147        self.file_type
148            .clone()
149            .map_err(|(kind, message)| io::Error::new(kind, message))
150    }
151
152    /// The entry's metadata, without following a link.
153    pub async fn metadata(&self) -> io::Result<Metadata> {
154        let entry = std::sync::Arc::clone(&self.inner);
155        blocking(move || entry.metadata()).await
156    }
157}