Update dtb patch to not use in-place modification

native/jni/magiskboot/dtb.cpp

@@ -6,6 +6,7 @@ extern "C" {
 }
 #include <utils.h>
 #include <bitset>
+#include <vector>
 
 #include "magiskboot.h"
 #include "format.h"
@@ -19,29 +20,17 @@ static void pretty_node(int depth) {
 	if (depth == 0)
 		return;
 
-	for (int i = 0; i < depth - 1; ++i) {
-		if (depth_set[i])
-			printf("│   ");
-		else
-			printf("    ");
-	}
-	if (depth_set[depth - 1])
-		printf("├── ");
-	else
-		printf("└── ");
+	for (int i = 0; i < depth - 1; ++i)
+		printf(depth_set[i] ? "│   " : "    ");
+
+	printf(depth_set[depth - 1] ? "├── " : "└── ");
 }
 
 static void pretty_prop(int depth) {
-	for (int i = 0; i < depth; ++i) {
-		if (depth_set[i])
-			printf("│   ");
-		else
-			printf("    ");
-	}
-	if (depth_set[depth])
-		printf("│  ");
-	else
-		printf("   ");
+	for (int i = 0; i < depth; ++i)
+		printf(depth_set[i] ? "│   " : "    ");
+
+	printf(depth_set[depth] ? "│  " : "   ");
 }
 
 static void print_node(const void *fdt, int node = 0, int depth = 0) {
@@ -94,13 +83,13 @@ static void print_node(const void *fdt, int node = 0, int depth = 0) {
 	}
 }
 
