From a46c6252c619f2e321630d889ef38b297074e240 Mon Sep 17 00:00:00 2001
From: topjohnwu <topjohnwu@gmail.com>
Date: Sat, 11 Aug 2018 15:56:12 +0800
Subject: [PATCH] Detect insufficient partition size

Close #388
---
 .../com/topjohnwu/magisk/asyncs/InstallMagisk.java | 14 ++++++++------
 app/src/full/res/raw/utils.sh                      |  5 +++++
 scripts/addon.d.sh                                 |  2 +-
 scripts/flash_script.sh                            |  2 +-
 scripts/magisk_uninstaller.sh                      |  2 +-
 scripts/util_functions.sh                          |  4 ++++
 6 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/app/src/full/java/com/topjohnwu/magisk/asyncs/InstallMagisk.java b/app/src/full/java/com/topjohnwu/magisk/asyncs/InstallMagisk.java
index ceccf23d3..afbfb8eb4 100644
--- a/app/src/full/java/com/topjohnwu/magisk/asyncs/InstallMagisk.java
+++ b/app/src/full/java/com/topjohnwu/magisk/asyncs/InstallMagisk.java
@@ -248,7 +248,7 @@ public class InstallMagisk extends ParallelTask<Void, Void, Boolean> {
         return patched;
     }
 
-    private void outputBoot(File patched) throws IOException {
+    private boolean outputBoot(File patched) throws IOException {
         switch (mode) {
             case PATCH_MODE:
                 String fmt = mm.prefs.getString(Const.Key.BOOT_FORMAT, ".img");
@@ -278,13 +278,14 @@ public class InstallMagisk extends ParallelTask<Void, Void, Boolean> {
                 break;
             case SECOND_SLOT_MODE:
             case DIRECT_MODE:
-                Shell.Job job = Shell.su(Utils.fmt("direct_install %s %s %s", patched, mBoot, installDir))
-                        .to(console, logs);
+                if (!Shell.su(Utils.fmt("direct_install %s %s %s", patched, mBoot, installDir))
+                        .to(console, logs).exec().isSuccess())
+                    return false;
                 if (!Data.keepVerity)
-                    job.add("find_dtbo_image", "patch_dtbo_image");
-                job.exec();
+                    Shell.su("find_dtbo_image", "patch_dtbo_image").to(console, logs).exec();
                 break;
         }
+        return true;
     }
 
     private void postOTA() {
@@ -374,7 +375,8 @@ public class InstallMagisk extends ParallelTask<Void, Void, Boolean> {
                 File patched = patchBoot();
                 if (patched == null)
                     return false;
-                outputBoot(patched);
+                if (!outputBoot(patched))
+                    return false;
                 if (mode == SECOND_SLOT_MODE)
                     postOTA();
                 console.add("- All done!");
diff --git a/app/src/full/res/raw/utils.sh b/app/src/full/res/raw/utils.sh
index b20cb62a5..89a89c272 100644
--- a/app/src/full/res/raw/utils.sh
+++ b/app/src/full/res/raw/utils.sh
@@ -60,8 +60,13 @@ direct_install() {
   cp -rf $3/* /data/adb/magisk
   echo "- Flashing new boot image"
   flash_image $1 $2
+  if [ $? -ne 0 ]; then
+    echo "! Insufficient partition size"
+    return 1
+  fi
   rm -f $1
   rm -rf $3
+  return 0
 }
 
 mm_patch_dtbo() {
diff --git a/scripts/addon.d.sh b/scripts/addon.d.sh
index 4e368d09f..844e2658b 100644
--- a/scripts/addon.d.sh
+++ b/scripts/addon.d.sh
@@ -65,7 +65,7 @@ installation() {
   . ./boot_patch.sh "$BOOTIMAGE"
 
   ui_print "- Flashing new boot image"
-  flash_image new-boot.img "$BOOTIMAGE"
+  flash_image new-boot.img "$BOOTIMAGE" || abort "! Insufficient partition size"
   rm -f new-boot.img
 
   if [ -f stock_boot* ]; then
diff --git a/scripts/flash_script.sh b/scripts/flash_script.sh
index 5f90834a8..7168c9a8c 100644
--- a/scripts/flash_script.sh
+++ b/scripts/flash_script.sh
@@ -115,7 +115,7 @@ cd $MAGISKBIN
 . ./boot_patch.sh "$BOOTIMAGE"
 
 ui_print "- Flashing new boot image"
-flash_image new-boot.img "$BOOTIMAGE"
+flash_image new-boot.img "$BOOTIMAGE" || abort "! Insufficient partition size"
 rm -f new-boot.img
 
 if [ -f stock_boot* ]; then
diff --git a/scripts/magisk_uninstaller.sh b/scripts/magisk_uninstaller.sh
index 84a00c496..494341be1 100644
--- a/scripts/magisk_uninstaller.sh
+++ b/scripts/magisk_uninstaller.sh
@@ -114,7 +114,7 @@ case $? in
       # Sign chromeos boot
       $CHROMEOS && sign_chromeos
       ui_print "- Flashing restored boot image"
-      flash_image new-boot.img $BOOTIMAGE
+      flash_image new-boot.img $BOOTIMAGE || abort "! Insufficient partition size"
     fi
     ;;
   2 ) # Other patched
diff --git a/scripts/util_functions.sh b/scripts/util_functions.sh
index 9282cc196..6c3127868 100644
--- a/scripts/util_functions.sh
+++ b/scripts/util_functions.sh
@@ -203,11 +203,15 @@ flash_image() {
     COM2="cat -"
   fi
   if [ -b "$2" ]; then
+    local s_size=`stat -c '%s' "$1"`
+    local t_size=`blockdev --getsize64 "$2"`
+    [ $s_size -gt $t_size ] && return 1
     eval $COM1 | eval $COM2 | cat - /dev/zero > "$2" 2>/dev/null
   else
     ui_print "- Not block device, storing image"
     eval $COM1 | eval $COM2 > "$2" 2>/dev/null
   fi
+  return 0
 }
 
 find_dtbo_image() {