Find boot image with MagiskInstaller

Fix #8211

app/apk/src/main/java/com/topjohnwu/magisk/dialog/EnvFixDialog.kt

@@ -1,13 +1,18 @@
 package com.topjohnwu.magisk.dialog
 
+import android.widget.Toast
+import androidx.core.os.postDelayed
 import androidx.lifecycle.lifecycleScope
 import com.topjohnwu.magisk.core.BuildConfig
 import com.topjohnwu.magisk.core.Info
 import com.topjohnwu.magisk.core.R
+import com.topjohnwu.magisk.core.ktx.reboot
+import com.topjohnwu.magisk.core.ktx.toast
 import com.topjohnwu.magisk.core.tasks.MagiskInstaller
 import com.topjohnwu.magisk.events.DialogBuilder
 import com.topjohnwu.magisk.ui.home.HomeViewModel
 import com.topjohnwu.magisk.view.MagiskDialog
+import com.topjohnwu.superuser.internal.UiThreadHandler
 import kotlinx.coroutines.launch
 
 class EnvFixDialog(private val vm: HomeViewModel, private val code: Int) : DialogBuilder {
@@ -27,9 +32,15 @@ class EnvFixDialog(private val vm: HomeViewModel, private val code: Int) : Dialo
                         setCancelable(false)
                     }
                     dialog.activity.lifecycleScope.launch {
-                        MagiskInstaller.FixEnv {
+                        MagiskInstaller.FixEnv().exec { success ->
                             dialog.dismiss()
-                        }.exec()
+                            context.toast(
+                                if (success) R.string.reboot_delay_toast else R.string.setup_fail,
+                                Toast.LENGTH_LONG
+                            )
+                            if (success)
+                                UiThreadHandler.handler.postDelayed(5000) { reboot() }
+                        }
                     }
                 }
             }

app/apk/src/main/java/com/topjohnwu/magisk/dialog/UninstallDialog.kt

@@ -1,15 +1,17 @@
 package com.topjohnwu.magisk.dialog
 
 import android.app.ProgressDialog
-import android.content.Context
 import android.widget.Toast
+import androidx.lifecycle.lifecycleScope
 import com.topjohnwu.magisk.arch.NavigationActivity
+import com.topjohnwu.magisk.arch.UIActivity
 import com.topjohnwu.magisk.core.R
 import com.topjohnwu.magisk.core.ktx.toast
+import com.topjohnwu.magisk.core.tasks.MagiskInstaller
 import com.topjohnwu.magisk.events.DialogBuilder
 import com.topjohnwu.magisk.ui.flash.FlashFragment
 import com.topjohnwu.magisk.view.MagiskDialog
