1use crate::client::Client;
5use crate::error::{Error, Result};
6use crate::wire::{DirEntry, FileType};
7
8const WALK_PAGE: u32 = 1000;
10
11pub const DEFAULT_MAX_DIR_ENTRIES: usize = 1_000_000;
13
14pub const DEFAULT_MAX_WALK_PENDING: usize = 65_536;
16
17fn join(dir: &str, name: &str) -> String {
18 format!("{}/{name}", dir.trim_end_matches('/'))
19}
20
21type Page = (Vec<DirEntry>, String);
22
23fn read_dir_all_with(
24 mut ls: impl FnMut(&str, u32, &str) -> Result<Page>,
25 path: &str,
26 limit: usize,
27) -> Result<Vec<DirEntry>> {
28 let mut out = Vec::new();
29 let mut cursor = String::new();
30 loop {
31 let (entries, next) = ls(path, WALK_PAGE, &cursor)?;
32 if out.len() + entries.len() > limit {
33 return Err(Error::TooManyEntries(format!("{path} has more than {limit} entries")));
34 }
35 out.extend(entries);
36 if next.is_empty() {
37 return Ok(out);
38 }
39 cursor = next;
40 }
41}
42
43fn walk_with(
44 mut ls: impl FnMut(&str, u32, &str) -> Result<Page>,
45 path: &str,
46 max_pending: usize,
47 mut f: impl FnMut(&str, &DirEntry) -> Result<()>,
48) -> Result<()> {
49 let mut stack = vec![path.to_string()];
50 while let Some(dir) = stack.pop() {
51 let mut subdirs = Vec::new();
52 let mut cursor = String::new();
53 loop {
54 let (entries, next) = ls(&dir, WALK_PAGE, &cursor)?;
55 for e in &entries {
56 f(&dir, e)?;
57 if e.r#type == FileType::Directory as i32 {
58 if stack.len() + subdirs.len() >= max_pending {
59 return Err(Error::TooManyEntries(format!(
60 "walk of {dir} has more than {max_pending} pending directories"
61 )));
62 }
63 subdirs.push(join(&dir, &e.name));
64 }
65 }
66 if next.is_empty() {
67 break;
68 }
69 cursor = next;
70 }
71 stack.extend(subdirs.into_iter().rev());
72 }
73 Ok(())
74}
75
76impl Client {
77 pub fn read_dir_all(&self, path: &str) -> Result<Vec<DirEntry>> {
80 self.read_dir_all_max(path, DEFAULT_MAX_DIR_ENTRIES)
81 }
82
83 pub fn read_dir_all_max(&self, path: &str, max_entries: usize) -> Result<Vec<DirEntry>> {
85 read_dir_all_with(|p, n, c| self.ls(p, n, c), path, max_entries)
86 }
87
88 pub fn walk(&self, path: &str, f: impl FnMut(&str, &DirEntry) -> Result<()>) -> Result<()> {
95 self.walk_max(path, DEFAULT_MAX_WALK_PENDING, f)
96 }
97
98 pub fn walk_max(
100 &self,
101 path: &str,
102 max_pending: usize,
103 f: impl FnMut(&str, &DirEntry) -> Result<()>,
104 ) -> Result<()> {
105 walk_with(|p, n, c| self.ls(p, n, c), path, max_pending, f)
106 }
107}
108
109#[cfg(test)]
110mod tests {
111 use super::*;
112 use std::cell::Cell;
113
114 fn entry(name: String, dir: bool) -> DirEntry {
115 DirEntry {
116 name,
117 r#type: if dir { FileType::Directory } else { FileType::Regular } as i32,
118 info: None,
119 }
120 }
121
122 fn endless<'a>(
124 calls: &'a Cell<usize>,
125 biggest: &'a Cell<usize>,
126 dirs: bool,
127 ) -> impl FnMut(&str, u32, &str) -> Result<Page> + 'a {
128 move |_p, n, c| {
129 calls.set(calls.get() + 1);
130 let from: usize = c.parse().unwrap_or(0);
131 let page: Vec<DirEntry> = (from..from + n as usize)
132 .map(|i| entry(format!("n{i}"), dirs))
133 .collect();
134 biggest.set(biggest.get().max(page.len()));
135 Ok((page, (from + n as usize).to_string()))
136 }
137 }
138
139 fn finite(total: usize, dirs: bool) -> impl FnMut(&str, u32, &str) -> Result<Page> {
140 move |_p, n, c| {
141 let from: usize = c.parse().unwrap_or(0);
142 let end = (from + n as usize).min(total);
143 let page: Vec<DirEntry> = (from..end).map(|i| entry(format!("n{i}"), dirs)).collect();
144 assert!(page.len() <= WALK_PAGE as usize);
145 Ok((page, if end < total { end.to_string() } else { String::new() }))
146 }
147 }
148
149 #[test]
150 fn read_dir_all_stops_at_the_limit() {
151 let (calls, biggest) = (Cell::new(0), Cell::new(0));
152 let e = read_dir_all_with(endless(&calls, &biggest, false), "/huge", 2500).unwrap_err();
153 assert!(matches!(e, Error::TooManyEntries(ref m) if m.contains("2500")));
154 assert_eq!(calls.get(), 3);
155 assert_eq!(biggest.get(), WALK_PAGE as usize);
156 }
157
158 #[test]
159 fn read_dir_all_at_the_limit_works() {
160 assert_eq!(read_dir_all_with(finite(2500, false), "/d", 2500).unwrap().len(), 2500);
161 assert!(read_dir_all_with(finite(2500, false), "/d", 2499).is_err());
162 }
163
164 #[test]
165 fn walk_bounds_pending_directories() {
166 let (calls, biggest) = (Cell::new(0), Cell::new(0));
167 let mut visited = 0;
168 let e = walk_with(endless(&calls, &biggest, true), "/", 3000, |_, _| {
169 visited += 1;
170 Ok(())
171 })
172 .unwrap_err();
173 assert!(matches!(e, Error::TooManyEntries(ref m) if m.contains("3000")));
174 assert_eq!(visited, 3001);
175 assert_eq!(biggest.get(), WALK_PAGE as usize);
176 assert_eq!(calls.get(), 4);
177 }
178
179 #[test]
180 fn walk_holds_one_page_of_a_flat_directory() {
181 let mut count = 0;
182 walk_with(finite(5003, false), "/f", 4, |_, _| {
183 count += 1;
184 Ok(())
185 })
186 .unwrap();
187 assert_eq!(count, 5003);
188 }
189}