feat: Implement bundle writing and file injecting

.gitignore

@@ -3,3 +3,4 @@
 .envrc
 liboo2corelinux64.so
 oo2core_8_win64.dll
+dictionary.csv

src/bin/cmd/bundle/inject.rs

@@ -0,0 +1,112 @@
+use std::{path::PathBuf, sync::Arc};
+
+use clap::{value_parser, Arg, ArgMatches, Command};
+use color_eyre::{
+    eyre::{self, Context, Result},
+    Help,
+};
+use dtmt::Bundle;
+use tokio::{fs::File, io::AsyncReadExt, sync::RwLock};
+
+pub(crate) fn command_definition() -> Command {
+    Command::new("inject")
+        .about("Inject a file into a bundle.")
+        .arg(
+            Arg::new("replace")
+                .help("The name of a file in the bundle whos content should be replaced.")
+                .short('r')
+                .long("replace"),
+        )
+        .arg(
+            Arg::new("output")
+                .help(
+                    "The path to write the changed bundle to. \
+                    If omitted, the input bundle will be overwritten.",
+                )
+                .short('o')
+                .long("output")
+                .value_parser(value_parser!(PathBuf)),
+        )
+        .arg(
+            Arg::new("bundle")
+                .help("Path to the bundle to inject the file into.")
+                .required(true)
+                .value_parser(value_parser!(PathBuf)),
+        )
+        .arg(
+            Arg::new("file")
+                .help("Path to the file to inject.")
+                .required(true)
+                .value_parser(value_parser!(PathBuf)),
+        )
+}
+
+#[tracing::instrument(skip_all)]
+pub(crate) async fn run(ctx: Arc<RwLock<dtmt::Context>>, matches: &ArgMatches) -> Result<()> {
+    let bundle_path = matches
+        .get_one::<PathBuf>("bundle")
+        .expect("required parameter not found");
+
+    let file_path = matches
+        .get_one::<PathBuf>("file")
+        .expect("required parameter not found");
+
+    tracing::trace!(bundle_path = %bundle_path.display(), file_path = %file_path.display());
+
+    let mut bundle = Bundle::open(ctx.clone(), bundle_path)
+        .await
+        .wrap_err("Failed to open bundle file")?;
+
+    if let Some(_name) = matches.get_one::<String>("replace") {
+        let mut file = File::open(&file_path)
+            .await
+            .wrap_err_with(|| format!("failed to open '{}'", file_path.display()))?;
+
+        if let Some(variant) = bundle
+            .files_mut()
+            .filter(|file| file.matches_name(_name))
+            // TODO: Handle file variants
+            .filter_map(|file| file.variants_mut().next())
+            .next()
+        {
+            let mut data = Vec::new();
+            file.read_to_end(&mut data)
+                .await
+                .wrap_err("failed to read input file")?;
+            variant.set_data(data);
+        } else {
+            let err = eyre::eyre!("No file '{}' in this bundle.", _name)
+                .with_suggestion(|| {
+                    format!(
+                        "Run '{} bundle list {}' to list the files in this bundle.",
+                        clap::crate_name!(),
+                        bundle_path.display()
+                    )
+                })
+                .with_suggestion(|| {
+                    format!(
+                        "Use '{} bundle inject --add {} {} {}' to add it as a new file",
+                        clap::crate_name!(),
+                        _name,
+                        bundle_path.display(),
+                        file_path.display()
+                    )
+                });
+
+            return Err(err);
+        }
+
+        let out_path = matches.get_one::<PathBuf>("output").unwrap_or(bundle_path);
+        let mut out_file = File::create(out_path)
+            .await
+            .wrap_err_with(|| format!("failed to open output file {}", out_path.display()))?;
+        bundle
+            .write(ctx.clone(), &mut out_file)
+            .await
+            .wrap_err("failed to write changed bundle to output")?;
+
+        Ok(())
+    } else {
+        eyre::bail!("Currently, only the '--replace' operation is supported.");
+    }
+}

src/bin/cmd/bundle/mod.rs

@@ -8,6 +8,7 @@ use dtmt::Oodle;
 
 mod decompress;
 mod extract;
+mod inject;
 mod list;
 
 #[cfg(target_os = "windows")]
@@ -33,6 +34,7 @@ pub(crate) fn command_definition() -> Command {
         )
         .subcommand(decompress::command_definition())
         .subcommand(extract::command_definition())
+        .subcommand(inject::command_definition())
         .subcommand(list::command_definition())
 }
 
