From 78c64d39ecc7f57105172b0f30dbbecd21a0c7d2 Mon Sep 17 00:00:00 2001
From: topjohnwu <topjohnwu@gmail.com>
Date: Mon, 30 Dec 2019 13:04:39 +0800
Subject: [PATCH] Add split command to magiskboot

Allow splitting image.*-dtb files to kernel and dtb
---
 native/jni/magiskboot/bootimg.cpp  | 50 +++++++++++++++++++++++-------
 native/jni/magiskboot/magiskboot.h |  1 +
 native/jni/magiskboot/main.cpp     | 24 +++++++-------
 3 files changed, 53 insertions(+), 22 deletions(-)

diff --git a/native/jni/magiskboot/bootimg.cpp b/native/jni/magiskboot/bootimg.cpp
index 3ab44d753..cf067e5ce 100644
--- a/native/jni/magiskboot/bootimg.cpp
+++ b/native/jni/magiskboot/bootimg.cpp
@@ -282,33 +282,61 @@ void boot_img::parse_image(uint8_t *addr) {
 	fprintf(stderr, "RAMDISK_FMT     [%s]\n", fmt2name[r_fmt]);
 }
 
-void boot_img::find_kernel_dtb() {
-	const int eof = static_cast<int>(hdr->kernel_size());
-	for (int i = 0; i < eof - (int) sizeof(fdt_header); ++i) {
-		auto fdt_hdr = reinterpret_cast<fdt_header *>(kernel + i);
+static int find_dtb_offset(uint8_t *buf, int sz) {
+	for (int off = 0; off < sz - (int) sizeof(fdt_header); ++off) {
+		auto fdt_hdr = reinterpret_cast<fdt_header *>(buf + off);
 		if (fdt32_to_cpu(fdt_hdr->magic) != FDT_MAGIC)
 			continue;
 
 		// Check that fdt_header.totalsize does not overflow kernel image size
 		uint32_t totalsize = fdt32_to_cpu(fdt_hdr->totalsize);
-		if (totalsize + i > eof)
+		if (totalsize + off > sz)
 			continue;
 
 		// Check that fdt_header.off_dt_struct does not overflow kernel image size
 		uint32_t off_dt_struct = fdt32_to_cpu(fdt_hdr->off_dt_struct);
-		if (off_dt_struct + i > eof)
+		if (off_dt_struct + off > sz)
 			continue;
 
 		// Check that fdt_node_header.tag of first node is FDT_BEGIN_NODE
-		auto fdt_node_hdr = reinterpret_cast<fdt_node_header *>(kernel + i + off_dt_struct);
+		auto fdt_node_hdr = reinterpret_cast<fdt_node_header *>(buf + off + off_dt_struct);
 		if (fdt32_to_cpu(fdt_node_hdr->tag) != FDT_BEGIN_NODE)
 			continue;
 
-		kernel_dtb = kernel + i;
-		kernel_dt_size = eof - i;
-		hdr->kernel_size() = i;
+		return off;
+	}
+	return -1;
+}
+
+void boot_img::find_kernel_dtb() {
+	if (int off = find_dtb_offset(kernel, hdr->kernel_size()); off > 0) {
+		kernel_dtb = kernel + off;
+		kernel_dt_size = hdr->kernel_size() - off;
+		hdr->kernel_size() = off;
 		fprintf(stderr, "KERNEL_DTB      [%u]\n", kernel_dt_size);
-		break;
+	}
+}
+
+int split_image_dtb(const char *filename) {
+	uint8_t *buf;
+	size_t sz;
+	mmap_ro(filename, buf, sz);
+	run_finally f([=]{ munmap(buf, sz); });
+
+	if (int off = find_dtb_offset(buf, sz); off > 0) {
+		format_t fmt = check_fmt(buf, sz);
+		if (COMPRESSED(fmt)) {
+			int fd = creat(KERNEL_FILE, 0644);
+			decompress(fmt, fd, buf, off);
+			close(fd);
+		} else {
+			dump(buf, off, KERNEL_FILE);
+		}
+		dump(buf + off, sz - off, KER_DTB_FILE);
+		return 0;
+	} else {
+		fprintf(stderr, "Cannot find DTB in %s\n", filename);
+		return 1;
 	}
 }
 
diff --git a/native/jni/magiskboot/magiskboot.h b/native/jni/magiskboot/magiskboot.h
index 580018e49..912b3657e 100644
--- a/native/jni/magiskboot/magiskboot.h
+++ b/native/jni/magiskboot/magiskboot.h
@@ -14,6 +14,7 @@
 
 int unpack(const char *image, bool nodecomp = false, bool hdr = false);
 void repack(const char* src_img, const char* out_img, bool nocomp = false);
+int split_image_dtb(const char *filename);
 int hexpatch(const char *image, const char *from, const char *to);
 int cpio_commands(int argc, char *argv[]);
 int dtb_commands(int argc, char *argv[]);
diff --git a/native/jni/magiskboot/main.cpp b/native/jni/magiskboot/main.cpp
index 2c5a89c5d..825d4d255 100644
--- a/native/jni/magiskboot/main.cpp
+++ b/native/jni/magiskboot/main.cpp
@@ -16,7 +16,7 @@ using namespace std;
 
 static void usage(char *arg0) {
 	fprintf(stderr,
-FULL_VER(MagiskBoot) R"EOF(  - Boot Image Modification Tool
+FULL_VER(MagiskBoot) R"EOF( - Boot Image Modification Tool
 
 Usage: %s <action> [args...]
 
@@ -84,6 +84,15 @@ Supported actions:
         If [OUT] is not specified, it will directly output to <input>
         Configure with env variables: KEEPVERITY TWOSTAGEINIT
 
+  split <input>
+    Split image.*-dtb into kernel + kernel_dtb
+
+  sha1 <file>
+    Print the SHA1 checksum for <file>
+
+  cleanup
+    Cleanup the current working directory
+
   compress[=method] <infile> [outfile]
     Compress <infile> with [method] (default: gzip), optionally to [outfile]
     <infile>/[outfile] can be '-' to be STDIN/STDOUT
@@ -102,16 +111,7 @@ Supported actions:
 	for (auto &it : name2fmt)
 		fprintf(stderr, "%s ", it.first.data());
 
-	fprintf(stderr, R"EOF(
-
-  sha1 <file>
-    Print the SHA1 checksum for <file>
-
-  cleanup
-    Cleanup the current working directory
-
-)EOF");
-
+	fprintf(stderr, "\n\n");
 	exit(1);
 }
 
@@ -147,6 +147,8 @@ int main(int argc, char *argv[]) {
 			printf("%02x", i);
 		printf("\n");
 		munmap(buf, size);
+	} else if (argc > 2 && action == "split") {
+		return split_image_dtb(argv[2]);
 	} else if (argc > 2 && action == "unpack") {
 		int idx = 2;
 		bool nodecomp = false;