From e139e8777b9839ac9659c37b79506c2501c095b3 Mon Sep 17 00:00:00 2001
From: topjohnwu <topjohnwu@gmail.com>
Date: Fri, 29 Dec 2017 04:25:03 +0800
Subject: [PATCH] Fix faulty magiskboot ramdisk patch code

---
 core/jni/magiskboot/main.c    |  8 ++---
 core/jni/magiskboot/ramdisk.c |  4 ++-
 core/jni/utils/pattern.c      | 63 +++++++++++++++++++----------------
 3 files changed, 41 insertions(+), 34 deletions(-)

diff --git a/core/jni/magiskboot/main.c b/core/jni/magiskboot/main.c
index 15bbdf0ba..8417173d9 100644
--- a/core/jni/magiskboot/main.c
+++ b/core/jni/magiskboot/main.c
@@ -73,12 +73,12 @@ static void usage(char *arg0) {
 		" --dtb-<cmd> <dtb>\n"
 		"  Do dtb related cmds to <dtb> (modifications are done directly)\n"
 		"  Supported commands:\n"
-		"    -dump\n"
+		"    dump\n"
 		"      Dump all contents from dtb for debugging\n"
-		"    -test\n"
+		"    test\n"
 		"      Check if fstab has verity/avb flags\n"
-		"      Return value: 0/no flags 1/flag exists"
-		"    -patch\n"
+		"      Return value: 0/no flags 1/flag exists\n"
+		"    patch\n"
 		"      Search for fstab and remove verity/avb\n"
 		"\n"
 		" --compress[=method] <infile> [outfile]\n"
diff --git a/core/jni/magiskboot/ramdisk.c b/core/jni/magiskboot/ramdisk.c
index 378ef3366..db4eddd2a 100644
--- a/core/jni/magiskboot/ramdisk.c
+++ b/core/jni/magiskboot/ramdisk.c
@@ -9,6 +9,8 @@
 #include "cpio.h"
 
 static void cpio_patch(struct vector *v, int keepverity, int keepforceencrypt) {
+	fprintf(stderr, "Patch with flag KEEPVERITY=[%s] KEEPFORCEENCRYPT=[%s]\n",
+			keepverity ? "true" : "false", keepforceencrypt ? "true" : "false");
 	cpio_entry *e;
 	vec_for_each(v, e) {
 		if (!e) continue;
@@ -134,7 +136,7 @@ static void cpio_backup(struct vector *v, struct vector *bak, const char *orig,
 			backup = 1;
 			fprintf(stderr, "Backup mismatch entry: ");
 		} else {
-			// Someting new in ramdisk, record in rem
+			// Something new in ramdisk, record in rem
 			++j;
 			rem->data = xrealloc(rem->data, rem->filesize + strlen(n->filename) + 1);
 			memcpy(rem->data + rem->filesize, n->filename, strlen(n->filename) + 1);
diff --git a/core/jni/utils/pattern.c b/core/jni/utils/pattern.c
index 36d62dc64..a3209146e 100644
--- a/core/jni/utils/pattern.c
+++ b/core/jni/utils/pattern.c
@@ -4,19 +4,19 @@
 #include "utils.h"
 
 static int check_verity_pattern(const char *s) {
-	int pos = 0;
-	if (s[0] == ',') ++pos;
-	if (strncmp(s + pos, "verify", 6) == 0)
-		pos += 6;
-	else if (strncmp(s + pos, "avb", 3) == 0)
-		pos += 3;
+	int skip = 0;
+	if (s[0] == ',') ++skip;
+	if (strncmp(s + skip, "verify", 6) == 0)
+		skip += 6;
+	else if (strncmp(s + skip, "avb", 3) == 0)
+		skip += 3;
 	else
 		return -1;
 
-	if (s[pos] == '=') {
-		while (s[pos] != '\0' && s[pos] != ' ' && s[pos] != '\n' && s[pos] != ',') ++pos;
+	if (s[skip] == '=') {
+		while (s[skip] != '\0' && s[skip] != ' ' && s[skip] != '\n' && s[skip] != ',') ++skip;
 	}
-	return pos;
+	return skip;
 }
 
 static int check_encryption_pattern(const char *s) {
@@ -59,34 +59,39 @@ void patch_init_rc(void **buf, size_t *size) {
 }
 
 int patch_verity(void **buf, uint32_t *size, int patch) {
-	int skip, found = 0;
-	for (int pos = 0; pos < *size; ++pos) {
-		if ((skip = check_verity_pattern(*buf + pos)) > 0) {
-			found = 1;
-			fprintf(stderr, "%s pattern [%.*s]\n", patch ? "Remove" : "Found", skip, (char *) *buf + pos);
-			if (patch) {
-				memcpy(*buf + pos, *buf + pos + skip, *size - pos - skip);
-				memset(*buf + *size - skip, '\0', skip);
-				*size -= skip;
-			} else {
-				pos += skip - 1;
-			}
+	int skip, src_size = *size;
+	char *src = *buf, *patched = patch ? xcalloc(src_size, 1) : NULL;
+	for (int read = 0, write = 0; read < src_size; ++read, ++write) {
+		if ((skip = check_verity_pattern(src + read)) > 0) {
+			if (!patch)
+				return 1;
+			fprintf(stderr, "Remove pattern [%.*s]\n", skip, src + read);
+			read += skip;
+			*size -= skip;
 		}
+		if (patch)
+			patched[write] = src[read];
 	}
-	return found;
+	free(*buf);
+	*buf = patched;
+	return 0;
 }
 
 void patch_encryption(void **buf, uint32_t *size) {
-	int skip;
-	for (int pos = 0; pos < *size; ++pos) {
-		if ((skip = check_encryption_pattern(*buf + pos)) > 0) {
-			fprintf(stderr, "Replace pattern [%.*s] with [encryptable]\n", skip, (char *) *buf + pos);
-			memcpy(*buf + pos, "encryptable", 11);
-			memcpy(*buf + pos + 11, *buf + pos + skip, *size - pos - skip);
-			memset(*buf + *size - skip + 11, '\0', skip - 11);
+	int skip, src_size = *size;
+	char *src = *buf, *patched = xcalloc(src_size, 1);
+	for (int read = 0, write = 0; read < src_size; ++read, ++write) {
+		if ((skip = check_encryption_pattern(src + read)) > 0) {
+			fprintf(stderr, "Replace pattern [%.*s] with [encryptable]\n", skip, src + read);
+			memcpy(patched + read, "encryptable", 11);
+			read += skip;
+			write += 11;
 			*size -= (skip - 11);
 		}
+		patched[write] = src[read];
 	}
+	free(*buf);
+	*buf = patched;
 }