From a0a149938b93818e0ff0d3837a89ee7c807d3018 Mon Sep 17 00:00:00 2001 From: cohaereo Date: Tue, 11 Jul 2023 22:21:26 +0200 Subject: [PATCH] Allow immutable access to package This helps with the API, especially with entry iteration --- src/bin/unpack.rs | 12 +++++++----- src/package.rs | 50 +++++++++++++++++++++++++++++------------------ 2 files changed, 38 insertions(+), 24 deletions(-) diff --git a/src/bin/unpack.rs b/src/bin/unpack.rs index 983b5e1..9782255 100644 --- a/src/bin/unpack.rs +++ b/src/bin/unpack.rs @@ -1,14 +1,12 @@ use destiny2_pkg::package::Package; -use destiny2_pkg::structs::EntryHeader; use std::fs::File; use std::io::Write; fn main() -> anyhow::Result<()> { - let mut package = Package::open(&std::env::args().nth(1).unwrap())?; - let entries: Vec = package.entries().cloned().collect(); + let package = Package::open(&std::env::args().nth(1).unwrap())?; std::fs::create_dir("./files/").ok(); - for (i, e) in entries.iter().enumerate() { + for (i, e) in package.entries().enumerate() { if e.reference != u32::MAX { print!( "{i} 0x{:x} - p={:x} f={} ", @@ -68,7 +66,11 @@ fn main() -> anyhow::Result<()> { let data = match package.read_entry(i) { Ok(data) => data, Err(e) => { - eprintln!("Failed to extract entry {}/{}: {e}", i, entries.len() - 1); + eprintln!( + "Failed to extract entry {}/{}: {e}", + i, + package.entries().count() - 1 + ); continue; } }; diff --git a/src/package.rs b/src/package.rs index 2a2a807..a4cfc50 100644 --- a/src/package.rs +++ b/src/package.rs @@ -5,8 +5,11 @@ use anyhow::{anyhow, Context}; use binrw::{BinReaderExt, VecArgs}; use nohash_hasher::IntMap; use std::borrow::Cow; +use std::cell::RefCell; +use std::collections::hash_map::Entry; use std::fs::File; use std::io::{BufReader, Read, Seek, SeekFrom}; +use std::rc::Rc; use std::slice::Iter; pub const BLOCK_SIZE: usize = 0x40000; @@ -15,16 +18,16 @@ pub trait ReadSeek: Read + Seek {} impl ReadSeek for R {} pub struct Package { - gcm: PkgGcmState, + gcm: RefCell, header: PackageHeader, entries: Vec, blocks: Vec, - reader: Box, + reader: RefCell>, path_base: String, - block_cache: IntMap>, + block_cache: RefCell>>>, } impl Package { @@ -60,8 +63,8 @@ impl Package { Ok(Package { path_base, - reader: Box::new(reader), - gcm: PkgGcmState::new(header.pkg_id), + reader: RefCell::new(Box::new(reader)), + gcm: RefCell::new(PkgGcmState::new(header.pkg_id)), header, entries, blocks, @@ -73,13 +76,15 @@ impl Package { self.entries.iter() } - fn get_block_raw(&mut self, block_index: usize) -> anyhow::Result> { + fn get_block_raw(&self, block_index: usize) -> anyhow::Result> { let bh = &self.blocks[block_index]; let mut data = vec![0u8; bh.size as usize]; if self.header.patch_id == bh.patch_id { - self.reader.seek(SeekFrom::Start(bh.offset as u64))?; - self.reader.read_exact(&mut data)?; + self.reader + .borrow_mut() + .seek(SeekFrom::Start(bh.offset as u64))?; + self.reader.borrow_mut().read_exact(&mut data)?; } else { let mut f = File::open(format!("{}_{}.pkg", self.path_base, bh.patch_id)).context(format!( @@ -94,16 +99,14 @@ impl Package { Ok(Cow::Owned(data)) } - pub fn get_block(&mut self, block_index: usize) -> anyhow::Result> { - if self.block_cache.contains_key(&block_index) { - return Ok(Cow::Borrowed(self.block_cache.get(&block_index).unwrap())); - } - + /// Reads, decrypts and decompresses the specified block + fn read_block(&self, block_index: usize) -> anyhow::Result> { let bh = self.blocks[block_index].clone(); let mut block_data = self.get_block_raw(block_index)?.to_vec(); if (bh.flags & 0x2) != 0 { self.gcm + .borrow_mut() .decrypt_block_in_place(bh.flags, &bh.gcm_tag, &mut block_data)?; }; @@ -115,24 +118,33 @@ impl Package { block_data }; - self.block_cache.insert(block_index, decompressed_data); - Ok(Cow::Borrowed(self.block_cache.get(&block_index).unwrap())) + Ok(decompressed_data) + } + + /// Gets the specified block from the cache or reads it + pub fn get_block(&self, block_index: usize) -> anyhow::Result>> { + Ok(match self.block_cache.borrow_mut().entry(block_index) { + Entry::Occupied(o) => o.get().clone(), + Entry::Vacant(v) => { + let block = self.read_block(*v.key())?; + v.insert(Rc::new(block)).clone() + } + }) } - pub fn read_entry(&mut self, index: usize) -> anyhow::Result> { + pub fn read_entry(&self, index: usize) -> anyhow::Result> { let entry = self .entries .get(index) - .ok_or(anyhow!("Entry index is out of range"))? - .clone(); + .ok_or(anyhow!("Entry index is out of range"))?; let mut buffer = Vec::with_capacity(entry.file_size as usize); let mut current_offset = 0usize; let mut current_block = entry.starting_block; while current_offset < entry.file_size as usize { - let block_data = self.get_block(current_block as usize)?; let remaining_bytes = entry.file_size as usize - current_offset; + let block_data = self.get_block(current_block as usize)?; if current_block == entry.starting_block { let block_start_offset = (entry.starting_block_offset * 16) as usize;