Refactor code for file injection

crates/dtmt/src/cmd/bundle/inject.rs

@@ -1,40 +1,68 @@
 use std::path::PathBuf;
 
 use clap::{value_parser, Arg, ArgMatches, Command};
-use color_eyre::eyre::{self, Context, Result};
+use color_eyre::eyre::{self, Context, OptionExt as _, Result};
 use color_eyre::Help;
+use path_slash::PathBufExt as _;
+use sdk::filetype::texture;
 use sdk::murmur::IdString64;
 use sdk::Bundle;
-use tokio::fs::{self, File};
+use tokio::fs;
-use tokio::io::AsyncReadExt;
 
 pub(crate) fn command_definition() -> Command {
     Command::new("inject")
+        .subcommand_required(true)
         .about("Inject a file into a bundle.\n\
             Raw binary data can be used to directly replace the file's variant data blob without affecting the metadata.\n\
             Alternatively, a compiler format may be specified, and a complete bundle file is created.")
-        .arg(
-            Arg::new("replace")
-                .help("The name of a file in the bundle whose content should be replaced.")
-                .short('r')
-                .long("replace"),
-        )
-        .arg(
-            Arg::new("compile")
-                .help("Compile the file with the given data format before injecting.")
-                .long("compile")
-                .short('c')
-        )
         .arg(
             Arg::new("output")
                 .help(
                     "The path to write the changed bundle to. \
-                    If omitted, the input bundle will be overwritten.",
+                    If omitted, the input bundle will be overwritten.\n\
+                    Remember to add a `.patch_<NUMBER>` suffix if you also use '--patch'.",
                 )
                 .short('o')
                 .long("output")
                 .value_parser(value_parser!(PathBuf)),
         )
+        .arg(
+            Arg::new("patch")
+                .help("Create a patch bundle. Optionally, a patch NUMBER may be specified as \
+                    '--patch=123'.\nThe maximum number is 1.")
+                .short('p')
+                .long("patch")
+                .num_args(0..=1)
+                .require_equals(true)
+                .default_missing_value("1")
+                .value_name("NUMBER")
+                .value_parser(value_parser!(u16))
+        )
+        .subcommand(
+            Command::new("replace")
+                .about("Replace an existing file in the bundle")
+                .arg(
+                    Arg::new("compile")
+                        .help("Compile the file as TYPE before injecting. \
+                            Metadata will also be overwritten")
+                        .long("compile")
+                        .short('c')
+                        .value_name("TYPE")
+                        .value_parser(["texture"])
+                )
+                .arg(
+                    Arg::new("variant")
+                        .help("Sepcify the variant index to replace.")
+                        .long("variant")
+                        .default_value("0")
+                        .value_parser(value_parser!(u8))
+                )
+                .arg(
+                    Arg::new("new-file")
+                        .help("Path to the file to inject.")
+                        .required(true)
+                        .value_parser(value_parser!(PathBuf)),
+                )
                 .arg(
                     Arg::new("bundle")
                         .help("Path to the bundle to inject the file into.")
@@ -42,85 +70,198 @@ pub(crate) fn command_definition() -> Command {
                         .value_parser(value_parser!(PathBuf)),
                 )
                 .arg(
-            Arg::new("file")
+                    Arg::new("bundle-file")
-                .help("Path to the file to inject.")
+                        .help("The name of a file in the bundle whose content should be replaced.")
-                .required(true)
+                        .required(true),
-                .value_parser(value_parser!(PathBuf)),
+                ),
         )
+    // .subcommand(
+    //     Command::new("add")
+    //         .about("Add a new file to the bundle")
+    //         .arg(
+    //             Arg::new("new-file")
+    //                 .help("Path to the file to inject.")
+    //                 .required(true)
+    //                 .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)),
+    //         ),
+    // )
 }
 
-#[tracing::instrument(skip_all)]
+#[tracing::instrument(
+    skip_all,
+    fields(
+        bundle_path = tracing::field::Empty,
+        in_file_path = tracing::field::Empty,
+        compile = tracing::field::Empty,
+        output_path = tracing::field::Empty,
+        file_name = tracing::field::Empty,
+    )
+)]
 pub(crate) async fn run(ctx: sdk::Context, matches: &ArgMatches) -> Result<()> {
-    let bundle_path = matches
+    let Some((op, sub_matches)) = matches.subcommand() else {
+        unreachable!("clap is configured to require a subcommand, and they're all handled above");
+    };
+
+    let compile = matches.get_one::<String>("compile");
+
+    let bundle_path = sub_matches
         .get_one::<PathBuf>("bundle")
         .expect("required parameter not found");
 
-    let file_path = matches
+    let in_file_path = sub_matches
         .get_one::<PathBuf>("file")
         .expect("required parameter not found");
 
-    tracing::trace!(bundle_path = %bundle_path.display(), file_path = %file_path.display());
+    let patch_number = matches
+        .get_one::<u16>("patch")
+        .map(|num| format!("{:03}", num));
 
-    let mut bundle = {
+    let output_path = matches
-        let binary = fs::read(bundle_path).await?;
+        .get_one::<PathBuf>("output")
-        let name = Bundle::get_name_from_path(&ctx, bundle_path);
+        .cloned()
-        Bundle::from_binary(&ctx, name, binary).wrap_err("Failed to open bundle file")?
+        .unwrap_or_else(|| {
-    };
+            let mut output_path = bundle_path.clone();
 
-    let name = match matches.get_one::<String>("replace") {
+            if let Some(patch_number) = patch_number.as_ref() {
-        Some(name) => match u64::from_str_radix(name, 16) {
+                output_path.set_extension(format!("patch_{:03}", patch_number));
+            }
+
+            output_path
+        });
+
+    let target_name = if op == "replace" {
+        sub_matches
+            .get_one::<String>("bundle-file")
+            .map(|name| match u64::from_str_radix(name, 16) {
                 Ok(id) => IdString64::from(id),
                 Err(_) => IdString64::String(name.clone()),
-        },
+            })
-        None => eyre::bail!("Currently, only the '--replace' operation is supported."),
+            .expect("argument is required")
+    } else {
+        let mut path = PathBuf::from(in_file_path);
+        path.set_extension("");
+        IdString64::from(path.to_slash_lossy().to_string())
     };
 
-    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.clone()))
-        // TODO: Handle file variants
-        .find_map(|file| file.variants_mut().next())
     {
-        let mut data = Vec::new();
+        let span = tracing::Span::current();
-        file.read_to_end(&mut data)
+        if !span.is_disabled() {
-            .await
+            span.record("bundle_path", bundle_path.display().to_string());
-            .wrap_err("Failed to read input file")?;
+            span.record("in_file_path", in_file_path.display().to_string());
-        variant.set_data(data);
+            span.record("output_path", output_path.display().to_string());
+            span.record("compile", compile);
+            span.record("file_name", target_name.display().to_string());
+        }
+    }
+
+    let mut bundle = {
+        let name = Bundle::get_name_from_path(&ctx, bundle_path);
+        if patch_number.is_some() {
+            unimplemented!("Create patch bundle");
         } else {
-        let err =
+            let binary = fs::read(bundle_path).await?;
-            eyre::eyre!("No file '{}' in this bundle.", name.display()).with_suggestion(|| {
+            Bundle::from_binary(&ctx, name, binary)
+                .wrap_err_with(|| format!("Failed to open bundle '{}'", bundle_path.display()))?
+        }
+    };
+
+    if op == "copy" {
+        unimplemented!("Implement copying a file from one bundle to the other.");
+    }
+
+    let file_data = tokio::fs::read(&in_file_path)
+        .await
+        .wrap_err_with(|| format!("Failed to read file '{}'", in_file_path.display()))?;
+
+    let data = if let Some(compile_type) = compile {
+        match compile_type.as_str() {
+            "texture" => {
+                let sjson = String::from_utf8(file_data).wrap_err_with(|| {
+                    format!("Invalid UTF8 data in '{}'", in_file_path.display())
+                })?;
+
+                let root = in_file_path
+                    .parent()
+                    .ok_or_eyre("File path has no parent")?;
+
+                let texture = texture::compile(target_name.clone(), sjson, root)
+                    .await
+                    .wrap_err_with(|| {
                         format!(
-                    "Run '{} bundle list {}' to list the files in this bundle.",
+                            "Failed to compile file as texture: {}",
+                            in_file_path.display()
+                        )
+                    })?;
+
+                texture.to_binary()?
+            }
+            _ => unreachable!("clap only allows file types that we implement"),
+        }
+    } else {
+        file_data
+    };
+
+    match op {
+        "replace" => {
+            let Some(file) = bundle
+                .files_mut()
+                .find(|file| file.matches_name(&target_name))
+            else {
+                let err = eyre::eyre!(
+                    "No file with name '{}' in bundle '{}'",
+                    target_name.display(),
+                    bundle_path.display()
+                );
+
+                return Err(err).with_suggestion(|| {
+                    format!(
+                        "Run '{} bundle list \"{}\"' to list the files in this bundle.",
                         clap::crate_name!(),
                         bundle_path.display()
                     )
-                // Not yet supported.
-                // })
-                // .with_suggestion(|| {
-                //     format!(
-                //         "Use '{} bundle inject --add {} {} {}' to add it as a new file",
-                //         clap::crate_name!(),
-                //         name.display(),
-                //         bundle_path.display(),
-                //         file_path.display()
-                //     )
                 });
+            };
 
-        return Err(err);
+            let variant_index = sub_matches
+                .get_one::<u8>("variant")
+                .expect("argument with default missing");
+
+            let Some(variant) = file.variants_mut().nth(*variant_index as usize) else {
+                let err = eyre::eyre!(
+                    "Variant index '{}' does not exist in '{}'",
+                    variant_index,
+                    target_name.display()
+                );
+
+                return Err(err).with_suggestion(|| {
+                    format!(
+                        "See '{} bundle inject add --help' if you want to add it as a new file",
+                        clap::crate_name!(),
+                    )
+                });
+            };
+
+            variant.set_data(data);
+        }
+        "add" => {
+            unimplemented!("Implement adding a new file to the bundle.");
+        }
+        _ => unreachable!("no other operations exist"),
     }
 
-    let out_path = matches.get_one::<PathBuf>("output").unwrap_or(bundle_path);
     let data = bundle
         .to_binary()
         .wrap_err("Failed to write changed bundle to output")?;
 
-    fs::write(out_path, &data)
+    fs::write(&output_path, &data)
         .await
-        .wrap_err("Failed to write data to output file")?;
+        .wrap_err_with(|| format!("Failed to write data to '{}'", output_path.display()))?;
 
     Ok(())
 }

lib/sdk/src/bundle/file.rs

@@ -335,14 +335,13 @@ impl BundleFile {
         s
     }
 
-    pub fn matches_name(&self, name: impl Into<IdString64>) -> bool {
+    pub fn matches_name(&self, name: &IdString64) -> bool {
-        let name = name.into();
+        if self.name == *name {
-        if self.name == name {
             return true;
         }
 
         if let IdString64::String(name) = name {
-            self.name(false, None) == name || self.name(true, None) == name
+            self.name(false, None) == *name || self.name(true, None) == *name
         } else {
             false
         }