1use core::fmt::Display;
2
3use alloc::borrow::ToOwned;
4use alloc::boxed::Box;
5use alloc::string::String;
6use alloc::vec::Vec;
7use maybe_async::maybe_async;
8
9use crate::blockdev::BlockDeviceWrite;
10
11use super::{FileEntry, FileType, Filesystem, PathVec};
12
13#[derive(Clone, Debug)]
14struct PathPrefixTree<T> {
15 children: [Option<Box<PathPrefixTree<T>>>; 256],
16 record: Option<(T, Box<PathPrefixTree<T>>)>,
17}
18
19struct PPTIter<'a, T> {
20 queue: Vec<(String, &'a PathPrefixTree<T>)>,
21}
22
23impl<'a, T> Iterator for PPTIter<'a, T> {
24 type Item = (String, &'a T);
25
26 fn next(&mut self) -> Option<Self::Item> {
27 use alloc::borrow::ToOwned;
28
29 while let Some(subtree) = self.queue.pop() {
31 let (name, node) = &subtree;
32 for (ch, child) in node.children.iter().enumerate() {
33 if let Some(child) = child {
34 let mut name = name.to_owned();
35 name.push(ch as u8 as char);
36 self.queue.push((name, child));
37 }
38 }
39
40 if let Some(record) = &node.record {
41 return Some((name.to_owned(), &record.0));
42 }
43 }
44
45 None
46 }
47}
48
49impl<T> Default for PathPrefixTree<T> {
50 fn default() -> Self {
51 Self {
52 children: [const { None }; 256],
53 record: None,
54 }
55 }
56}
57
58impl<T> PathPrefixTree<T> {
59 fn lookup_node(&self, path: &PathVec) -> Option<&Self> {
61 let mut node = self;
62
63 let mut component_iter = path.iter().peekable();
64 while let Some(component) = component_iter.next() {
65 for ch in component.chars() {
66 let next = &node.children[ch as usize];
67 node = next.as_ref()?;
68 }
69
70 if component_iter.peek().is_some() {
71 let record = &node.record;
72 let (_, subtree) = record.as_ref()?;
73 node = subtree;
74 }
75 }
76
77 Some(node)
78 }
79
80 fn lookup_subdir(&self, path: &PathVec) -> Option<&Self> {
82 let mut node = self;
83
84 for component in path.iter() {
85 for ch in component.chars() {
86 let next = &node.children[ch as usize];
87 node = next.as_ref()?;
88 }
89
90 let record = &node.record;
91 let (_, subtree) = record.as_ref()?;
92 node = subtree;
93 }
94
95 Some(node)
96 }
97
98 fn insert_tail(&mut self, tail: &str, val: T) -> &mut Self {
99 let mut node = self;
100
101 for ch in tail.chars() {
102 let next = &mut node.children[ch as usize];
103 if next.is_none() {
104 *next = Some(Box::new(Self::default()));
105 }
106
107 node = next.as_mut().unwrap().as_mut();
109 }
110
111 if let Some(ref mut record) = node.record {
112 return record.1.as_mut();
113 }
114
115 node.record = Some((val, Box::new(Self::default())));
116 node.record.as_mut().map(|x| x.1.as_mut()).unwrap()
118 }
119
120 fn get(&self, path: &PathVec) -> Option<&T> {
121 let node = self.lookup_node(path)?;
122 node.record.as_ref().map(|v| &v.0)
123 }
124
125 fn iter(&self) -> PPTIter<'_, T> {
126 PPTIter {
127 queue: alloc::vec![(String::new(), self)],
128 }
129 }
130}
131
132#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
133pub struct RemapOverlayConfig {
134 pub map_rules: Vec<(String, String)>,
136}
137
138#[derive(Clone, Debug)]
139struct MapEntry {
140 host_path: PathVec,
141 host_entry: FileEntry,
142 is_prefix_directory: bool,
143}
144
145#[derive(Debug)]
146pub enum InvalidRewriteSubstitutionKind {
147 NonDigitCharacter,
148 UnclosedBrace,
149}
150
151#[derive(Debug)]
152pub enum RemapOverlayFilesystemBuildingError<E> {
153 GlobBuildingError(wax::BuildError),
154 InvalidRewriteSubstitution(usize, String, InvalidRewriteSubstitutionKind),
155 FilesystemError(E),
156}
157
158impl<E: Display> RemapOverlayFilesystemBuildingError<E> {
159 pub fn as_string(&self) -> String {
160 match self {
161 Self::FilesystemError(e) => alloc::format!("error in underlying filesystem: {}", e),
162 Self::GlobBuildingError(e) => alloc::format!("failed to build glob pattern: {}", e),
163 Self::InvalidRewriteSubstitution(idx, rewrite, kind) => alloc::format!(
164 "invalid rewrite substitution \"{}\" (at {}): {}",
165 rewrite,
166 idx,
167 match kind {
168 InvalidRewriteSubstitutionKind::NonDigitCharacter => "expected digit character",
169 InvalidRewriteSubstitutionKind::UnclosedBrace => "unclosed brace",
170 }
171 ),
172 }
173 }
174}
175
176impl<E: Display> Display for RemapOverlayFilesystemBuildingError<E> {
177 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
178 f.write_str(self.as_string().as_str())
179 }
180}
181
182impl<E: Display + core::fmt::Debug> std::error::Error for RemapOverlayFilesystemBuildingError<E> {}
183
184#[non_exhaustive]
185#[derive(Clone, Debug)]
186pub enum RemapOverlayError<E> {
187 NoSuchFile(String),
188 UnderlyingError(E),
189}
190
191impl<E> From<E> for RemapOverlayError<E> {
192 fn from(value: E) -> Self {
193 Self::UnderlyingError(value)
194 }
195}
196
197impl<E: Display> RemapOverlayError<E> {
198 pub fn as_string(&self) -> String {
199 match self {
200 Self::NoSuchFile(image) => alloc::format!("no host mapping for image path: {}", image),
201 Self::UnderlyingError(e) => alloc::format!("error in underlying filesystem: {}", e),
202 }
203 }
204}
205
206impl<E: Display> Display for RemapOverlayError<E> {
207 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
208 f.write_str(self.as_string().as_str())
209 }
210}
211
212impl<E: Display + core::fmt::Debug> std::error::Error for RemapOverlayError<E> {}
213
214#[derive(Clone, Debug)]
215pub struct RemapOverlayFilesystem<BDE, BD: BlockDeviceWrite<BDE>, FS: Filesystem<BD, BDE>> {
216 img_to_host: PathPrefixTree<MapEntry>,
217 fs: FS,
218
219 bde_type: core::marker::PhantomData<BDE>,
220 bd_type: core::marker::PhantomData<BD>,
221 fs_type: core::marker::PhantomData<FS>,
222}
223
224impl<BDE, BD, FS> RemapOverlayFilesystem<BDE, BD, FS>
225where
226 BDE: Into<RemapOverlayError<BDE>> + Send + Sync,
227 BD: BlockDeviceWrite<BDE>,
228 FS: Filesystem<BD, BDE>,
229{
230 fn find_match_indices(
231 rewrite: &str,
232 ) -> Result<Vec<usize>, RemapOverlayFilesystemBuildingError<BDE>> {
233 let mut match_indices: Vec<usize> = Vec::new();
234 let mut match_index = 0;
235 let mut matching = false;
236
237 for (idx, c) in rewrite.chars().enumerate() {
238 if c == '{' {
240 matching = true;
241 continue;
242 }
243
244 if !matching {
245 continue;
246 }
247
248 if c == '}' {
249 matching = false;
250 match_indices.push(match_index);
251 match_index = 0;
252 continue;
253 }
254
255 if let Some(digit) = c.to_digit(10) {
256 match_index *= 10;
257 match_index += digit as usize;
258 continue;
259 }
260
261 return Err(
262 RemapOverlayFilesystemBuildingError::InvalidRewriteSubstitution(
263 idx,
264 rewrite.to_owned(),
265 InvalidRewriteSubstitutionKind::NonDigitCharacter,
266 ),
267 );
268 }
269
270 if matching {
271 return Err(
272 RemapOverlayFilesystemBuildingError::InvalidRewriteSubstitution(
273 rewrite.len() - 1,
274 rewrite.to_owned(),
275 InvalidRewriteSubstitutionKind::UnclosedBrace,
276 ),
277 );
278 }
279
280 Ok(match_indices)
281 }
282
283 #[maybe_async]
284 pub async fn new(
285 mut fs: FS,
286 cfg: RemapOverlayConfig,
287 ) -> Result<Self, RemapOverlayFilesystemBuildingError<BDE>> {
288 use wax::Pattern;
289
290 let glob_keys: Result<Vec<wax::Glob>, _> = cfg
291 .map_rules
292 .iter()
293 .map(|(from, _)| wax::Glob::new(from.trim_start_matches('!')))
294 .collect();
295 let glob_keys =
296 glob_keys.map_err(|e| RemapOverlayFilesystemBuildingError::GlobBuildingError(e))?;
297 let all_globs = wax::any(glob_keys.clone().into_iter())
298 .map_err(|e| RemapOverlayFilesystemBuildingError::GlobBuildingError(e))?;
299
300 let mut host_dirs = alloc::vec![(PathVec::default(), None)];
302 let mut matches: Vec<(PathVec, FileEntry, PathVec)> = Vec::new();
303 while let Some((dir, parent_match_prefix)) = host_dirs.pop() {
304 let listing = fs
305 .read_dir(&dir)
306 .await
307 .map_err(|e| RemapOverlayFilesystemBuildingError::FilesystemError(e))?;
308 for entry in listing.iter() {
309 let path = PathVec::from_base(&dir, &entry.name);
310 let match_prefix = if all_globs.is_match(path.as_string().trim_start_matches('/')) {
311 Some(path.clone())
312 } else if parent_match_prefix.is_some() {
313 parent_match_prefix.clone()
314 } else {
315 None
316 };
317
318 if let FileType::Directory = entry.file_type {
319 host_dirs.push((PathVec::from_base(&dir, &entry.name), match_prefix.clone()));
320 }
321
322 if let Some(prefix) = match_prefix {
323 matches.push((path.clone(), entry.clone(), prefix));
324 }
325 }
326 }
327
328 let mut img_to_host = PathPrefixTree::default();
329 for (path, entry, prefix) in matches {
330 let suffix = path.suffix(&prefix);
331 let mut rewritten_path: Option<PathVec> = None;
332
333 for (idx, glob) in glob_keys.iter().enumerate() {
334 let path_str = prefix.as_string();
335
336 let cand_path = wax::CandidatePath::from(path_str.trim_start_matches('/'));
338 let matched = glob.matched(&cand_path);
339 let Some(matched) = matched else {
340 continue;
341 };
342
343 if cfg.map_rules[idx].0.starts_with('!') {
345 rewritten_path = None;
346 continue;
347 }
348
349 if rewritten_path.is_some() {
351 continue;
352 }
353
354 let mut rewrite = cfg.map_rules[idx].1.clone();
355 let match_indices = Self::find_match_indices(&rewrite)?;
356 for index in match_indices {
357 let replace = matched.get(index).unwrap_or("");
358 rewrite = rewrite.replace(&alloc::format!("{{{index}}}"), replace);
359 }
360
361 if !suffix.is_root() {
364 rewrite =
365 alloc::format!("{}{}", rewrite.trim_end_matches('/'), suffix.as_string());
366 }
367
368 let rewrite = PathVec::from_iter(
369 rewrite
370 .trim_start_matches('.')
371 .trim_start_matches('/')
372 .split('/'),
373 );
374 rewritten_path = Some(rewrite);
375 }
376
377 if let Some(rewrite) = rewritten_path {
379 let mut rewrite = rewrite.iter().peekable();
380 let mut node = &mut img_to_host;
381
382 while let Some(component) = rewrite.next() {
383 let is_prefix_directory = rewrite.peek().is_some();
384 let entry = if !is_prefix_directory {
385 entry.clone()
386 } else {
387 FileEntry {
388 name: component.to_owned(),
389 file_type: FileType::Directory,
390 len: 0,
391 }
392 };
393
394 node = node.insert_tail(
395 component,
396 MapEntry {
397 host_entry: entry,
398 host_path: path.clone(),
399 is_prefix_directory,
400 },
401 );
402 }
403 }
404 }
405
406 Ok(Self {
407 img_to_host,
408 fs,
409
410 bde_type: core::marker::PhantomData,
411 bd_type: core::marker::PhantomData,
412 fs_type: core::marker::PhantomData,
413 })
414 }
415
416 pub fn dump(&self) -> Vec<(PathVec, PathVec)> {
417 let mut queue = alloc::vec![PathVec::default()];
418 let mut output: Vec<(PathVec, PathVec)> = Vec::new();
419
420 while let Some(path) = queue.pop() {
421 let listing = self
422 .img_to_host
423 .lookup_subdir(&path)
424 .expect("failed trie lookup for vfs directory");
425 for (name, entry) in listing.iter() {
426 let path = PathVec::from_base(&path, &name);
427
428 if !entry.is_prefix_directory {
429 output.push((entry.host_path.clone(), path.clone()));
430 }
431
432 if let FileType::Directory = entry.host_entry.file_type {
433 queue.push(path);
434 }
435 }
436 }
437
438 output
439 }
440}
441
442#[maybe_async]
443impl<BDE, BD, FS> Filesystem<BD, RemapOverlayError<BDE>, BDE>
444 for RemapOverlayFilesystem<BDE, BD, FS>
445where
446 BDE: Into<RemapOverlayError<BDE>> + Send + Sync,
447 BD: BlockDeviceWrite<BDE>,
448 FS: Filesystem<BD, BDE>,
449{
450 async fn read_dir(&mut self, path: &PathVec) -> Result<Vec<FileEntry>, RemapOverlayError<BDE>> {
451 let dir = self
452 .img_to_host
453 .lookup_subdir(path)
454 .expect("failed trie lookup for virtual filesystem directory");
455 let entries: Vec<FileEntry> = dir
456 .iter()
457 .map(|(name, entry)| FileEntry {
458 name,
459 file_type: entry.host_entry.file_type,
460 len: entry.host_entry.len,
461 })
462 .collect();
463
464 Ok(entries)
465 }
466
467 async fn copy_file_in(
468 &mut self,
469 src: &PathVec,
470 dest: &mut BD,
471 offset: u64,
472 size: u64,
473 ) -> Result<u64, RemapOverlayError<BDE>> {
474 let entry = self
475 .img_to_host
476 .get(src)
477 .ok_or_else(|| RemapOverlayError::NoSuchFile(src.as_string()))?;
478 self.fs
479 .copy_file_in(&entry.host_path, dest, offset, size)
480 .await
481 .map_err(|e| e.into())
482 }
483
484 async fn copy_file_buf(
485 &mut self,
486 src: &PathVec,
487 buf: &mut [u8],
488 offset: u64,
489 ) -> Result<u64, RemapOverlayError<BDE>> {
490 let entry = self
491 .img_to_host
492 .get(src)
493 .ok_or_else(|| RemapOverlayError::NoSuchFile(src.as_string()))?;
494 self.fs
495 .copy_file_buf(&entry.host_path, buf, offset)
496 .await
497 .map_err(|e| e.into())
498 }
499}