From 653225bb5b3d27e559947e69bd572e95cf1ab03e Mon Sep 17 00:00:00 2001
From: weishu <twsxtd@gmail.com>
Date: Mon, 15 Jan 2024 20:17:10 +0800
Subject: [PATCH] ksud: Add support for boot patch

---
 userspace/ksud/Cargo.lock        |  75 +++++++++++-
 userspace/ksud/Cargo.toml        |   4 +-
 userspace/ksud/src/boot_patch.rs | 202 +++++++++++++++++++++++++++++++
 userspace/ksud/src/cli.rs        |  46 +++++++
 userspace/ksud/src/main.rs       |   1 +
 5 files changed, 323 insertions(+), 5 deletions(-)
 create mode 100644 userspace/ksud/src/boot_patch.rs

diff --git a/userspace/ksud/Cargo.lock b/userspace/ksud/Cargo.lock
index 889422738929..0d11992cfbbc 100644
--- a/userspace/ksud/Cargo.lock
+++ b/userspace/ksud/Cargo.lock
@@ -231,8 +231,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "16b0a3d9ed01224b22057780a37bb8c5dbfe1be8ba48678e7bf57ec4b385411f"
 dependencies = [
  "iana-time-zone",
+ "js-sys",
  "num-integer",
  "num-traits",
+ "time 0.1.45",
+ "wasm-bindgen",
  "winapi",
 ]
 
@@ -576,6 +579,12 @@ dependencies = [
  "miniz_oxide",
 ]
 
+[[package]]
+name = "fuchsia-cprng"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a06f77d526c1a601b7c4cdd98f54b5eaabffc14d5f2f0296febdc7f357c6d3ba"
+
 [[package]]
 name = "generic-array"
 version = "0.14.6"
