feat: more efficient zipalign

src/main/kotlin/app/revanced/utils/signing/align/ZipAligner.kt

@@ -1,63 +1,29 @@
 package app.revanced.utils.signing.align
 
-import app.revanced.utils.signing.align.stream.MultiOutputStream
-import app.revanced.utils.signing.align.stream.PeekingFakeStream
-import java.io.BufferedOutputStream
+import app.revanced.utils.signing.align.zip.ZipFile
 import java.io.File
 import java.util.*
-import java.util.zip.ZipEntry
-import java.util.zip.ZipFile
-import java.util.zip.ZipOutputStream
 
 internal object ZipAligner {
-    fun align(input: File, output: File, alignment: Int = 4) {
-        val zipFile = ZipFile(input)
+    const val DEFAULT_ALIGNMENT = 4
+    const val LIBRARY_ALIGNEMNT = 4096
 
-        val entries: Enumeration<out ZipEntry?> = zipFile.entries()
+    fun align(input: File, output: File) {
+        val inputZip = ZipFile(input)
+        val outputZip = ZipFile(output)
 
-        // fake
-        val peekingFakeStream = PeekingFakeStream()
-        val fakeOutputStream = ZipOutputStream(peekingFakeStream)
-        // real
-        val zipOutputStream = ZipOutputStream(BufferedOutputStream(output.outputStream()))
+        for (entry in inputZip.entries) {
+            val data = inputZip.getDataForEntry(entry)
 
-        val multiOutputStream = MultiOutputStream(
-            listOf(
-                fakeOutputStream, // fake, used to add the data to the fake stream
-                zipOutputStream // real
-            )
-        )
+            if (entry.compression == 0.toUShort()) {
+                val alignment = if (entry.fileName.endsWith(".so")) LIBRARY_ALIGNEMNT else DEFAULT_ALIGNMENT
 
-        var bias = 0
-        while (entries.hasMoreElements()) {
-            var padding = 0
-
-            val entry: ZipEntry = entries.nextElement()!!
-            // fake, used to calculate the file offset of the entry
-            fakeOutputStream.putNextEntry(entry)
-
-            if (entry.size == entry.compressedSize) {
-                val fileOffset = peekingFakeStream.peek()
-                val newOffset = fileOffset + bias
-                padding = ((alignment - (newOffset % alignment)) % alignment).toInt()
-
-                // real
-                entry.extra = if (entry.extra == null) ByteArray(padding)
-                else Arrays.copyOf(entry.extra, entry.extra.size + padding)
+                outputZip.addEntryAligned(entry, data, alignment)
+            } else {
+                outputZip.addEntry(entry, data)
             }
-
-            zipOutputStream.putNextEntry(entry)
-            zipFile.getInputStream(entry).copyTo(multiOutputStream)
-
-            // fake, used to add remaining bytes
-            fakeOutputStream.closeEntry()
-            // real
-            zipOutputStream.closeEntry()
-
-            bias += padding
         }
 
-        zipFile.close()
-        zipOutputStream.close()
+        outputZip.finish()
     }
-}
+}

src/main/kotlin/app/revanced/utils/signing/align/zip/Extensions.kt

@@ -0,0 +1,33 @@
+package app.revanced.utils.signing.align.zip
+
+import java.io.DataInput
+import java.io.DataOutput
+import java.nio.ByteBuffer
+
+fun UInt.toLittleEndian() =
+    (((this.toInt() and 0xff000000.toInt()) shr 24) or ((this.toInt() and 0x00ff0000) shr 8) or ((this.toInt() and 0x0000ff00) shl 8) or (this.toInt() shl 24)).toUInt()
+
+fun UShort.toLittleEndian() = (this.toUInt() shl 16).toLittleEndian().toUShort()
+
+fun UInt.toBigEndian() = (((this.toInt() and 0xff) shl 24) or ((this.toInt() and 0xff00) shl 8)
+        or ((this.toInt() and 0x00ff0000) ushr 8) or (this.toInt() ushr 24)).toUInt()
+
+fun UShort.toBigEndian() = (this.toUInt() shl 16).toBigEndian().toUShort()
+
+fun ByteBuffer.getUShort() = this.getShort().toUShort()
+fun ByteBuffer.getUInt() = this.getInt().toUInt()
+
+fun ByteBuffer.putUShort(ushort: UShort) = this.putShort(ushort.toShort())
+fun ByteBuffer.putUInt(uint: UInt) = this.putInt(uint.toInt())
+
+fun DataInput.readUShort() = this.readShort().toUShort()
+fun DataInput.readUInt() = this.readInt().toUInt()
+
+fun DataOutput.writeUShort(ushort: UShort) = this.writeShort(ushort.toInt())
+fun DataOutput.writeUInt(uint: UInt) = this.writeInt(uint.toInt())
+
+fun DataInput.readUShortLE() = this.readUShort().toBigEndian()
+fun DataInput.readUIntLE() = this.readUInt().toBigEndian()
+
+fun DataOutput.writeUShortLE(ushort: UShort) = this.writeUShort(ushort.toLittleEndian())
+fun DataOutput.writeUIntLE(uint: UInt) = this.writeUInt(uint.toLittleEndian())

src/main/kotlin/app/revanced/utils/signing/align/zip/ZipFile.kt

@@ -0,0 +1,128 @@
+package app.revanced.utils.signing.align.zip
+
+import app.revanced.utils.signing.align.zip.structures.ZipEndRecord
+import app.revanced.utils.signing.align.zip.structures.ZipEntry
+import java.io.File
+import java.io.RandomAccessFile
+import java.nio.ByteBuffer
+import java.nio.channels.FileChannel
+
+class ZipFile(val file: File) {
+    var entries: MutableList<ZipEntry> = mutableListOf()
+
+    private val filePointer: RandomAccessFile = RandomAccessFile(file, "rw")
+
+    init {
+        //if file isn't empty try to load entries
+        if (file.length() > 0) {
+            val endRecord = findEndRecord()
+
+            if (endRecord.diskNumber > 0u || endRecord.totalEntries != endRecord.diskEntries)
+                throw IllegalArgumentException("Multi-file archives are not supported")
+
+            entries = readEntries(endRecord).toMutableList()
+        }
+
+        //seek back to start for writing
+        filePointer.seek(0)
+    }
+
+    private fun findEndRecord(): ZipEndRecord {
+        //look from end to start since end record is at the end
+        for (i in filePointer.length() - 1 downTo 0) {
+            filePointer.seek(i)
+            //possible beginning of signature
+            if (filePointer.readByte() == 0x50.toByte()) {
+                //seek back to get the full int
+                filePointer.seek(i)
+                val possibleSignature = filePointer.readUIntLE()
+                if (possibleSignature == ZipEndRecord.ECD_SIGNATURE) {
+                    filePointer.seek(i)
+                    return ZipEndRecord.fromECD(filePointer)
+                }
+            }
+        }
+
+        throw Exception("Couldn't find end record")
+    }
+
+    private fun readEntries(endRecord: ZipEndRecord): List<ZipEntry> {
+        filePointer.seek(endRecord.centralDirectoryStartOffset.toLong())
+
+        val numberOfEntries = endRecord.diskEntries.toInt()
+
+        return buildList(numberOfEntries) {
+            for (i in 1..numberOfEntries) {
+                add(ZipEntry.fromCDE(filePointer).also
+                {
+                    //for some reason the local extra field can be different from the central one
+                    it.readLocalExtra(
+                        filePointer.channel.map(
+                            FileChannel.MapMode.READ_ONLY,
+                            it.localHeaderOffset.toLong() + 28,
+                            2
+                        )
+                    )
+                })
+            }
+        }
+    }
+
+    private fun writeCDE() {
+        val CDEStart = filePointer.channel.position().toUInt()
+
+        entries.forEach {
+            filePointer.channel.write(it.toCDE())
+        }
+
+        val endRecord = ZipEndRecord(
+            0u,
+            0u,
+            entries.count().toUShort(),
+            entries.count().toUShort(),
+            filePointer.channel.position().toUInt() - CDEStart,
+            CDEStart,
+            ""
+        )
+
+        filePointer.channel.write(endRecord.toECD())
+    }
+
+    fun addEntry(entry: ZipEntry, data: ByteBuffer) {
+        entry.localHeaderOffset = filePointer.channel.position().toUInt()
+
+        filePointer.channel.write(entry.toLFH())
+        filePointer.channel.write(data)
+
+        entries.add(entry)
+    }
+
+    fun addEntryAligned(entry: ZipEntry, data: ByteBuffer, alignment: Int) {
+        //calculate where data would end up
+        val dataOffset = filePointer.filePointer + entry.LFHSize
+
+        val mod = dataOffset % alignment
+
+        //wrong alignment
+        if (mod != 0L) {
+            //add padding at end of extra field
+            entry.localExtraField =
+                entry.localExtraField.copyOf((entry.localExtraField.size + (alignment - mod)).toInt())
+        }
+
+        addEntry(entry, data)
+    }
+
+    fun getDataForEntry(entry: ZipEntry): ByteBuffer {
+        return filePointer.channel.map(
+            FileChannel.MapMode.READ_ONLY,
+            entry.dataOffset.toLong(),
+            entry.compressedSize.toLong()
+        )
+    }
+
+    fun finish() {
+        writeCDE()
+        filePointer.close()
+    }
+}

src/main/kotlin/app/revanced/utils/signing/align/zip/structures/ZipEndRecord.kt

@@ -0,0 +1,77 @@
+package app.revanced.utils.signing.align.zip.structures
+
+import app.revanced.utils.signing.align.zip.putUInt
+import app.revanced.utils.signing.align.zip.putUShort
+import app.revanced.utils.signing.align.zip.readUIntLE
+import app.revanced.utils.signing.align.zip.readUShortLE
+import java.io.DataInput
+import java.nio.ByteBuffer
+import java.nio.ByteOrder
+
+data class ZipEndRecord(
+    val diskNumber: UShort,
+    val startingDiskNumber: UShort,
+    val diskEntries: UShort,
+    val totalEntries: UShort,
+    val centralDirectorySize: UInt,
+    val centralDirectoryStartOffset: UInt,
+    val fileComment: String,
+) {
+
+    companion object {
+        const val ECD_HEADER_SIZE = 22
+        const val ECD_SIGNATURE = 0x06054b50u
+
+        fun fromECD(input: DataInput): ZipEndRecord {
+            val signature = input.readUIntLE()
+
+            if (signature != ECD_SIGNATURE)
+                throw IllegalArgumentException("Input doesn't start with end record signature")
+
+            val diskNumber = input.readUShortLE()
+            val startingDiskNumber = input.readUShortLE()
+            val diskEntries = input.readUShortLE()
+            val totalEntries = input.readUShortLE()
+            val centralDirectorySize = input.readUIntLE()
+            val centralDirectoryStartOffset = input.readUIntLE()
+            val fileCommentLength = input.readUShortLE()
+            var fileComment = ""
+
+            if (fileCommentLength > 0u) {
+                val fileCommentBytes = ByteArray(fileCommentLength.toInt())
+                input.readFully(fileCommentBytes)
+                fileComment = fileCommentBytes.toString(Charsets.UTF_8)
+            }
+
+            return ZipEndRecord(
+                diskNumber,
+                startingDiskNumber,
+                diskEntries,
+                totalEntries,
+                centralDirectorySize,
+                centralDirectoryStartOffset,
+                fileComment
+            )
+        }
+    }
+
+    fun toECD(): ByteBuffer {
+        val commentBytes = fileComment.toByteArray(Charsets.UTF_8)
+
+        val buffer = ByteBuffer.allocate(ECD_HEADER_SIZE + commentBytes.size).also { it.order(ByteOrder.LITTLE_ENDIAN) }
+
+        buffer.putUInt(ECD_SIGNATURE)
+        buffer.putUShort(diskNumber)
+        buffer.putUShort(startingDiskNumber)
+        buffer.putUShort(diskEntries)
+        buffer.putUShort(totalEntries)
+        buffer.putUInt(centralDirectorySize)
+        buffer.putUInt(centralDirectoryStartOffset)
+        buffer.putUShort(commentBytes.size.toUShort())
+
+        buffer.put(commentBytes)
+
+        buffer.flip()
+        return buffer
+    }
+}