ksud: create our own new ramdisk when no compatible ramdisk found

Co-authored-by: rifsxd <rifat.44.azad.rifs@gmail.com>

userspace/ksud/src/boot_patch.rs

@@ -153,40 +153,126 @@ fn parse_kmi_from_boot(magiskboot: &Path, image: &PathBuf, workdir: &Path) -> Re
     parse_kmi_from_kernel(&image_path, workdir)
 }
 
-fn do_cpio_cmd(magiskboot: &Path, workdir: &Path, cpio_path: &Path, cmd: &str) -> Result<()> {
+fn do_cpio_cmd(magiskboot: &Path, workdir: &Path, cmd: &str) -> Result<()> {
     let status = Command::new(magiskboot)
         .current_dir(workdir)
         .stdout(Stdio::null())
         .stderr(Stdio::null())
         .arg("cpio")
-        .arg(cpio_path)
+        .arg("ramdisk.cpio")
         .arg(cmd)
         .status()?;
+
     ensure!(status.success(), "magiskboot cpio {} failed", cmd);
     Ok(())
 }
 
-fn is_magisk_patched(magiskboot: &Path, workdir: &Path, cpio_path: &Path) -> Result<bool> {
+fn do_vendor_init_boot_cpio_cmd(magiskboot: &Path, workdir: &Path, cmd: &str) -> Result<()> {
+    let vendor_init_boot_cpio = workdir.join("vendor_ramdisk").join("init_boot.cpio");
     let status = Command::new(magiskboot)
         .current_dir(workdir)
         .stdout(Stdio::null())
         .stderr(Stdio::null())
         .arg("cpio")
-        .arg(cpio_path)
+        .arg(vendor_init_boot_cpio)
-        .arg("test")
+        .arg(cmd)
         .status()?;
+
+    ensure!(status.success(), "magiskboot cpio {} failed", cmd);
+    Ok(())
+}
+
+fn do_vendor_ramdisk_cpio_cmd(magiskboot: &Path, workdir: &Path, cmd: &str) -> Result<()> {
+    let vendor_ramdisk_cpio = workdir.join("vendor_ramdisk").join("ramdisk.cpio");
+    let status = Command::new(magiskboot)
+        .current_dir(workdir)
+        .stdout(Stdio::null())
+        .stderr(Stdio::null())
+        .arg("cpio")
+        .arg(vendor_ramdisk_cpio)
+        .arg(cmd)
+        .status()?;
+
+    ensure!(status.success(), "magiskboot cpio {} failed", cmd);
+    Ok(())
+}
+
+fn is_magisk_patched(magiskboot: &Path, workdir: &Path) -> Result<bool> {
+    let status = Command::new(magiskboot)
+        .current_dir(workdir)
+        .stdout(Stdio::null())
+        .stderr(Stdio::null())
+        .args(["cpio", "ramdisk.cpio", "test"])
+        .status()?;
+
     // 0: stock, 1: magisk
     Ok(status.code() == Some(1))
 }
 
-fn is_kernelsu_patched(magiskboot: &Path, workdir: &Path, cpio_path: &Path) -> Result<bool> {
+fn is_magisk_patched_vendor_init_boot(magiskboot: &Path, workdir: &Path) -> Result<bool> {
+    let vendor_init_boot_cpio = workdir.join("vendor_ramdisk").join("init_boot.cpio");
     let status = Command::new(magiskboot)
         .current_dir(workdir)
         .stdout(Stdio::null())
         .stderr(Stdio::null())
-        .arg("cpio")
+        .args(["cpio", vendor_init_boot_cpio.to_str().unwrap(), "test"])
-        .arg(cpio_path)
+        .status()?;
-        .arg("exists kernelsu.ko")
+
+    // 0: stock, 1: magisk
+    Ok(status.code() == Some(1))
+}
+
+fn is_magisk_patched_vendor_ramdisk(magiskboot: &Path, workdir: &Path) -> Result<bool> {
+    let vendor_ramdisk_cpio = workdir.join("vendor_ramdisk").join("ramdisk.cpio");
+    let status = Command::new(magiskboot)
+        .current_dir(workdir)
+        .stdout(Stdio::null())
+        .stderr(Stdio::null())
+        .args(["cpio", vendor_ramdisk_cpio.to_str().unwrap(), "test"])
+        .status()?;
+
+    // 0: stock, 1: magisk
+    Ok(status.code() == Some(1))
+}
+
+fn is_kernelsu_patched(magiskboot: &Path, workdir: &Path) -> Result<bool> {
+    let status = Command::new(magiskboot)
+        .current_dir(workdir)
+        .stdout(Stdio::null())
+        .stderr(Stdio::null())
+        .args(["cpio", "ramdisk.cpio", "exists kernelsu.ko"])
+        .status()?;
+
+    Ok(status.success())
+}
+
+fn is_kernelsu_patched_vendor_init_boot(magiskboot: &Path, workdir: &Path) -> Result<bool> {
+    let vendor_ramdisk_cpio = workdir.join("vendor_ramdisk").join("init_boot.cpio");
+    let status = Command::new(magiskboot)
+        .current_dir(workdir)
+        .stdout(Stdio::null())
+        .stderr(Stdio::null())
+        .args([
+            "cpio",
+            vendor_ramdisk_cpio.to_str().unwrap(),
+            "exists kernelsu.ko",
+        ])
+        .status()?;
+
+    Ok(status.success())
+}
+
+fn is_kernelsu_patched_vendor_ramdisk(magiskboot: &Path, workdir: &Path) -> Result<bool> {
+    let vendor_ramdisk_cpio = workdir.join("vendor_ramdisk").join("ramdisk.cpio");
+    let status = Command::new(magiskboot)
+        .current_dir(workdir)
+        .stdout(Stdio::null())
+        .stderr(Stdio::null())
+        .args([
+            "cpio",
+            vendor_ramdisk_cpio.to_str().unwrap(),
+            "exists kernelsu.ko",
+        ])
         .status()?;
 
     Ok(status.success())
@@ -224,7 +310,8 @@ pub fn restore(
 
     let skip_init = kmi.starts_with("android12-");
 
-    let (bootimage, bootdevice) = find_boot_image(&image, skip_init, false, false, workdir)?;
+    let (bootimage, bootdevice) =
+        find_boot_image(&image, skip_init, false, false, workdir, &magiskboot)?;
 
     println!("- Unpacking boot image");
     let status = Command::new(&magiskboot)
@@ -236,36 +323,33 @@ pub fn restore(
         .status()?;
     ensure!(status.success(), "magiskboot unpack failed");
 
-    let mut ramdisk = workdir.join("ramdisk.cpio");
+    let no_ramdisk = !workdir.join("ramdisk.cpio").exists();
-    if !ramdisk.exists() {
+    // let no_ramdisk = !workdir.join("ramdisk.cpio").exists();
-        ramdisk = workdir.join("vendor_ramdisk").join("init_boot.cpio")
+    let no_vendor_init_boot = !workdir
-    }
+        .join("vendor_ramdisk")
-    if !ramdisk.exists() {
+        .join("init_boot.cpio")
-        ramdisk = workdir.join("vendor_ramdisk").join("ramdisk.cpio");
+        .exists();
-    }
+    let no_vendor_ramdisk = !workdir.join("vendor_ramdisk").join("ramdisk.cpio").exists();
-    if !ramdisk.exists() {
+    let is_kernelsu_patched = is_kernelsu_patched(&magiskboot, workdir)?;
-        bail!("No compatible ramdisk found.")
+    let is_kernelsu_patched_vendor_init_boot =
-    }
+        is_kernelsu_patched_vendor_init_boot(&magiskboot, workdir)?;
-    let ramdisk = ramdisk.as_path();
+    let is_kernelsu_patched_vendor_ramdisk =
-    let is_kernelsu_patched = is_kernelsu_patched(&magiskboot, workdir, ramdisk)?;
+        is_kernelsu_patched_vendor_ramdisk(&magiskboot, workdir)?;
-    ensure!(is_kernelsu_patched, "boot image is not patched by KernelSU");
+    ensure!(
+        is_kernelsu_patched
+            || is_kernelsu_patched_vendor_init_boot
+            || is_kernelsu_patched_vendor_ramdisk,
+        "boot image is not patched by KernelSU"
+    );
 
     let mut new_boot = None;
     let mut from_backup = false;
 
     #[cfg(target_os = "android")]
-    if do_cpio_cmd(
+    if do_cpio_cmd(&magiskboot, workdir, &format!("exists {BACKUP_FILENAME}")).is_ok() {
-        &magiskboot,
-        workdir,
-        ramdisk,
-        &format!("exists {BACKUP_FILENAME}"),
-    )
-    .is_ok()
-    {
         do_cpio_cmd(
             &magiskboot,
             workdir,
-            ramdisk,
             &format!("extract {BACKUP_FILENAME} {BACKUP_FILENAME}"),
         )?;
         let sha = std::fs::read(workdir.join(BACKUP_FILENAME))?;
@@ -288,13 +372,81 @@ pub fn restore(
     }
 
     if new_boot.is_none() {
+        if !no_ramdisk {
+            println!("- Restoring /ramdisk");
+            println!("- Removing /ramdisk/kernelsu.ko");
             // remove kernelsu.ko
-        do_cpio_cmd(&magiskboot, workdir, ramdisk, "rm kernelsu.ko")?;
+            do_cpio_cmd(&magiskboot, workdir, "rm kernelsu.ko")?;
 
             // if init.real exists, restore it
-        let status = do_cpio_cmd(&magiskboot, workdir, ramdisk, "exists init.real").is_ok();
+            println!("- Checking if init.real exists");
+            let status = do_cpio_cmd(&magiskboot, workdir, "exists init.real").is_ok();
             if status {
-            do_cpio_cmd(&magiskboot, workdir, ramdisk, "mv init.real init")?;
+                println!("- /ramdisk/init.real exists");
+                println!("- Restoring /ramdisk/init.real to init");
+                do_cpio_cmd(&magiskboot, workdir, "mv init.real init")?;
+            } else {
+                println!("- /ramdisk/init.real not found");
+                println!("- Removing ramdisk.cpio");
+                let ramdisk = workdir.join("ramdisk.cpio");
+                std::fs::remove_file(ramdisk)?;
+            }
+        } else if !no_vendor_init_boot {
+            println!("- Restoring /vendor_ramdisk/init_boot");
+            println!("- Removing /vendor_ramdisk/init_boot/kernelsu.ko");
+            // vendor init_boot restore
+            do_vendor_init_boot_cpio_cmd(&magiskboot, workdir, "rm kernelsu.ko")?;
+
+            println!("- Checking if init.real exists");
+            let status =
+                do_vendor_init_boot_cpio_cmd(&magiskboot, workdir, "exists init.real").is_ok();
+            if status {
+                println!("- /vendor_ramdisk/init_boot/init.real exists");
+                println!("- Restoring /vendor_ramdisk/init_boot/init.real to init");
+                do_vendor_init_boot_cpio_cmd(&magiskboot, workdir, "mv init.real init")?;
+            } else {
+                println!("- /vendor_ramdisk/init_boot/init.real not found");
+                println!("- Removing vendor_ramdisk/init_boot.cpio");
+                let vendor_init_boot = workdir.join("vendor_ramdisk").join("init_boot.cpio");
+                std::fs::remove_file(vendor_init_boot)?;
+            }
+        } else if !no_vendor_ramdisk {
+            println!("- Restoring /vendor_ramdisk/ramdisk");
+            println!("- Removing /vendor_ramdisk/ramdisk/kernelsu.ko");
+            // vendor ramdisk restore
+            do_vendor_ramdisk_cpio_cmd(&magiskboot, workdir, "rm kernelsu.ko")?;
+
+            let status =
+                do_vendor_ramdisk_cpio_cmd(&magiskboot, workdir, "exists init.real").is_ok();
+            if status {
+                println!("- /vendor_ramdisk/ramdisk/init.real exists");
+                println!("- Restoring /vendor_ramdisk/ramdisk/init.real to init");
+                do_vendor_ramdisk_cpio_cmd(&magiskboot, workdir, "mv init.real init")?;
+            } else {
+                println!("- /vendor_ramdisk/ramdisk/init.real not found");
+                println!("- Removing vendor_ramdisk/ramdisk.cpio");
+                let vendor_ramdisk = workdir.join("vendor_ramdisk").join("ramdisk.cpio");
+                std::fs::remove_file(vendor_ramdisk)?;
+            }
+        } else {
+            println!("- Restoring /ramdisk");
+            println!("- Removing /ramdisk/kernelsu.ko");
+            // remove kernelsu.ko
+            do_cpio_cmd(&magiskboot, workdir, "rm kernelsu.ko")?;
+
+            // if init.real exists, restore it
+            println!("- Checking if init.real exists");
+            let status = do_cpio_cmd(&magiskboot, workdir, "exists init.real").is_ok();
+            if status {
+                println!("- /ramdisk/init.real exists");
+                println!("- Restoring /ramdisk/init.real to init");
+                do_cpio_cmd(&magiskboot, workdir, "mv init.real init")?;
+            } else {
+                println!("- /ramdisk/init.real not found");
+                println!("- Removing ramdisk.cpio");
+                let ramdisk = workdir.join("ramdisk.cpio");
+                std::fs::remove_file(ramdisk)?;
+            }
         }
 
         println!("- Repacking boot image");
@@ -303,7 +455,7 @@ pub fn restore(
             .stdout(Stdio::null())
             .stderr(Stdio::null())
             .arg("repack")
-            .arg(&bootimage)
+            .arg(bootimage.display().to_string())
             .status()?;
         ensure!(status.success(), "magiskboot repack failed");
         new_boot = Some(workdir.join("new-boot.img"));
@@ -316,7 +468,7 @@ pub fn restore(
         let output_dir = std::env::current_dir()?;
         let now = chrono::Utc::now();
         let output_image = output_dir.join(format!(
-            "kernelsu_restore_{}.img",
+            "kernelsu_next_restore_{}.img",
             now.format("%Y%m%d_%H%M%S")
         ));
 
@@ -424,10 +576,16 @@ fn do_patch(
 
     let skip_init = kmi.starts_with("android12-");
 
-    let (bootimage, bootdevice) =
+    let (bootimage, bootdevice) = find_boot_image(
-        find_boot_image(&image, skip_init, ota, is_replace_kernel, workdir)?;
+        &image,
+        skip_init,
+        ota,
+        is_replace_kernel,
+        workdir,
+        &magiskboot,
+    )?;
 
-    let bootimage = bootimage.as_path();
+    let bootimage = bootimage.display().to_string();
 
     // try extract magiskboot/bootctl
     let _ = assets::ensure_binaries(false);
@@ -462,48 +620,99 @@ fn do_patch(
         .stdout(Stdio::null())
         .stderr(Stdio::null())
         .arg("unpack")
-        .arg(bootimage)
+        .arg(&bootimage)
         .status()?;
     ensure!(status.success(), "magiskboot unpack failed");
 
-    let mut ramdisk = workdir.join("ramdisk.cpio");
+    let no_ramdisk = !workdir.join("ramdisk.cpio").exists();
-    if !ramdisk.exists() {
+    let no_vendor_init_boot = !workdir
-        ramdisk = workdir.join("vendor_ramdisk").join("init_boot.cpio")
+        .join("vendor_ramdisk")
+        .join("init_boot.cpio")
+        .exists();
+    let no_vendor_ramdisk = !workdir.join("vendor_ramdisk").join("ramdisk.cpio").exists();
+    if no_ramdisk && no_vendor_init_boot && no_vendor_ramdisk {
+        println!("- No compatible ramdisk found.");
+        println!("- Will create our own ramdisk!");
     }
-    if !ramdisk.exists() {
+    let is_magisk_patched = is_magisk_patched(&magiskboot, workdir)?;
-        ramdisk = workdir.join("vendor_ramdisk").join("ramdisk.cpio");
+    let is_magisk_patched_vendor_init_boot =
-    }
+        is_magisk_patched_vendor_init_boot(&magiskboot, workdir)?;
-    if !ramdisk.exists() {
+    let is_magisk_patched_vendor_ramdisk = is_magisk_patched_vendor_ramdisk(&magiskboot, workdir)?;
-        bail!("No compatible ramdisk found.");
+    ensure!(
-    }
+        !is_magisk_patched
-    let ramdisk = ramdisk.as_path();
+            || !is_magisk_patched_vendor_init_boot
-    let is_magisk_patched = is_magisk_patched(&magiskboot, workdir, ramdisk)?;
+            || !is_magisk_patched_vendor_ramdisk,
-    ensure!(!is_magisk_patched, "Cannot work with Magisk patched image");
+        "Cannot work with Magisk patched image"
+    );
 
     println!("- Adding KernelSU LKM");
-    let is_kernelsu_patched = is_kernelsu_patched(&magiskboot, workdir, ramdisk)?;
+    let is_kernelsu_patched = is_kernelsu_patched(&magiskboot, workdir)?;
+    let is_kernelsu_patched_vendor_init_boot =
+        is_kernelsu_patched_vendor_init_boot(&magiskboot, workdir)?;
+    let is_kernelsu_patched_vendor_ramdisk =
+        is_kernelsu_patched_vendor_ramdisk(&magiskboot, workdir)?;
 
     let mut need_backup = false;
-    if !is_kernelsu_patched {
+    if (no_ramdisk && !is_kernelsu_patched_vendor_init_boot)
-        // kernelsu.ko is not exist, backup init if necessary
+        || (no_ramdisk && no_vendor_init_boot && !is_kernelsu_patched_vendor_ramdisk)
-        let status = do_cpio_cmd(&magiskboot, workdir, ramdisk, "exists init");
+        || !is_kernelsu_patched
+    {
+        if !no_ramdisk {
+            println!("- Checking if /ramdisk/init exists");
+            let status = do_cpio_cmd(&magiskboot, workdir, "exists init");
             if status.is_ok() {
-            do_cpio_cmd(&magiskboot, workdir, ramdisk, "mv init init.real")?;
+                println!("- Backing up ramdisk/init");
+                do_cpio_cmd(&magiskboot, workdir, "mv init init.real")?;
+            }
+            need_backup = flash;
+        } else if !no_vendor_init_boot {
+            println!("- Checking if /vendor_ramdisk/init_boot/init exists");
+            let status = do_vendor_init_boot_cpio_cmd(&magiskboot, workdir, "exists init");
+            if status.is_ok() {
+                println!("- Backing up vendor_ramdisk/init_boot/init");
+                do_vendor_init_boot_cpio_cmd(&magiskboot, workdir, "mv init init.real")?;
+            }
+            need_backup = flash;
+        } else if !no_vendor_ramdisk {
+            println!("- Checking if /vendor_ramdisk/ramdisk/init exists");
+            let status = do_vendor_ramdisk_cpio_cmd(&magiskboot, workdir, "exists init");
+            if status.is_ok() {
+                println!("- Backing up vendor_ramdisk/ramdisk/init");
+                do_vendor_ramdisk_cpio_cmd(&magiskboot, workdir, "mv init init.real")?;
+            }
+            need_backup = flash;
+        } else {
+            println!("- Checking if /ramdisk/init exists");
+            let status = do_cpio_cmd(&magiskboot, workdir, "exists init");
+            if status.is_ok() {
+                println!("- Backing up ramdisk/init");
+                do_cpio_cmd(&magiskboot, workdir, "mv init init.real")?;
             }
             need_backup = flash;
         }
+    }
 
-    do_cpio_cmd(&magiskboot, workdir, ramdisk, "add 0755 init init")?;
+    if !no_ramdisk {
-    do_cpio_cmd(
+        println!("- Patching /ramdisk");
-        &magiskboot,
+        do_cpio_cmd(&magiskboot, workdir, "add 0755 init init")?;
-        workdir,
+        do_cpio_cmd(&magiskboot, workdir, "add 0755 kernelsu.ko kernelsu.ko")?;
-        ramdisk,
+    } else if !no_vendor_init_boot {
-        "add 0755 kernelsu.ko kernelsu.ko",
+        println!("- Patching /vendor_ramdisk/init_boot");
-    )?;
+        do_vendor_init_boot_cpio_cmd(&magiskboot, workdir, "add 0755 init init")?;
+        do_vendor_init_boot_cpio_cmd(&magiskboot, workdir, "add 0755 kernelsu.ko kernelsu.ko")?;
+    } else if !no_vendor_ramdisk {
+        println!("- Patching /vendor_ramdisk/ramdisk");
+        do_vendor_ramdisk_cpio_cmd(&magiskboot, workdir, "add 0750 init init")?;
+        do_vendor_ramdisk_cpio_cmd(&magiskboot, workdir, "add 0750 kernelsu.ko kernelsu.ko")?;
+    } else {
+        println!("- Creating and Patching /ramdisk");
+        do_cpio_cmd(&magiskboot, workdir, "add 0755 init init")?;
+        do_cpio_cmd(&magiskboot, workdir, "add 0755 kernelsu.ko kernelsu.ko")?;
+    }
 
     #[cfg(target_os = "android")]
     if need_backup {
-        if let Err(e) = do_backup(&magiskboot, workdir, ramdisk, bootimage) {
+        if let Err(e) = do_backup(&magiskboot, workdir, &bootimage) {
             println!("- Backup stock image failed: {e}");
         }
     }
@@ -515,7 +724,7 @@ fn do_patch(
         .stdout(Stdio::null())
         .stderr(Stdio::null())
         .arg("repack")
-        .arg(bootimage)
+        .arg(&bootimage)
         .status()?;
     ensure!(status.success(), "magiskboot repack failed");
     let new_boot = workdir.join("new-boot.img");
@@ -525,7 +734,7 @@ fn do_patch(
         let output_dir = out.unwrap_or(std::env::current_dir()?);
         let now = chrono::Utc::now();
         let output_image = output_dir.join(format!(
-            "kernelsu_patched_{}.img",
+            "kernelsu_next_patched_{}.img",
             now.format("%Y%m%d_%H%M%S")
         ));
 
@@ -570,7 +779,7 @@ fn calculate_sha1(file_path: impl AsRef<Path>) -> Result<String> {
 }
 
 #[cfg(target_os = "android")]
-fn do_backup(magiskboot: &Path, workdir: &Path, cpio_path: &Path, image: &Path) -> Result<()> {
+fn do_backup(magiskboot: &Path, workdir: &Path, image: &str) -> Result<()> {
     let sha1 = calculate_sha1(image)?;
     let filename = format!("{KSU_BACKUP_FILE_PREFIX}{sha1}");
 
@@ -582,7 +791,6 @@ fn do_backup(magiskboot: &Path, workdir: &Path, cpio_path: &Path, image: &Path)
     do_cpio_cmd(
         magiskboot,
         workdir,
-        cpio_path,
         &format!("add 0755 {BACKUP_FILENAME} {BACKUP_FILENAME}"),
     )?;
     println!("- Stock image has been backup to");
@@ -656,16 +864,17 @@ fn find_boot_image(
     ota: bool,
     is_replace_kernel: bool,
     workdir: &Path,
+    magiskboot: &Path,
 ) -> Result<(PathBuf, Option<String>)> {
     let bootimage;
     let mut bootdevice = None;
     if let Some(ref image) = *image {
-        ensure!(image.exists(), "boot image not found");
+        ensure!(image.exists(), "- Boot image not found");
         bootimage = std::fs::canonicalize(image)?;
     } else {
         if cfg!(not(target_os = "android")) {
             println!("- Current OS is not android, refusing auto bootimage/bootdevice detection");
-            bail!("Please specify a boot image");
+            bail!("- Please specify a boot image");
         }
         let mut slot_suffix =
             utils::getprop("ro.boot.slot_suffix").unwrap_or_else(|| String::from(""));
@@ -678,27 +887,112 @@ fn find_boot_image(
             }
         };
 
-        let init_boot_exist =
+        let init_boot_partition = format!("/dev/block/by-name/init_boot{slot_suffix}");
-            Path::new(&format!("/dev/block/by-name/init_boot{slot_suffix}")).exists();
+        let vendor_boot_partition = format!("/dev/block/by-name/vendor_boot{slot_suffix}");
-        let vendor_boot_exist =
+        let boot_partition = format!("/dev/block/by-name/boot{slot_suffix}");
-            Path::new(&format!("/dev/block/by-name/vendor_boot{slot_suffix}")).exists();
-        let boot_partition = if !is_replace_kernel && init_boot_exist && !skip_init {
-            format!("/dev/block/by-name/init_boot{slot_suffix}")
-        } else if !is_replace_kernel && vendor_boot_exist && !skip_init {
-            format!("/dev/block/by-name/vendor_boot{slot_suffix}")
-        } else {
-            format!("/dev/block/by-name/boot{slot_suffix}")
-        };
 
-        println!("- Bootdevice: {boot_partition}");
+        let init_boot_exist = Path::new(&init_boot_partition).exists();
+        let vendor_boot_exist = Path::new(&vendor_boot_partition).exists();
+
+        // helper: unpack a partition and check for a ramdisk and init
+        fn unpack_and_check_init(
+            magiskboot: &Path,
+            workdir: &Path,
+            partition: &str,
+            ramdisk_cpio: &str,
+        ) -> Result<bool> {
+            let tmp_img = workdir.join("probe.img");
+            dd(partition, &tmp_img)?;
+            let status = Command::new(magiskboot)
+                .current_dir(workdir)
+                .stdout(Stdio::null())
+                .stderr(Stdio::null())
+                .arg("unpack")
+                .arg(&tmp_img)
+                .status()?;
+            if !status.success() {
+                let _ = std::fs::remove_file(&tmp_img);
+                return Ok(false);
+            }
+            let ramdisk_path = workdir.join(ramdisk_cpio);
+            let has_init = if ramdisk_path.exists() {
+                Command::new(magiskboot)
+                    .current_dir(workdir)
+                    .stdout(Stdio::null())
+                    .stderr(Stdio::null())
+                    .arg("cpio")
+                    .arg(ramdisk_cpio)
+                    .arg("exists init")
+                    .status()
+                    .map(|s| s.success())
+                    .unwrap_or(false)
+            } else {
+                false
+            };
+            // Clean up
+            let _ = std::fs::remove_file(&tmp_img);
+            let _ = std::fs::remove_file(workdir.join("ramdisk.cpio"));
+            let _ = std::fs::remove_dir_all(workdir.join("vendor_ramdisk"));
+            Ok(has_init)
+        }
+
+        let mut selected_partition = &boot_partition;
+
+        if !is_replace_kernel && init_boot_exist && !skip_init {
+            // try init_boot/ramdisk.cpio
+            if unpack_and_check_init(magiskboot, workdir, &init_boot_partition, "ramdisk.cpio")? {
+                println!("- Using init_boot partition (ramdisk.cpio).");
+                selected_partition = &init_boot_partition;
+            }
+        }
+
+        // try vendor_boot/vendor_ramdisk/init_boot.cpio
+        if selected_partition == &boot_partition
+            && !is_replace_kernel
+            && vendor_boot_exist
+            && !skip_init
+        {
+            if unpack_and_check_init(
+                magiskboot,
+                workdir,
+                &vendor_boot_partition,
+                "vendor_ramdisk/init_boot.cpio",
+            )? {
+                println!("- Using vendor_boot partition (vendor_ramdisk/init_boot.cpio).");
+                selected_partition = &vendor_boot_partition;
+            }
+        }
+
+        // try vendor_boot/vendor_ramdisk/ramdisk.cpio
+        if selected_partition == &boot_partition
+            && !is_replace_kernel
+            && vendor_boot_exist
+            && !skip_init
+        {
+            if unpack_and_check_init(
+                magiskboot,
+                workdir,
+                &vendor_boot_partition,
+                "vendor_ramdisk/ramdisk.cpio",
+            )? {
+                println!("- Using vendor_boot partition (vendor_ramdisk/ramdisk.cpio).");
+                selected_partition = &vendor_boot_partition;
+            }
+        }
+
+        if selected_partition == &boot_partition {
+            println!("- Using boot partition (ramdisk.cpio).");
+        }
+
+        println!("- Bootdevice: {selected_partition}");
         let tmp_boot_path = workdir.join("boot.img");
 
-        dd(&boot_partition, &tmp_boot_path)?;
+        dd(selected_partition, &tmp_boot_path)?;
 
-        ensure!(tmp_boot_path.exists(), "boot image not found");
+        ensure!(tmp_boot_path.exists(), "- Tmp boot image not found");
 
         bootimage = tmp_boot_path;
-        bootdevice = Some(boot_partition);
+        bootdevice = Some(selected_partition.to_string());
     };
     Ok((bootimage, bootdevice))
 }