@@ -813,6 +822,7 @@ dependencies = [
  "android-properties",
  "android_logger",
  "anyhow",
+ "chrono",
  "clap",
  "const_format",
  "derive-new",
@@ -835,6 +845,7 @@ dependencies = [
  "serde_json",
  "sha256",
  "sys-mount",
+ "tempdir",
  "which",
  "zip 0.6.4",
  "zip-extensions",
@@ -1037,7 +1048,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7676374caaee8a325c9e7a2ae557f216c5563a171d6997b0ef8a65af35147700"
 dependencies = [
  "base64ct",
- "rand_core",
+ "rand_core 0.6.4",
  "subtle",
 ]
 
@@ -1151,6 +1162,19 @@ dependencies = [
  "proc-macro2",
 ]
 
+[[package]]
+name = "rand"
+version = "0.4.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "552840b97013b1a26992c11eac34bdd778e464601a4c2054b5f0bff7c6761293"
+dependencies = [
+ "fuchsia-cprng",
+ "libc",
+ "rand_core 0.3.1",
+ "rdrand",
+ "winapi",
+]
+
 [[package]]
 name = "rand"
 version = "0.8.5"
@@ -1159,7 +1183,7 @@ checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404"
 dependencies = [
  "libc",
  "rand_chacha",
- "rand_core",
+ "rand_core 0.6.4",
 ]
 
 [[package]]
@@ -1169,9 +1193,24 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88"
 dependencies = [
  "ppv-lite86",
- "rand_core",
+ "rand_core 0.6.4",
+]
+
+[[package]]
+name = "rand_core"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7a6fdeb83b075e8266dcc8762c22776f6877a63111121f5f8c7411e5be7eed4b"
+dependencies = [
+ "rand_core 0.4.2",
 ]
 
+[[package]]
+name = "rand_core"
+version = "0.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9c33a3c44ca05fa6f1807d8e6743f3824e8509beca625669633be0acbdf509dc"
+
 [[package]]
 name = "rand_core"
 version = "0.6.4"
@@ -1203,6 +1242,15 @@ dependencies = [
  "num_cpus",
 ]
 
+[[package]]
+name = "rdrand"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "678054eb77286b51581ba43620cc911abf02758c91f93f479767aed0f90458b2"
+dependencies = [
+ "rand_core 0.3.1",
+]
+
 [[package]]
 name = "regex"
 version = "1.7.1"
@@ -1220,13 +1268,22 @@ version = "0.6.28"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "456c603be3e8d448b072f410900c09faf164fbce2d480456f50eea6e25f9c848"
 
+[[package]]
+name = "remove_dir_all"
+version = "0.5.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3acd125665422973a33ac9d3dd2df85edad0f4ae9b00dafb1a05e43a9f5ef8e7"
+dependencies = [
+ "winapi",
+]
+
 [[package]]
 name = "retry"
 version = "2.0.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9166d72162de3575f950507683fac47e30f6f2c3836b71b7fbc61aa517c9c5f4"
 dependencies = [
- "rand",
+ "rand 0.8.5",
 ]
 
 [[package]]
@@ -1452,6 +1509,16 @@ dependencies = [
  "tracing",
 ]
 
+[[package]]
+name = "tempdir"
+version = "0.3.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "15f2b5fb00ccdf689e0149d1b1b3c03fead81c2b37735d812fa8bddbbf41b6d8"
+dependencies = [
+ "rand 0.4.6",
+ "remove_dir_all",
+]
+
 [[package]]
 name = "termcolor"
 version = "1.2.0"
diff --git a/userspace/ksud/Cargo.toml b/userspace/ksud/Cargo.toml
index f4c6422591d9..4a4f938c6dda 100644
--- a/userspace/ksud/Cargo.toml
+++ b/userspace/ksud/Cargo.toml
@@ -44,7 +44,9 @@ procfs = "0.16"
 [target.'cfg(target_os = "android")'.dependencies]
 android_logger = "0.13"
 
+tempdir = "0.3"
+chrono = "0.4"
 [profile.release]
 strip = true
 opt-level = "z"
-lto = true
+#lto = true
diff --git a/userspace/ksud/src/boot_patch.rs b/userspace/ksud/src/boot_patch.rs
new file mode 100644
index 000000000000..45ba3144ef46
--- /dev/null
+++ b/userspace/ksud/src/boot_patch.rs
@@ -0,0 +1,202 @@
+use std::os::unix::fs::PermissionsExt;
+
+use anyhow::bail;
+use anyhow::ensure;
+use anyhow::Context;
+use anyhow::Result;
+use is_executable::IsExecutable;
+use std::path::Path;
+use std::path::PathBuf;
+use std::process::Command;
+use std::process::Stdio;
+
+use crate::utils;
+
+fn ensure_gki_kernel() -> Result<()> {
+    let version =
+        procfs::sys::kernel::Version::current().with_context(|| "get kernel version failed")?;
+    let is_gki = version.major == 5 && version.minor >= 10 || version.major > 5;
+    ensure!(is_gki, "only support GKI kernel");
+    Ok(())
+}
+
+fn do_cpio_cmd(magiskboot: &Path, workding_dir: &Path, cmd: &str) -> Result<()> {
+    let status = Command::new(magiskboot)
+        .current_dir(workding_dir)
+        .stdout(Stdio::null())
+        .stderr(Stdio::null())
+        .arg("cpio")
+        .arg("ramdisk.cpio")
+        .arg(cmd)
+        .status()?;
+
+    ensure!(status.success(), "magiskboot cpio {} failed", cmd);
+    Ok(())
+}
+
+fn dd<P: AsRef<Path> + std::fmt::Debug, Q: AsRef<Path> + std::fmt::Debug>(
+    ifile: P,
+    ofile: Q,
+) -> Result<()> {
+    let status = Command::new("dd")
+        .stdout(Stdio::null())
+        .stderr(Stdio::null())
+        .arg(format!("if={ifile:?}"))
+        .arg(format!("of={ofile:?}"))
+        .status()?;
+    ensure!(status.success(), "dd if={:?} of={:?} failed", ifile, ofile);
+    Ok(())
+}
+
+#[allow(clippy::too_many_arguments)]
+pub fn patch(
+    image: Option<PathBuf>,
+    kernel: Option<PathBuf>,
+    kmod: Option<PathBuf>,
+    init: Option<PathBuf>,
+    ota: bool,
+    flash: bool,
+    out: Option<PathBuf>,
+    magiskboot_path: Option<PathBuf>,
+) -> Result<()> {
+    ensure_gki_kernel()?;
+
+    if kernel.is_some() {
+        ensure!(
+            init.is_none() && kmod.is_none(),
+            "init and module must not be specified."
+        );
+    } else {
+        ensure!(
+            init.is_some() && kmod.is_some(),
+            "init and module must be specified"
+        );
+    }
+
+    let workding_dir = tempdir::TempDir::new("KernelSU")?;
+
+    let bootimage;
+
+    let mut bootdevice = None;
+
+    if let Some(image) = image {
+        ensure!(image.exists(), "boot image not found");
+        bootimage = image;
+    } else {
+        let mut slot_suffix =
+            utils::getprop("ro.boot.slot_suffix").unwrap_or_else(|| String::from(""));
+
+        if !slot_suffix.is_empty() && ota {
+            if slot_suffix == "_a" {
+                slot_suffix = "_b".to_string()
+            } else {
+                slot_suffix = "_a".to_string()
+            }
+        };
+
+        let init_boot_exist =
+            Path::new(&format!("/dev/block/by-name/init_boot{slot_suffix}")).exists();
+        let boot_partition = if init_boot_exist {
+            format!("/dev/block/by-name/init_boot{slot_suffix}")
+        } else {
+            format!("/dev/block/by-name/boot{slot_suffix}")
+        };
+
+        println!("bootdevice: {boot_partition}");
+        let tmp_boot_path = workding_dir.path().join("boot.img");
+
+        dd(&boot_partition, &tmp_boot_path)?;
+
+        ensure!(tmp_boot_path.exists(), "boot image not found");
+
+        bootimage = tmp_boot_path;
+        bootdevice = Some(boot_partition);
+    };
+
+    println!("boot image: {bootimage:?}");
+
+    let magiskboot = magiskboot_path
+        .map(std::fs::canonicalize)
+        .transpose()?
+        .unwrap_or_else(|| "magiskboot".into());
+
+    if !magiskboot.is_executable() {
+        std::fs::set_permissions(&magiskboot, std::fs::Permissions::from_mode(0o755))
+            .with_context(|| "set magiskboot executable failed".to_string())?;
+    }
+
+    ensure!(magiskboot.exists(), "magiskboot not found");
+
+    if let Some(kernel) = kernel {
+        std::fs::copy(kernel, workding_dir.path().join("kernel"))
+            .with_context(|| "copy kernel from failed".to_string())?;
+    }
+
+    if let (Some(kmod), Some(init)) = (kmod, init) {
+        std::fs::copy(kmod, workding_dir.path().join("kernelsu.ko"))
+            .with_context(|| "copy kernel module failed".to_string())?;
+        std::fs::copy(init, workding_dir.path().join("init"))
+            .with_context(|| "copy init failed".to_string())?;
+
+        // magiskboot unpack boot.img
+        // magiskboot cpio ramdisk.cpio 'cp init init.real'
+        // magiskboot cpio ramdisk.cpio 'add 0755 ksuinit init'
+        // magiskboot cpio ramdisk.cpio 'add 0755 <kmod> kernelsu.ko'
+
+        let status = Command::new(&magiskboot)
+            .current_dir(workding_dir.path())
+            .stdout(Stdio::null())
+            .stderr(Stdio::null())
+            .arg("unpack")
+            .arg(bootimage.display().to_string())
+            .status()?;
+        ensure!(status.success(), "magiskboot unpack failed");
+
+        let status = do_cpio_cmd(&magiskboot, workding_dir.path(), "exists init");
+        if status.is_ok() {
+            // init exist, backup it.
+            do_cpio_cmd(&magiskboot, workding_dir.path(), "mv init init.real")?;
+        }
+
+        do_cpio_cmd(&magiskboot, workding_dir.path(), "add 0755 init init")?;
+        do_cpio_cmd(
+            &magiskboot,
+            workding_dir.path(),
+            "add 0755 kernelsu.ko kernelsu.ko",
+        )?;
+    }
+
+    // magiskboot repack boot.img
+    let status = Command::new(&magiskboot)
+        .current_dir(workding_dir.path())
+        .stdout(Stdio::null())
+        .stderr(Stdio::null())
+        .arg("repack")
+        .arg(bootimage.display().to_string())
+        .status()?;
+    ensure!(status.success(), "magiskboot repack failed");
+
+    let out = out.unwrap_or(std::env::current_dir()?);
+
+    let now = chrono::Utc::now();
+    let output_image = out.join(format!(
+        "kernelsu_patched_boot_{}.img",
+        now.format("%Y%m%d_%H%M%S")
+    ));
+    std::fs::copy(workding_dir.path().join("new-boot.img"), &output_image)
+        .with_context(|| "copy out new boot failed".to_string())?;
+
+    if flash {
+        let Some(bootdevice) = bootdevice else {
+            bail!("boot device not found")
+        };
+        let status = Command::new("blockdev")
+            .arg("--setrw")
+            .arg(&bootdevice)
+            .status()?;
+        ensure!(status.success(), "set boot device rw failed");
+
+        dd(&output_image, &bootdevice).with_context(|| "flash boot failed")?;
+    }
+    Ok(())
+}
diff --git a/userspace/ksud/src/cli.rs b/userspace/ksud/src/cli.rs
index 9c8f156c61d5..5e941522dabf 100644
--- a/userspace/ksud/src/cli.rs
+++ b/userspace/ksud/src/cli.rs
@@ -1,5 +1,6 @@
 use anyhow::{Ok, Result};
 use clap::Parser;
+use std::path::PathBuf;
 
 #[cfg(target_os = "android")]
 use android_logger::Config;
@@ -48,6 +49,40 @@ enum Commands {
         command: Profile,
     },
 
+    /// Patch boot or init_boot images to apply KernelSU
+    BootPatch {
+        /// boot image path, if not specified, will try to find the boot image automatically
+        #[arg(short, long)]
+        boot: Option<PathBuf>,
+
+        /// kernel image path to replace
+        #[arg(short, long)]
+        kernel: Option<PathBuf>,
+
+        /// LKM module path to replace
+        #[arg(short, long, requires("init"))]
+        module: Option<PathBuf>,
+
+        /// init to be replaced, if use LKM, this must be specified
+        #[arg(short, long, requires("module"))]
+        init: Option<PathBuf>,
+
+        /// will use another slot when boot image is not specified
+        #[arg(short = 'u', long, default_value = "false")]
+        ota: bool,
+
+        /// Flash it to boot partition after patch
+        #[arg(short, long, default_value = "false")]
+        flash: bool,
+
+        /// output path, if not specified, will use current directory
+        #[arg(short, long, default_value = None)]
+        out: Option<PathBuf>,
+
+        /// magiskboot path, if not specified, will use builtin one
+        #[arg(long, default_value = None)]
+        magiskboot: Option<PathBuf>,
+    },
     /// For developers
     Debug {
         #[command(subcommand)]
@@ -244,6 +279,17 @@ pub fn run() -> Result<()> {
             Debug::Mount => event::mount_systemlessly(defs::MODULE_DIR),
             Debug::Test => todo!(),
         },
+
+        Commands::BootPatch {
+            boot,
+            init,
+            kernel,
+            module,
+            ota,
+            flash,
+            out,
+            magiskboot,
+        } => crate::boot_patch::patch(boot, kernel, module, init, ota, flash, out, magiskboot),
     };
 
     if let Err(e) = &result {
diff --git a/userspace/ksud/src/main.rs b/userspace/ksud/src/main.rs
index d12bbf9395bb..f95ab0371e83 100644
--- a/userspace/ksud/src/main.rs
+++ b/userspace/ksud/src/main.rs
@@ -1,5 +1,6 @@
 mod apk_sign;
 mod assets;
+mod boot_patch;
 mod cli;
 mod debug;
 mod defs;