diff --git a/src/io/detect.rs b/src/io/detect.rs index 28a3a23..4ca65d4 100644 --- a/src/io/detect.rs +++ b/src/io/detect.rs @@ -23,6 +23,8 @@ pub enum VersionSlice<'b> { #[derive(thiserror::Error, Debug)] pub enum DetectError { + #[error("Not a BDAT file")] + NotBdat, #[error("Can't determine legacy platform: no tables found")] LegacyNoTables, } @@ -135,6 +137,8 @@ pub fn detect_bytes_version(bytes: &[u8]) -> Result { /// An error ([`BdatError::VersionDetect`]) might be returned if the version couldn't be detected /// because of ambiguous details. /// +/// **Note**: the state of the reader will be modified after the call. +/// /// [`BdatError::VersionDetect`]: crate::BdatError::VersionDetect pub fn detect_file_version(reader: R) -> Result { detect_version(reader) @@ -145,7 +149,6 @@ fn detect_version(mut reader: R) -> Result { reader.read_exact(&mut magic)?; if magic == BDAT_MAGIC { // XC3 BDAT files start with "BDAT" - reader.seek(SeekFrom::Start(0))?; return Ok(BdatVersion::Modern); } @@ -187,11 +190,14 @@ fn detect_version(mut reader: R) -> Result { reader.seek(SeekFrom::Start( SwitchEndian::read_u32(&first_offset.to_be_bytes()) as u64, ))?; - if reader.read_u32::()? == MAGIC_INT { + reader.read_exact(&mut new_magic)?; + if WiiEndian::read_u32(&new_magic) == MAGIC_INT { // Table magic in big endian, this is a 3DS file. return Ok(LegacyVersion::New3ds.into()); + } else if SwitchEndian::read_u32(&new_magic) == MAGIC_INT { + return Ok(LegacyVersion::Switch.into()); } - return Ok(LegacyVersion::Switch.into()); + return Err(DetectError::NotBdat.into()); } // If we've reached this point, we either have a XC1 (Wii) file or a XCX file, which are both @@ -217,7 +223,12 @@ fn detect_version(mut reader: R) -> Result { // offset), then the table is from XCX. // In any other case, it's the XC1 format. - reader.seek(SeekFrom::Start(first_offset as u64 + 32 - 4 * 2))?; + reader.seek(SeekFrom::Start(first_offset as u64))?; + // Magic is always BDAT for non-3DS games + if reader.read_u32::()? != MAGIC_INT { + return Err(DetectError::NotBdat.into()); + } + reader.seek(SeekFrom::Current(32 - 4 * 3))?; let string_table_offset = reader.read_u32::()?; let string_table_len = reader.read_u32::()?; let final_offset = string_table_offset + string_table_len; diff --git a/toolset/src/convert/mod.rs b/toolset/src/convert/mod.rs index ca05fcd..b8b463a 100644 --- a/toolset/src/convert/mod.rs +++ b/toolset/src/convert/mod.rs @@ -13,7 +13,7 @@ use rayon::prelude::*; use crate::{ error::Error, - filter::{Filter, FilterArg}, + filter::{BdatFileFilter, Filter, FilterArg, SchemaFileFilter}, util::hash::HashNameTable, InputData, }; @@ -115,7 +115,7 @@ pub fn run_serialization(args: ConvertArgs, hash_table: HashNameTable) -> Result let files = args .input - .list_files("bdat", false)? + .list_files(BdatFileFilter, false)? .into_iter() .collect::>>()?; let base_path = crate::util::get_common_denominator(&files); @@ -204,7 +204,7 @@ pub fn run_serialization(args: ConvertArgs, hash_table: HashNameTable) -> Result fn run_deserialization(args: ConvertArgs) -> Result<()> { let schema_files = args .input - .list_files("bschema", false)? + .list_files(SchemaFileFilter, false)? .into_iter() .collect::>>()?; if schema_files.is_empty() { diff --git a/toolset/src/diff.rs b/toolset/src/diff.rs index 9e7d660..33d4eb5 100644 --- a/toolset/src/diff.rs +++ b/toolset/src/diff.rs @@ -14,7 +14,7 @@ use rayon::{iter::Either, prelude::*}; use bdat::{BdatFile, Cell, CompatRef, CompatRowRef, CompatTable, Label, RowId}; -use crate::{util::hash::MurmurHashSet, InputData}; +use crate::{filter::BdatFileFilter, util::hash::MurmurHashSet, InputData}; #[derive(Args)] pub struct DiffArgs { @@ -71,14 +71,14 @@ pub fn run_diff(args: DiffArgs) -> Result<()> { .with_message(" (Reading files)"); let new_files = args .input - .list_files("bdat", !args.no_file_names)? + .list_files(BdatFileFilter, !args.no_file_names)? .into_iter(); let old_files = InputData { files: args.old_files, ..Default::default() }; let old_files = old_files - .list_files("bdat", !args.no_file_names)? + .list_files(BdatFileFilter, !args.no_file_names)? .into_iter(); let hash_table = args.input.load_hashes()?; diff --git a/toolset/src/filter.rs b/toolset/src/filter.rs index 15ca76e..ab7c03b 100644 --- a/toolset/src/filter.rs +++ b/toolset/src/filter.rs @@ -1,3 +1,5 @@ +use std::{fs::File, io::BufReader, path::Path}; + use bdat::Label; #[derive(Debug)] @@ -7,6 +9,17 @@ pub struct Filter { pub struct FilterArg(pub String); +pub trait FileFilter: Clone { + /// This function does not fail: we only care about BDAT files, so if there is an error + /// in parsing a BDAT file for the purpose of file type discovery, it should panic instead. + fn filter_file(&self, path: impl AsRef, extension: Option<&str>) -> bool; +} + +#[derive(Clone, Copy)] +pub struct BdatFileFilter; +#[derive(Clone, Copy)] +pub struct SchemaFileFilter; + impl Filter { pub fn contains(&self, label: &Label) -> bool { if self.hashes.is_empty() { @@ -25,6 +38,27 @@ impl Filter { } } +impl FileFilter for BdatFileFilter { + fn filter_file(&self, path: impl AsRef, extension: Option<&str>) -> bool { + if extension.is_some_and(|e| e == "bdat") { + // This also makes sure to throw errors if a ".bdat" file failed + // version detection + return true; + } + // Accept non-".bdat" files that actually appear to be BDAT files + File::open(path) + .map_err(|_| ()) + .and_then(|f| bdat::detect_file_version(BufReader::new(f)).map_err(|_| ())) + .is_ok() + } +} + +impl FileFilter for SchemaFileFilter { + fn filter_file(&self, _: impl AsRef, extension: Option<&str>) -> bool { + extension.is_some_and(|e| e == "bschema") + } +} + impl FromIterator for Filter { fn from_iter>(iter: T) -> Self { Self::from_iter(iter.into_iter().flat_map(|s| { diff --git a/toolset/src/info.rs b/toolset/src/info.rs index f4ad04f..50e16fe 100644 --- a/toolset/src/info.rs +++ b/toolset/src/info.rs @@ -1,5 +1,5 @@ use crate::{ - filter::{Filter, FilterArg}, + filter::{BdatFileFilter, Filter, FilterArg}, util::hash::HashNameTable, InputData, }; @@ -26,7 +26,7 @@ pub fn get_info(args: InfoArgs) -> Result<()> { let table_filter: Filter = args.tables.into_iter().map(FilterArg).collect(); let column_filter: Filter = args.columns.into_iter().map(FilterArg).collect(); - for file in args.input.list_files("bdat", false)? { + for file in args.input.list_files(BdatFileFilter, false)? { let path = file?; let mut file = std::fs::read(&path)?; let tables = args diff --git a/toolset/src/main.rs b/toolset/src/main.rs index a83f7c1..1cab8b7 100644 --- a/toolset/src/main.rs +++ b/toolset/src/main.rs @@ -9,6 +9,7 @@ use anyhow::{Context, Result}; use clap::{Args, Parser, Subcommand}; use convert::ConvertArgs; use diff::DiffArgs; +use filter::FileFilter; use hash::HashArgs; use info::InfoArgs; use itertools::Itertools; @@ -89,12 +90,11 @@ fn main() -> anyhow::Result<()> { } impl InputData { - pub fn list_files<'a, 'b: 'a, E: Into>>( - &'a self, - extension: E, + pub fn list_files( + &self, + filter: impl FileFilter + 'static, canonical_paths: bool, - ) -> Result> + 'a> { - let extension = extension.into(); + ) -> Result> + '_> { let paths: Vec<_> = self .files .iter() @@ -108,17 +108,18 @@ impl InputData { .try_collect()?; Ok(paths.into_iter().flat_map(move |name| { + let filter = filter.clone(); WalkDir::new(name) .into_iter() - .filter_map(move |p| match (p, extension) { + .filter_map(move |p| match (p, Some(&filter)) { (Err(e), _) => Some(Err(e)), + // might want this later (Ok(e), None) => Some(Ok(e.path().to_owned())), - (Ok(e), Some(ext)) => { + (Ok(e), Some(filter)) => { let path = e.path(); - if let Some(path_ext) = path.extension() { - if matches!(path_ext.to_str(), Some(p) if p == ext) { - return Some(Ok(path.to_owned())); - } + let ext = path.extension().and_then(|p| p.to_str()); + if filter.filter_file(path, ext) { + return Some(Ok(path.to_owned())); } None } diff --git a/toolset/src/scramble.rs b/toolset/src/scramble.rs index c64ab27..eb13610 100644 --- a/toolset/src/scramble.rs +++ b/toolset/src/scramble.rs @@ -1,4 +1,5 @@ use crate::error::Error; +use crate::filter::BdatFileFilter; use crate::util::{ProgressBarState, RayonPoolJobs}; use crate::InputData; use anyhow::{Context, Result}; @@ -41,7 +42,7 @@ fn run( let files = args .input - .list_files("bdat", false)? + .list_files(BdatFileFilter, false)? .into_iter() .collect::>>()?;