Fix magiskboot cli parsing

native/src/base/misc.rs

@@ -1,5 +1,5 @@
 use crate::{Utf8CStr, cstr, ffi};
-use argh::EarlyExit;
+use argh::{EarlyExit, MissingRequirements};
 use libc::c_char;
 use std::{
     fmt,
@@ -102,6 +102,51 @@ impl<T> EarlyExitExt<T> for Result<T, EarlyExit> {
     }
 }
 
+pub struct PositionalArgParser<'a>(pub slice::Iter<'a, &'a str>);
+
+impl PositionalArgParser<'_> {
+    pub fn required(&mut self, field_name: &'static str) -> Result<String, EarlyExit> {
+        if let Some(next) = self.0.next() {
+            Ok(next.to_string())
+        } else {
+            let mut missing = MissingRequirements::default();
+            missing.missing_positional_arg(field_name);
+            missing.err_on_any()?;
+            unreachable!()
+        }
+    }
+
+    pub fn optional(&mut self) -> Option<String> {
+        self.0.next().map(ToString::to_string)
+    }
+
+    pub fn last_required(&mut self, field_name: &'static str) -> Result<String, EarlyExit> {
+        let r = self.required(field_name)?;
+        self.ensure_end()?;
+        Ok(r)
+    }
+
+    pub fn last_optional(&mut self) -> Result<Option<String>, EarlyExit> {
+        let r = self.optional();
+        if r.is_none() {
+            return Ok(r);
+        }
+        self.ensure_end()?;
+        Ok(r)
+    }
+
+    fn ensure_end(&mut self) -> Result<(), EarlyExit> {
+        if self.0.len() == 0 {
+            Ok(())
+        } else {
+            Err(EarlyExit::from(format!(
+                "Unrecognized argument: {}\n",
+                self.0.next().unwrap()
+            )))
+        }
+    }
+}
+
 pub struct FmtAdaptor<'a, T>(pub &'a mut T)
 where
     T: Write;

native/src/boot/cli.rs

@@ -5,10 +5,10 @@ use crate::ffi::{BootImage, FileFormat, cleanup, repack, split_image_dtb, unpack
 use crate::patch::hexpatch;
 use crate::payload::extract_boot_from_payload;
 use crate::sign::{sha1_hash, sign_boot_image};
-use argh::FromArgs;
+use argh::{CommandInfo, EarlyExit, FromArgs, SubCommand};
 use base::{
-    CmdArgs, EarlyExitExt, LoggedResult, MappedFile, ResultExt, Utf8CStr, WriteExt,
-    cmdline_logging, cstr, libc, libc::umask, log_err,
+    CmdArgs, EarlyExitExt, LoggedResult, MappedFile, PositionalArgParser, ResultExt, Utf8CStr,
+    WriteExt, cmdline_logging, cstr, libc, libc::umask, log_err,
 };
 use std::ffi::c_char;
 use std::io::{Seek, SeekFrom, Write};
@@ -78,13 +78,28 @@ struct Sign {
     args: Vec<String>,
 }
 
-#[derive(FromArgs)]
-#[argh(subcommand, name = "extract")]
 struct Extract {
-    #[argh(positional)]
     payload: String,
-    #[argh(positional)]
-    args: Vec<String>,
+    partition: Option<String>,
+    outfile: Option<String>,
+}
+
+impl FromArgs for Extract {
+    fn from_args(_command_name: &[&str], args: &[&str]) -> Result<Self, EarlyExit> {
+        let mut parse = PositionalArgParser(args.iter());
+        Ok(Extract {
+            payload: parse.required("payload.bin")?,
+            partition: parse.optional(),
+            outfile: parse.last_optional()?,
+        })
+    }
+}
+
+impl SubCommand for Extract {
+    const COMMAND: &'static CommandInfo = &CommandInfo {
+        name: "extract",
+        description: "",
+    };
 }
 
 #[derive(FromArgs)]
@@ -136,26 +151,61 @@ struct Sha1 {
 #[argh(subcommand, name = "cleanup")]
 struct Cleanup {}
 
-#[derive(FromArgs)]
-#[argh(subcommand, name = "compress")]
 struct Compress {
-    #[argh(option, short = 'f', default = r#""gzip".to_string()"#)]
-    format: String,
-    #[argh(positional)]
+    format: FileFormat,
     file: String,
-    #[argh(positional)]
     out: Option<String>,
 }
 
-#[derive(FromArgs)]
-#[argh(subcommand, name = "decompress")]
+impl FromArgs for Compress {
+    fn from_args(command_name: &[&str], args: &[&str]) -> Result<Self, EarlyExit> {
+        let cmd = command_name.last().copied().unwrap_or_default();
+        let fmt = cmd.strip_prefix("compress=").unwrap_or("gzip");
+
+        let Ok(fmt) = FileFormat::from_str(fmt) else {
+            return Err(EarlyExit::from(format!(
+                "Unsupported or unknown compression format: {fmt}\n"
+            )));
+        };
+
+        let mut iter = PositionalArgParser(args.iter());
+        Ok(Compress {
+            format: fmt,
+            file: iter.required("infile")?,
+            out: iter.last_optional()?,
+        })
+    }
+}
+
+impl SubCommand for Compress {
+    const COMMAND: &'static CommandInfo = &CommandInfo {
+        name: "compress",
+        description: "",
+    };
+}
+
 struct Decompress {
-    #[argh(positional)]
     file: String,
-    #[argh(positional)]
     out: Option<String>,
 }
 
