toolset: allow reading bdat files without bdat extension

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<BdatVersion> {
 /// 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<R: Read + Seek>(reader: R) -> Result<BdatVersion> {
     detect_version(reader)
@@ -145,7 +149,6 @@ fn detect_version<R: Read + Seek>(mut reader: R) -> Result<BdatVersion> {
     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<R: Read + Seek>(mut reader: R) -> Result<BdatVersion> {
         reader.seek(SeekFrom::Start(
             SwitchEndian::read_u32(&first_offset.to_be_bytes()) as u64,
         ))?;
-        if reader.read_u32::<WiiEndian>()? == 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<R: Read + Seek>(mut reader: R) -> Result<BdatVersion> {
     // 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::<SwitchEndian>()? != MAGIC_INT {
+        return Err(DetectError::NotBdat.into());
+    }
+    reader.seek(SeekFrom::Current(32 - 4 * 3))?;
     let string_table_offset = reader.read_u32::<WiiEndian>()?;
     let string_table_len = reader.read_u32::<WiiEndian>()?;
     let final_offset = string_table_offset + string_table_len;

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::<walkdir::Result<Vec<_>>>()?;
     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::<walkdir::Result<Vec<_>>>()?;
     if schema_files.is_empty() {

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()?;
 

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<Path>, 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<Path>, 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<Path>, extension: Option<&str>) -> bool {
+        extension.is_some_and(|e| e == "bschema")
+    }
+}
+
 impl FromIterator<FilterArg> for Filter {
     fn from_iter<T: IntoIterator<Item = FilterArg>>(iter: T) -> Self {
         Self::from_iter(iter.into_iter().flat_map(|s| {

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

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<Option<&'b str>>>(
-        &'a self,
-        extension: E,
+    pub fn list_files(
+        &self,
+        filter: impl FileFilter + 'static,
         canonical_paths: bool,
-    ) -> Result<impl IntoIterator<Item = walkdir::Result<PathBuf>> + 'a> {
-        let extension = extension.into();
+    ) -> Result<impl IntoIterator<Item = walkdir::Result<PathBuf>> + '_> {
         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
                     }

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::<walkdir::Result<Vec<_>>>()?;