diff --git a/crates/league-toolkit/Cargo.toml b/crates/league-toolkit/Cargo.toml index aeac964e..5261601b 100644 --- a/crates/league-toolkit/Cargo.toml +++ b/crates/league-toolkit/Cargo.toml @@ -29,7 +29,8 @@ io-ext = { path = "../io-ext" } league-primitives = { path = "../league-primitives" } zstd = { version = "0.13", default-features = false, optional = true } ruzstd = { version = "0.7", optional = true } - +xxhash-rust = { version = "0.8.15", features = ["xxh64", "xxh3"] } +itertools = "0.14.0" serde = { version = "1.0.204", features = ["derive"], optional = true } paste = "1.0.15" miette = "7.2.0" diff --git a/crates/league-toolkit/src/core/wad/builder.rs b/crates/league-toolkit/src/core/wad/builder.rs new file mode 100644 index 00000000..59aa7d9e --- /dev/null +++ b/crates/league-toolkit/src/core/wad/builder.rs @@ -0,0 +1,262 @@ +use std::io::{self, BufReader, BufWriter, Cursor, Read, Seek, SeekFrom, Write}; + +use byteorder::{WriteBytesExt, LE}; +use flate2::read::GzEncoder; +use io_ext::measure; +use itertools::Itertools; +use xxhash_rust::{xxh3, xxh64}; + +use crate::league_file::LeagueFileKind; + +use super::{WadChunk, WadChunkCompression, WadError}; + +#[derive(Debug, thiserror::Error)] +pub enum WadBuilderError { + #[error("wad error")] + WadError(#[from] WadError), + + #[error("io error")] + IoError(#[from] io::Error), + + #[error("unsupported compression type: {0}")] + UnsupportedCompressionType(WadChunkCompression), +} + +/// Implements a builder interface for creating WAD files. +/// +/// ## This example builds a WAD file in memory +/// ``` +/// # use league_toolkit::core::wad::*; +/// # use std::io::{Cursor, Write}; +/// +/// let mut builder = WadBuilder::default(); +/// let scratch = Vec::new(); +/// let mut wad_cursor = Cursor::new(scratch); +/// +/// builder = builder.with_chunk(WadChunkBuilder::default().with_path("path/to/chunk")); +/// builder.build_to_writer(&mut wad_cursor, |path, cursor| { +/// cursor.write_all(&[0xAA; 100])?; +/// +/// Ok(()) +/// }) +/// .expect("Failed to build WAD"); +/// ``` +#[derive(Debug, Default)] +pub struct WadBuilder { + chunk_builders: Vec, +} + +impl WadBuilder { + pub fn with_chunk(mut self, chunk: WadChunkBuilder) -> Self { + self.chunk_builders.push(chunk); + self + } + + /// Build the WAD file and write it to the given writer. + /// + /// * `writer` - The writer to write the WAD file to. + /// * `provide_chunk_data` - A function that provides the rawdata for each chunk. + pub fn build_to_writer< + TWriter: io::Write + io::Seek, + TChunkDataProvider: Fn(u64, &mut Cursor>) -> Result<(), WadBuilderError>, + >( + self, + writer: &mut TWriter, + provide_chunk_data: TChunkDataProvider, + ) -> Result<(), WadBuilderError> { + // First we need to write a dummy header and TOC, so we can calculate from where to start writing the chunks + let mut writer = BufWriter::new(writer); + + let (_, toc_offset) = self.write_dummy_toc::(&mut writer)?; + + // Sort the chunks by path hash, otherwise League wont load the WAD + let ordered_chunks = self + .chunk_builders + .iter() + .sorted_by_key(|chunk| chunk.path) + .collect::>(); + + let mut final_chunks = Vec::new(); + + for chunk in ordered_chunks { + let mut cursor = Cursor::new(Vec::new()); + provide_chunk_data(chunk.path, &mut cursor)?; + + let chunk_data_size = cursor.get_ref().len(); + let (compressed_data, compression) = + Self::compress_chunk_data(cursor.get_ref(), chunk.force_compression)?; + let compressed_data_size = compressed_data.len(); + let compressed_checksum = xxh3::xxh3_64(&compressed_data); + + let chunk_data_offset = writer.stream_position()?; + writer.write_all(&compressed_data)?; + + final_chunks.push(WadChunk { + path_hash: chunk.path, + data_offset: chunk_data_offset as usize, + compressed_size: compressed_data_size, + uncompressed_size: chunk_data_size, + compression_type: compression, + is_duplicated: false, + frame_count: 0, + start_frame: 0, + checksum: compressed_checksum, + }); + } + + writer.seek(SeekFrom::Start(toc_offset))?; + + for chunk in final_chunks { + chunk.write_v3_4(&mut writer)?; + } + + Ok(()) + } + + fn write_dummy_toc( + &self, + writer: &mut BufWriter<&mut W>, + ) -> Result<(u64, u64), WadBuilderError> { + let (header_toc_size, toc_offset) = measure(writer, |writer| { + // Write the header + writer.write_u16::(0x5752)?; + writer.write_u8(3)?; // major + writer.write_u8(4)?; // minor + + // Write dummy ECDSA signature + writer.write_all(&[0; 256])?; + writer.write_u64::(0)?; + + // Write dummy TOC + writer.write_u32::(self.chunk_builders.len() as u32)?; + let toc_offset = writer.stream_position()?; + for _ in self.chunk_builders.iter() { + writer.write_all(&[0; 32])?; + } + + Ok::<_, WadBuilderError>(toc_offset) + })?; + + Ok((header_toc_size, toc_offset)) + } + + fn compress_chunk_data( + data: &[u8], + force_compression: Option, + ) -> Result<(Vec, WadChunkCompression), WadBuilderError> { + let (compressed_data, compression) = match force_compression { + Some(compression) => ( + Self::compress_chunk_data_by_compression(data, compression)?, + compression, + ), + None => { + let kind = LeagueFileKind::identify_from_bytes(data); + let compression = kind.ideal_compression(); + let compressed_data = Self::compress_chunk_data_by_compression(data, compression)?; + + (compressed_data, compression) + } + }; + + Ok((compressed_data, compression)) + } + + fn compress_chunk_data_by_compression( + data: &[u8], + compression: WadChunkCompression, + ) -> Result, WadBuilderError> { + let mut compressed_data = Vec::new(); + match compression { + WadChunkCompression::None => { + compressed_data = data.to_vec(); + } + WadChunkCompression::GZip => { + let reader = BufReader::new(data); + let mut encoder = GzEncoder::new(reader, flate2::Compression::default()); + + encoder.read_to_end(&mut compressed_data)?; + } + WadChunkCompression::Zstd => { + let mut encoder = zstd::Encoder::new(BufWriter::new(&mut compressed_data), 3)?; + encoder.write_all(data)?; + encoder.finish()?; + } + WadChunkCompression::Satellite => { + return Err(WadBuilderError::UnsupportedCompressionType(compression)); + } + WadChunkCompression::ZstdMulti => { + return Err(WadBuilderError::UnsupportedCompressionType(compression)); + } + } + + Ok(compressed_data) + } +} + +/// Implements a builder interface for creating WAD chunks. +/// +/// # Examples +/// ``` +/// # use league_toolkit::core::wad::*; +/// # +/// let builder = WadChunkBuilder::default(); +/// builder.with_path("path/to/chunk"); +/// builder.with_force_compression(WadChunkCompression::Zstd); +/// ``` +#[derive(Debug, Clone, Copy, Default)] +pub struct WadChunkBuilder { + /// The path hash of the chunk. Hashed using xxhash64. + path: u64, + + /// If provided, the chunk will be compressed using the given compression type, otherwise the ideal compression will be used. + force_compression: Option, +} + +impl WadChunkBuilder { + pub fn with_path(mut self, path: impl AsRef) -> Self { + self.path = xxh64::xxh64(path.as_ref().to_lowercase().as_bytes(), 0); + self + } + + pub fn with_force_compression(mut self, compression: WadChunkCompression) -> Self { + self.force_compression = Some(compression); + self + } +} + +#[cfg(test)] +mod tests { + use crate::core::wad::Wad; + + use super::*; + + #[test] + fn test_wad_builder() { + let scratch = Vec::new(); + let mut cursor = Cursor::new(scratch); + + let mut builder = WadBuilder::default(); + builder = builder.with_chunk(WadChunkBuilder::default().with_path("test1")); + builder = builder.with_chunk(WadChunkBuilder::default().with_path("test2")); + builder = builder.with_chunk(WadChunkBuilder::default().with_path("test3")); + + builder + .build_to_writer(&mut cursor, |path, cursor| { + cursor.write_all(&[0xAA; 100])?; + + Ok(()) + }) + .expect("Failed to build WAD"); + + cursor.set_position(0); + + let wad = Wad::mount(cursor).expect("Failed to mount WAD"); + assert_eq!(wad.chunks().len(), 3); + + let chunk = wad.chunks.get(&xxh64::xxh64(b"test1", 0)).unwrap(); + assert_eq!(chunk.path_hash, xxh64::xxh64(b"test1", 0)); + assert_eq!(chunk.compressed_size, 17); + assert_eq!(chunk.uncompressed_size, 100); + assert_eq!(chunk.compression_type, WadChunkCompression::Zstd); + } +} diff --git a/crates/league-toolkit/src/core/wad/chunk.rs b/crates/league-toolkit/src/core/wad/chunk.rs index 268d3721..a6dcd2da 100644 --- a/crates/league-toolkit/src/core/wad/chunk.rs +++ b/crates/league-toolkit/src/core/wad/chunk.rs @@ -1,6 +1,7 @@ use std::io::{BufReader, Read}; +use std::{fmt, io}; -use byteorder::{ReadBytesExt as _, LE}; +use byteorder::{ReadBytesExt as _, WriteBytesExt as _, LE}; use num_enum::{IntoPrimitive, TryFromPrimitive}; use super::WadError; @@ -16,6 +17,18 @@ pub enum WadChunkCompression { ZstdMulti = 4, } +impl fmt::Display for WadChunkCompression { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + WadChunkCompression::None => write!(f, "None"), + WadChunkCompression::GZip => write!(f, "GZip"), + WadChunkCompression::Satellite => write!(f, "Satellite"), + WadChunkCompression::Zstd => write!(f, "Zstd"), + WadChunkCompression::ZstdMulti => write!(f, "ZstdMulti"), + } + } +} + #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] #[derive(Clone, Copy, Debug, PartialEq, Eq)] /// A single wad chunk @@ -27,12 +40,12 @@ pub struct WadChunk { pub compression_type: WadChunkCompression, pub is_duplicated: bool, pub frame_count: u8, - pub start_frame: u16, + pub start_frame: u32, pub checksum: u64, } impl WadChunk { - pub(crate) fn read(reader: &mut BufReader) -> Result { + pub fn read_v3_1(reader: &mut BufReader) -> Result { let path_hash = reader.read_u64::()?; let data_offset = reader.read_u32::()? as usize; let compressed_size = reader.read_i32::()? as usize; @@ -41,7 +54,9 @@ impl WadChunk { let type_frame_count = reader.read_u8()?; let frame_count = type_frame_count >> 4; let compression_type = WadChunkCompression::try_from_primitive(type_frame_count & 0xF) - .expect("failed to read chunk compression"); + .map_err(|_| WadError::InvalidChunkCompression { + compression: type_frame_count & 0xF, + })?; let is_duplicated = reader.read_u8()? == 1; let start_frame = reader.read_u16::()?; @@ -55,11 +70,57 @@ impl WadChunk { compression_type, is_duplicated, frame_count, - start_frame, + start_frame: start_frame as u32, checksum, }) } + pub fn read_v3_4(reader: &mut BufReader) -> Result { + let path_hash = reader.read_u64::()?; + let data_offset = reader.read_u32::()? as usize; + let compressed_size = reader.read_u32::()? as usize; + let uncompressed_size = reader.read_u32::()? as usize; + + let type_frame_count = reader.read_u8()?; + let frame_count = type_frame_count >> 4; + let compression_type = WadChunkCompression::try_from_primitive(type_frame_count & 0xF) + .map_err(|_| WadError::InvalidChunkCompression { + compression: type_frame_count & 0xF, + })?; + + let start_frame = read_24_bit_subchunk_start_frame(reader)?; + + let checksum = reader.read_u64::()?; + + Ok(WadChunk { + path_hash, + data_offset, + compressed_size, + uncompressed_size, + compression_type, + is_duplicated: false, // v3.4 always has is_duplicated = false + frame_count, + start_frame: start_frame as u32, + checksum, + }) + } + + pub fn write_v3_4(&self, writer: &mut W) -> Result<(), WadError> { + writer.write_u64::(self.path_hash)?; + writer.write_u32::(self.data_offset as u32)?; + writer.write_u32::(self.compressed_size as u32)?; + writer.write_u32::(self.uncompressed_size as u32)?; + + let type_frame_count = (self.frame_count << 4) | (self.compression_type as u8 & 0xF); + writer.write_u8(type_frame_count)?; + + write_24_bit_subchunk_start_frame(writer, self.start_frame)?; + + writer.write_u64::(self.checksum)?; + + Ok(()) + } + pub fn path_hash(&self) -> u64 { self.path_hash } @@ -79,3 +140,46 @@ impl WadChunk { self.checksum } } + +pub(crate) fn read_24_bit_subchunk_start_frame( + reader: &mut BufReader, +) -> Result { + let start_frame_hi = reader.read_u8()? as u32; + let start_frame_lo = reader.read_u8()? as u32; + let start_frame_mi = reader.read_u8()? as u32; + let start_frame = start_frame_hi << 16 | start_frame_mi << 8 | start_frame_lo; + + Ok(start_frame) +} + +pub(crate) fn write_24_bit_subchunk_start_frame( + writer: &mut W, + start_frame: u32, +) -> Result<(), WadError> { + writer.write_u8(((start_frame >> 16) & 0xFF) as u8)?; + writer.write_u8((start_frame & 0xFF) as u8)?; + writer.write_u8(((start_frame >> 8) & 0xFF) as u8)?; + + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::io::{BufReader, Cursor}; + + use crate::core::wad::{read_24_bit_subchunk_start_frame, write_24_bit_subchunk_start_frame}; + + #[test] + fn test_read_24_bit_subchunk_start_frame() { + let mut reader = BufReader::new(Cursor::new([0x01, 0x03, 0x02])); + let start_frame = read_24_bit_subchunk_start_frame(&mut reader).unwrap(); + assert_eq!(start_frame, 0x010203); + } + + #[test] + fn test_write_24_bit_subchunk_start_frame() { + let mut writer = Vec::new(); + write_24_bit_subchunk_start_frame(&mut writer, 0x010302).unwrap(); + assert_eq!(writer, [0x01, 0x02, 0x03]); + } +} diff --git a/crates/league-toolkit/src/core/wad/mod.rs b/crates/league-toolkit/src/core/wad/mod.rs index 1e2beb5e..2f34bdc5 100644 --- a/crates/league-toolkit/src/core/wad/mod.rs +++ b/crates/league-toolkit/src/core/wad/mod.rs @@ -1,8 +1,11 @@ //! Wad file handling + +mod builder; mod chunk; mod decoder; mod error; +pub use builder::*; pub use chunk::*; pub use decoder::*; pub use error::*; @@ -63,7 +66,12 @@ impl Wad { let chunk_count = reader.read_i32::()? as usize; let mut chunks = HashMap::::with_capacity(chunk_count); for _ in 0..chunk_count { - let chunk = WadChunk::read(&mut reader)?; + let chunk = match (major, minor) { + (3, 1) => WadChunk::read_v3_1(&mut reader), + (3, 4) => WadChunk::read_v3_4(&mut reader), + _ => Err(WadError::InvalidVersion { major, minor }), + }?; + chunks .insert(chunk.path_hash(), chunk) .map_or(Ok(()), |chunk| { diff --git a/crates/league-toolkit/src/league_file/kind.rs b/crates/league-toolkit/src/league_file/kind.rs index 66271c4b..05ae937e 100644 --- a/crates/league-toolkit/src/league_file/kind.rs +++ b/crates/league-toolkit/src/league_file/kind.rs @@ -1,3 +1,5 @@ +use crate::core::wad::WadChunkCompression; + use super::pattern::LEAGUE_FILE_MAGIC_BYTES; #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] @@ -142,4 +144,23 @@ impl LeagueFileKind { LeagueFileKind::Unknown } + + /// Get the ideal compression for this file type. + /// + /// # Examples + /// ``` + /// # use league_toolkit::league_file::*; + /// # use league_toolkit::core::wad::WadChunkCompression; + /// # + /// assert_eq!(LeagueFileKind::Animation.ideal_compression(), WadChunkCompression::Zstd); + /// assert_eq!(LeagueFileKind::WwisePackage.ideal_compression(), WadChunkCompression::None); + /// ``` + pub fn ideal_compression(&self) -> WadChunkCompression { + // TODO: Maybe we should move this into the wad module ? + + match self { + LeagueFileKind::WwisePackage | LeagueFileKind::WwiseBank => WadChunkCompression::None, + _ => WadChunkCompression::Zstd, + } + } } diff --git a/crates/league-toolkit/src/league_file/pattern.rs b/crates/league-toolkit/src/league_file/pattern.rs index 211a294b..27381d13 100644 --- a/crates/league-toolkit/src/league_file/pattern.rs +++ b/crates/league-toolkit/src/league_file/pattern.rs @@ -1,5 +1,3 @@ -use lazy_static::lazy_static; - use super::LeagueFileKind; pub static LEAGUE_FILE_MAGIC_BYTES: &[LeagueFilePattern] = &[ @@ -83,6 +81,7 @@ impl LeagueFilePattern { } } +#[cfg(test)] mod tests { use super::*;