-static int find_fstab(const void *fdt, int parent = 0) {
-	int node, fstab;
-	fdt_for_each_subnode(node, fdt, parent) {
-		if (strcmp(fdt_get_name(fdt, node, nullptr), "fstab") == 0)
-			return node;
-		fstab = find_fstab(fdt, node);
-		if (fstab != -1)
+static int find_fstab(const void *fdt, int node = 0) {
+	if (fdt_get_name(fdt, node, nullptr) == "fstab"sv)
+		return node;
+	int child;
+	fdt_for_each_subnode(child, fdt, node) {
+		int fstab = find_fstab(fdt, child);
+		if (fstab >= 0)
 			return fstab;
 	}
 	return -1;
@@ -126,7 +115,7 @@ static void dtb_print(const char *file, bool fstab) {
 				fprintf(stderr, "Printing dtb.%04d\n", dtb_num);
 				print_node(fdt);
 			}
-			dtb_num++;
+			++dtb_num;
 		}
 	}
 	fprintf(stderr, "\n");
@@ -134,44 +123,50 @@ static void dtb_print(const char *file, bool fstab) {
 	exit(0);
 }
 
-static void dtb_patch(const char *file, int patch) {
-	size_t size ;
-	uint8_t *dtb, *fdt;
-	fprintf(stderr, "Loading dtbs from [%s]\n", file);
-	if (patch)
-		mmap_rw(file, dtb, size);
-	else
-		mmap_ro(file, dtb, size);
-	// Loop through all the dtbs
-	int dtb_num = 0;
-	bool found = false;
-	for (int i = 0; i < size; ++i) {
+static void dtb_patch(const char *in, const char *out) {
+	vector<uint8_t *> fdt_list;
+	size_t dtb_sz ;
+	uint8_t *dtb;
+	fprintf(stderr, "Loading dtbs from [%s]\n", in);
+	mmap_ro(in, dtb, dtb_sz);
+	bool modified = false;
+	for (int i = 0; i < dtb_sz; ++i) {
 		if (memcmp(dtb + i, DTB_MAGIC, 4) == 0) {
-			fdt = dtb + i;
-			int fstab = find_fstab(fdt, 0);
-			if (fstab > 0) {
-				fprintf(stderr, "Found fstab in dtb.%04d\n", dtb_num++);
-				int block;
-				fdt_for_each_subnode(block, fdt, fstab) {
-					fprintf(stderr, "Found block [%s] in fstab\n", fdt_get_name(fdt, block, nullptr));
-					uint32_t value_size;
-					void *value = (void *) fdt_getprop(fdt, block, "fsmgr_flags", (int *)&value_size);
-					if (patch) {
-						void *dup = xmalloc(value_size);
-						memcpy(dup, value, value_size);
-						memset(value, 0, value_size);
-						found |= patch_verity(&dup, &value_size);
-						memcpy(value, dup, value_size);
-						free(dup);
-					} else {
-						found |= patch_verity(&value, &value_size, false);
-					}
+			// Patched will only be smaller
+			int len = fdt_totalsize(dtb + i);
+			auto fdt = static_cast<uint8_t *>(xmalloc(len));
+			memcpy(fdt, dtb + i, len);
+			int fstab = find_fstab(fdt);
+			if (fstab < 0)
+				continue;
+			fprintf(stderr, "Found fstab in dtb.%04d\n", fdt_list.size());
+			int block;
+			fdt_for_each_subnode(block, fdt, fstab) {
+				fprintf(stderr, "Found entry [%s] in fstab\n", fdt_get_name(fdt, block, nullptr));
+				uint32_t size;
+				auto value = static_cast<const char *>(
+						fdt_getprop(fdt, block, "fsmgr_flags", reinterpret_cast<int *>(&size)));
+				char *pval = patch_verity(value, size);
+				if (pval) {
+					modified = true;
+					fdt_setprop_string(fdt, block, "fsmgr_flags", pval);
 				}
 			}
+			fdt_list.push_back(fdt);
 		}
 	}
-	munmap(dtb, size);
-	exit(!found);
+	munmap(dtb, dtb_sz);
+	if (modified) {
+		if (!out)
+			out = in;
+		int fd = xopen(out, O_WRONLY | O_CREAT | O_CLOEXEC);
+		for (auto fdt : fdt_list) {
+			fdt_pack(fdt);
+			xwrite(fd, fdt, fdt_totalsize(fdt));
+		}
+		close(fd);
+	}
+	exit(!modified);
 }
 
 int dtb_commands(int argc, char *argv[]) {
@@ -182,9 +177,7 @@ int dtb_commands(int argc, char *argv[]) {
 	if (argv[0] == "print"sv) {
 		dtb_print(dtb, argc > 1 && argv[1] == "-f"sv);
 	} else if (argv[0] == "patch"sv) {
-		dtb_patch(dtb, 1);
-	} else if (argv[0] == "test"sv) {
-		dtb_patch(dtb, 0);
+		dtb_patch(dtb, argv[1]);
 	} else {
 		return 1;
 	}

native/jni/magiskboot/magiskboot.h

@@ -20,5 +20,5 @@ int cpio_commands(int argc, char *argv[]);
 int dtb_commands(int argc, char *argv[]);
 
 // Pattern
-bool patch_verity(void **buf, uint32_t *size, bool patch = true);
+char *patch_verity(const void *buf, uint32_t &size);
 void patch_encryption(void **buf, uint32_t *size);

native/jni/magiskboot/main.cpp

@@ -70,18 +70,15 @@ Supported actions:
       sha1
         Print stock boot SHA1 if previously backed up in ramdisk
 
-  dtb <dtb> <command> [args...]
-    Do commands to <dtb> (modifications are done directly)
-    Supported commands:
+  dtb <input> <action> [args...]
+    Do dtb related actions to <input>
+    Supported actions:
       print [-f]
-        Print all contents from dtb for debugging
+        Print all contents of dtb for debugging
         Specify [-f] to only print fstab nodes
-      test
-        Check if fstab has verity/avb flags
-        Return values:
-        0:flag exists    1:no flags
-      patch
+      patch [OUT]
         Search for fstab and remove verity/avb
+        If [OUT] is not specified, it will directly output to <input>
 
   compress[=method] <infile> [outfile]
     Compress <infile> with [method] (default: gzip), optionally to [outfile]

native/jni/magiskboot/pattern.cpp

@@ -22,7 +22,7 @@ static int check_verity_pattern(const char *s) {
 }
 
 static int check_encryption_pattern(const char *s) {
-	const char *encrypt_list[] = { "forceencrypt", "forcefdeorfbe", NULL };
+	const char *encrypt_list[] = { "forceencrypt", "forcefdeorfbe", nullptr };
 	for (int i = 0 ; encrypt_list[i]; ++i) {
 		int len = strlen(encrypt_list[i]);
 		if (strncmp(s, encrypt_list[i], len) == 0)
@@ -31,29 +31,23 @@ static int check_encryption_pattern(const char *s) {
 	return -1;
 }
 
-bool patch_verity(void **buf, uint32_t *size, bool patch) {
-	int skip, src_size = *size;
+char *patch_verity(const void *buf, uint32_t &size) {
+	auto src = static_cast<const char *>(buf);
+	int src_size = size;
 	bool found = false;
-	char *src = (char *) *buf, *patched = patch ? (char *) xcalloc(src_size, 1) : nullptr;
-	for (int read = 0, write = 0; read < src_size; ++read, ++write) {
-		if ((skip = check_verity_pattern(src + read)) > 0) {
-			if (patch) {
-				fprintf(stderr, "Remove pattern [%.*s]\n", skip, src + read);
-				*size -= skip;
-			} else {
-				fprintf(stderr, "Found pattern [%.*s]\n", skip, src + read);
-			}
+	char patched[4096];
+	int write = 0;
+	for (int read = 0; read < src_size; ++read, ++write) {
+		if (int skip; (skip = check_verity_pattern(src + read)) > 0) {
+			fprintf(stderr, "Found pattern [%.*s]\n", skip, src + read);
+			size -= skip;
 			read += skip;
 			found = true;
 		}
-		if (patch)
-			patched[write] = src[read];
+		patched[write] = src[read];
 	}
-	if (patch) {
-		free(*buf);
-		*buf = patched;
-	}
-	return found;
+	patched[write] = '\0';
+	return found ? strdup(patched) : nullptr;
 }
 
 void patch_encryption(void **buf, uint32_t *size) {

native/jni/magiskboot/ramdisk.cpp

@@ -46,7 +46,9 @@ void magisk_cpio::patch(bool keepverity, bool keepforceencrypt) {
 					 str_contains(cur->first, "fstab") && S_ISREG(cur->second->mode);
 		if (!keepverity) {
 			if (fstab) {
-				patch_verity(&cur->second->data, &cur->second->filesize);
+				auto buf = patch_verity(cur->second->data, cur->second->filesize);
+				free(cur->second->data);
+				cur->second->data = buf;
 			} else if (cur->first == "verity_key") {
 				rm(cur);
 				continue;

scripts/boot_patch.sh

@@ -158,7 +158,7 @@ rm -f ramdisk.cpio.orig config
 
 if ! $KEEPVERITY; then
   for dt in dtb kernel_dtb extra recovery_dtbo; do
-    [ -f $dt ] && ./magiskboot dtb-patch $dt && ui_print "- Removing dm(avb)-verity in $dt"
+    [ -f $dt ] && ./magiskboot dtb $dt patch && ui_print "- Removing dm(avb)-verity in $dt"
   done
 fi
 

scripts/util_functions.sh

@@ -280,11 +280,13 @@ patch_dtbo_image() {
   find_dtbo_image
   if [ ! -z $DTBOIMAGE ]; then
     ui_print "- DTBO image: $DTBOIMAGE"
-    if $MAGISKBIN/magiskboot --dtb-test $DTBOIMAGE; then
+    local PATCHED=$TMPDIR/dtbo
+    if $MAGISKBIN/magiskboot dtb $DTBOIMAGE patch $PATCHED; then
       ui_print "- Backing up stock DTBO image"
-      $MAGISKBIN/magiskboot --compress $DTBOIMAGE $MAGISKBIN/stock_dtbo.img.gz
+      $MAGISKBIN/magiskboot compress $DTBOIMAGE $MAGISKBIN/stock_dtbo.img.gz
       ui_print "- Patching DTBO to remove avb-verity"
-      $MAGISKBIN/magiskboot --dtb-patch $DTBOIMAGE
+      cat $PATCHED /dev/zero > $DTBOIMAGE
+      rm -f $PATCHED
       return 0
     fi
   fi