add method for shredding cache database file

Author: snshn

src/cache.rs

@@ -1,4 +1,6 @@
 use std::collections::HashMap;
+use std::fs::File;
+use std::io::{BufWriter, Write};
 use std::path::Path;
 
 use redb::{Database, Error, TableDefinition};
@@ -15,8 +17,10 @@ pub struct Cache {
     metadata: HashMap<String, CacheMetadataItem>, // Dictionary of metadata (and occasionally data [mostly for very small files])
     db: Option<Database>, // Pointer to database instance; None if not yet initialized or if failed to initialize
     db_ok: Option<bool>, // None by default, Some(true) if was able to initialize database, Some (false) if an error occured
+    db_file_path: Option<String>, // Filesystem path to file used for storing database
 }
 
+const FILE_WRITE_BUF_LEN: usize = 1024 * 100; // On-disk cache file write buffer size (in bytes)
 const TABLE: TableDefinition<&str, &[u8]> = TableDefinition::new("_");
 
 impl Cache {
@@ -26,6 +30,7 @@ impl Cache {
             metadata: HashMap::new(),
             db: None,
             db_ok: None,
+            db_file_path: db_file_path.clone(),
         };
 
         if db_file_path.is_some() {
@@ -52,14 +57,14 @@ impl Cache {
             data: if self.db_ok.is_some() && self.db_ok.unwrap() {
                 None
             } else {
-                Some((data.to_owned()).to_vec())
+                Some(data.to_owned().to_vec())
             },
             media_type: Some(media_type.to_owned()),
             charset: Some(charset),
         };
 
         if (self.db_ok.is_none() || !self.db_ok.unwrap()) || data.len() <= self.min_file_size {
-            cache_metadata_item.data = Some((data.to_owned()).to_vec());
+            cache_metadata_item.data = Some(data.to_owned().to_vec());
         } else {
             match self.db.as_ref().unwrap().begin_write() {
                 Ok(write_txn) => {
@@ -71,7 +76,7 @@ impl Cache {
                 }
                 Err(..) => {
                     // Fall back to caching everything in memory
-                    cache_metadata_item.data = Some((data.to_owned()).to_vec());
+                    cache_metadata_item.data = Some(data.to_owned().to_vec());
                 }
             }
         }
@@ -110,4 +115,40 @@ impl Cache {
     pub fn contains_key(&self, key: &str) -> bool {
         self.metadata.contains_key(key)
     }
+
+    pub fn destroy_database_file(&mut self) {
+        if self.db_ok.is_none() || !self.db_ok.unwrap() {
+            return;
+        }
+
+        // Destroy database instance (prevents writes into file)
+        self.db = None;
+        self.db_ok = Some(false);
+
+        // Wipe database file
+        if let Some(db_file_path) = self.db_file_path.to_owned() {
+            // Overwrite file with zeroes
+            if let Ok(temp_file) = File::options()
+                .read(true)
+                .write(true)
+                .open(db_file_path.clone())
+            {
+                let mut buffer = [0; FILE_WRITE_BUF_LEN];
+                let mut remaining_size: usize = temp_file.metadata().unwrap().len() as usize;
+                let mut writer = BufWriter::new(temp_file);
+
+                while remaining_size > 0 {
+                    let bytes_to_write: usize = if remaining_size < FILE_WRITE_BUF_LEN {
+                        remaining_size
+                    } else {
+                        FILE_WRITE_BUF_LEN
+                    };
+                    let buffer = &mut buffer[..bytes_to_write];
+                    writer.write(buffer).unwrap();
+
+                    remaining_size -= bytes_to_write;
+                }
+            }
+        }
+    }
 }

src/main.rs

@@ -54,14 +54,15 @@ const ASCII: &str = " \
 |   | \\___/ |          | \\                      |   |   |   |   |  |   |
 |___|       |__________|  \\_____________________|   |___|   |___|  |___|
 ";
-const CACHE_ASSET_FILE_SIZE_THRESHOLD: usize = 1024 * 50; // Minimum file size for on-disk caching (in bytes)
+const CACHE_ASSET_FILE_SIZE_THRESHOLD: usize = 1024 * 10; // Minimum file size for on-disk caching (in bytes)
 const DEFAULT_NETWORK_TIMEOUT: u64 = 120;
 const DEFAULT_USER_AGENT: &str =
     "Mozilla/5.0 (X11; Ubuntu; Linux x86_64; rv:135.0) Gecko/20100101 Firefox/135.0";
 
 fn main() {
     // Process CLI flags and options
     let mut cookie_file_path: Option<String> = None;
+    let mut exit_code = 0;
     let mut options: Options = Options::default();
     let source;
     let destination;
@@ -164,7 +165,7 @@ fn main() {
         Ok(tempfile) => Some(tempfile),
         Err(_) => None,
     };
-    let cache = Cache::new(
+    let mut cache = Some(Cache::new(
         CACHE_ASSET_FILE_SIZE_THRESHOLD,
         if temp_cache_file.is_some() {
             Some(
@@ -178,7 +179,7 @@ fn main() {
         } else {
             None
         },
-    );
+    ));
 
     // Read and parse cookie file
     if let Some(opt_cookie_file) = cookie_file_path.clone() {
@@ -211,7 +212,7 @@ fn main() {
         }
     }
 
-    match create_monolithic_document(source, &options, &mut Some(cache)) {
+    match create_monolithic_document(source, &options, &mut cache) {
         Ok(result) => {
             // Define output
             let mut output = Output::new(&destination).expect("could not prepare output");
@@ -222,7 +223,14 @@ fn main() {
         Err(error) => {
             print_error_message(&format!("Error: {}", error), &options);
 
-            process::exit(1);
+            exit_code = 1;
         }
     }
+
+    // Clean up (shred database file)
+    cache.unwrap().destroy_database_file();
+
+    if exit_code > 0 {
+        process::exit(exit_code);
+    }
 }