mod naming;
mod report;
mod resolver;
mod writer;
#[cfg(test)]
mod tests;
use std::{
collections::HashSet,
fmt,
io::{Read, Seek},
num::NonZeroUsize,
sync::{
atomic::{AtomicBool, Ordering},
mpsc, Mutex, MutexGuard, PoisonError,
},
thread,
};
use camino::{Utf8Path, Utf8PathBuf};
use ltk_file::LeagueFileKind;
use ltk_hash::WadHash;
use crate::{ChunkDecoder, NameRecovery, RecoveredNames, Wad, WadChunk, WadError};
pub use self::{
naming::{chunk_hash_of, strip_ltk_suffix, NamingPolicy},
report::{DisplacedChunk, ExtractProgress, ExtractReport, ExtractResult, PathIssue},
resolver::{hex_chunk_hash, hex_name, is_hex_chunk_path, NoResolver, PathResolver},
writer::{ExistingFilePolicy, ExtractLayout},
};
use self::{
naming::{is_evil, DirectoryPaths},
report::ChunkOutcome,
writer::ChunkWriter,
};
pub(crate) use self::resolver::resolve_all_checked;
const DEFAULT_WORKER_CAP: usize = 8;
type PathFilter<'a> = Box<dyn Fn(&str) -> bool + 'a>;
type ProgressCallback<'a> = Box<dyn FnMut(&ExtractProgress<'_>) + 'a>;
pub struct WadExtractor<'a> {
resolver: &'a dyn PathResolver,
filter: Option<PathFilter<'a>>,
type_filter: Option<Vec<LeagueFileKind>>,
progress: Option<ProgressCallback<'a>>,
layout: ExtractLayout,
existing: ExistingFilePolicy,
naming: NamingPolicy,
cancel: Option<&'a AtomicBool>,
workers: Option<NonZeroUsize>,
recover_names: bool,
}
impl fmt::Debug for WadExtractor<'_> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("WadExtractor")
.field("layout", &self.layout)
.field("existing", &self.existing)
.field("naming", &self.naming)
.field("type_filter", &self.type_filter)
.field("has_filter", &self.filter.is_some())
.field("has_progress", &self.progress.is_some())
.field("workers", &self.workers)
.field("recover_names", &self.recover_names)
.finish_non_exhaustive()
}
}
impl<'a> WadExtractor<'a> {
pub fn new(resolver: &'a dyn PathResolver) -> Self {
Self {
resolver,
filter: None,
type_filter: None,
progress: None,
layout: ExtractLayout::default(),
existing: ExistingFilePolicy::default(),
naming: NamingPolicy::default(),
cancel: None,
workers: None,
recover_names: false,
}
}
pub fn with_filter(mut self, filter: impl Fn(&str) -> bool + 'a) -> Self {
self.filter = Some(Box::new(filter));
self
}
pub fn with_type_filter(mut self, kinds: impl IntoIterator<Item = LeagueFileKind>) -> Self {
self.type_filter = Some(kinds.into_iter().collect());
self
}
pub fn on_progress(mut self, callback: impl FnMut(&ExtractProgress<'_>) + 'a) -> Self {
self.progress = Some(Box::new(callback));
self
}
pub fn with_layout(mut self, layout: ExtractLayout) -> Self {
self.layout = layout;
self
}
pub fn with_existing_file_policy(mut self, policy: ExistingFilePolicy) -> Self {
self.existing = policy;
self
}
pub fn with_naming_policy(mut self, policy: NamingPolicy) -> Self {
self.naming = policy;
self
}
pub fn with_cancel_flag(mut self, flag: &'a AtomicBool) -> Self {
self.cancel = Some(flag);
self
}
pub fn with_workers(mut self, workers: NonZeroUsize) -> Self {
self.workers = Some(workers);
self
}
pub fn with_name_recovery(mut self) -> Self {
self.recover_names = true;
self
}
pub fn extract_all<S: Read + Seek>(
&mut self,
wad: &mut Wad<S>,
output_dir: impl AsRef<Utf8Path>,
) -> Result<ExtractReport, WadError> {
let chunks: Vec<WadChunk> = wad.chunks().iter().copied().collect();
self.run(wad, chunks, Vec::new(), output_dir.as_ref())
}
pub fn extract_chunks<S: Read + Seek>(
&mut self,
wad: &mut Wad<S>,
path_hashes: impl IntoIterator<Item = WadHash>,
output_dir: impl AsRef<Utf8Path>,
) -> Result<ExtractReport, WadError> {
let mut seen = HashSet::new();
let mut chunks = Vec::new();
let mut missing = Vec::new();
for path_hash in path_hashes {
if !seen.insert(path_hash) {
continue;
}
match wad.chunks().get(path_hash) {
Some(chunk) => chunks.push(*chunk),
None => missing.push(path_hash),
}
}
self.run(wad, chunks, missing, output_dir.as_ref())
}
fn resolve_chunks(&self, chunks: Vec<WadChunk>, resolver: &dyn PathResolver) -> Vec<Named> {
let path_hashes: Vec<WadHash> = chunks.iter().map(|chunk| chunk.path_hash).collect();
let resolved = resolve_all_checked(resolver, &path_hashes);
chunks
.into_iter()
.zip(resolved)
.map(|(chunk, resolved)| {
let (path, named) = match resolved {
Some(path) => (path, true),
None => (hex_name(chunk.path_hash), false),
};
let refused = if is_evil(&path) {
Some(Refusal::Rejected)
} else if self
.filter
.as_ref()
.is_some_and(|filter| !filter(path.as_str()))
{
Some(Refusal::Filtered)
} else {
None
};
Named {
chunk,
path,
named,
refused,
}
})
.collect()
}
fn run<S: Read + Seek>(
&mut self,
wad: &mut Wad<S>,
mut chunks: Vec<WadChunk>,
missing: Vec<WadHash>,
output_dir: &Utf8Path,
) -> Result<ExtractReport, WadError> {
let workers = self.workers.map_or_else(default_workers, NonZeroUsize::get);
chunks.sort_unstable_by_key(|chunk| chunk.data_offset);
let recovered = if self.recover_names {
NameRecovery {
workers: self.workers,
cancel: self.cancel,
}
.run(wad, self.resolver)?
} else {
RecoveredNames::default()
};
let resolver = recovered.over(self.resolver);
let total = chunks.len();
let chunks = self.resolve_chunks(chunks, &resolver);
let directories = match self.layout {
ExtractLayout::Paths => directory_paths(&chunks),
ExtractLayout::Flat => DirectoryPaths::default(),
};
let shared = Shared {
writer: ChunkWriter {
layout: self.layout,
existing: self.existing,
naming: self.naming,
type_filter: self.type_filter.as_deref(),
output_dir,
directories,
output_occupied: std::fs::read_dir(output_dir.as_std_path())
.map(|mut entries| entries.next().is_some())
.unwrap_or(false),
claimed: Mutex::default(),
created: Mutex::default(),
subchunk_toc: wad.subchunk_toc().cloned(),
},
report: Mutex::default(),
failure: Mutex::default(),
};
let (sender, receiver) = mpsc::sync_channel::<Job>(workers);
let receiver = Mutex::new(receiver);
let (done_sender, done) = mpsc::channel::<Done>();
let mut reader = Reader {
cancel: self.cancel,
progress: Progress {
callback: self.progress.as_deref_mut(),
done: 0,
total,
},
report: ExtractReport::default(),
};
let cancelled = thread::scope(|scope| {
let shared = &shared;
let receiver = &receiver;
for _ in 0..workers {
let done_sender = done_sender.clone();
scope.spawn(move || shared.run_worker(receiver, &done_sender));
}
drop(done_sender);
let result = reader.read_chunks(wad, chunks, &sender, shared, &done);
drop(sender);
for finished in done.iter() {
reader.progress.report(&finished);
}
result
})?;
if let Some(error) = shared
.failure
.into_inner()
.unwrap_or_else(PoisonError::into_inner)
{
return Err(error);
}
let mut report = shared
.report
.into_inner()
.unwrap_or_else(PoisonError::into_inner);
report.merge(reader.report);
report.cancelled = cancelled;
report.missing = missing;
report.recovered = recovered;
Ok(report)
}
}
pub(crate) fn default_workers() -> usize {
thread::available_parallelism().map_or(1, |count| count.get().min(DEFAULT_WORKER_CAP))
}
pub(crate) fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
mutex.lock().unwrap_or_else(PoisonError::into_inner)
}
struct Named {
chunk: WadChunk,
path: String,
named: bool,
refused: Option<Refusal>,
}
#[derive(Debug, Clone, Copy)]
enum Refusal {
Rejected,
Filtered,
}
fn directory_paths(chunks: &[Named]) -> DirectoryPaths {
DirectoryPaths::of(chunks.iter().filter_map(|resolved| {
(resolved.named && resolved.refused.is_none()).then_some(resolved.path.as_str())
}))
}
struct Job {
chunk: WadChunk,
path: String,
named: bool,
raw: Box<[u8]>,
}
struct Done {
path_hash: WadHash,
path: String,
named: bool,
output_path: Option<Utf8PathBuf>,
outcome: ChunkOutcome,
}
impl Done {
fn displaced(&self, issue: Option<PathIssue>) -> Option<DisplacedChunk> {
Some(DisplacedChunk {
path_hash: self.path_hash,
path: self.path.clone(),
issue: issue?,
})
}
}
struct Progress<'c, 'a> {
callback: Option<&'c mut (dyn FnMut(&ExtractProgress<'_>) + 'a)>,
done: usize,
total: usize,
}
impl Progress<'_, '_> {
fn report(&mut self, finished: &Done) {
self.done += 1;
if let Some(callback) = &mut self.callback {
callback(&ExtractProgress {
done: self.done,
total: self.total,
path_hash: finished.path_hash,
path: &finished.path,
named: finished.named,
output_path: finished.output_path.as_deref(),
result: finished.outcome.into(),
bytes: finished.outcome.bytes(),
});
}
}
}
struct Reader<'r, 'a> {
cancel: Option<&'r AtomicBool>,
progress: Progress<'r, 'a>,
report: ExtractReport,
}
impl Reader<'_, '_> {
fn read_chunks<S: Read + Seek>(
&mut self,
wad: &mut Wad<S>,
chunks: Vec<Named>,
sender: &mpsc::SyncSender<Job>,
shared: &Shared<'_>,
done: &mpsc::Receiver<Done>,
) -> Result<bool, WadError> {
for Named {
chunk,
path,
named,
refused,
} in chunks
{
if self.cancel.is_some_and(|flag| flag.load(Ordering::Relaxed)) {
return Ok(true);
}
if shared.failed() {
return Ok(false);
}
if let Some(refusal) = refused {
let (outcome, issue) = match refusal {
Refusal::Rejected => {
(ChunkOutcome::SkippedRejectedPath, Some(PathIssue::Rejected))
}
Refusal::Filtered => (ChunkOutcome::SkippedByPath, None),
};
let finished = Done {
path_hash: chunk.path_hash,
path,
named,
output_path: None,
outcome,
};
self.report
.record_chunk(finished.outcome, finished.displaced(issue));
self.progress.report(&finished);
continue;
}
if let Some(written) = shared.writer.probe_skip(&chunk, &path, named) {
let finished = Done {
path_hash: chunk.path_hash,
path,
named,
output_path: written.path,
outcome: written.outcome,
};
self.report
.record_chunk(finished.outcome, finished.displaced(written.issue));
self.progress.report(&finished);
continue;
}
let raw = wad
.load_chunk_raw(&chunk)
.map_err(|error| WadError::chunk(chunk.path_hash, &path, error))?;
let job = Job {
chunk,
path,
named,
raw,
};
if sender.send(job).is_err() {
break;
}
for finished in done.try_iter() {
self.progress.report(&finished);
}
}
Ok(false)
}
}
struct Shared<'s> {
writer: ChunkWriter<'s>,
report: Mutex<ExtractReport>,
failure: Mutex<Option<WadError>>,
}
impl Shared<'_> {
fn failed(&self) -> bool {
lock(&self.failure).is_some()
}
fn run_worker(&self, receiver: &Mutex<mpsc::Receiver<Job>>, done: &mpsc::Sender<Done>) {
let mut decoder = ChunkDecoder::new();
let mut report = ExtractReport::default();
loop {
let Ok(job) = lock(receiver).recv() else {
break;
};
if self.failed() {
continue;
}
match self.writer.write(&job, &mut decoder) {
Ok(written) => {
let finished = Done {
path_hash: job.chunk.path_hash,
named: job.named,
path: job.path,
output_path: written.path,
outcome: written.outcome,
};
report.record_chunk(finished.outcome, finished.displaced(written.issue));
let _ = done.send(finished);
}
Err(error) => {
let mut failure = lock(&self.failure);
if failure.is_none() {
*failure = Some(WadError::chunk(job.chunk.path_hash, &job.path, error));
}
}
}
}
lock(&self.report).merge(report);
}
}