From 32ee8e462c287ad53c69f0d414c4097d2c782eaa Mon Sep 17 00:00:00 2001
From: topjohnwu <topjohnwu@gmail.com>
Date: Wed, 15 Nov 2017 04:41:30 +0800
Subject: [PATCH] Properly and fully support dtbo patching

---
 java                      |  2 +-
 jni/core/bootstages.c     |  5 +++--
 scripts/flash_script.sh   |  2 +-
 scripts/util_functions.sh | 27 +++++++++++++++------------
 4 files changed, 20 insertions(+), 16 deletions(-)

diff --git a/java b/java
index 99c74b31b..389299afd 160000
--- a/java
+++ b/java
@@ -1 +1 @@
-Subproject commit 99c74b31be15a239fdffe4a727d2d83feb52b8bb
+Subproject commit 389299afd103337ec3ab787278ac56ac249f9195
diff --git a/jni/core/bootstages.c b/jni/core/bootstages.c
index b11ea74ed..a7b1c2d8d 100644
--- a/jni/core/bootstages.c
+++ b/jni/core/bootstages.c
@@ -621,10 +621,11 @@ void post_fs_data(int client) {
 		rm_rf(DATABIN);
 		cp_afc(bin_path, DATABIN);
 		rm_rf(bin_path);
-		// Lazy.... use shell blob to match files
-		exec_command_sync("sh", "-c", "mv /data/magisk/stock_boot* /data", NULL);
 	}
 
+	// Lazy.... use shell blob to match files
+	exec_command_sync("sh", "-c", "mv -f /data/magisk/stock_*.img.gz /data", NULL);
+
 	// Initialize
 	daemon_init();
 
diff --git a/scripts/flash_script.sh b/scripts/flash_script.sh
index 7e83271f9..6e87299f8 100644
--- a/scripts/flash_script.sh
+++ b/scripts/flash_script.sh
@@ -102,7 +102,7 @@ $BOOTMODE || recovery_actions
 # Boot patching
 ##########################################################################################
 
-find_boot_image
+[ -z $BOOTIMAGE ] && find_boot_image
 [ -z $BOOTIMAGE ] && abort "! Unable to detect boot image"
 ui_print "- Found boot image: $BOOTIMAGE"
 
diff --git a/scripts/util_functions.sh b/scripts/util_functions.sh
index dc2a54286..3d6f5fb49 100644
--- a/scripts/util_functions.sh
+++ b/scripts/util_functions.sh
@@ -100,15 +100,14 @@ resolve_link() {
 }
 
 find_boot_image() {
-  if [ -z "$BOOTIMAGE" ]; then
-    if [ ! -z $SLOT ]; then
-      BOOTIMAGE=`find /dev/block -iname boot$SLOT | head -n 1` 2>/dev/null
-    else
-      for BLOCK in boot_a kern-a android_boot kernel boot lnx bootimg; do
-        BOOTIMAGE=`find /dev/block -iname $BLOCK | head -n 1` 2>/dev/null
-        [ ! -z $BOOTIMAGE ] && break
-      done
-    fi
+  BOOTIMAGE=
+  if [ ! -z $SLOT ]; then
+    BOOTIMAGE=`find /dev/block -iname boot$SLOT | head -n 1` 2>/dev/null
+  else
+    for BLOCK in boot_a kern-a android_boot kernel boot lnx bootimg; do
+      BOOTIMAGE=`find /dev/block -iname $BLOCK | head -n 1` 2>/dev/null
+      [ ! -z $BOOTIMAGE ] && break
+    done
   fi
   # Recovery fallback
   if [ -z "$BOOTIMAGE" ]; then
@@ -123,8 +122,7 @@ find_boot_image() {
 migrate_boot_backup() {
   # Update the broken boot backup
   if [ -f /data/stock_boot_.img.gz ]; then
-    $MAGISKBIN/magiskboot --decompress /data/stock_boot_.img.gz
-    mv /data/stock_boot_.img /data/stock_boot.img
+    $MAGISKBIN/magiskboot --decompress /data/stock_boot_.img.gz /data/stock_boot.img
   fi
   # Update our previous backup to new format if exists
   if [ -f /data/stock_boot.img ]; then
@@ -171,13 +169,18 @@ patch_dtbo_image() {
       $MAGISKBIN/magiskboot --compress $DTBOIMAGE $MAGISKBIN/stock_dtbo.img.gz
       ui_print "- Patching fstab in dtbo to remove avb-verity"
       $MAGISKBIN/magiskboot --dtb-patch $DTBOIMAGE
+      return 0
     fi
   fi
+  return 1
 }
 
 restore_imgs() {
   STOCKBOOT=/data/stock_boot_${1}.img.gz
   STOCKDTBO=/data/stock_dtbo.img.gz
+
+  # Make sure all blocks are writable
+  $MAGISKBIN/magisk --unlock-blocks 2>/dev/null
   find_dtbo_image
   if [ ! -z "$DTBOIMAGE" -a -f "$STOCKDTBO" ]; then
     ui_print "- Restoring stock dtbo image"
@@ -187,7 +190,7 @@ restore_imgs() {
   find_boot_image
   if [ ! -z "$BOOTIMAGE" -a -f "$STOCKBOOT" ]; then
     ui_print "- Restoring stock boot image"
-    flash_boot_image $STOCKBOOT "$BOOTIMAGE"
+    gzip -d < $STOCKBOOT | cat - /dev/zero 2>/dev/null | dd of="$BOOTIMAGE" bs=4096 2>/dev/null
     return 0
   fi
   return 1