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(&self, path: &str, max_pending: usize, f: impl FnMut(&str, &DirEntry) -> Result<()>) -> Result<()> {
100 walk_with(|p, n, c| self.ls(p, n, c), path, max_pending, f)
101 }
102}
103
104#[cfg(test)]
105mod tests {
106 use super::*;
107 use std::cell::Cell;
108
109 fn entry(name: String, dir: bool) -> DirEntry {
110 DirEntry {
111 name,
112 r#type: if dir { FileType::Directory } else { FileType::Regular } as i32,
113 info: None,
114 }
115 }
116
117 fn endless<'a>(
119 calls: &'a Cell<usize>,
120 biggest: &'a Cell<usize>,
121 dirs: bool,
122 ) -> impl FnMut(&str, u32, &str) -> Result<Page> + 'a {
123 move |_p, n, c| {
124 calls.set(calls.get() + 1);
125 let from: usize = c.parse().unwrap_or(0);
126 let page: Vec<DirEntry> = (from..from + n as usize)
127 .map(|i| entry(format!("n{i}"), dirs))
128 .collect();
129 biggest.set(biggest.get().max(page.len()));
130 Ok((page, (from + n as usize).to_string()))
131 }
132 }
133
134 fn finite(total: usize, dirs: bool) -> impl FnMut(&str, u32, &str) -> Result<Page> {
135 move |_p, n, c| {
136 let from: usize = c.parse().unwrap_or(0);
137 let end = (from + n as usize).min(total);
138 let page: Vec<DirEntry> = (from..end).map(|i| entry(format!("n{i}"), dirs)).collect();
139 assert!(page.len() <= WALK_PAGE as usize);
140 Ok((page, if end < total { end.to_string() } else { String::new() }))
141 }
142 }
143
144 #[test]
145 fn read_dir_all_stops_at_the_limit() {
146 let (calls, biggest) = (Cell::new(0), Cell::new(0));
147 let e = read_dir_all_with(endless(&calls, &biggest, false), "/huge", 2500).unwrap_err();
148 assert!(matches!(e, Error::TooManyEntries(ref m) if m.contains("2500")));
149 assert_eq!(calls.get(), 3);
150 assert_eq!(biggest.get(), WALK_PAGE as usize);
151 }
152
153 #[test]
154 fn read_dir_all_at_the_limit_works() {
155 assert_eq!(read_dir_all_with(finite(2500, false), "/d", 2500).unwrap().len(), 2500);
156 assert!(read_dir_all_with(finite(2500, false), "/d", 2499).is_err());
157 }
158
159 #[test]
160 fn walk_bounds_pending_directories() {
161 let (calls, biggest) = (Cell::new(0), Cell::new(0));
162 let mut visited = 0;
163 let e = walk_with(endless(&calls, &biggest, true), "/", 3000, |_, _| {
164 visited += 1;
165 Ok(())
166 })
167 .unwrap_err();
168 assert!(matches!(e, Error::TooManyEntries(ref m) if m.contains("3000")));
169 assert_eq!(visited, 3001);
170 assert_eq!(biggest.get(), WALK_PAGE as usize);
171 assert_eq!(calls.get(), 4);
172 }
173
174 #[test]
175 fn walk_holds_one_page_of_a_flat_directory() {
176 let mut count = 0;
177 walk_with(finite(5003, false), "/f", 4, |_, _| {
178 count += 1;
179 Ok(())
180 })
181 .unwrap();
182 assert_eq!(count, 5003);
183 }
184}