Skip to main content

monty_fs/
overlay_state.rs

1//! Internal storage for in-memory overlay mounts.
2//!
3//! This module keeps the overlay data structures separate from the public
4//! [`MountMode`](super::MountMode) definition so the public API stays easy to
5//! scan while the storage internals can evolve independently.
6
7use std::{
8    collections::BTreeMap,
9    fs, mem,
10    ops::Bound,
11    path::{Path, PathBuf},
12    time::SystemTime,
13};
14
15use super::{MountError, common::as_u64};
16
17/// Conservative bookkeeping charge for each overlay map entry.
18///
19/// This covers the map node, key allocation, and entry metadata. Variable-size
20/// file contents and host paths are charged separately.
21pub(super) const ENTRY_MEMORY_USAGE: u64 = 256;
22
23/// In-memory overlay state for [`super::MountMode::OverlayMemory`].
24///
25/// A single [`BTreeMap`] stores relative mount paths and the overlay entry that
26/// currently shadows or extends the underlying real filesystem.
27#[derive(Debug, Default)]
28pub struct OverlayState {
29    /// Entries keyed by forward-slash-separated relative path (e.g.
30    /// `"subdir/file.txt"`). The mount root is represented by `""`.
31    ///
32    /// [`BTreeMap`] is used so prefix walks for directory operations can stay
33    /// `O(log n + k)` rather than scanning the entire overlay.
34    entries: BTreeMap<String, OverlayEntry>,
35    /// Estimated live bytes retained by `entries`.
36    memory_usage: u64,
37}
38
39impl OverlayState {
40    /// Creates a new empty overlay state.
41    #[must_use]
42    pub fn new() -> Self {
43        Self::default()
44    }
45
46    /// Looks up the overlay entry for `relative_path`.
47    #[must_use]
48    pub(super) fn get(&self, relative_path: &str) -> Option<&OverlayEntry> {
49        self.entries.get(relative_path)
50    }
51
52    /// Returns the estimated live memory retained by this overlay.
53    #[must_use]
54    pub(super) fn memory_usage(&self) -> u64 {
55        self.memory_usage
56    }
57
58    /// Removes and returns the entry for `relative_path`.
59    pub(super) fn remove(&mut self, relative_path: &str) -> Option<OverlayEntry> {
60        let entry = self.entries.remove(relative_path)?;
61        self.memory_usage = self
62            .memory_usage
63            .saturating_sub(entry_memory_usage(relative_path, &entry));
64        Some(entry)
65    }
66
67    /// Inserts an entry if the resulting overlay stays within `limit`.
68    pub(super) fn insert(&mut self, relative_path: String, entry: OverlayEntry, limit: u64) -> Result<(), MountError> {
69        let projected = self.projected_usage(&relative_path, &entry);
70        if projected > limit {
71            Err(MountError::MemoryUsageLimitExceeded(limit))
72        } else {
73            self.entries.insert(relative_path, entry);
74            self.memory_usage = projected;
75            Ok(())
76        }
77    }
78
79    /// Inserts an entry after the caller has preflighted a multi-entry update.
80    pub(super) fn insert_unchecked(&mut self, relative_path: String, entry: OverlayEntry) {
81        self.memory_usage = self.projected_usage(&relative_path, &entry);
82        self.entries.insert(relative_path, entry);
83    }
84
85    /// Returns total retained usage as if `entry` replaced `relative_path`.
86    fn projected_usage(&self, relative_path: &str, entry: &OverlayEntry) -> u64 {
87        let old_usage = self
88            .entries
89            .get(relative_path)
90            .map_or(0, |old| entry_memory_usage(relative_path, old));
91        let new_usage = entry_memory_usage(relative_path, entry);
92        self.memory_usage.saturating_sub(old_usage).saturating_add(new_usage)
93    }
94
95    /// Appends bytes to an overlay file while accounting for retained content.
96    pub(super) fn append_file(
97        &mut self,
98        relative_path: &str,
99        data: &[u8],
100        mtime: f64,
101        limit: u64,
102    ) -> Result<bool, MountError> {
103        let Some(OverlayEntry::File(file)) = self.entries.get_mut(relative_path) else {
104            return Ok(false);
105        };
106        let projected = self.memory_usage.saturating_add(as_u64(data.len()));
107        if projected > limit {
108            Err(MountError::MemoryUsageLimitExceeded(limit))
109        } else {
110            file.content.extend_from_slice(data);
111            file.mtime = mtime;
112            self.memory_usage = projected;
113            Ok(true)
114        }
115    }
116
117    /// Checks replacing `relative_path` with a file of `content_len` bytes.
118    pub(super) fn check_file_replacement(
119        &self,
120        relative_path: &str,
121        content_len: usize,
122        limit: u64,
123    ) -> Result<(), MountError> {
124        let old_usage = self
125            .entries
126            .get(relative_path)
127            .map_or(0, |old| entry_memory_usage(relative_path, old));
128        let new_usage = base_entry_memory_usage(relative_path).saturating_add(as_u64(content_len));
129        let projected = self.memory_usage.saturating_sub(old_usage).saturating_add(new_usage);
130        if projected > limit {
131            Err(MountError::MemoryUsageLimitExceeded(limit))
132        } else {
133            Ok(())
134        }
135    }
136
137    /// Checks a sequence of replacements as one atomic overlay update.
138    pub(super) fn check_replacements<'a>(
139        &self,
140        replacements: impl IntoIterator<Item = (&'a str, &'a OverlayEntry)>,
141        limit: u64,
142    ) -> Result<(), MountError> {
143        let mut projected = self.memory_usage;
144        let mut replaced = BTreeMap::new();
145        for (path, entry) in replacements {
146            let old_usage = replaced
147                .get(path)
148                .copied()
149                .unwrap_or_else(|| self.entries.get(path).map_or(0, |old| entry_memory_usage(path, old)));
150            let new_usage = entry_memory_usage(path, entry);
151            projected = projected.saturating_sub(old_usage).saturating_add(new_usage);
152            replaced.insert(path, new_usage);
153        }
154        if projected > limit {
155            Err(MountError::MemoryUsageLimitExceeded(limit))
156        } else {
157            Ok(())
158        }
159    }
160
161    /// Iterates over overlay entries whose keys start with `prefix`.
162    ///
163    /// `prefix` must be either `""` or end with `'/'`. The upper bound uses a
164    /// lexical successor so the range query stays tight without scanning the
165    /// whole map.
166    pub(super) fn prefix_iter(&self, prefix: &str) -> impl Iterator<Item = (&str, &OverlayEntry)> {
167        debug_assert!(prefix.is_empty() || prefix.ends_with('/'));
168
169        let upper_storage;
170        let bounds: (Bound<&str>, Bound<&str>) = if prefix.is_empty() {
171            (Bound::Unbounded, Bound::Unbounded)
172        } else {
173            upper_storage = {
174                let mut upper = prefix.to_owned();
175                upper.pop();
176                upper.push('0');
177                upper
178            };
179            (Bound::Included(prefix), Bound::Excluded(upper_storage.as_str()))
180        };
181
182        self.entries
183            .range::<str, _>(bounds)
184            .map(|(key, value)| (key.as_str(), value))
185    }
186}
187
188/// Estimates retained heap bytes for one overlay entry.
189fn entry_memory_usage(relative_path: &str, entry: &OverlayEntry) -> u64 {
190    let variable = match entry {
191        OverlayEntry::File(file) => file.content.len(),
192        OverlayEntry::RealFileRef(file_ref) => file_ref.host_path.as_os_str().len(),
193        OverlayEntry::Directory { .. } | OverlayEntry::Deleted => 0,
194    };
195    base_entry_memory_usage(relative_path).saturating_add(as_u64(variable))
196}
197
198/// Returns the fixed and key-dependent charge for an overlay entry.
199fn base_entry_memory_usage(relative_path: &str) -> u64 {
200    ENTRY_MEMORY_USAGE
201        .saturating_add(as_u64(relative_path.len()))
202        .saturating_add(as_u64(mem::size_of::<OverlayEntry>()))
203}
204
205/// An entry stored in an overlay mount.
206#[derive(Debug)]
207pub(super) enum OverlayEntry {
208    /// A file written by sandbox code and stored directly in memory.
209    File(OverlayFile),
210
211    /// A lazily-read reference to a real host file that has been renamed into
212    /// the overlay without eagerly loading its contents.
213    RealFileRef(OverlayFileRef),
214
215    /// A directory that exists only in the overlay.
216    Directory {
217        /// Modification time recorded for synthetic stat results.
218        mtime: f64,
219    },
220
221    /// A tombstone hiding a real or previously-overlay entry.
222    Deleted,
223}
224
225/// In-memory contents of a file owned by the overlay.
226#[derive(Debug)]
227pub(super) struct OverlayFile {
228    /// Raw file contents.
229    pub content: Vec<u8>,
230    /// Modification time recorded for synthetic stat results.
231    pub mtime: f64,
232}
233
234/// A lazy reference to a real host file preserved during overlay rename.
235#[derive(Debug)]
236pub(super) struct OverlayFileRef {
237    /// Canonical host path for the original file contents.
238    pub host_path: PathBuf,
239    /// Modification time copied from the original file.
240    pub mtime: f64,
241    /// File size in bytes.
242    pub size: i64,
243}
244
245impl OverlayFileRef {
246    /// Builds a lazy file reference from a host path if metadata can be read.
247    ///
248    /// Uses `fs::metadata` which follows symlinks, so the size and mtime
249    /// reflect the target file. Use [`from_lstat`](Self::from_lstat) when
250    /// the path itself is a symlink that should be preserved as-is.
251    #[must_use]
252    pub fn from_host_path(path: &Path) -> Option<Self> {
253        let metadata = fs::metadata(path).ok()?;
254        let mtime = metadata
255            .modified()
256            .unwrap_or(SystemTime::UNIX_EPOCH)
257            .duration_since(SystemTime::UNIX_EPOCH)
258            .map_or(0.0, |duration| duration.as_secs_f64());
259        let size = i64::try_from(metadata.len()).unwrap_or(i64::MAX);
260        Some(Self {
261            host_path: path.to_path_buf(),
262            mtime,
263            size,
264        })
265    }
266
267    /// Builds a lazy file reference using `symlink_metadata` (lstat).
268    ///
269    /// Unlike [`from_host_path`](Self::from_host_path), this does not follow
270    /// symlinks. The stored `host_path` is the symlink itself, preserving
271    /// symlink identity across overlay renames.
272    #[must_use]
273    pub fn from_lstat(path: &Path) -> Option<Self> {
274        let metadata = fs::symlink_metadata(path).ok()?;
275        let mtime = metadata
276            .modified()
277            .unwrap_or(SystemTime::UNIX_EPOCH)
278            .duration_since(SystemTime::UNIX_EPOCH)
279            .map_or(0.0, |duration| duration.as_secs_f64());
280        let size = i64::try_from(metadata.len()).unwrap_or(i64::MAX);
281        Some(Self {
282            host_path: path.to_path_buf(),
283            mtime,
284            size,
285        })
286    }
287}