Skip to main content

xdvdfs/write/fs/
remap.rs

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        // Expand until we find a node with a record
30        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    /// Looks up a node, only descending into subdirs if the path is not consumed
60    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    /// Looks up a subdir, returning its subtree
81    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            // Unwrap safe, set above
108            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        // Unwrap safe, set above
117        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    // host path regex -> image path rewrite
135    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            // TODO: Allow these to be escaped. Are {} characters valid anyway?
239            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        // Walk the host and store any paths that match the mapping rules
301        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                // Find which specific glob was matched by this path
337                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                // Negating patterns erase any rewritten_path we have come across
344                if cfg.map_rules[idx].0.starts_with('!') {
345                    rewritten_path = None;
346                    continue;
347                }
348
349                // Prefer previously matched patterns, if any
350                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 this path matched a prefix (e.g. the rule "bin") and has a suffix (e.g.
362                // "/default.xbe"), then we need to re-add the suffix to the rewritten prefix
363                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 we have a valid rewritten path, we can insert it into the new filesystem
378            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}