Skip to main content

microsandbox_image/tree/
model.rs

1use std::collections::{BTreeMap, HashMap, HashSet};
2use std::ffi::{OsStr, OsString};
3use std::fmt;
4use std::io::{Read, Seek, SeekFrom};
5use std::os::unix::ffi::OsStrExt;
6use std::path::{Path, PathBuf};
7use std::sync::Arc;
8use std::sync::atomic::{AtomicU64, Ordering};
9
10//--------------------------------------------------------------------------------------------------
11// Constants
12//--------------------------------------------------------------------------------------------------
13
14const DEFAULT_MAX_TOTAL_SIZE: u64 = 10 * 1024 * 1024 * 1024; // 10 GiB
15const DEFAULT_MAX_FILE_SIZE: u64 = 5 * 1024 * 1024 * 1024; // 5 GiB
16const DEFAULT_MAX_ENTRY_COUNT: u64 = 1_000_000;
17const DEFAULT_MAX_PATH_LENGTH: usize = 4096;
18const DEFAULT_MAX_PATH_DEPTH: usize = 128;
19const DEFAULT_MAX_SYMLINK_TARGET: usize = 4096;
20
21const DEFAULT_DIR_MODE: u16 = 0o755;
22
23/// Overlayfs whiteout: char device with major=0, minor=0 signals deletion.
24pub(crate) const WHITEOUT_MAJOR: u32 = 0;
25pub(crate) const WHITEOUT_MINOR: u32 = 0;
26
27/// Overlayfs opaque directory xattr: hides all lower-layer entries.
28pub(crate) const OPAQUE_XATTR_NAME: &[u8] = b"trusted.overlay.opaque";
29pub(crate) const OPAQUE_XATTR_VALUE: &[u8] = b"y";
30
31//--------------------------------------------------------------------------------------------------
32// Types
33//--------------------------------------------------------------------------------------------------
34
35/// File content storage — either in-memory for small files or spooled to
36/// disk for large files to keep memory usage bounded.
37#[derive(Clone)]
38pub enum FileData {
39    /// Small file content held in memory.
40    Memory(Vec<u8>),
41    /// In-memory content shared by hardlinked regular-file aliases.
42    SharedMemory(Arc<[u8]>),
43    /// Large file content written to a shared spool file on disk.
44    /// Multiple `FileData::Spool` entries can reference different regions
45    /// of the same underlying spool file via `Arc`.
46    Spool {
47        spool: Arc<std::sync::Mutex<std::fs::File>>,
48        offset: u64,
49        len: u64,
50    },
51}
52
53impl std::fmt::Debug for FileData {
54    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
55        match self {
56            FileData::Memory(data) => f.debug_tuple("Memory").field(&data.len()).finish(),
57            FileData::SharedMemory(data) => {
58                f.debug_tuple("SharedMemory").field(&data.len()).finish()
59            }
60            FileData::Spool { offset, len, .. } => f
61                .debug_struct("Spool")
62                .field("offset", offset)
63                .field("len", len)
64                .finish(),
65        }
66    }
67}
68
69impl PartialEq for FileData {
70    fn eq(&self, other: &Self) -> bool {
71        match (self, other) {
72            (FileData::Memory(a), FileData::Memory(b)) => a == b,
73            (FileData::Memory(a), FileData::SharedMemory(b))
74            | (FileData::SharedMemory(b), FileData::Memory(a)) => a.as_slice() == b.as_ref(),
75            (FileData::SharedMemory(a), FileData::SharedMemory(b)) => a.as_ref() == b.as_ref(),
76            _ => false,
77        }
78    }
79}
80
81/// Threshold below which file data is kept in memory (64 KiB).
82/// Files at or above this size are spooled to disk during tar ingestion.
83pub const SPOOL_THRESHOLD: u64 = 64 * 1024;
84
85/// A writable spool file for large file data during tar ingestion.
86pub struct DataSpool {
87    file: std::fs::File,
88    shared: Arc<std::sync::Mutex<std::fs::File>>,
89    offset: u64,
90}
91
92pub struct ResourceLimits {
93    pub max_total_size: u64,
94    pub max_file_size: u64,
95    pub max_entry_count: u64,
96    pub max_path_length: usize,
97    pub max_path_depth: usize,
98    pub max_symlink_target: usize,
99}
100
101#[derive(Clone)]
102pub struct InodeMetadata {
103    pub uid: u32,
104    pub gid: u32,
105    pub mode: u16,
106    pub mtime: u64,
107    pub mtime_nsec: u32,
108}
109
110#[derive(Clone)]
111pub struct Xattr {
112    pub name: Vec<u8>,
113    pub value: Vec<u8>,
114}
115
116#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
117pub struct RegularFileId(u64);
118
119#[derive(Clone)]
120pub enum TreeNode {
121    RegularFile(RegularFileNode),
122    Directory(DirectoryNode),
123    Symlink(SymlinkNode),
124    CharDevice(DeviceNode),
125    BlockDevice(DeviceNode),
126    Fifo(InodeMetadata),
127    Socket(InodeMetadata),
128}
129
130#[derive(Clone)]
131pub struct RegularFileNode {
132    pub id: RegularFileId,
133    pub metadata: InodeMetadata,
134    pub xattrs: Vec<Xattr>,
135    pub data: FileData,
136    pub nlink: u32,
137}
138
139#[derive(Clone)]
140pub struct DirectoryNode {
141    pub metadata: InodeMetadata,
142    pub xattrs: Vec<Xattr>,
143    pub entries: BTreeMap<OsString, TreeNode>,
144}
145
146#[derive(Clone)]
147pub struct SymlinkNode {
148    pub metadata: InodeMetadata,
149    pub target: Vec<u8>,
150}
151
152#[derive(Clone)]
153pub struct DeviceNode {
154    pub metadata: InodeMetadata,
155    pub major: u32,
156    pub minor: u32,
157}
158
159#[derive(Clone)]
160pub struct FileTree {
161    pub root: DirectoryNode,
162}
163
164#[derive(Debug)]
165pub enum FileTreeError {
166    PathEmpty,
167    PathTraversal(String),
168    NotADirectory(String),
169    EntryExists(String),
170}
171
172//--------------------------------------------------------------------------------------------------
173// Methods
174//--------------------------------------------------------------------------------------------------
175
176impl FileData {
177    /// Total byte length of the file content.
178    pub fn len(&self) -> usize {
179        match self {
180            FileData::Memory(v) => v.len(),
181            FileData::SharedMemory(v) => v.len(),
182            FileData::Spool { len, .. } => *len as usize,
183        }
184    }
185
186    pub fn is_empty(&self) -> bool {
187        self.len() == 0
188    }
189
190    /// Read the entire content into memory. For `Memory` variant this
191    /// clones; for `Spool` this reads from disk.
192    pub fn read_all(&self) -> std::io::Result<Vec<u8>> {
193        match self {
194            FileData::Memory(v) => Ok(v.clone()),
195            FileData::SharedMemory(v) => Ok(v.to_vec()),
196            FileData::Spool { spool, offset, len } => {
197                let mut buf = vec![0u8; *len as usize];
198                let mut file = spool
199                    .lock()
200                    .map_err(|_| std::io::Error::other("spool lock poisoned"))?;
201                file.seek(SeekFrom::Start(*offset))?;
202                file.read_exact(&mut buf)?;
203                Ok(buf)
204            }
205        }
206    }
207
208    /// Borrow the in-memory bytes directly (only for `Memory` variant).
209    pub fn as_bytes(&self) -> Option<&[u8]> {
210        match self {
211            FileData::Memory(v) => Some(v),
212            FileData::SharedMemory(v) => Some(v),
213            FileData::Spool { .. } => None,
214        }
215    }
216
217    /// Write content to an output writer, reading from spool if needed.
218    /// Avoids loading the entire file into memory for large spooled files.
219    pub fn write_to(&self, out: &mut impl std::io::Write) -> std::io::Result<()> {
220        self.write_range(0, self.len(), out)
221    }
222
223    /// Write a byte range of the content to an output writer.
224    pub fn write_range(
225        &self,
226        start: usize,
227        len: usize,
228        out: &mut impl std::io::Write,
229    ) -> std::io::Result<()> {
230        match self {
231            FileData::Memory(v) => out.write_all(&v[start..start + len]),
232            FileData::SharedMemory(v) => out.write_all(&v[start..start + len]),
233            FileData::Spool { spool, offset, .. } => {
234                let mut file = spool
235                    .lock()
236                    .map_err(|_| std::io::Error::other("spool lock poisoned"))?;
237                file.seek(SeekFrom::Start(*offset + start as u64))?;
238                let mut remaining = len;
239                let mut buf = [0u8; 65536];
240                while remaining > 0 {
241                    let to_read = remaining.min(buf.len());
242                    file.read_exact(&mut buf[..to_read])?;
243                    out.write_all(&buf[..to_read])?;
244                    remaining -= to_read;
245                }
246                Ok(())
247            }
248        }
249    }
250
251    /// Return a shared reference suitable for a hardlinked alias.
252    pub fn clone_ref(&mut self) -> FileData {
253        match self {
254            FileData::Memory(data) => {
255                let shared: Arc<[u8]> = std::mem::take(data).into();
256                *self = FileData::SharedMemory(Arc::clone(&shared));
257                FileData::SharedMemory(shared)
258            }
259            FileData::SharedMemory(data) => FileData::SharedMemory(Arc::clone(data)),
260            FileData::Spool { spool, offset, len } => FileData::Spool {
261                spool: Arc::clone(spool),
262                offset: *offset,
263                len: *len,
264            },
265        }
266    }
267}
268
269impl RegularFileId {
270    pub fn new() -> Self {
271        static NEXT_ID: AtomicU64 = AtomicU64::new(1);
272        Self(NEXT_ID.fetch_add(1, Ordering::Relaxed))
273    }
274}
275
276impl Default for RegularFileId {
277    fn default() -> Self {
278        Self::new()
279    }
280}
281
282impl DataSpool {
283    /// Create a new spool file at the given path.
284    pub fn new(path: &std::path::Path) -> std::io::Result<Self> {
285        let file = std::fs::OpenOptions::new()
286            .create(true)
287            .truncate(true)
288            .read(true)
289            .write(true)
290            .open(path)?;
291        let shared = Arc::new(std::sync::Mutex::new(file.try_clone()?));
292        Ok(Self {
293            file,
294            shared,
295            offset: 0,
296        })
297    }
298
299    /// Write data to the spool and return a `FileData::Spool` reference.
300    pub fn write_data(&mut self, data: &[u8]) -> std::io::Result<FileData> {
301        use std::io::Write;
302        let offset = self.offset;
303        self.file.write_all(data)?;
304        self.offset += data.len() as u64;
305        Ok(FileData::Spool {
306            spool: Arc::clone(&self.shared),
307            offset,
308            len: data.len() as u64,
309        })
310    }
311
312    pub fn current_offset(&self) -> u64 {
313        self.offset
314    }
315
316    pub fn write_chunk(&mut self, data: &[u8]) -> std::io::Result<()> {
317        use std::io::Write;
318        self.file.write_all(data)?;
319        self.offset += data.len() as u64;
320        Ok(())
321    }
322
323    pub fn data_ref(&self, offset: u64, len: u64) -> FileData {
324        FileData::Spool {
325            spool: Arc::clone(&self.shared),
326            offset,
327            len,
328        }
329    }
330}
331
332impl DirectoryNode {
333    pub fn new(metadata: InodeMetadata) -> Self {
334        Self {
335            metadata,
336            xattrs: Vec::new(),
337            entries: BTreeMap::new(),
338        }
339    }
340
341    pub fn entry_count(&self) -> usize {
342        self.entries.len()
343    }
344}
345
346impl Default for FileTree {
347    fn default() -> Self {
348        Self::new()
349    }
350}
351
352impl FileTree {
353    pub fn new() -> Self {
354        Self {
355            root: DirectoryNode::new(InodeMetadata::default()),
356        }
357    }
358
359    pub fn insert(&mut self, path: &[u8], node: TreeNode) -> Result<(), FileTreeError> {
360        use std::collections::btree_map::Entry;
361
362        let components = split_path(path)?;
363        if components.is_empty() {
364            return Err(FileTreeError::PathEmpty);
365        }
366
367        let (parent_components, file_name) = components.split_at(components.len() - 1);
368
369        // Traverse to the parent directory, creating missing intermediates.
370        // Uses the BTreeMap entry API to do a single lookup per component
371        // instead of contains_key + insert + get_mut (3 lookups).
372        let mut current = &mut self.root;
373        for component in parent_components {
374            let key = OsStr::from_bytes(component).to_os_string();
375            current = match current.entries.entry(key) {
376                Entry::Vacant(e) => {
377                    let dir = TreeNode::Directory(DirectoryNode::new(InodeMetadata::default()));
378                    match e.insert(dir) {
379                        TreeNode::Directory(d) => d,
380                        _ => unreachable!(),
381                    }
382                }
383                Entry::Occupied(e) => match e.into_mut() {
384                    TreeNode::Directory(d) => d,
385                    _ => {
386                        let path_str = String::from_utf8_lossy(component).into_owned();
387                        return Err(FileTreeError::NotADirectory(path_str));
388                    }
389                },
390            };
391        }
392
393        // Insert the final node. Directory-over-directory merges metadata
394        // but keeps existing entries. Non-directory replaces non-directory.
395        let key = OsStr::from_bytes(file_name[0]).to_os_string();
396        match current.entries.entry(key) {
397            Entry::Vacant(e) => {
398                e.insert(node);
399            }
400            Entry::Occupied(mut e) => match (e.get(), &node) {
401                (TreeNode::Directory(_), TreeNode::Directory(_)) => {
402                    if let TreeNode::Directory(existing) = e.get_mut()
403                        && let TreeNode::Directory(new_dir) = node
404                    {
405                        existing.metadata = new_dir.metadata;
406                        existing.xattrs = new_dir.xattrs;
407                    }
408                }
409                (TreeNode::Directory(_), _) => {
410                    let path_str = String::from_utf8_lossy(file_name[0]).into_owned();
411                    return Err(FileTreeError::EntryExists(path_str));
412                }
413                _ => {
414                    e.insert(node);
415                }
416            },
417        }
418
419        Ok(())
420    }
421
422    pub fn get(&self, path: &[u8]) -> Option<&TreeNode> {
423        let components = split_path(path).ok()?;
424        if components.is_empty() {
425            return None;
426        }
427
428        let (parent_components, file_name) = components.split_at(components.len() - 1);
429
430        let mut current = &self.root;
431        for component in parent_components {
432            let key = OsStr::from_bytes(component);
433            match current.entries.get(key) {
434                Some(TreeNode::Directory(dir)) => {
435                    current = dir;
436                }
437                _ => return None,
438            }
439        }
440
441        current.entries.get(OsStr::from_bytes(file_name[0]))
442    }
443
444    pub fn get_mut(&mut self, path: &[u8]) -> Option<&mut TreeNode> {
445        let components = split_path(path).ok()?;
446        if components.is_empty() {
447            return None;
448        }
449
450        let (parent_components, file_name) = components.split_at(components.len() - 1);
451
452        let mut current = &mut self.root;
453        for component in parent_components {
454            let key = OsStr::from_bytes(component);
455            match current.entries.get_mut(key) {
456                Some(TreeNode::Directory(dir)) => {
457                    current = dir;
458                }
459                _ => return None,
460            }
461        }
462
463        current.entries.get_mut(OsStr::from_bytes(file_name[0]))
464    }
465
466    pub fn remove(&mut self, path: &[u8]) -> Option<TreeNode> {
467        let components = split_path(path).ok()?;
468        if components.is_empty() {
469            return None;
470        }
471
472        let (parent_components, file_name) = components.split_at(components.len() - 1);
473
474        let mut current = &mut self.root;
475        for component in parent_components {
476            let key = OsStr::from_bytes(component);
477            match current.entries.get_mut(key) {
478                Some(TreeNode::Directory(dir)) => {
479                    current = dir;
480                }
481                _ => return None,
482            }
483        }
484
485        current.entries.remove(OsStr::from_bytes(file_name[0]))
486    }
487
488    pub fn node_count(&self) -> u64 {
489        count_nodes_in_dir(&self.root)
490    }
491
492    pub fn total_data_size(&self) -> u64 {
493        data_size_in_dir(&self.root)
494    }
495
496    pub(crate) fn regular_file_link_counts(&self) -> HashMap<RegularFileId, u32> {
497        let mut counts = HashMap::new();
498        count_regular_links_in_dir(&self.root, &mut counts);
499        counts
500    }
501
502    pub(crate) fn refresh_regular_nlinks(&mut self) {
503        let counts = self.regular_file_link_counts();
504        refresh_regular_nlinks_in_dir(&mut self.root, &counts);
505    }
506
507    pub fn merge_layer(&mut self, layer: FileTree) {
508        merge_directory(&mut self.root, layer.root);
509        self.refresh_regular_nlinks();
510    }
511
512    /// Strip file data from this tree, keeping only directory structure and metadata.
513    ///
514    /// After calling this, all `RegularFile` nodes have empty `FileData::Memory(Vec::new())`.
515    /// Used to reduce memory after writing a per-layer EROFS while retaining the tree
516    /// for fsmeta merge.
517    pub fn strip_file_data(&mut self) {
518        strip_data_in_dir(&mut self.root);
519    }
520}
521
522//--------------------------------------------------------------------------------------------------
523// Trait Implementations
524//--------------------------------------------------------------------------------------------------
525
526impl Default for ResourceLimits {
527    fn default() -> Self {
528        Self {
529            max_total_size: DEFAULT_MAX_TOTAL_SIZE,
530            max_file_size: DEFAULT_MAX_FILE_SIZE,
531            max_entry_count: DEFAULT_MAX_ENTRY_COUNT,
532            max_path_length: DEFAULT_MAX_PATH_LENGTH,
533            max_path_depth: DEFAULT_MAX_PATH_DEPTH,
534            max_symlink_target: DEFAULT_MAX_SYMLINK_TARGET,
535        }
536    }
537}
538
539impl Default for InodeMetadata {
540    fn default() -> Self {
541        Self {
542            uid: 0,
543            gid: 0,
544            mode: DEFAULT_DIR_MODE,
545            mtime: 0,
546            mtime_nsec: 0,
547        }
548    }
549}
550
551impl fmt::Display for FileTreeError {
552    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
553        match self {
554            FileTreeError::PathEmpty => write!(f, "path is empty"),
555            FileTreeError::PathTraversal(p) => {
556                write!(f, "path traversal attempt: \"..\" in path \"{p}\"")
557            }
558            FileTreeError::NotADirectory(p) => {
559                write!(f, "not a directory: \"{p}\"")
560            }
561            FileTreeError::EntryExists(p) => {
562                write!(f, "entry already exists: \"{p}\"")
563            }
564        }
565    }
566}
567
568impl std::error::Error for FileTreeError {}
569
570//--------------------------------------------------------------------------------------------------
571// Functions
572//--------------------------------------------------------------------------------------------------
573
574fn split_path(path: &[u8]) -> Result<Vec<&[u8]>, FileTreeError> {
575    let components: Vec<&[u8]> = path
576        .split(|&b| b == b'/')
577        .filter(|c| !c.is_empty())
578        .collect();
579
580    if components.is_empty() {
581        return Err(FileTreeError::PathEmpty);
582    }
583
584    for component in &components {
585        if *component == b".." {
586            let path_str = String::from_utf8_lossy(path).into_owned();
587            return Err(FileTreeError::PathTraversal(path_str));
588        }
589    }
590
591    Ok(components)
592}
593
594fn count_nodes_in_dir(dir: &DirectoryNode) -> u64 {
595    let mut count = 0u64;
596    for node in dir.entries.values() {
597        count += 1;
598        if let TreeNode::Directory(child_dir) = node {
599            count += count_nodes_in_dir(child_dir);
600        }
601    }
602    count
603}
604
605fn data_size_in_dir(dir: &DirectoryNode) -> u64 {
606    let mut seen = HashSet::new();
607    data_size_in_dir_once(dir, &mut seen)
608}
609
610fn data_size_in_dir_once(dir: &DirectoryNode, seen: &mut HashSet<RegularFileId>) -> u64 {
611    let mut size = 0u64;
612    for node in dir.entries.values() {
613        match node {
614            TreeNode::RegularFile(file) if seen.insert(file.id) => {
615                size += file.data.len() as u64;
616            }
617            TreeNode::Directory(child_dir) => {
618                size += data_size_in_dir_once(child_dir, seen);
619            }
620            _ => {}
621        }
622    }
623    size
624}
625
626fn count_regular_links_in_dir(dir: &DirectoryNode, counts: &mut HashMap<RegularFileId, u32>) {
627    for node in dir.entries.values() {
628        match node {
629            TreeNode::RegularFile(file) => {
630                *counts.entry(file.id).or_insert(0) += 1;
631            }
632            TreeNode::Directory(child_dir) => {
633                count_regular_links_in_dir(child_dir, counts);
634            }
635            _ => {}
636        }
637    }
638}
639
640fn refresh_regular_nlinks_in_dir(dir: &mut DirectoryNode, counts: &HashMap<RegularFileId, u32>) {
641    for node in dir.entries.values_mut() {
642        match node {
643            TreeNode::RegularFile(file) => {
644                file.nlink = counts.get(&file.id).copied().unwrap_or(1);
645            }
646            TreeNode::Directory(child_dir) => {
647                refresh_regular_nlinks_in_dir(child_dir, counts);
648            }
649            _ => {}
650        }
651    }
652}
653
654fn strip_data_in_dir(dir: &mut DirectoryNode) {
655    for node in dir.entries.values_mut() {
656        match node {
657            TreeNode::RegularFile(f) => {
658                f.data = FileData::Memory(Vec::new());
659            }
660            TreeNode::Directory(d) => {
661                strip_data_in_dir(d);
662            }
663            _ => {}
664        }
665    }
666}
667
668/// Merge multiple layer trees into a single tree with provenance tracking.
669///
670/// Applies layers bottom-to-top, handling whiteouts and opaque directories.
671/// Unlike `merge_layer()`, whiteout entries and opaque xattrs are consumed
672/// and NOT propagated — the output is the clean final state.
673///
674/// Returns the merged tree and a map from file path to source layer index
675/// (0-based, bottom-to-top).
676pub fn merge_layers_with_provenance(layers: Vec<FileTree>) -> (FileTree, HashMap<PathBuf, usize>) {
677    let mut merged = FileTree::new();
678    let mut provenance: HashMap<PathBuf, usize> = HashMap::new();
679
680    for (layer_idx, layer) in layers.into_iter().enumerate() {
681        let path = PathBuf::new();
682        merge_directory_with_provenance(
683            &mut merged.root,
684            layer.root,
685            layer_idx,
686            &path,
687            &mut provenance,
688        );
689    }
690
691    // Strip opaque xattrs from the final merged tree — they are overlayfs
692    // directives consumed by the merge, not meaningful in fsmeta.
693    strip_opaque_xattrs(&mut merged.root);
694    merged.refresh_regular_nlinks();
695
696    (merged, provenance)
697}
698
699fn merge_directory_with_provenance(
700    base: &mut DirectoryNode,
701    layer: DirectoryNode,
702    layer_idx: usize,
703    current_path: &Path,
704    provenance: &mut HashMap<PathBuf, usize>,
705) {
706    for (name, layer_node) in layer.entries {
707        let child_path = current_path.join(&name);
708
709        // Whiteout: remove target from merged tree and its provenance.
710        if is_whiteout_device(&layer_node) {
711            // Remove provenance entries for the deleted item and all its descendants.
712            base.entries.remove(&name);
713            provenance.retain(|k, _| !k.starts_with(&child_path));
714            continue;
715        }
716
717        match layer_node {
718            TreeNode::Directory(layer_dir) => {
719                let opaque = has_opaque_xattr(&layer_dir);
720
721                match base.entries.get_mut(&name) {
722                    Some(TreeNode::Directory(base_dir)) => {
723                        if opaque {
724                            // Remove all provenance for entries under this directory.
725                            provenance.retain(|k, _| !k.starts_with(&child_path));
726                            base_dir.entries.clear();
727                        }
728                        base_dir.metadata = layer_dir.metadata;
729                        base_dir.xattrs = layer_dir.xattrs;
730                        merge_directory_with_provenance(
731                            base_dir,
732                            DirectoryNode {
733                                metadata: InodeMetadata::default(),
734                                xattrs: Vec::new(),
735                                entries: layer_dir.entries,
736                            },
737                            layer_idx,
738                            &child_path,
739                            provenance,
740                        );
741                    }
742                    _ => {
743                        // New directory replaces whatever was there.
744                        provenance.retain(|k, _| !k.starts_with(&child_path));
745                        // Record provenance for all entries in the new directory.
746                        record_provenance_recursive(&layer_dir, layer_idx, &child_path, provenance);
747                        base.entries.insert(name, TreeNode::Directory(layer_dir));
748                    }
749                }
750            }
751            other => {
752                // Non-directory entry: record provenance.
753                provenance.insert(child_path, layer_idx);
754                base.entries.insert(name, other);
755            }
756        }
757    }
758}
759
760fn record_provenance_recursive(
761    dir: &DirectoryNode,
762    layer_idx: usize,
763    current_path: &Path,
764    provenance: &mut HashMap<PathBuf, usize>,
765) {
766    for (name, child) in &dir.entries {
767        let child_path = current_path.join(name);
768        match child {
769            TreeNode::Directory(child_dir) => {
770                record_provenance_recursive(child_dir, layer_idx, &child_path, provenance);
771            }
772            _ => {
773                provenance.insert(child_path, layer_idx);
774            }
775        }
776    }
777}
778
779fn strip_opaque_xattrs(dir: &mut DirectoryNode) {
780    dir.xattrs
781        .retain(|x| !(x.name == OPAQUE_XATTR_NAME && x.value == OPAQUE_XATTR_VALUE));
782    for node in dir.entries.values_mut() {
783        if let TreeNode::Directory(child_dir) = node {
784            strip_opaque_xattrs(child_dir);
785        }
786    }
787}
788
789fn is_whiteout_device(node: &TreeNode) -> bool {
790    matches!(node, TreeNode::CharDevice(dev) if dev.major == WHITEOUT_MAJOR && dev.minor == WHITEOUT_MINOR)
791}
792
793fn has_opaque_xattr(dir: &DirectoryNode) -> bool {
794    dir.xattrs
795        .iter()
796        .any(|x| x.name == OPAQUE_XATTR_NAME && x.value == OPAQUE_XATTR_VALUE)
797}
798
799fn merge_directory(base: &mut DirectoryNode, layer: DirectoryNode) {
800    for (name, layer_node) in layer.entries {
801        if is_whiteout_device(&layer_node) {
802            base.entries.remove(&name);
803            continue;
804        }
805
806        match layer_node {
807            TreeNode::Directory(layer_dir) => {
808                let opaque = has_opaque_xattr(&layer_dir);
809
810                match base.entries.get_mut(&name) {
811                    Some(TreeNode::Directory(base_dir)) => {
812                        if opaque {
813                            base_dir.entries.clear();
814                        }
815                        base_dir.metadata = layer_dir.metadata;
816                        base_dir.xattrs = layer_dir.xattrs;
817                        merge_directory(
818                            base_dir,
819                            DirectoryNode {
820                                metadata: InodeMetadata::default(),
821                                xattrs: Vec::new(),
822                                entries: layer_dir.entries,
823                            },
824                        );
825                    }
826                    _ => {
827                        base.entries.insert(name, TreeNode::Directory(layer_dir));
828                    }
829                }
830            }
831            other => {
832                base.entries.insert(name, other);
833            }
834        }
835    }
836}
837
838//--------------------------------------------------------------------------------------------------
839// Tests
840//--------------------------------------------------------------------------------------------------
841
842#[cfg(test)]
843mod tests {
844    use super::*;
845
846    fn make_regular_file(data: &[u8]) -> TreeNode {
847        make_regular_file_with_id(data, RegularFileId::new())
848    }
849
850    fn make_regular_file_with_id(data: &[u8], id: RegularFileId) -> TreeNode {
851        TreeNode::RegularFile(RegularFileNode {
852            id,
853            metadata: InodeMetadata::default(),
854            xattrs: Vec::new(),
855            data: FileData::Memory(data.to_vec()),
856            nlink: 1,
857        })
858    }
859
860    fn make_directory() -> TreeNode {
861        TreeNode::Directory(DirectoryNode::new(InodeMetadata::default()))
862    }
863
864    fn make_whiteout() -> TreeNode {
865        TreeNode::CharDevice(DeviceNode {
866            metadata: InodeMetadata::default(),
867            major: 0,
868            minor: 0,
869        })
870    }
871
872    fn make_opaque_directory() -> DirectoryNode {
873        DirectoryNode {
874            metadata: InodeMetadata::default(),
875            xattrs: vec![Xattr {
876                name: OPAQUE_XATTR_NAME.to_vec(),
877                value: OPAQUE_XATTR_VALUE.to_vec(),
878            }],
879            entries: BTreeMap::new(),
880        }
881    }
882
883    #[test]
884    fn insert_and_get_file() {
885        let mut tree = FileTree::new();
886        tree.insert(b"hello.txt", make_regular_file(b"hello world"))
887            .unwrap();
888
889        let node = tree.get(b"hello.txt").unwrap();
890        match node {
891            TreeNode::RegularFile(f) => {
892                assert_eq!(f.data, FileData::Memory(b"hello world".to_vec()))
893            }
894            _ => panic!("expected regular file"),
895        }
896    }
897
898    #[test]
899    fn insert_with_missing_parents_creates_them() {
900        let mut tree = FileTree::new();
901        tree.insert(b"a/b/c/file.txt", make_regular_file(b"deep"))
902            .unwrap();
903
904        // Intermediate directories should exist.
905        let node = tree.get(b"a").unwrap();
906        assert!(matches!(node, TreeNode::Directory(_)));
907
908        let node = tree.get(b"a/b").unwrap();
909        assert!(matches!(node, TreeNode::Directory(_)));
910
911        let node = tree.get(b"a/b/c").unwrap();
912        assert!(matches!(node, TreeNode::Directory(_)));
913
914        let node = tree.get(b"a/b/c/file.txt").unwrap();
915        assert!(matches!(node, TreeNode::RegularFile(_)));
916    }
917
918    #[test]
919    fn reject_dotdot_in_path() {
920        let mut tree = FileTree::new();
921        let result = tree.insert(b"a/../etc/passwd", make_regular_file(b"bad"));
922        assert!(matches!(result, Err(FileTreeError::PathTraversal(_))));
923    }
924
925    #[test]
926    fn merge_layer_replaces_file() {
927        let mut base = FileTree::new();
928        base.insert(b"config.txt", make_regular_file(b"old"))
929            .unwrap();
930
931        let mut layer = FileTree::new();
932        layer
933            .insert(b"config.txt", make_regular_file(b"new"))
934            .unwrap();
935
936        base.merge_layer(layer);
937
938        match base.get(b"config.txt").unwrap() {
939            TreeNode::RegularFile(f) => assert_eq!(f.data, FileData::Memory(b"new".to_vec())),
940            _ => panic!("expected regular file"),
941        }
942    }
943
944    #[test]
945    fn merge_layer_whiteout_removes_file() {
946        let mut base = FileTree::new();
947        base.insert(b"dir/secret.txt", make_regular_file(b"sensitive"))
948            .unwrap();
949
950        let mut layer = FileTree::new();
951        layer.insert(b"dir", make_directory()).unwrap();
952        layer.insert(b"dir/secret.txt", make_whiteout()).unwrap();
953
954        base.merge_layer(layer);
955
956        assert!(base.get(b"dir/secret.txt").is_none());
957        // The parent directory should still exist.
958        assert!(base.get(b"dir").is_some());
959    }
960
961    #[test]
962    fn merge_layer_opaque_dir_clears_existing_entries() {
963        let mut base = FileTree::new();
964        base.insert(b"dir/a.txt", make_regular_file(b"a")).unwrap();
965        base.insert(b"dir/b.txt", make_regular_file(b"b")).unwrap();
966
967        let mut layer = FileTree::new();
968        let mut opaque_dir = make_opaque_directory();
969        opaque_dir
970            .entries
971            .insert(OsString::from("c.txt"), make_regular_file(b"c"));
972        layer
973            .root
974            .entries
975            .insert(OsString::from("dir"), TreeNode::Directory(opaque_dir));
976
977        base.merge_layer(layer);
978
979        // Old entries should be gone.
980        assert!(base.get(b"dir/a.txt").is_none());
981        assert!(base.get(b"dir/b.txt").is_none());
982        // New entry should be present.
983        match base.get(b"dir/c.txt").unwrap() {
984            TreeNode::RegularFile(f) => assert_eq!(f.data, FileData::Memory(b"c".to_vec())),
985            _ => panic!("expected regular file"),
986        }
987    }
988
989    #[test]
990    fn node_count_and_data_size() {
991        let mut tree = FileTree::new();
992        tree.insert(b"a/file1.txt", make_regular_file(b"hello"))
993            .unwrap();
994        tree.insert(b"a/file2.txt", make_regular_file(b"world!"))
995            .unwrap();
996        tree.insert(b"b/nested/file3.txt", make_regular_file(b"!"))
997            .unwrap();
998
999        // a, a/file1.txt, a/file2.txt, b, b/nested, b/nested/file3.txt = 6
1000        assert_eq!(tree.node_count(), 6);
1001        // 5 + 6 + 1 = 12
1002        assert_eq!(tree.total_data_size(), 12);
1003    }
1004
1005    #[test]
1006    fn data_size_counts_hardlinked_regular_file_once() {
1007        let mut tree = FileTree::new();
1008        let file_id = RegularFileId::new();
1009
1010        tree.insert(b"a.txt", make_regular_file_with_id(b"shared", file_id))
1011            .unwrap();
1012        tree.insert(b"b.txt", make_regular_file_with_id(b"shared", file_id))
1013            .unwrap();
1014
1015        tree.refresh_regular_nlinks();
1016
1017        assert_eq!(tree.total_data_size(), b"shared".len() as u64);
1018        for path in [b"a.txt".as_slice(), b"b.txt".as_slice()] {
1019            match tree.get(path).unwrap() {
1020                TreeNode::RegularFile(file) => assert_eq!(file.nlink, 2),
1021                _ => panic!("expected regular file"),
1022            }
1023        }
1024    }
1025
1026    #[test]
1027    fn remove_node() {
1028        let mut tree = FileTree::new();
1029        tree.insert(b"a/b.txt", make_regular_file(b"data")).unwrap();
1030        assert!(tree.get(b"a/b.txt").is_some());
1031
1032        let removed = tree.remove(b"a/b.txt");
1033        assert!(removed.is_some());
1034        assert!(tree.get(b"a/b.txt").is_none());
1035    }
1036
1037    #[test]
1038    fn empty_path_is_rejected() {
1039        let mut tree = FileTree::new();
1040        let result = tree.insert(b"", make_regular_file(b"data"));
1041        assert!(matches!(result, Err(FileTreeError::PathEmpty)));
1042    }
1043
1044    #[test]
1045    fn not_a_directory_error() {
1046        let mut tree = FileTree::new();
1047        tree.insert(b"a", make_regular_file(b"file")).unwrap();
1048
1049        let result = tree.insert(b"a/b", make_regular_file(b"nested"));
1050        assert!(matches!(result, Err(FileTreeError::NotADirectory(_))));
1051    }
1052
1053    #[test]
1054    fn resource_limits_default() {
1055        let limits = ResourceLimits::default();
1056        assert_eq!(limits.max_total_size, 10 * 1024 * 1024 * 1024);
1057        assert_eq!(limits.max_file_size, 5 * 1024 * 1024 * 1024);
1058        assert_eq!(limits.max_entry_count, 1_000_000);
1059        assert_eq!(limits.max_path_length, 4096);
1060        assert_eq!(limits.max_path_depth, 128);
1061        assert_eq!(limits.max_symlink_target, 4096);
1062    }
1063}