Simplify magiskboot FFI

native/src/boot/cli.rs

@@ -1,16 +1,17 @@
 use crate::compress::{compress, decompress};
 use crate::cpio::{cpio_commands, print_cpio_usage};
 use crate::dtb::{DtbAction, dtb_commands, print_dtb_usage};
-use crate::ffi::{FileFormat, cleanup, repack, sign, split_image_dtb, unpack, verify};
+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;
+use crate::sign::{sha1_hash, sign_boot_image};
 use argh::FromArgs;
 use base::{
-    CmdArgs, EarlyExitExt, LoggedResult, MappedFile, ResultExt, Utf8CStr, cmdline_logging, cstr,
-    libc::umask, log_err,
+    CmdArgs, EarlyExitExt, LoggedResult, MappedFile, ResultExt, Utf8CStr, WriteExt,
+    cmdline_logging, cstr, libc, libc::umask, log_err,
 };
 use std::ffi::c_char;
+use std::io::{Seek, SeekFrom, Write};
 use std::str::FromStr;
 
 #[derive(FromArgs)]
@@ -159,7 +160,7 @@ fn print_usage(cmd: &str) {
     eprintln!(
         r#"MagiskBoot - Boot Image Modification Tool
 
-Usage: {} <action> [args...]
+Usage: {0} <action> [args...]
 
 Supported actions:
   unpack [-n] [-h] <bootimg>
@@ -259,6 +260,44 @@ Supported actions:
     );
 }
 
+fn verify_cmd(image: &Utf8CStr, cert: Option<&Utf8CStr>) -> bool {
+    let image = BootImage::new(image);
+    match cert {
+        None => {
+            // Boot image parsing already checks if the image is signed
+            image.is_signed()
+        }
+        Some(_) => {
+            // Provide a custom certificate and re-verify
+            image.verify(cert).is_ok()
+        }
+    }
+}
+
+fn sign_cmd(
+    image: &Utf8CStr,
+    name: Option<&Utf8CStr>,
+    cert: Option<&Utf8CStr>,
+    key: Option<&Utf8CStr>,
+) -> LoggedResult<()> {
+    let img = BootImage::new(image);
+    let name = name.unwrap_or(cstr!("/boot"));
+    let sig = sign_boot_image(img.payload(), name, cert, key)?;
+    let tail_off = img.tail_off();
+    drop(img);
+    let mut fd = image.open(libc::O_WRONLY | libc::O_CLOEXEC)?;
+    fd.seek(SeekFrom::Start(tail_off))?;
+    fd.write_all(&sig)?;
+    let current = fd.stream_position()?;
+    let eof = fd.seek(SeekFrom::End(0))?;
+    if eof > current {
+        // Zero out rest of the file
+        fd.seek(SeekFrom::Start(current))?;
+        fd.write_zeros((eof - current) as usize)?;
+    }
+    Ok(())
+}
+
 #[unsafe(no_mangle)]
 pub extern "C" fn main(argc: i32, argv: *const *const c_char, _envp: *const *const c_char) -> i32 {
     cmdline_logging();
@@ -302,40 +341,24 @@ pub extern "C" fn main(argc: i32, argv: *const *const c_char, _envp: *const *con
                     no_compress,
                 );
             }
-            Action::Verify(Verify {
-                ref mut img,
-                ref mut cert,
-            }) => {
-                return unsafe {
-                    verify(
-                        Utf8CStr::from_string(img),
-                        cert.as_mut()
-                            .map(|x| Utf8CStr::from_string(x).as_ptr())
-                            .unwrap_or(std::ptr::null()),
-                    )
+            Action::Verify(Verify { mut img, mut cert }) => {
+                return if verify_cmd(
+                    Utf8CStr::from_string(&mut img),
+                    cert.as_mut().map(Utf8CStr::from_string),
+                ) {
+                    0
+                } else {
+                    1
                 };
             }
-            Action::Sign(Sign {
-                ref mut img,
-                ref mut args,
-            }) => {
-                let (pem, pk8) = match args.get_mut(1..=2) {
-                    Some([pem, pk8]) => (
-                        Utf8CStr::from_string(pem).as_ptr(),
-                        Utf8CStr::from_string(pk8).as_ptr(),
-                    ),
-                    _ => (std::ptr::null(), std::ptr::null()),
-                };
-                return unsafe {
-                    sign(
-                        Utf8CStr::from_string(img),
-                        args.first_mut()
-                            .map(Utf8CStr::from_string)
-                            .unwrap_or(cstr!("/boot")),
-                        pem,
-                        pk8,
-                    )
-                };
+            Action::Sign(Sign { mut img, mut args }) => {
+                let mut iter = args.iter_mut();
+                sign_cmd(
+                    Utf8CStr::from_string(&mut img),
+                    iter.next().map(Utf8CStr::from_string),
+                    iter.next().map(Utf8CStr::from_string),
+                    iter.next().map(Utf8CStr::from_string),
+                )?;
             }
             Action::Extract(Extract {
                 ref payload,