Refine cpio and dtb subcommands

native/src/boot/cli.rs

@@ -1,5 +1,5 @@
 use crate::cpio::{cpio_commands, print_cpio_usage};
-use crate::dtb::{dtb_commands, print_dtb_usage};
+use crate::dtb::{dtb_commands, print_dtb_usage, DtbAction};
 use crate::ffi::{
     cleanup, compress, decompress_raw, formats, repack, sign, split_image_dtb, unpack, verify
 };
@@ -100,6 +100,8 @@ struct HexPatch {
 #[derive(FromArgs)]
 #[argh(subcommand, name = "cpio")]
 struct Cpio {
+    #[argh(positional)]
+    file: String,
     #[argh(positional)]
     cmds: Vec<String>,
 }
@@ -108,7 +110,9 @@ struct Cpio {
 #[argh(subcommand, name = "dtb")]
 struct Dtb {
     #[argh(positional)]
-    cmds: Vec<String>,
+    file: String,
+    #[argh(subcommand)]
+    action: DtbAction,
 }
 
 #[derive(FromArgs)]
@@ -353,15 +357,27 @@ pub unsafe extern "C" fn main(
                     Err(log_err!("Failed to patch"))?;
                 }
             }
-            Action::Cpio(Cpio { ref cmds }) => {
-                return if cpio_commands(&cmds.iter().map(|x| x.as_str()).collect::<Vec<_>>()) {
+            Action::Cpio(Cpio {
+                ref mut file,
+                ref mut cmds,
+            }) => {
+                return if cpio_commands(file, cmds)
+                    .log_with_msg(|w| w.write_str("Failed to process cpio"))
+                    .is_ok()
+                {
                     0
                 } else {
                     1
                 }
             }
-            Action::Dtb(Dtb { ref cmds }) => {
-                return if dtb_commands(&cmds.iter().map(|x| x.as_str()).collect::<Vec<_>>()) {
+            Action::Dtb(Dtb {
+                ref mut file,
+                ref action,
+            }) => {
+                return if dtb_commands(file, action)
+                    .log_with_msg(|w| w.write_str("Failed to process dtb"))
+                    .unwrap_or(false)
+                {
                     0
                 } else {
                     1

native/src/boot/cpio.rs

@@ -28,14 +28,6 @@ use crate::check_env;
 use crate::ffi::{unxz, xz};
 use crate::patch::{patch_encryption, patch_verity};
 
-#[derive(FromArgs)]
-struct CpioCli {
-    #[argh(positional)]
-    file: String,
-    #[argh(positional)]
-    commands: Vec<String>,
-}
-
 #[derive(FromArgs)]
 struct CpioCommand {
     #[argh(subcommand)]
@@ -753,68 +745,58 @@ impl Display for CpioEntry {
     }
 }
 
-pub fn cpio_commands(cmds: &Vec<&str>) -> bool {
-    let res: LoggedResult<()> = try {
-        let mut cli =
-            CpioCli::from_args(&["magiskboot", "cpio"], cmds).on_early_exit(print_cpio_usage);
-
-        let file = Utf8CStr::from_string(&mut cli.file);
-        let mut cpio = if FsPath::from(file).exists() {
-            Cpio::load_from_file(file)?
-        } else {
-            Cpio::new()
-        };
-
-        for cmd in cli.commands {
-            if cmd.starts_with('#') {
-                continue;
-            }
-            let mut cli = CpioCommand::from_args(
-                &["magiskboot", "cpio", file],
-                cmd.split(' ')
-                    .filter(|x| !x.is_empty())
-                    .collect::<Vec<_>>()
-                    .as_slice(),
-            )
-            .on_early_exit(print_cpio_usage);
-
-            match &mut cli.action {
-                CpioAction::Test(_) => exit(cpio.test()),
-                CpioAction::Restore(_) => cpio.restore()?,
-                CpioAction::Patch(_) => cpio.patch(),
-                CpioAction::Exists(Exists { path }) => {
-                    if cpio.exists(path) {
-                        exit(0);
-                    } else {
-                        exit(1);
-                    }
-                }
-                CpioAction::Backup(Backup {
-                    origin,
-                    skip_compress,
-                }) => cpio.backup(origin, *skip_compress)?,
-                CpioAction::Remove(Remove { path, recursive }) => cpio.rm(path, *recursive),
-                CpioAction::Move(Move { from, to }) => cpio.mv(from, to)?,
-                CpioAction::MakeDir(MakeDir { mode, dir }) => cpio.mkdir(*mode, dir),
-                CpioAction::Link(Link { src, dst }) => cpio.ln(src, dst),
-                CpioAction::Add(Add { mode, path, file }) => cpio.add(*mode, path, file)?,
-                CpioAction::Extract(Extract { paths }) => {
-                    if !paths.is_empty() && paths.len() != 2 {
-                        Err(log_err!("invalid arguments"))?;
-                    }
-                    let mut it = paths.iter_mut();
-                    cpio.extract(it.next(), it.next())?;
-                }
-                CpioAction::List(List { path, recursive }) => {
-                    cpio.ls(path.as_str(), *recursive);
-                    exit(0);
-                }
-            };
-        }
-        cpio.dump(file)?;
+pub(crate) fn cpio_commands(file: &mut String, cmds: &mut Vec<String>) -> LoggedResult<bool> {
+    let file = Utf8CStr::from_string(file);
+    let mut cpio = if FsPath::from(file).exists() {
+        Cpio::load_from_file(file)?
+    } else {
+        Cpio::new()
     };
-    res.log_with_msg(|w| w.write_str("Failed to process cpio"))
-        .is_ok()
+
+    for cmd in cmds {
+        if cmd.starts_with('#') {
+            continue;
+        }
+        let mut cmd = CpioCommand::from_args(
+            &["magiskboot", "cpio", file],
+            cmd.split(' ')
+                .filter(|x| !x.is_empty())
+                .collect::<Vec<_>>()
+                .as_slice(),
+        )
+        .on_early_exit(print_cpio_usage);
+
+        match &mut cmd.action {
+            CpioAction::Test(_) => exit(cpio.test()),
+            CpioAction::Restore(_) => cpio.restore()?,
+            CpioAction::Patch(_) => cpio.patch(),
+            CpioAction::Exists(Exists { path }) => {
+                return Ok(cpio.exists(path));
+            }
+            CpioAction::Backup(Backup {
+                origin,
+                skip_compress,
+            }) => cpio.backup(origin, *skip_compress)?,
+            CpioAction::Remove(Remove { path, recursive }) => cpio.rm(path, *recursive),
+            CpioAction::Move(Move { from, to }) => cpio.mv(from, to)?,
+            CpioAction::MakeDir(MakeDir { mode, dir }) => cpio.mkdir(*mode, dir),
+            CpioAction::Link(Link { src, dst }) => cpio.ln(src, dst),
+            CpioAction::Add(Add { mode, path, file }) => cpio.add(*mode, path, file)?,
+            CpioAction::Extract(Extract { paths }) => {
+                if !paths.is_empty() && paths.len() != 2 {
+                    Err(log_err!("invalid arguments"))?;
+                }
+                let mut it = paths.iter_mut();
+                cpio.extract(it.next(), it.next())?;
+            }
+            CpioAction::List(List { path, recursive }) => {
+                cpio.ls(path.as_str(), *recursive);
+                return Ok(true);
+            }
+        };
+    }
+    cpio.dump(file)?;
+    Ok(true)
 }
 
 fn x8u(x: &[u8; 8]) -> LoggedResult<u32> {

native/src/boot/dtb.rs

@@ -1,4 +1,4 @@
-use std::{cell::UnsafeCell, process::exit};
+use std::cell::UnsafeCell;
 
 use argh::FromArgs;
 use fdt::{
@@ -6,21 +6,14 @@ use fdt::{
     Fdt, FdtError,
 };
 
-use base::{EarlyExitExt, LoggedResult, MappedFile, ResultExt, Utf8CStr};
+use base::{LoggedResult, MappedFile, Utf8CStr};
 
 use crate::{check_env, patch::patch_verity};
 
-#[derive(FromArgs)]
-struct DtbCli {
-    #[argh(positional)]
-    file: String,
-    #[argh(subcommand)]
-    action: DtbAction,
-}
-
 #[derive(FromArgs)]
 #[argh(subcommand)]
-enum DtbAction {
+#[allow(private_interfaces)]
+pub(crate) enum DtbAction {
     Print(Print),
     Patch(Patch),
     Test(Test),
@@ -272,29 +265,18 @@ fn dtb_patch(file: &Utf8CStr) -> LoggedResult<bool> {
     Ok(patched)
 }
 
-pub fn dtb_commands(cmds: &Vec<&str>) -> bool {
-    let res: LoggedResult<()> = try {
-        let mut cli =
-            DtbCli::from_args(&["magiskboot", "dtb"], &cmds).on_early_exit(print_dtb_usage);
-
-        let file = Utf8CStr::from_string(&mut cli.file);
-
-        match cli.action {
-            DtbAction::Print(Print { fstab }) => {
-                dtb_print(file, fstab)?;
-            }
-            DtbAction::Test(_) => {
-                if !dtb_test(file)? {
-                    exit(1);
-                }
-            }
-            DtbAction::Patch(_) => {
-                if !dtb_patch(file)? {
-                    exit(1);
-                }
-            }
+pub(crate) fn dtb_commands(file: &mut String, action: &DtbAction) -> LoggedResult<bool> {
+    let file = Utf8CStr::from_string(file);
+    match action {
+        DtbAction::Print(Print { fstab }) => {
+            dtb_print(file, *fstab)?;
+            Ok(true)
         }
-    };
-    res.log_with_msg(|w| w.write_str("Failed to process dtb"))
-        .is_ok()
+        DtbAction::Test(_) => {
+            Ok(dtb_test(file)?)
+        }
+        DtbAction::Patch(_) => {
+            Ok(dtb_patch(file)?)
+        }
+    }
 }