gravityfile_scan/
snapshot.rs

1//! Reading saved scans: snapshots written by `grav export` and dumps written by
2//! `ncdu -o` (or `gdu -o`), either of them optionally zstd-compressed.
3
4use std::io::{self, BufRead, BufReader};
5use std::path::{Path, PathBuf};
6
7use gravityfile_core::{FileNode, FileTree};
8use serde::Deserialize;
9
10/// Newest snapshot schema this reader understands.
11pub const SCHEMA_VERSION: u32 = 2;
12
13/// First bytes of a zstd frame.
14const ZSTD_MAGIC: [u8; 4] = [0x28, 0xb5, 0x2f, 0xfd];
15
16/// How to measure an imported tree.
17#[derive(Debug, Clone, Default)]
18pub struct ImportOptions {
19    /// Measure apparent lengths instead of allocated blocks.
20    pub apparent: bool,
21    /// Count every hard link to a file, not only the first.
22    pub count_hard_links: bool,
23    /// Fail if the snapshot has more than this many entries below the root.
24    pub max_entries: Option<u64>,
25}
26
27#[derive(Debug, thiserror::Error)]
28pub enum SnapshotError {
29    #[error("cannot read snapshot {}: {source}", gravityfile_core::escape_filename(path.as_os_str()))]
30    Open { path: PathBuf, source: io::Error },
31
32    #[error(transparent)]
33    Io(#[from] io::Error),
34
35    #[error("not an ncdu dump: {message} (at byte {offset})")]
36    Syntax { offset: u64, message: String },
37
38    #[error("not a gravityfile snapshot: {0}")]
39    Json(#[from] serde_json::Error),
40
41    #[error("unsupported snapshot: {0}")]
42    Unsupported(String),
43
44    #[error("damaged snapshot: {0}")]
45    Damaged(String),
46
47    #[error("snapshot has more than {limit} entries")]
48    EntryLimitExceeded { limit: u64 },
49}
50
51/// Which format a saved scan is in.
52#[derive(Debug, Clone, Copy, PartialEq, Eq)]
53pub enum Format {
54    /// `grav export`
55    Gravityfile,
56    /// `ncdu -o`
57    Ncdu,
58}
59
60/// Read the saved scan at `path`, measured as `options` asks.
61pub fn read(path: &Path, options: &ImportOptions) -> Result<FileTree, SnapshotError> {
62    let open = |source| SnapshotError::Open {
63        path: path.to_path_buf(),
64        source,
65    };
66    let file = std::fs::File::open(path).map_err(open)?;
67    read_from(BufReader::with_capacity(256 * 1024, file), options)
68}
69
70/// Read a saved scan from `input`, decompressing it if it is zstd.
71pub fn read_from(
72    mut input: impl BufRead,
73    options: &ImportOptions,
74) -> Result<FileTree, SnapshotError> {
75    if input.fill_buf()?.starts_with(&ZSTD_MAGIC) {
76        let decoder = zstd::stream::read::Decoder::with_buffer(input)?;
77        return read_plain(BufReader::with_capacity(256 * 1024, decoder), options);
78    }
79    read_plain(input, options)
80}
81
82fn read_plain(mut input: impl BufRead, options: &ImportOptions) -> Result<FileTree, SnapshotError> {
83    match detect(&mut input)? {
84        Format::Ncdu => crate::ncdu::read(input, options),
85        Format::Gravityfile => {
86            let mut tree = read_gravityfile(input)?;
87            check_structure(&tree.root)?;
88            // Totals are derived from the entries, as for an ncdu dump, never
89            // taken on trust: an edited or damaged file could claim anything.
90            // `grav export` writes complete trees, so nothing is lost.
91            tree.root.update_counts();
92            tree.remeasure(options.apparent, options.count_hard_links);
93            if let Some(limit) = options.max_entries {
94                let entries = tree.stats.total_files
95                    + tree.stats.total_dirs.saturating_sub(1)
96                    + tree.stats.total_symlinks;
97                if entries > limit {
98                    return Err(SnapshotError::EntryLimitExceeded { limit });
99                }
100            }
101            Ok(tree)
102        }
103    }
104}
105
106/// Only directories have entries. Walks without recursion, as trees can be
107/// as deep as the filesystem.
108fn check_structure(root: &FileNode) -> Result<(), SnapshotError> {
109    let mut pending = vec![root];
110    while let Some(node) = pending.pop() {
111        if !node.is_dir() && !node.children.is_empty() {
112            return Err(SnapshotError::Damaged(format!(
113                "{} is not a directory but has entries",
114                node.display_name()
115            )));
116        }
117        pending.extend(&node.children);
118    }
119    Ok(())
120}
121
122/// Tell the formats apart by their first significant byte.
123fn detect(input: &mut impl BufRead) -> Result<Format, SnapshotError> {
124    loop {
125        let buffer = input.fill_buf()?;
126        let Some(&first) = buffer.first() else {
127            return Err(SnapshotError::Unsupported("the file is empty".into()));
128        };
129        if first.is_ascii_whitespace() {
130            input.consume(1);
131            continue;
132        }
133        return match first {
134            b'[' => Ok(Format::Ncdu),
135            b'{' => Ok(Format::Gravityfile),
136            _ => Err(SnapshotError::Unsupported(
137                "expected a grav export or an ncdu dump".into(),
138            )),
139        };
140    }
141}
142
143/// A `grav export` document: the versioned envelope, or the bare tree that
144/// `grav export` wrote before schema 2. The two are told apart by their first
145/// key so the document streams straight into the tree without being buffered.
146fn read_gravityfile(mut input: impl BufRead) -> Result<FileTree, SnapshotError> {
147    #[derive(Deserialize)]
148    struct Envelope {
149        schema_version: u32,
150        kind: String,
151        data: Option<Box<FileTree>>,
152    }
153    let bare = first_key(&mut input)?.as_deref() == Some(b"root".as_slice());
154    let mut json = serde_json::Deserializer::from_reader(input);
155    json.disable_recursion_limit();
156    // Trees are as deep as the filesystem; grow the stack instead of
157    // overflowing it.
158    let deserializer = serde_stacker::Deserializer::new(&mut json);
159    let tree = if bare {
160        FileTree::deserialize(deserializer)?
161    } else {
162        let envelope = Envelope::deserialize(deserializer)?;
163        if envelope.schema_version > SCHEMA_VERSION {
164            return Err(SnapshotError::Unsupported(format!(
165                "schema version {} is newer than this grav understands ({SCHEMA_VERSION})",
166                envelope.schema_version
167            )));
168        }
169        if envelope.kind != "scan" {
170            return Err(SnapshotError::Unsupported(format!(
171                "a {:?} report, not a scan; save one with `grav export`",
172                envelope.kind
173            )));
174        }
175        *envelope
176            .data
177            .ok_or_else(|| SnapshotError::Unsupported("the document has no data".into()))?
178    };
179    json.end()?;
180    Ok(tree)
181}
182
183/// The first key of the object at the start of `input`, if it is visible in
184/// the buffered bytes and has no escapes. Nothing is consumed.
185fn first_key(input: &mut impl BufRead) -> io::Result<Option<Vec<u8>>> {
186    let buffer = input.fill_buf()?;
187    let Some(rest) = buffer.strip_prefix(b"{") else {
188        return Ok(None);
189    };
190    let start = rest.iter().position(|byte| !byte.is_ascii_whitespace());
191    let Some(rest) = start.and_then(|start| rest[start..].strip_prefix(b"\"")) else {
192        return Ok(None);
193    };
194    Ok(rest
195        .iter()
196        .position(|&byte| byte == b'"' || byte == b'\\')
197        .filter(|&end| rest[end] == b'"')
198        .map(|end| rest[..end].to_vec()))
199}
200
201#[cfg(test)]
202mod tests {
203    use super::*;
204    use gravityfile_core::{FileNode, ScanConfig, Timestamps, TreeStats};
205
206    fn tree() -> FileTree {
207        let mut root =
208            FileNode::new_directory("r", Timestamps::with_modified(std::time::UNIX_EPOCH));
209        root.children.push(FileNode::new_file(
210            "f",
211            1000,
212            8,
213            Timestamps::with_modified(std::time::UNIX_EPOCH),
214            false,
215        ));
216        root.update_counts();
217        let mut tree = FileTree::new(
218            root,
219            PathBuf::from("/r"),
220            ScanConfig::new("/r"),
221            TreeStats::new(),
222            std::time::Duration::ZERO,
223            Vec::new(),
224        );
225        tree.remeasure(false, false);
226        tree
227    }
228
229    fn envelope(tree: &FileTree) -> Vec<u8> {
230        serde_json::to_vec(&serde_json::json!({
231            "schema_version": 2,
232            "kind": "scan",
233            "root": "/r",
234            "data": tree,
235        }))
236        .unwrap()
237    }
238
239    #[test]
240    fn reads_envelopes_bare_trees_and_ncdu_dumps_alike() {
241        let tree = tree();
242        let mut dump = Vec::new();
243        crate::ncdu::write(&tree, &mut dump).unwrap();
244        for document in [envelope(&tree), serde_json::to_vec(&tree).unwrap(), dump] {
245            let read = read_from(document.as_slice(), &ImportOptions::default()).unwrap();
246            assert_eq!(read.root.size, 4096);
247            assert_eq!(read.stats.total_files, 1);
248        }
249    }
250
251    #[test]
252    fn compressed_snapshots_are_detected() {
253        let compressed = zstd::encode_all(envelope(&tree()).as_slice(), 3).unwrap();
254        let read = read_from(compressed.as_slice(), &ImportOptions::default()).unwrap();
255        assert_eq!(read.root.size, 4096);
256    }
257
258    #[test]
259    fn imports_can_switch_the_measure() {
260        let options = ImportOptions {
261            apparent: true,
262            ..ImportOptions::default()
263        };
264        let read = read_from(envelope(&tree()).as_slice(), &options).unwrap();
265        assert_eq!(read.root.size, 1000);
266        assert!(read.config.apparent_size);
267    }
268
269    #[test]
270    fn totals_come_from_the_entries_not_the_file() {
271        let mut edited = tree();
272        edited.root.size = 1 << 40;
273        edited.stats.total_files = 1_000_000;
274        let read = read_from(envelope(&edited).as_slice(), &ImportOptions::default()).unwrap();
275        assert_eq!(read.root.size, 4096);
276        assert_eq!(read.stats.total_files, 1);
277        // The entry limit sees the real count, not the claimed one.
278        let limited = ImportOptions {
279            max_entries: Some(1),
280            ..ImportOptions::default()
281        };
282        assert!(read_from(envelope(&edited).as_slice(), &limited).is_ok());
283    }
284
285    /// Found by fuzzing: such a tree read differently after a trip through
286    /// an ncdu dump.
287    #[test]
288    fn a_file_with_entries_is_damage() {
289        let mut damaged = tree();
290        let stray = damaged.root.children[0].clone();
291        damaged.root.children[0].children.push(stray);
292        let error = read_from(envelope(&damaged).as_slice(), &ImportOptions::default());
293        assert!(
294            matches!(&error, Err(SnapshotError::Damaged(message)) if message.contains('f')),
295            "{error:?}"
296        );
297    }
298
299    #[test]
300    fn deep_snapshots_load_without_overflowing_the_stack() {
301        let mut node =
302            FileNode::new_directory("d", Timestamps::with_modified(std::time::UNIX_EPOCH));
303        for _ in 0..5000 {
304            let mut parent =
305                FileNode::new_directory("d", Timestamps::with_modified(std::time::UNIX_EPOCH));
306            parent.children.push(node);
307            node = parent;
308        }
309        let mut deep = tree();
310        deep.root = node;
311        let mut bytes = Vec::new();
312        {
313            // Serialize on a large stack: serde's writer is recursive.
314            let deep = &deep;
315            std::thread::scope(|scope| {
316                std::thread::Builder::new()
317                    .stack_size(256 * 1024 * 1024)
318                    .spawn_scoped(scope, || bytes = envelope(deep))
319                    .unwrap();
320            });
321        }
322        let read = read_from(bytes.as_slice(), &ImportOptions::default()).unwrap();
323        let mut depth = 0;
324        let mut node = &read.root;
325        while let Some(child) = node.children.first() {
326            depth += 1;
327            node = child;
328        }
329        assert_eq!(depth, 5000);
330    }
331
332    #[test]
333    fn rejects_other_reports_and_future_versions() {
334        for document in [
335            r#"{"schema_version":2,"kind":"age","data":{}}"#,
336            r#"{"schema_version":99,"kind":"scan","data":{}}"#,
337            "",
338            "garbage",
339            "{}",
340        ] {
341            assert!(
342                read_from(document.as_bytes(), &ImportOptions::default()).is_err(),
343                "{document:?}"
344            );
345        }
346    }
347}
gravityfile_scan/snapshot.rs

gravityfile_scan/
snapshot.rs
