Merge pull request #36 from jorpic/inline-pdf-storage

Inline bits of pdf crate for better performance

crates/cracker/Cargo.toml

@@ -8,4 +8,5 @@ authors.workspace = true
 
 [dependencies]
 log.workspace = true
-pdf.workspace = true
+pdf.workspace = true
+anyhow.workspace = true

crates/cracker/src/lib.rs

@@ -1,6 +1,11 @@
-use std::{fs, io};
+use std::{fs, io, cell::RefCell, sync::Arc};
+use std::collections::hash_map::HashMap;
 
-use pdf::file::FileOptions;
+use anyhow::anyhow;
+use pdf::PdfError;
+use pdf::any::AnySync;
+use pdf::file::{Cache, Storage};
+use pdf::object::{ParseOptions, PlainRef};
 
 #[derive(Clone)]
 pub struct PDFCracker(Vec<u8>);
@@ -12,12 +17,57 @@ impl PDFCracker {
     }
 }
 
-impl PDFCracker {
+type ObjectCache = SimpleCache<Result<AnySync, Arc<PdfError>>>;
+type StreamCache = SimpleCache<Result<Arc<[u8]>, Arc<PdfError>>>;
+
+pub struct PDFCrackerState(
+    Storage<Vec<u8>, ObjectCache, StreamCache>
+);
+
+impl PDFCrackerState {
+    /// Init `pdf::file::Storage` so we can reuse it on each `attempt`.
+    pub fn from_cracker(pdf_file: &PDFCracker) -> anyhow::Result<Self> {
+        let res = Storage::with_cache(
+            pdf_file.0.clone(),
+            ParseOptions::strict(),
+            SimpleCache::new(),
+            SimpleCache::new(),
+        );
+
+        match res {
+            Ok(storage) => Ok(Self(storage)),
+            Err(err) => Err(anyhow!(err).context("Failed to init cracker")),
+        }
+    }
+
     /// Attempt to crack the cryptography using the password, return true on success.
-    pub fn attempt(&self, password: &[u8]) -> bool {
-        FileOptions::cached()
-            .password(password)
-            .load(self.0.as_ref())
+    pub fn attempt(&mut self, password: &[u8]) -> bool {
+        self.0.load_storage_and_trailer_password(password)
             .is_ok()
     }
 }
+
+/// Single-threaded adapter to use `HashMap` as a `pdf::file::Cache`.
+struct SimpleCache<T>(
+    RefCell<HashMap<PlainRef, T>>
+);
+
+impl<T: Clone> SimpleCache<T> {
+    fn new() -> Self {
+        Self(RefCell::new(HashMap::new()))
+    }
+}
+
+impl<T: Clone> Cache<T> for SimpleCache<T> {
+    fn get_or_compute(&self, key: PlainRef, compute: impl FnOnce() -> T) -> T {
+        let mut hash = self.0.borrow_mut();
+        match hash.get(&key) {
+            Some(value) => value.clone(),
+            None => {
+                let value = compute();
+                hash.insert(key, value.clone());
+                value
+            }
+        }
+    }
+}

crates/engine/src/lib.rs

@@ -8,7 +8,7 @@ pub mod producers {
 
 /// Expose our available crackers
 pub mod crackers {
-    pub use cracker::PDFCracker;
+    pub use cracker::{PDFCracker, PDFCrackerState};
 }
 
 // We will run a SPMC layout where a single producer produces passwords
@@ -22,7 +22,7 @@ use crossbeam::channel::{Receiver, Sender, TryRecvError};
 
 use producer::Producer;
 
-use cracker::PDFCracker;
+use cracker::{PDFCracker, PDFCrackerState};
 
 /// Returns Ok(Some(<Password in bytes>)) if it successfully cracked the file.
 /// Returns Ok(None) if it did not find the password.
@@ -46,8 +46,12 @@ pub fn crack_file(
         let r2 = r.clone();
         let c2 = cracker_handle.clone();
         let id: std::thread::JoinHandle<()> = std::thread::spawn(move || {
+            let Ok(mut cracker) = PDFCrackerState::from_cracker(&c2) else {
+                return
+            };
+
             while let Ok(passwd) = r2.recv() {
-                if c2.attempt(&passwd) {
+                if cracker.attempt(&passwd) {
                     // inform main thread we found a good password then die
                     success.send(passwd).unwrap_or_default();
                     return;