+impl FromArgs for Decompress {
+    fn from_args(_command_name: &[&str], args: &[&str]) -> Result<Self, EarlyExit> {
+        let mut iter = PositionalArgParser(args.iter());
+        Ok(Decompress {
+            file: iter.required("infile")?,
+            out: iter.last_optional()?,
+        })
+    }
+}
+
+impl SubCommand for Decompress {
+    const COMMAND: &'static CommandInfo = &CommandInfo {
+        name: "decompress",
+        description: "",
+    };
+}
+
 fn print_usage(cmd: &str) {
     eprintln!(
         r#"MagiskBoot - Boot Image Modification Tool
@@ -307,17 +357,20 @@ fn boot_main(cmds: CmdArgs) -> LoggedResult<i32> {
         cmds[1] = &cmds[1][2..];
     }
 
-    if let Some(fmt) = str::strip_prefix(cmds[1], "compress=") {
-        cmds.insert(1, "compress");
-        cmds.insert(2, "-f");
-        cmds[3] = fmt;
+    let mut cli = if cmds[1].starts_with("compress=") {
+        // Skip the main parser, directly parse the subcommand
+        Compress::from_args(&cmds[..2], &cmds[2..]).map(|compress| Cli {
+            action: Action::Compress(compress),
+        })
+    } else {
+        Cli::from_args(&[cmds[0]], &cmds[1..])
     }
-
-    let mut cli = Cli::from_args(&[cmds[0]], &cmds[1..]).on_early_exit(|| match cmds.get(1) {
-        Some(&"dtb") => print_dtb_usage(),
-        Some(&"cpio") => print_cpio_usage(),
+    .on_early_exit(|| match cmds[1] {
+        "dtb" => print_dtb_usage(),
+        "cpio" => print_cpio_usage(),
         _ => print_usage(cmds[0]),
     });
+
     match cli.action {
         Action::Unpack(Unpack {
             no_decompress,
@@ -358,16 +411,13 @@ fn boot_main(cmds: CmdArgs) -> LoggedResult<i32> {
                 iter.next().map(Utf8CStr::from_string),
             )?;
         }
-        Action::Extract(Extract { payload, args }) => {
-            if args.len() > 2 {
-                log_err!("Too many arguments")?;
-            }
-            extract_boot_from_payload(
-                &payload,
-                args.first().map(|x| x.as_str()),
-                args.get(1).map(|x| x.as_str()),
-            )
-            .log_with_msg(|w| w.write_str("Failed to extract from payload"))?;
+        Action::Extract(Extract {
+            payload,
+            partition,
+            outfile,
+        }) => {
+            extract_boot_from_payload(&payload, partition.as_deref(), outfile.as_deref())
+                .log_with_msg(|w| w.write_str("Failed to extract from payload"))?;
         }
         Action::HexPatch(HexPatch {
             mut file,
@@ -417,15 +467,11 @@ fn boot_main(cmds: CmdArgs) -> LoggedResult<i32> {
             decompress_cmd(&mut file, out.as_mut())?;
         }
         Action::Compress(Compress {
-            ref mut file,
-            ref format,
-            ref mut out,
+            format,
+            mut file,
+            mut out,
         }) => {
-            compress_cmd(
-                FileFormat::from_str(format).unwrap_or(FileFormat::UNKNOWN),
-                file,
-                out.as_mut(),
-            )?;
+            compress_cmd(format, &mut file, out.as_mut())?;
         }
     }
     Ok(0)

native/src/boot/compress.rs

@@ -1,6 +1,6 @@
 use crate::ffi::{FileFormat, check_fmt};
 use base::libc::{O_RDONLY, O_TRUNC, O_WRONLY};
-use base::{Chunker, FileOrStd, LoggedResult, Utf8CStr, Utf8CString, WriteExt, error, log_err};
+use base::{Chunker, FileOrStd, LoggedResult, Utf8CStr, Utf8CString, WriteExt, log_err};
 use bytemuck::bytes_of_mut;
 use bzip2::{Compression as BzCompression, write::BzDecoder, write::BzEncoder};
 use flate2::{Compression as GzCompression, write::GzEncoder, write::MultiGzDecoder};
@@ -480,10 +480,6 @@ pub(crate) fn compress_cmd(
     infile: &mut String,
     outfile: Option<&mut String>,
 ) -> LoggedResult<()> {
-    if method == FileFormat::UNKNOWN {
-        error!("Unsupported compression format");
-    }
-
     let infile = Utf8CStr::from_string(infile);
     let outfile = outfile.map(Utf8CStr::from_string);