Support modern Samsung AP.tar patching

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

@@ -37,11 +37,7 @@ import org.koin.core.KoinComponent
 import org.koin.core.get
 import org.koin.core.inject
 import timber.log.Timber
-import java.io.File
+import java.io.*
-import java.io.FileOutputStream
-import java.io.IOException
-import java.io.InputStream
-import java.nio.ByteBuffer
 import java.util.zip.ZipEntry
 import java.util.zip.ZipInputStream
 
@@ -49,7 +45,6 @@ abstract class MagiskInstallImpl : KoinComponent {
 
     protected lateinit var installDir: File
     private lateinit var srcBoot: String
-    private lateinit var destFile: MediaStoreUtils.UriFile
     private lateinit var zipUri: Uri
 
     protected val console: MutableList<String>
@@ -163,15 +158,14 @@ abstract class MagiskInstallImpl : KoinComponent {
     }
 
     @Throws(IOException::class)
-    private fun handleTar(input: InputStream) {
+    private fun handleTar(input: InputStream, output: OutputStream): OutputStream {
         console.add("- Processing tar file")
-        var vbmeta = false
+        val tarOut = TarOutputStream(output)
-        val tarOut = TarOutputStream(destFile.uri.outputStream())
-        this.tarOut = tarOut
         TarInputStream(input).use { tarIn ->
             lateinit var entry: TarEntry
             while (tarIn.nextEntry?.let { entry = it } != null) {
-                if (entry.name.contains("boot.img") || entry.name.contains("recovery.img")) {
+                if (entry.name.contains("boot.img") ||
+                    (Config.recovery && entry.name.contains("recovery.img"))) {
                     val name = entry.name
                     console.add("-- Extracting: $name")
                     val extract = File(installDir, name)
@@ -180,25 +174,15 @@ abstract class MagiskInstallImpl : KoinComponent {
                         console.add("-- Decompressing: $name")
                         "./magiskboot decompress $extract".sh()
                     }
-                } else if (entry.name.contains("vbmeta.img")) {
-                    vbmeta = true
-                    val buf = ByteBuffer.allocate(256)
-                    buf.put("AVB0".toByteArray())   // magic
-                    buf.putInt(1)                   // required_libavb_version_major
-                    buf.putInt(120, 2)              // flags
-                    buf.position(128)               // release_string
-                    buf.put("avbtool 1.1.0".toByteArray())
-                    tarOut.putNextEntry(newEntry("vbmeta.img", 256))
-                    tarOut.write(buf.array())
                 } else {
-                    console.add("-- Writing: " + entry.name)
+                    console.add("-- Copying: " + entry.name)
                     tarOut.putNextEntry(entry)
                     tarIn.copyTo(tarOut)
                 }
             }
             val boot = SuFile.open(installDir, "boot.img")
             val recovery = SuFile.open(installDir, "recovery.img")
-            if (vbmeta && recovery.exists() && boot.exists()) {
+            if (recovery.exists() && boot.exists()) {
                 // Install Magisk to recovery
                 srcBoot = recovery.path
                 // Repack boot image to prevent restore
@@ -220,27 +204,33 @@ abstract class MagiskInstallImpl : KoinComponent {
                 srcBoot = boot.path
             }
         }
+        return tarOut
     }
 
     private fun handleFile(uri: Uri): Boolean {
+        val outStream: OutputStream
+        val outFile: MediaStoreUtils.UriFile
+
+        // Process input file
         try {
-            uri.inputStream().buffered().use {
+            uri.inputStream().buffered().use { src ->
-                it.mark(500)
+                src.mark(500)
                 val magic = ByteArray(5)
-                if (it.skip(257) != 257L || it.read(magic) != magic.size) {
+                if (src.skip(257) != 257L || src.read(magic) != magic.size) {
-                    console.add("! Invalid file")
+                    console.add("! Invalid input file")
                     return false
                 }
-                it.reset()
+                src.reset()
-                if (magic.contentEquals("ustar".toByteArray())) {
+                outStream = if (magic.contentEquals("ustar".toByteArray())) {
-                    destFile = MediaStoreUtils.getFile("magisk_patched.tar")
+                    outFile = MediaStoreUtils.getFile("magisk_patched.tar")
-                    handleTar(it)
+                    handleTar(src, outFile.uri.outputStream())
                 } else {
                     // Raw image
                     srcBoot = File(installDir, "boot.img").path
-                    destFile = MediaStoreUtils.getFile("magisk_patched.img")
                     console.add("- Copying image to cache")
-                    FileOutputStream(srcBoot).use { out -> it.copyTo(out) }
+                    FileOutputStream(srcBoot).use { src.copyTo(it) }
+                    outFile = MediaStoreUtils.getFile("magisk_patched.img")
+                    outFile.uri.outputStream()
                 }
             }
         } catch (e: IOException) {
@@ -249,6 +239,31 @@ abstract class MagiskInstallImpl : KoinComponent {
             return false
         }
 
+        // Patch file
+        if (!patchBoot())
+            return false
+
+        // Output file
+        try {
+            val patched = SuFile.open(installDir, "new-boot.img")
+            if (outStream is TarOutputStream) {
+                val name = if (srcBoot.contains("recovery")) "recovery.img" else "boot.img"
+                outStream.putNextEntry(newEntry(name, patched.length()))
+            }
+            withStreams(SuFileInputStream(patched), outStream) { src, out -> src.copyTo(out) }
+            patched.delete()
+
+            console.add("")
+            console.add("****************************")
+            console.add(" Output file is written to ")
+            console.add(" $outFile ")
+            console.add("****************************")
+        } catch (e: IOException) {
+            console.add("! Failed to output to $outFile")
+            Timber.e(e)
+            return false
+        }
+
         return true
     }
 
@@ -311,33 +326,7 @@ abstract class MagiskInstallImpl : KoinComponent {
     private fun flashBoot(): Boolean {
         if (!"direct_install $installDir $srcBoot".sh().isSuccess)
             return false
-        arrayOf("run_migrations").sh()
+        "run_migrations".sh()
-        return true
-    }
-
-    private fun storeBoot(): Boolean {
-        val patched = SuFile.open(installDir, "new-boot.img")
-        try {
-            val os = tarOut?.let {
-                it.putNextEntry(newEntry(
-                    if (srcBoot.contains("recovery")) "recovery.img" else "boot.img",
-                    patched.length()))
-                tarOut = null
-                it
-            } ?: destFile.uri.outputStream()
-            SuFileInputStream(patched).use { it.copyTo(os); os.close() }
-        } catch (e: IOException) {
-            console.add("! Failed to output to $destFile")
-            Timber.e(e)
-            return false
-        }
-
-        patched.delete()
-        console.add("")
-        console.add("****************************")
-        console.add(" Output file is placed in ")
-        console.add(" $destFile ")
-        console.add("****************************")
         return true
     }
 
@@ -366,8 +355,7 @@ abstract class MagiskInstallImpl : KoinComponent {
     private fun String.fsh() = ShellUtils.fastCmd(this)
     private fun Array<String>.fsh() = ShellUtils.fastCmd(*this)
 
-    protected fun doPatchFile(patchFile: Uri) =
+    protected fun doPatchFile(patchFile: Uri) = extractZip() && handleFile(patchFile)
-        extractZip() && handleFile(patchFile) && patchBoot() && storeBoot()
 
     protected fun direct() = findImage() && extractZip() && patchBoot() && flashBoot()