@@ -47,6 +49,7 @@ pub(crate) async fn run(ctx: Arc<RwLock<dtmt::Context>>, matches: &ArgMatches) -
     match matches.subcommand() {
         Some(("decompress", sub_matches)) => decompress::run(ctx, sub_matches).await,
         Some(("extract", sub_matches)) => extract::run(ctx, sub_matches).await,
+        Some(("inject", sub_matches)) => inject::run(ctx, sub_matches).await,
         Some(("list", sub_matches)) => list::run(ctx, sub_matches).await,
         _ => unreachable!(
             "clap is configured to require a subcommand, and they're all handled above"

src/bin/cmd/util.rs

@@ -18,7 +18,7 @@ where
         }
         Err(err) => {
             if err.kind() != io::ErrorKind::NotADirectory {
-                tracing::error!(%err, "Failed to read path");
+                tracing::error!("Failed to read path: {:?}", err);
             }
             let paths = vec![PathBuf::from(path.as_ref())];
             tracing::debug!(is_dir = false, resolved_paths = ?paths);

src/bin/dtmt.rs

@@ -88,7 +88,7 @@ async fn main() -> Result<()> {
                     if is_default {
                         return;
                     }
-                    tracing::error!("{}", err);
+                    tracing::error!("{:#}", err);
 
                     return;
                 }

src/binary.rs

@@ -35,11 +35,11 @@ macro_rules! make_read {
 
 macro_rules! make_write {
     ($func:ident, $op:ident, $type:ty) => {
-        pub(crate) async fn $func<W>(r: &mut W, val: $type) -> Result<()>
+        pub(crate) async fn $func<W>(w: &mut W, val: $type) -> Result<()>
         where
             W: AsyncWrite + AsyncSeek + std::marker::Unpin,
         {
-            let res = r
+            let res = w
                 .$op(val)
                 .await
                 .wrap_err(concat!("failed to write ", stringify!($type)));
@@ -48,7 +48,7 @@ macro_rules! make_write {
                 return res;
             }
 
-            let pos = r.stream_position().await;
+            let pos = w.stream_position().await;
             if pos.is_ok() {
                 res.with_section(|| {
                     format!("{pos:#X} ({pos})", pos = pos.unwrap()).header("Position: ")
@@ -61,7 +61,7 @@ macro_rules! make_write {
 }
 
 macro_rules! make_skip {
-    ($func:ident, $read:ident, $op:ident, $type:ty) => {
+    ($func:ident, $read:ident, $type:ty) => {
         pub(crate) async fn $func<R>(r: &mut R, cmp: $type) -> Result<()>
         where
             R: AsyncRead + AsyncSeek + std::marker::Unpin,
@@ -92,8 +92,8 @@ make_write!(write_u8, write_u8, u8);
 make_write!(write_u32, write_u32_le, u32);
 make_write!(write_u64, write_u64_le, u64);
 
-make_skip!(skip_u8, read_u8, read_u8, u8);
+make_skip!(skip_u8, read_u8, u8);
-make_skip!(skip_u32, read_u32, read_u32_le, u32);
+make_skip!(skip_u32, read_u32, u32);
 
 pub(crate) async fn skip_padding<S>(stream: &mut S) -> Result<()>
 where
@@ -112,7 +112,7 @@ where
     Ok(())
 }
 
-pub(crate) async fn read_up_to<R>(r: &mut R, buf: &mut Vec<u8>) -> Result<usize>
+pub(crate) async fn _read_up_to<R>(r: &mut R, buf: &mut Vec<u8>) -> Result<usize>
 where
     R: AsyncRead + AsyncSeek + std::marker::Unpin,
 {
@@ -142,6 +142,8 @@ where
     let pos = w.stream_position().await?;
     let size = 16 - (pos % 16) as usize;
 
+    tracing::trace!(padding_size = size, "Writing padding");
+
     if size > 0 && size < 16 {
         let buf = vec![0; size];
         w.write_all(&buf).await?;

src/bundle/file.rs

@@ -1,4 +1,3 @@
-use std::ops::Deref;
 use std::sync::Arc;
 
 use color_eyre::{Help, Result, SectionExt};
@@ -158,8 +157,8 @@ impl BundleFileType {
         }
     }
 
-    pub fn hash(&self) -> u64 {
+    pub fn hash(&self) -> Murmur64 {
-        *Murmur64::from(*self).deref()
+        Murmur64::from(*self)
     }
 }
 
@@ -171,7 +170,7 @@ impl From<u64> for BundleFileType {
 
 impl From<Murmur64> for BundleFileType {
     fn from(hash: Murmur64) -> BundleFileType {
-        match hash.deref() {
+        match *hash {
             0x931e336d7646cc26 => BundleFileType::Animation,
             0xdcfb9e18fff13984 => BundleFileType::AnimationCurves,
             0x3eed05ba83af5090 => BundleFileType::Apb,
@@ -361,6 +360,11 @@ impl BundleFileVariant {
     pub fn data(&self) -> &Vec<u8> {
         &self.data
     }
+
+    pub fn set_data(&mut self, data: Vec<u8>) {
+        self.header.size = data.len();
+        self.data = data;
+    }
 }
 
 pub struct BundleFile {
@@ -418,7 +422,7 @@ impl BundleFile {
     where
         W: AsyncWrite + AsyncSeek + std::marker::Unpin,
     {
-        write_u64(w, self.file_type.hash()).await?;
+        write_u64(w, *self.file_type.hash()).await?;
         write_u64(w, *self.hash).await?;
 
         let header_count = self.variants.len();
@@ -459,6 +463,14 @@ impl BundleFile {
         }
     }
 
+    pub fn matches_name<S>(&self, name: S) -> bool
+    where
+        S: AsRef<str>,
+    {
+        let name = name.as_ref();
+        self.name == name || self.name(false) == name || self.name(true) == name
+    }
+
     pub fn hash(&self) -> Murmur64 {
         self.hash
     }
@@ -471,6 +483,10 @@ impl BundleFile {
         &self.variants
     }
 
+    pub fn variants_mut(&mut self) -> impl Iterator<Item = &mut BundleFileVariant> {
+        self.variants.iter_mut()
+    }
+
     pub fn raw(&self) -> Result<Vec<UserFile>> {
         let files = self
             .variants

src/bundle/mod.rs

@@ -56,8 +56,8 @@ impl EntryHeader {
     where
         R: AsyncRead + AsyncSeek + std::marker::Unpin,
     {
-        let extension_hash = r.read_u64().await?;
+        let extension_hash = read_u64(r).await?;
-        let name_hash = r.read_u64().await?;
+        let name_hash = read_u64(r).await?;
         let flags = read_u32(r).await?;
 
         // NOTE: Known values so far:
@@ -67,7 +67,7 @@ impl EntryHeader {
         if flags != 0x0 {
             tracing::debug!(
                 flags,
-                "Unexpected meta flags for file {:08X}.{:08X}",
+                "Unexpected meta flags for file {:016X}.{:016X}",
                 name_hash,
                 extension_hash
             );
@@ -113,17 +113,20 @@ impl Bundle {
         // `AsyncRead` and the bundle name separately.
         let path = path.as_ref();
         let bundle_name = if let Some(name) = path.file_name() {
-            let hash = Murmur64::try_from(name.to_string_lossy().as_ref())
+            match Murmur64::try_from(name.to_string_lossy().as_ref()) {
-                .wrap_err_with(|| format!("failed to turn string into hash: {:?}", name))?;
+                Ok(hash) => ctx.read().await.lookup_hash(hash, HashGroup::Filename),
-            ctx.read().await.lookup_hash(hash, HashGroup::Filename)
+                Err(err) => {
+                    tracing::debug!("failed to turn bundle name into hash: {}", err);
+                    name.to_string_lossy().to_string()
+                }
+            }
         } else {
             eyre::bail!("Invalid path to bundle file: {}", path.display());
         };
 
         let f = fs::File::open(path)
             .await
-            .wrap_err("Failed to open bundle file")
+            .wrap_err_with(|| format!("failed to open bundle file {}", path.display()))?;
-            .with_section(|| path.display().to_string().header("Path"))?;
 
         let mut r = BufReader::new(f);
 
@@ -171,6 +174,7 @@ impl Bundle {
         r.seek(SeekFrom::Current(4)).await?;
 
         let mut decompressed = Vec::with_capacity(unpacked_size);
+        let mut unpacked_size_tracked = unpacked_size;
 
         for (chunk_index, chunk_size) in chunk_sizes.into_iter().enumerate() {
             let span = tracing::debug_span!("Decompressing chunk", chunk_index, chunk_size);
@@ -200,6 +204,14 @@ impl Bundle {
                     OodleLZ_CheckCRC::No,
                 )?;
 
+                if unpacked_size_tracked < CHUNK_SIZE {
+                    raw_buffer.resize(unpacked_size_tracked, 0);
+                } else {
+                    unpacked_size_tracked -= CHUNK_SIZE;
+                }
+
+                tracing::trace!(raw_size = raw_buffer.len());
+
                 decompressed.append(&mut raw_buffer);
                 Ok(())
             }
@@ -254,6 +266,8 @@ impl Bundle {
             let buf = Vec::new();
             let mut c = Cursor::new(buf);
 
+            tracing::trace!(num_files = self.files.len());
+
             async {
                 for file in self.files.iter() {
                     file.write(ctx.clone(), &mut c).await?;
@@ -267,19 +281,48 @@ impl Bundle {
             c.into_inner()
         };
 
+        // Ceiling division (or division toward infinity) to calculate
+        // the number of chunks required to fit the unpacked data.
+        let num_chunks = (unpacked_data.len() + CHUNK_SIZE - 1) / CHUNK_SIZE;
+        tracing::trace!(num_chunks);
+        write_u32(w, num_chunks as u32).await?;
+
+        let chunk_sizes_start = w.stream_position().await?;
+        tracing::trace!(chunk_sizes_start);
+        w.seek(SeekFrom::Current(num_chunks as i64 * 4)).await?;
+
+        write_padding(w).await?;
+
+        tracing::trace!(unpacked_size = unpacked_data.len());
+        write_u32(w, unpacked_data.len() as u32).await?;
+        // NOTE: Unknown u32 that's always been 0 so far
+        write_u32(w, 0).await?;
+
         let chunks = unpacked_data.chunks(CHUNK_SIZE);
 
         let ctx = ctx.read().await;
         let oodle_lib = ctx.oodle.as_ref().unwrap();
+        let mut chunk_sizes = Vec::with_capacity(num_chunks);
 
         for chunk in chunks {
             let compressed = oodle_lib.compress(chunk)?;
+            tracing::trace!(
+                raw_chunk_size = chunk.len(),
+                compressed_chunk_size = compressed.len()
+            );
+            chunk_sizes.push(compressed.len());
             write_u32(w, compressed.len() as u32).await?;
             write_padding(w).await?;
             w.write_all(&compressed).await?;
         }
 
-        todo!("compress data and count chunks");
+        w.seek(SeekFrom::Start(chunk_sizes_start)).await?;
+
+        for size in chunk_sizes {
+            write_u32(w, size as u32).await?;
+        }
+
+        Ok(())
     }
 
     pub fn name(&self) -> &String {
@@ -289,6 +332,10 @@ impl Bundle {
     pub fn files(&self) -> &Vec<BundleFile> {
         &self.files
     }
+
+    pub fn files_mut(&mut self) -> impl Iterator<Item = &mut BundleFile> {
+        self.files.iter_mut()
+    }
 }
 
 /// Returns a decompressed version of the bundle data.

src/murmur/dictionary.rs

@@ -15,6 +15,12 @@ pub enum HashGroup {
     Other,
 }
 
+impl HashGroup {
+    pub fn all() -> [Self; 3] {
+        [Self::Filename, Self::Filetype, Self::Other]
+    }
+}
+
 impl std::fmt::Display for HashGroup {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {

src/murmur/mod.rs

@@ -19,6 +19,14 @@ pub use murmurhash64::hash;
 pub use murmurhash64::hash32;
 pub use murmurhash64::hash_inverse as inverse;
 
+fn _swap_bytes_u32(value: u32) -> u32 {
+    u32::from_le_bytes(value.to_be_bytes())
+}
+
+fn _swap_bytes_u64(value: u64) -> u64 {
+    u64::from_le_bytes(value.to_be_bytes())
+}
+
 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
 pub struct Murmur64(u64);
 
@@ -74,7 +82,7 @@ impl<'de> Visitor<'de> for Murmur64 {
         E: serde::de::Error,
     {
         let bytes = value.to_le_bytes();
-        self.visit_u64(u64::from_le_bytes(bytes))
+        Ok(Self::from(u64::from_le_bytes(bytes)))
     }
 
     fn visit_u64<E>(self, value: u64) -> Result<Self::Value, E>

src/oodle/mod.rs

@@ -79,8 +79,6 @@ impl Oodle {
             )
         };
 
-        tracing::debug!(uncompressed_size = ret, "Decompressed chunk");
-
         if ret == 0 {
             eyre::bail!("Failed to decompress chunk.");
         }
@@ -93,7 +91,9 @@ impl Oodle {
     where
         I: AsRef<[u8]>,
     {
-        let raw = data.as_ref();
+        let mut raw = Vec::from(data.as_ref());
+        raw.resize(CHUNK_SIZE, 0);
+
         // TODO: Query oodle for buffer size
         let mut out = vec![0u8; CHUNK_SIZE];
 
@@ -123,6 +123,8 @@ impl Oodle {
             eyre::bail!("Failed to compress chunk.");
         }
 
+        out.resize(ret as usize, 0);
+
         Ok(out)
     }
 }