-import com.topjohnwu.superuser.Shell
+import kotlinx.coroutines.launch
 
 class UninstallDialog : DialogBuilder {
 
@@ -19,7 +21,7 @@ class UninstallDialog : DialogBuilder {
             setMessage(R.string.uninstall_magisk_msg)
             setButton(MagiskDialog.ButtonType.POSITIVE) {
                 text = R.string.restore_img
-                onClick { restore(dialog.context) }
+                onClick { restore(dialog.activity) }
             }
             setButton(MagiskDialog.ButtonType.NEGATIVE) {
                 text = R.string.complete_uninstall
@@ -29,18 +31,20 @@ class UninstallDialog : DialogBuilder {
     }
 
     @Suppress("DEPRECATION")
-    private fun restore(context: Context) {
-        val dialog = ProgressDialog(context).apply {
-            setMessage(context.getString(R.string.restore_img_msg))
+    private fun restore(activity: UIActivity<*>) {
+        val dialog = ProgressDialog(activity).apply {
+            setMessage(activity.getString(R.string.restore_img_msg))
             show()
         }
 
-        Shell.cmd("restore_imgs").submit { result ->
-            dialog.dismiss()
-            if (result.isSuccess) {
-                context.toast(R.string.restore_done, Toast.LENGTH_SHORT)
-            } else {
-                context.toast(R.string.restore_fail, Toast.LENGTH_LONG)
+        activity.lifecycleScope.launch {
+            MagiskInstaller.Restore().exec { success ->
+                dialog.dismiss()
+                if (success) {
+                    activity.toast(R.string.restore_done, Toast.LENGTH_SHORT)
+                } else {
+                    activity.toast(R.string.restore_fail, Toast.LENGTH_LONG)
+                }
             }
         }
     }

app/core/src/main/java/com/topjohnwu/magisk/core/tasks/MagiskInstaller.kt

@@ -6,7 +6,6 @@ import android.system.ErrnoException
 import android.system.Os
 import android.system.OsConstants
 import android.system.OsConstants.O_WRONLY
-import android.widget.Toast
 import androidx.annotation.WorkerThread
 import androidx.core.os.postDelayed
 import com.topjohnwu.magisk.StubApk
@@ -15,13 +14,10 @@ import com.topjohnwu.magisk.core.BuildConfig
 import com.topjohnwu.magisk.core.Config
 import com.topjohnwu.magisk.core.Const
 import com.topjohnwu.magisk.core.Info
-import com.topjohnwu.magisk.core.R
 import com.topjohnwu.magisk.core.di.ServiceLocator
 import com.topjohnwu.magisk.core.isRunningAsStub
 import com.topjohnwu.magisk.core.ktx.copyAll
 import com.topjohnwu.magisk.core.ktx.copyAndClose
-import com.topjohnwu.magisk.core.ktx.reboot
-import com.topjohnwu.magisk.core.ktx.toast
 import com.topjohnwu.magisk.core.ktx.writeTo
 import com.topjohnwu.magisk.core.utils.DummyList
 import com.topjohnwu.magisk.core.utils.MediaStoreUtils
@@ -585,6 +581,8 @@ abstract class MagiskInstallImpl protected constructor(
 
     protected suspend fun fixEnv() = extractFiles() && "fix_env $installDir".sh().isSuccess
 
+    protected fun restore() = findImage() && "restore_imgs $srcBoot".sh().isSuccess
+
     protected fun uninstall() = "run_uninstaller $AppApkPath".sh().isSuccess
 
     @WorkerThread
@@ -608,11 +606,10 @@ abstract class MagiskInstallImpl protected constructor(
     }
 }
 
-abstract class MagiskInstaller(
+abstract class ConsoleInstaller(
     console: MutableList<String>,
     logs: MutableList<String>
 ) : MagiskInstallImpl(console, logs) {
-
     override suspend fun exec(): Boolean {
         val success = super.exec()
         if (success) {
@@ -622,40 +619,51 @@ abstract class MagiskInstaller(
         }
         return success
     }
+}
+
+abstract class CallBackInstaller : MagiskInstallImpl(DummyList, DummyList) {
+    suspend fun exec(callback: (Boolean) -> Unit): Boolean {
+        val success = exec()
+        callback(success)
+        return success
+    }
+}
+
+class MagiskInstaller {
 
     class Patch(
         private val uri: Uri,
         console: MutableList<String>,
         logs: MutableList<String>
-    ) : MagiskInstaller(console, logs) {
+    ) : ConsoleInstaller(console, logs) {
         override suspend fun operations() = patchFile(uri)
     }
 
     class SecondSlot(
         console: MutableList<String>,
         logs: MutableList<String>
-    ) : MagiskInstaller(console, logs) {
+    ) : ConsoleInstaller(console, logs) {
         override suspend fun operations() = secondSlot()
     }
 
     class Direct(
         console: MutableList<String>,
         logs: MutableList<String>
-    ) : MagiskInstaller(console, logs) {
+    ) : ConsoleInstaller(console, logs) {
         override suspend fun operations() = direct()
     }
 
     class Emulator(
         console: MutableList<String>,
         logs: MutableList<String>
-    ) : MagiskInstaller(console, logs) {
+    ) : ConsoleInstaller(console, logs) {
         override suspend fun operations() = fixEnv()
     }
 
     class Uninstall(
         console: MutableList<String>,
         logs: MutableList<String>
-    ) : MagiskInstallImpl(console, logs) {
+    ) : ConsoleInstaller(console, logs) {
         override suspend fun operations() = uninstall()
 
         override suspend fun exec(): Boolean {
@@ -669,19 +677,11 @@ abstract class MagiskInstaller(
         }
     }
 
-    class FixEnv(private val callback: () -> Unit) : MagiskInstallImpl(DummyList, DummyList) {
-        override suspend fun operations() = fixEnv()
+    class Restore : CallBackInstaller() {
+        override suspend fun operations() = restore()
+    }
 
-        override suspend fun exec(): Boolean {
-            val success = super.exec()
-            callback()
-            context.toast(
-                if (success) R.string.reboot_delay_toast else R.string.setup_fail,
-                Toast.LENGTH_LONG
-            )
-            if (success)
-                UiThreadHandler.handler.postDelayed(5000) { reboot() }
-            return success
-        }
+    class FixEnv : CallBackInstaller() {
+        override suspend fun operations() = fixEnv()
     }
 }

scripts/app_functions.sh

@@ -2,10 +2,14 @@
 # Magisk app internal scripts
 ##################################
 
+# $1 = delay
+# $2 = command
 run_delay() {
   (sleep $1; $2)&
 }
 
+# $1 = version string
+# $2 = version code
 env_check() {
   for file in busybox magiskboot magiskinit util_functions.sh boot_patch.sh; do
     [ -f "$MAGISKBIN/$file" ] || return 1
@@ -21,6 +25,8 @@ env_check() {
   return 0
 }
 
+# $1 = dir to copy
+# $2 = destination (optional)
 cp_readlink() {
   if [ -z $2 ]; then
     cd $1
@@ -39,6 +45,7 @@ cp_readlink() {
   cd /
 }
 
+# $1 = install dir
 fix_env() {
   # Cleanup and make dirs
   rm -rf $MAGISKBIN/*
@@ -49,6 +56,8 @@ fix_env() {
   chown -R 0:0 $MAGISKBIN
 }
 
+# $1 = install dir
+# $2 = boot partition
 direct_install() {
   echo "- Flashing new boot image"
   flash_image $1/new-boot.img $2
@@ -70,6 +79,7 @@ direct_install() {
   return 0
 }
 
+# $1 = uninstaller zip
 run_uninstaller() {
   rm -rf /dev/tmp
   mkdir -p /dev/tmp/install
@@ -77,24 +87,16 @@ run_uninstaller() {
   INSTALLER=/dev/tmp/install sh /dev/tmp/install/assets/uninstaller.sh dummy 1 "$1"
 }
 
+# $1 = boot partition
 restore_imgs() {
   local SHA1=$(grep_prop SHA1 $MAGISKTMP/.magisk/config)
   local BACKUPDIR=/data/magisk_backup_$SHA1
   [ -d $BACKUPDIR ] || return 1
-
-  get_flags
-  find_boot_image
-
-  for name in dtb dtbo; do
-    [ -f $BACKUPDIR/${name}.img.gz ] || continue
-    local IMAGE=$(find_block $name$SLOT)
-    [ -z $IMAGE ] && continue
-    flash_image $BACKUPDIR/${name}.img.gz $IMAGE
-  done
   [ -f $BACKUPDIR/boot.img.gz ] || return 1
-  flash_image $BACKUPDIR/boot.img.gz $BOOTIMAGE
+  flash_image $BACKUPDIR/boot.img.gz $1
 }
 
+# $1 = path to bootctl executable
 post_ota() {
   cd /data/adb
   cp -f $1 bootctl
@@ -131,6 +133,8 @@ EOF
   cd /
 }
 
+# $1 = APK
+# $2 = package name
 adb_pm_install() {
   local tmp=/data/local/tmp/temp.apk
   cp -f "$1" $tmp