From e9d0f615ba30f8d4a9281eafc7a10a736c635123 Mon Sep 17 00:00:00 2001
From: topjohnwu <topjohnwu@gmail.com>
Date: Sat, 11 Nov 2017 01:30:33 +0800
Subject: [PATCH] Add dtb test command

---
 jni/magiskboot/dtb.c        | 37 ++++++++++++++++++++++++++++---------
 jni/magiskboot/magiskboot.h |  5 ++---
 jni/magiskboot/main.c       | 18 +++++++++---------
 3 files changed, 39 insertions(+), 21 deletions(-)

diff --git a/jni/magiskboot/dtb.c b/jni/magiskboot/dtb.c
index a68313793..68bd8c14f 100644
--- a/jni/magiskboot/dtb.c
+++ b/jni/magiskboot/dtb.c
@@ -41,7 +41,7 @@ static int find_fstab(const void *fdt, int parent) {
 	return -1;
 }
 
-void dtb_dump(const char *file) {
+static void dtb_dump(const char *file) {
 	size_t size ;
 	void *dtb, *fdt;
 	fprintf(stderr, "Loading dtbs from [%s]\n", file);
@@ -60,13 +60,16 @@ void dtb_dump(const char *file) {
 	exit(0);
 }
 
-void dtb_patch(const char *file) {
+static void dtb_patch(const char *file, int patch) {
 	size_t size ;
 	void *dtb, *fdt;
 	fprintf(stderr, "Loading dtbs from [%s]\n\n", file);
-	mmap_rw(file, &dtb, &size);
+	if (patch)
+		mmap_rw(file, &dtb, &size);
+	else
+		mmap_ro(file, &dtb, &size);
 	// Loop through all the dtbs
-	int dtb_num = 0, patched = 0;
+	int dtb_num = 0, found = 0;
 	for (int i = 0; i < size; ++i) {
 		if (memcmp(dtb + i, DTB_MAGIC, 4) == 0) {
 			fdt = dtb + i;
@@ -80,10 +83,15 @@ void dtb_patch(const char *file) {
 					char *value = (char *) fdt_getprop(fdt, block, "fsmgr_flags", &value_size);
 					for (int i = 0; i < value_size; ++i) {
 						if ((skip = check_verity_pattern(value + i)) > 0) {
-							fprintf(stderr, "Remove pattern [%.*s] in [fsmgr_flags]\n", skip, value + i);
-							memcpy(value + i, value + i + skip, value_size - i - skip);
-							memset(value + value_size - skip, '\0', skip);
-							patched = 1;
+							if (patch) {
+								fprintf(stderr, "Remove pattern [%.*s] in [fsmgr_flags]\n", skip, value + i);
+								memcpy(value + i, value + i + skip, value_size - i - skip);
+								memset(value + value_size - skip, '\0', skip);
+							} else {
+								fprintf(stderr, "Found pattern [%.*s] in [fsmgr_flags]\n", skip, value + i);
+								i += skip - 1;
+							}
+							found = 1;
 						}
 					}
 				}
@@ -92,7 +100,18 @@ void dtb_patch(const char *file) {
 	}
 	fprintf(stderr, "\n");
 	munmap(dtb, size);
-	exit(!patched);
+	exit(!found);
+}
+
+int dtb_commands(const char *cmd, int argc, char *argv[]) {
+	if (argc == 0) return 1;
+	if (strcmp(cmd, "dump") == 0)
+		dtb_dump(argv[0]);
+	else if (strcmp(cmd, "patch") == 0)
+		dtb_patch(argv[0], 1);
+	else if (strcmp(cmd, "test") == 0)
+		dtb_patch(argv[0], 0);
+	return 0;
 }
 
 
diff --git a/jni/magiskboot/magiskboot.h b/jni/magiskboot/magiskboot.h
index d4e0259d4..0cc6b116a 100644
--- a/jni/magiskboot/magiskboot.h
+++ b/jni/magiskboot/magiskboot.h
@@ -18,11 +18,10 @@ void unpack(const char *image);
 void repack(const char* orig_image, const char* out_image);
 void hexpatch(const char *image, const char *from, const char *to);
 int parse_img(void *orig, size_t size, boot_img *boot);
-int cpio_commands(const char *command, int argc, char *argv[]);
+int cpio_commands(const char *cmd, int argc, char *argv[]);
 void comp_file(const char *method, const char *from, const char *to);
 void decomp_file(char *from, const char *to);
-void dtb_dump(const char *file);
-void dtb_patch(const char *file);
+int dtb_commands(const char *cmd, int argc, char *argv[]);
 
 // Compressions
 size_t gzip(int mode, int fd, const void *buf, size_t size);
diff --git a/jni/magiskboot/main.c b/jni/magiskboot/main.c
index b1f2dff62..f42511648 100644
--- a/jni/magiskboot/main.c
+++ b/jni/magiskboot/main.c
@@ -124,10 +124,6 @@ int main(int argc, char *argv[]) {
 		repack(argv[2], argc > 3 ? argv[3] : NEW_BOOT);
 	} else if (argc > 2 && strcmp(argv[1], "--decompress") == 0) {
 		decomp_file(argv[2], argc > 3 ? argv[3] : NULL);
-	} else if (argc > 2 && strcmp(argv[1], "--dtb-dump") == 0) {
-		dtb_dump(argv[2]);
-	} else if (argc > 2 && strcmp(argv[1], "--dtb-patch") == 0) {
-		dtb_patch(argv[2]);
 	} else if (argc > 2 && strncmp(argv[1], "--compress", 10) == 0) {
 		char *method;
 		method = strchr(argv[1], '=');
@@ -137,11 +133,15 @@ int main(int argc, char *argv[]) {
 	} else if (argc > 4 && strcmp(argv[1], "--hexpatch") == 0) {
 		hexpatch(argv[2], argv[3], argv[4]);
 	} else if (argc > 2 && strncmp(argv[1], "--cpio", 6) == 0) {
-		char *command;
-		command = strchr(argv[1] + 2, '-');
-		if (command == NULL) usage(argv[0]);
-		else ++command;
-		if (cpio_commands(command, argc - 2, argv + 2)) usage(argv[0]);
+		char *cmd = argv[1] + 6;
+		if (*cmd == '\0') usage(argv[0]);
+		else ++cmd;
+		if (cpio_commands(cmd, argc - 2, argv + 2)) usage(argv[0]);
+	} else if (argc > 2 && strncmp(argv[1], "--dtb", 5) == 0) {
+		char *cmd = argv[1] + 5;
+		if (*cmd == '\0') usage(argv[0]);
+		else ++cmd;
+		if (dtb_commands(cmd, argc - 2, argv + 2)) usage(argv[0]);
 	} else {
 		usage(argv[0]);
 	}