Rework PakReader to allow for parallel reads

repak/src/pak.rs

@@ -6,9 +6,8 @@ use std::collections::BTreeMap;
 use std::io::{self, Read, Seek, Write};
 
 #[derive(Debug)]
-pub struct PakReader<R: Read + Seek> {
+pub struct PakReader {
     pak: Pak,
-    reader: R,
     key: Option<aes::Aes256Dec>,
 }
 
@@ -58,6 +57,13 @@ impl Index {
         }
     }
 
+    fn into_entries(self) -> BTreeMap<String, super::entry::Entry> {
+        match self {
+            Index::V1(index) => index.entries,
+            Index::V2(index) => index.entries,
+        }
+    }
+
     fn add_entry(&mut self, path: &str, entry: super::entry::Entry) {
         match self {
             Index::V1(index) => index.entries.insert(path.to_string(), entry),
@@ -89,12 +95,15 @@ fn decrypt(key: &Option<aes::Aes256Dec>, bytes: &mut [u8]) -> Result<(), super::
     }
 }
 
-impl<R: Read + Seek> PakReader<R> {
-    pub fn new_any(mut reader: R, key: Option<aes::Aes256Dec>) -> Result<Self, super::Error> {
+impl PakReader {
+    pub fn new_any<R: Read + Seek>(
+        mut reader: R,
+        key: Option<aes::Aes256Dec>,
+    ) -> Result<Self, super::Error> {
         for ver in Version::iter() {
             match Pak::read(&mut reader, ver, key.clone()) {
                 Ok(pak) => {
-                    return Ok(PakReader { pak, reader, key });
+                    return Ok(PakReader { pak, key });
                 }
                 _ => continue,
             }
@@ -102,10 +111,6 @@ impl<R: Read + Seek> PakReader<R> {
         Err(super::Error::Other("version unsupported"))
     }
 
-    pub fn into_reader(self) -> R {
-        self.reader
-    }
-
     pub fn version(&self) -> super::Version {
         self.pak.version
     }
@@ -114,36 +119,30 @@ impl<R: Read + Seek> PakReader<R> {
         &self.pak.mount_point
     }
 
-    pub fn get(&mut self, path: &str) -> Result<Vec<u8>, super::Error> {
+    pub fn get<R: Read + Seek>(
+        &mut self,
+        path: &str,
+        reader: &mut R,
+    ) -> Result<Vec<u8>, super::Error> {
         let mut data = Vec::new();
-        self.read_file(path, &mut data)?;
+        self.read_file(path, reader, &mut data)?;
         Ok(data)
     }
 
-    pub fn read_file<W: io::Write>(
-        &mut self,
+    pub fn read_file<R: Read + Seek, W: io::Write>(
+        &self,
         path: &str,
+        reader: &mut R,
         writer: &mut W,
     ) -> Result<(), super::Error> {
         match self.pak.index.entries().get(path) {
-            Some(entry) => entry.read_file(
-                &mut self.reader,
-                self.pak.version,
-                self.key.as_ref(),
-                writer,
-            ),
+            Some(entry) => entry.read_file(reader, self.pak.version, self.key.as_ref(), writer),
             None => Err(super::Error::Other("no file found at given path")),
         }
     }
 
-    pub fn files(&self) -> std::vec::IntoIter<String> {
-        self.pak
-            .index
-            .entries()
-            .keys()
-            .cloned()
-            .collect::<Vec<String>>()
-            .into_iter()
+    pub fn files(&self) -> Vec<String> {
+        self.pak.index.entries().keys().cloned().collect()
     }
 }
 
@@ -588,17 +587,18 @@ mod test {
         use std::io::Cursor;
         let bytes = include_bytes!("../tests/packs/pack_v8b.pak");
 
-        let mut reader = super::PakReader::new_any(Cursor::new(bytes), None).unwrap();
+        let mut reader = Cursor::new(bytes);
+        let mut pak_reader = super::PakReader::new_any(&mut reader, None).unwrap();
         let writer = Cursor::new(vec![]);
         let mut pak_writer = super::PakWriter::new(
             writer,
             None,
             super::Version::V8B,
-            reader.mount_point().to_owned(),
+            pak_reader.mount_point().to_owned(),
         );
 
-        for path in reader.files() {
-            let data = reader.get(&path).unwrap();
+        for path in pak_reader.files() {
+            let data = pak_reader.get(&path, &mut reader).unwrap();
             pak_writer
                 .write_file(&path, &mut std::io::Cursor::new(data))
                 .unwrap();
@@ -613,17 +613,18 @@ mod test {
         use std::io::Cursor;
         let bytes = include_bytes!("../tests/packs/pack_v11.pak");
 
-        let mut reader = super::PakReader::new_any(Cursor::new(bytes), None).unwrap();
+        let mut reader = Cursor::new(bytes);
+        let mut pak_reader = super::PakReader::new_any(&mut reader, None).unwrap();
         let writer = Cursor::new(vec![]);
         let mut pak_writer = super::PakWriter::new(
             writer,
             None,
             super::Version::V11,
-            reader.mount_point().to_owned(),
+            pak_reader.mount_point().to_owned(),
         );
 
-        for path in reader.files() {
-            let data = reader.get(&path).unwrap();
+        for path in pak_reader.files() {
+            let data = pak_reader.get(&path, &mut reader).unwrap();
             pak_writer
                 .write_file(&path, &mut std::io::Cursor::new(data))
                 .unwrap();

repak/tests/test.rs

@@ -119,11 +119,9 @@ macro_rules! encryptindex {
 
                     let mut inner_reader = std::io::Cursor::new(include_bytes!(concat!("packs/pack_", $version, $compress, $encrypt, $encryptindex, ".pak")));
                     let len = inner_reader.seek(SeekFrom::End(0)).unwrap();
+                    let mut reader = ReadCounter::new_size(inner_reader, len as usize);
 
-                    let mut pak = repak::PakReader::new_any(
-                        ReadCounter::new_size(inner_reader, len as usize),
-                        Some(key),
-                    ).unwrap();
+                    let pak = repak::PakReader::new_any(&mut reader, Some(key)).unwrap();
 
                     assert_eq!(pak.mount_point(), "../mount/point/root/");
                     assert_eq!(pak.version(), $exp_version);
@@ -134,7 +132,7 @@ macro_rules! encryptindex {
                     for file in files {
                         let mut buf = vec![];
                         let mut writer = std::io::Cursor::new(&mut buf);
-                        pak.read_file(&file, &mut writer).unwrap();
+                        pak.read_file(&file, &mut reader, &mut writer).unwrap();
                         match file.as_str() {
                             "test.txt" => assert_eq!(buf, include_bytes!("pack/root/test.txt"), "test.txt incorrect contents"),
                             "test.png" => assert_eq!(buf, include_bytes!("pack/root/test.png"), "test.png incorrect contents"),
@@ -144,7 +142,7 @@ macro_rules! encryptindex {
                         }
                     }
 
-                    for r in pak.into_reader().into_reads() {
+                    for r in reader.into_reads() {
                         // sanity check. a pak file can be constructed with a lot of dead space
                         // which wouldn't have to be read, but so far all bytes in paks generated
                         // by UnrealPak are meaningful

repak_cli/Cargo.toml

@@ -15,4 +15,5 @@ aes = { workspace = true }
 base64 = { workspace = true }
 clap = { version = "4.1.4", features = ["derive"] }
 path-clean = "0.1.0"
+rayon = "1.6.1"
 repak = { version = "0.1.0", path = "../repak" }

repak_cli/src/main.rs

@@ -4,6 +4,7 @@ use std::path::{Path, PathBuf};
 
 use clap::{Parser, Subcommand};
 use path_clean::PathClean;
+use rayon::prelude::*;
 
 #[derive(Parser, Debug)]
 struct ActionInfo {
@@ -133,7 +134,7 @@ fn list(args: ActionInfo) -> Result<(), repak::Error> {
 }
 
 fn unpack(args: ActionUnpack) -> Result<(), repak::Error> {
-    let mut pak = repak::PakReader::new_any(
+    let pak = repak::PakReader::new_any(
         BufReader::new(File::open(&args.input)?),
         args.aes_key.map(|k| aes_key(k.as_str())).transpose()?,
     )?;
@@ -151,25 +152,31 @@ fn unpack(args: ActionUnpack) -> Result<(), repak::Error> {
     }
     let mount_point = PathBuf::from(pak.mount_point());
     let prefix = Path::new(&args.strip_prefix);
-    for file in pak.files() {
-        if args.verbose {
-            println!("extracting {}", &file);
-        }
-        let file_path = output.join(
-            mount_point
-                .join(&file)
-                .strip_prefix(prefix)
-                .map_err(|_| repak::Error::Other("prefix does not match"))?,
-        );
-        if !file_path.clean().starts_with(&output) {
-            return Err(repak::Error::Other(
-                "tried to write file outside of output directory",
-            ));
-        }
-        fs::create_dir_all(file_path.parent().expect("will be a file"))?;
-        pak.read_file(&file, &mut fs::File::create(file_path)?)?;
-    }
-    Ok(())
+    pak.files().into_par_iter().try_for_each_init(
+        || File::open(&args.input),
+        |file, path| -> Result<(), repak::Error> {
+            if args.verbose {
+                println!("extracting {path}");
+            }
+            let file_path = output.join(
+                mount_point
+                    .join(&path)
+                    .strip_prefix(prefix)
+                    .map_err(|_| repak::Error::Other("prefix does not match"))?,
+            );
+            if !file_path.clean().starts_with(&output) {
+                return Err(repak::Error::Other(
+                    "tried to write file outside of output directory",
+                ));
+            }
+            fs::create_dir_all(file_path.parent().expect("will be a file"))?;
+            pak.read_file(
+                &path,
+                &mut BufReader::new(file.as_ref().unwrap()), // TODO: avoid this unwrap
+                &mut fs::File::create(file_path)?,
+            )
+        },
+    )
 }
 
 fn pack(args: ActionPack) -> Result<(), repak::Error> {