Merge pull request #31 from LeagueToolkit/feat/wad-builder

feat: wad builder

Author: Crauzer

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"

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<WadChunkBuilder>,
+}
+
+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<Vec<u8>>) -> 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::<TWriter>(&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::<Vec<_>>();
+
+        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<W: io::Write + io::Seek>(
+        &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::<LE>(0x5752)?;
+            writer.write_u8(3)?; // major
+            writer.write_u8(4)?; // minor
+
+            // Write dummy ECDSA signature
+            writer.write_all(&[0; 256])?;
+            writer.write_u64::<LE>(0)?;
+
+            // Write dummy TOC
+            writer.write_u32::<LE>(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<WadChunkCompression>,
+    ) -> Result<(Vec<u8>, 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<Vec<u8>, 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<WadChunkCompression>,
+}
+
+impl WadChunkBuilder {
+    pub fn with_path(mut self, path: impl AsRef<str>) -> 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);
+    }
+}