use a sized cache to prevent memory explosion

4 files changed, 32 insertions, 18 deletions

diff --git a/Cargo.toml b/Cargo.toml
index 37df112..398370b 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -23,6 +23,7 @@ snafu = {version="0.3.1", features = ["backtraces", "rust_1_30"]}
 num-traits = "0.2"
 num-derive = "0.2"
 winstructs = {git = "https://github.com/omerbenamram/winstructs.git", branch = "master"}
+cached = "0.8.0"
 
 [dependencies.chrono]
 version = "0.4.6"
diff --git a/src/bin/mft_dump.rs b/src/bin/mft_dump.rs
index 7362b2b..a3fe3dd 100644
--- a/src/bin/mft_dump.rs
+++ b/src/bin/mft_dump.rs
@@ -1,10 +1,15 @@
 use clap::{App, Arg, ArgMatches};
 use env_logger;
-use log::{info};
+use log::info;
 
 use mft::mft::MftParser;
 use std::path::PathBuf;
 
+enum OutputFormat {
+    JSON,
+    CSV,
+}
+
 struct MftDump {
     filepath: PathBuf,
     indent: bool,
diff --git a/src/entry.rs b/src/entry.rs
index 6cf4f2c..c20af87 100644
--- a/src/entry.rs
+++ b/src/entry.rs
@@ -1,7 +1,6 @@
 use crate::err::{self, Result};
 use crate::impl_serialize_for_bitflags;
 
-
 use log::trace;
 use snafu::{ensure, ResultExt};
 
@@ -81,7 +80,6 @@ pub struct EntryHeader {
     pub next_attribute_id: u16,
     #[serde(skip_serializing)]
     pub record_number: u64,
-    pub update_sequence_value: u32,
     pub entry_reference: MftReference,
 }
 bitflags! {
@@ -153,7 +151,6 @@ impl EntryHeader {
             base_reference,
             next_attribute_id,
             record_number,
-            update_sequence_value: 0,
             entry_reference,
         })
     }
diff --git a/src/mft.rs b/src/mft.rs
index d6871bf..90beaf5 100644
--- a/src/mft.rs
+++ b/src/mft.rs
@@ -2,12 +2,13 @@ use crate::entry::MftEntry;
 use crate::err::{self, Result};
 
 use crate::{EntryHeader, ReadSeek};
-use log::debug;
+use log::{debug, trace};
 use snafu::ResultExt;
 
 use crate::attribute::MftAttributeContent::AttrX30;
 
-use std::collections::HashMap;
+use cached::stores::SizedCache;
+use cached::Cached;
 use std::fs::{self, File};
 use std::io::{BufReader, Cursor, SeekFrom};
 use std::path::{Path, PathBuf};
@@ -18,7 +19,7 @@ pub struct MftParser<T: ReadSeek> {
     /// Instead this will be guessed by the entry size of the first entry.
     entry_size: u32,
     size: u64,
-    entries_cache: HashMap<u64, PathBuf>,
+    entries_cache: SizedCache<u64, PathBuf>,
 }
 
 impl MftParser<BufReader<File>> {
@@ -55,7 +56,7 @@ impl<T: ReadSeek> MftParser<T> {
             data,
             entry_size: first_entry.total_entry_size,
             size,
-            entries_cache: HashMap::new(),
+            entries_cache: SizedCache::with_size(1000),
         })
     }
 
@@ -100,13 +101,21 @@ impl<T: ReadSeek> MftParser<T> {
             if let AttrX30(filename_header) = attribute.data {
                 let parent_entry_id = filename_header.parent.entry;
 
-                if parent_entry_id > 0 {
-                    // If i'm my own parent, I'm the root path.
-                    if parent_entry_id == entry_id {
-                        return Ok(Some(PathBuf::from(filename_header.name)));
-                    }
+                // MFT entry 5 is the root path.
+                if parent_entry_id == 5 {
+                    return Ok(Some(PathBuf::from(filename_header.name)));
+                }
 
-                    let cached_entry = self.entries_cache.get(&parent_entry_id);
+                if parent_entry_id == entry_id {
+                    trace!(
+                        "Found self-referential file path, for entry ID {}",
+                        entry_id
+                    );
+                    return Ok(Some(PathBuf::from("[Orphaned]").join(filename_header.name)));
+                }
+
+                if parent_entry_id > 0 {
+                    let cached_entry = self.entries_cache.cache_get(&parent_entry_id);
 
                     // If my parent path is known, then my path is parent's full path + my name.
                     // Else, retrieve and cache my parent's path.
@@ -128,15 +137,17 @@ impl<T: ReadSeek> MftParser<T> {
                             None => PathBuf::from("[Unknown]"),
                         };
 
-                        self.entries_cache.insert(parent_entry_id, path.clone());
+                        self.entries_cache.cache_set(parent_entry_id, path.clone());
                         return Ok(Some(path.join(filename_header.name)));
                     }
                 } else {
-                    let root = PathBuf::from(filename_header.name);
+                    trace!("Found orphaned entry ID {}", entry_id);
+
+                    let orphan = PathBuf::from("[Orphaned]").join(filename_header.name);
 
                     self.entries_cache
-                        .insert(entry.header.entry_reference.entry, root.clone());
-                    return Ok(Some(root));
+                        .cache_set(entry.header.entry_reference.entry, orphan.clone());
+                    return Ok(Some(orphan));
                 }
             }
         }