Provide upgrade path for stubs

app/src/main/java/com/topjohnwu/magisk/model/download/DownloadService.kt

@@ -7,20 +7,17 @@ import android.content.Intent
 import android.os.Build
 import android.webkit.MimeTypeMap
 import androidx.core.app.NotificationCompat
-import com.topjohnwu.magisk.ProcessPhoenix
 import com.topjohnwu.magisk.R
 import com.topjohnwu.magisk.extensions.chooser
 import com.topjohnwu.magisk.extensions.exists
 import com.topjohnwu.magisk.extensions.provide
 import com.topjohnwu.magisk.intent
-import com.topjohnwu.magisk.isRunningAsStub
 import com.topjohnwu.magisk.model.entity.internal.Configuration.*
 import com.topjohnwu.magisk.model.entity.internal.Configuration.Flash.Secondary
 import com.topjohnwu.magisk.model.entity.internal.DownloadSubject
 import com.topjohnwu.magisk.model.entity.internal.DownloadSubject.*
 import com.topjohnwu.magisk.ui.flash.FlashActivity
 import com.topjohnwu.magisk.utils.APKInstall
-import com.topjohnwu.magisk.utils.DynAPK
 import org.koin.core.get
 import java.io.File
 import kotlin.random.Random.Default.nextInt
@@ -65,15 +62,7 @@ open class DownloadService : RemoteFileService() {
     ) {
         remove(id)
         when (subject.configuration)  {
-            is APK.Upgrade -> {
-                if (isRunningAsStub) {
-                    subject.file.copyTo(DynAPK.update(this), overwrite = true)
-                    subject.file.delete()
-                    ProcessPhoenix.triggerRebirth(this)
-                } else {
-                    APKInstall.install(this, subject.file)
-                }
-            }
+            is APK.Upgrade -> APKInstall.install(this, subject.file)
             is APK.Restore -> Unit
         }
     }

app/src/main/java/com/topjohnwu/magisk/model/download/ManagerUpgrade.kt

@@ -1,18 +1,19 @@
 package com.topjohnwu.magisk.model.download
 
-import com.topjohnwu.magisk.BuildConfig
-import com.topjohnwu.magisk.Config
-import com.topjohnwu.magisk.R
-import com.topjohnwu.magisk.isRunningAsStub
+import com.topjohnwu.magisk.*
+import com.topjohnwu.magisk.extensions.writeTo
 import com.topjohnwu.magisk.model.entity.internal.Configuration.APK.Restore
 import com.topjohnwu.magisk.model.entity.internal.Configuration.APK.Upgrade
 import com.topjohnwu.magisk.model.entity.internal.DownloadSubject
+import com.topjohnwu.magisk.utils.DynAPK
 import com.topjohnwu.magisk.utils.PatchAPK
 import com.topjohnwu.superuser.Shell
 import java.io.File
 
-private fun RemoteFileService.patchPackage(apk: File, id: Int) {
-    if (!isRunningAsStub && packageName != BuildConfig.APPLICATION_ID) {
+private fun RemoteFileService.patch(apk: File, id: Int) {
+    if (packageName == BuildConfig.APPLICATION_ID)
+        return
+
     update(id) { notification ->
         notification.setProgress(0, 0, true)
             .setProgress(0, 0, true)
@@ -23,6 +24,25 @@ private fun RemoteFileService.patchPackage(apk: File, id: Int) {
     PatchAPK.patch(apk, patched, packageName, applicationInfo.nonLocalizedLabel.toString())
     apk.delete()
     patched.renameTo(apk)
+}
+
+private fun RemoteFileService.upgrade(apk: File, id: Int) {
+    if (isRunningAsStub) {
+        // Move to upgrade location
+        apk.copyTo(DynAPK.update(this), overwrite = true)
+        apk.delete()
+        if (ClassMap.data!!.version < Info.remote.stub.versionCode) {
+            // We also want to upgrade stub
+            service.fetchFile(Info.remote.stub.link).blockingGet().byteStream().use {
+                it.writeTo(apk)
+            }
+            patch(apk, id)
+        } else {
+            // Simply relaunch the app
+            ProcessPhoenix.triggerRebirth(this)
+        }
+    } else {
+        patch(apk, id)
     }
 }
 
@@ -41,6 +61,6 @@ private fun RemoteFileService.restore(apk: File, id: Int) {
 
 fun RemoteFileService.handleAPK(subject: DownloadSubject.Manager)
     = when (subject.configuration) {
-        is Upgrade -> patchPackage(subject.file, subject.hashCode())
+        is Upgrade -> upgrade(subject.file, subject.hashCode())
         is Restore -> restore(subject.file, subject.hashCode())
     }

app/src/main/java/com/topjohnwu/magisk/model/download/RemoteFileService.kt

@@ -22,7 +22,7 @@ import java.io.InputStream
 
 abstract class RemoteFileService : NotificationService() {
 
-    private val service: GithubRawServices by inject()
+    val service: GithubRawServices by inject()
 
     override val defaultNotification: NotificationCompat.Builder
         get() = Notifications.progress(this, "")

app/src/main/java/com/topjohnwu/magisk/utils/PatchAPK.kt

@@ -138,6 +138,11 @@ object PatchAPK {
 
     fun patch(apk: File, out: File, pkg: String, label: String): Boolean {
         try {
+            if (apk.length() < 1 shl 18) {
+                // APK is smaller than 256K, must be stub
+                return patch(apk.path, out.path, pkg, label)
+            }
+
             // Try using the new APK to patch itself
             val loader = DynamicClassLoader(apk)
             val cls = loader.loadClass("a.a")
@@ -152,9 +157,8 @@ object PatchAPK {
         } catch (e: Exception) {
             Timber.e(e)
             // Fallback to use the current implementation
-            patch(apk.path, out.path, pkg, label)
+            return patch(apk.path, out.path, pkg, label)
         }
-        return false
     }
 
     fun hideManager(context: Context, label: String) {