fix: Fixed writer & signature resolver, improved tests & speed, minor refactoring

src/main/kotlin/app/revanced/patcher/Patcher.kt

@@ -2,7 +2,7 @@ package app.revanced.patcher
 
 import app.revanced.patcher.cache.Cache
 import app.revanced.patcher.patch.Patch
-import app.revanced.patcher.resolver.MethodResolver
+import app.revanced.patcher.resolver.SignatureResolver
 import app.revanced.patcher.signature.MethodSignature
 import lanchon.multidexlib2.BasicDexFileNamer
 import lanchon.multidexlib2.MultiDexIO
@@ -21,23 +21,20 @@ class Patcher(
     private val patches = mutableSetOf<Patch>()
 
     init {
-        // TODO: find a way to load all dex classes, the code below only loads the first .dex file
         val dexFile = MultiDexIO.readDexFile(true, input, BasicDexFileNamer(), Opcodes.getDefault(), null)
-        cache = Cache(dexFile.classes, MethodResolver(dexFile.classes, signatures).resolve())
+        cache = Cache(dexFile.classes, SignatureResolver(dexFile.classes, signatures).resolve())
     }
 
     fun save() {
         val newDexFile = object : DexFile {
-            override fun getClasses(): MutableSet<out ClassDef> {
+            override fun getClasses(): Set<ClassDef> {
-                // TODO: find a way to return a set with a custom iterator
+                // this is a slow workaround for now
-                // TODO: the iterator would return the proxied class matching the current index of the list
+                val mutableClassList = cache.classes.toMutableList()
-                // TODO: instead of the original class
+                cache.classProxy
-                for (classProxy in cache.classProxy) {
+                    .filter { it.proxyUsed }.forEach { proxy ->
-                    if (!classProxy.proxyUsed) continue
+                        mutableClassList[proxy.originalIndex] = proxy.mutatedClass
-                    // TODO: merge this class with cache.classes somehow in an iterator
+                    }
-                    classProxy.mutatedClass
+                return mutableClassList.toSet()
-                }
-                return cache.classes.toMutableSet()
             }
 
             override fun getOpcodes(): Opcodes {
@@ -46,8 +43,8 @@ class Patcher(
             }
         }
 
-        // TODO: not sure about maxDexPoolSize & we should use the multithreading overload for writeDexFile
+        // TODO: we should use the multithreading capable overload for writeDexFile
-        MultiDexIO.writeDexFile(true, output, BasicDexFileNamer(), newDexFile, 10, null)
+        MultiDexIO.writeDexFile(true, output, BasicDexFileNamer(), newDexFile, 50000, null)
     }
 
     fun addPatches(vararg patches: Patch) {
@@ -56,6 +53,7 @@ class Patcher(
 
     fun applyPatches(stopOnError: Boolean = false): Map<String, Result<Nothing?>> {
         return buildMap {
+            // TODO: after each patch execution we could clear left overs like proxied classes to safe memory
             for (patch in patches) {
                 val result: Result<Nothing?> = try {
                     val pr = patch.execute(cache)

src/main/kotlin/app/revanced/patcher/cache/Cache.kt

@@ -1,22 +1,25 @@
 package app.revanced.patcher.cache
 
 import app.revanced.patcher.proxy.ClassProxy
-import app.revanced.patcher.signature.MethodSignatureScanResult
+import app.revanced.patcher.signature.SignatureResolverResult
 import org.jf.dexlib2.iface.ClassDef
 
 class Cache(
     internal val classes: Set<ClassDef>,
     val resolvedMethods: MethodMap
 ) {
+    // TODO: currently we create ClassProxies at multiple places, which is why we could have merge conflicts
+    //  this can be solved by creating a dedicated method for creating class proxies,
+    //  if the class proxy already exists in the cached proxy list below
     internal val classProxy = mutableSetOf<ClassProxy>()
 
     fun findClass(predicate: (ClassDef) -> Boolean): ClassProxy? {
         // if we already proxied the class matching the predicate,
-        val proxiedClass = classProxy.singleOrNull{classProxy -> predicate(classProxy.immutableClass)}
+        val proxiedClass = classProxy.singleOrNull { classProxy -> predicate(classProxy.immutableClass) }
         // return that proxy
         if (proxiedClass != null) return proxiedClass
         // else search the original class list
-        val foundClass =  classes.singleOrNull(predicate) ?: return null
+        val foundClass = classes.singleOrNull(predicate) ?: return null
         // create a class proxy with the index of the class in the classes list
         // TODO: There might be a more elegant way to the comment above
         val classProxy = ClassProxy(foundClass, classes.indexOf(foundClass))
@@ -27,8 +30,8 @@ class Cache(
     }
 }
 
-class MethodMap : LinkedHashMap<String, MethodSignatureScanResult>() {
+class MethodMap : LinkedHashMap<String, SignatureResolverResult>() {
-    override fun get(key: String): MethodSignatureScanResult {
+    override fun get(key: String): SignatureResolverResult {
         return super.get(key) ?: throw MethodNotFoundException("Method $key was not found in the method cache")
     }
 }

src/main/kotlin/app/revanced/patcher/proxy/ClassProxy.kt

@@ -3,10 +3,9 @@ package app.revanced.patcher.proxy
 import app.revanced.patcher.proxy.mutableTypes.MutableClass
 import org.jf.dexlib2.iface.ClassDef
 
-
 class ClassProxy(
     val immutableClass: ClassDef,
-    val originalClassIndex: Int,
+    val originalIndex: Int,
 ) {
     internal var proxyUsed = false
     internal lateinit var mutatedClass: MutableClass

src/main/kotlin/app/revanced/patcher/resolver/PatternScanData.kt

@@ -1,6 +1,6 @@
 package app.revanced.patcher.resolver
 
-internal data class MethodResolverScanResult(
+internal data class PatternScanData(
     val found: Boolean,
     val startIndex: Int? = 0,
     val endIndex: Int? = 0

src/main/kotlin/app/revanced/patcher/resolver/SignatureResolver.kt

@@ -1,43 +1,51 @@
 package app.revanced.patcher.resolver
 
 import app.revanced.patcher.cache.MethodMap
-import app.revanced.patcher.signature.MethodSignatureScanResult
+import app.revanced.patcher.proxy.ClassProxy
-import app.revanced.patcher.signature.PatternScanData
 import app.revanced.patcher.signature.MethodSignature
+import app.revanced.patcher.signature.PatternScanResult
+import app.revanced.patcher.signature.SignatureResolverResult
 import org.jf.dexlib2.Opcode
 import org.jf.dexlib2.iface.ClassDef
 import org.jf.dexlib2.iface.Method
 
-// TODO: add logger
+// TODO: add logger back
-internal class MethodResolver(private val classes: Set<ClassDef>, private val signatures: Array<MethodSignature>) {
+internal class SignatureResolver(
+    private val classes: Set<ClassDef>,
+    private val methodSignatures: Array<MethodSignature>
+) {
     fun resolve(): MethodMap {
         val methodMap = MethodMap()
 
-        for (classDef in classes) {
+        for ((index, classDef) in classes.withIndex()) {
-            for (method in classDef.methods) {
+            for (signature in methodSignatures) {
-                for (methodSignature in signatures) {
+                if (methodMap.containsKey(signature.name)) {
-                    if (methodMap.containsKey(methodSignature.name)) { // method already found for this sig
+                    continue
+                }
+
+                for (method in classDef.methods) {
+                    val (isMatch, patternScanData) = compareSignatureToMethod(signature, method)
+
+                    if (!isMatch || patternScanData == null) {
                         continue
                     }
 
-                    val (r, sr) = cmp(method, methodSignature)
+                    // create class proxy, in case a patch needs mutability
-                    if (!r || sr == null) {
+                    val classProxy = ClassProxy(classDef, index)
-                        continue
+                    methodMap[signature.name] = SignatureResolverResult(
-                    }
+                        classProxy,
-
+                        method.name,
-                    methodMap[methodSignature.name] = MethodSignatureScanResult(
+                        PatternScanResult(
-                        method,
+                            patternScanData.startIndex!!,
-                        PatternScanData(
+                            patternScanData.endIndex!!
-                            // sadly we cannot create contracts for a data class, so we must assert
-                            sr.startIndex!!,
-                            sr.endIndex!!
                         )
                     )
                 }
             }
         }
 
-        for (signature in signatures) {
+        // TODO: remove?
+        for (signature in methodSignatures) {
             if (methodMap.containsKey(signature.name)) continue
         }
 
@@ -46,45 +54,49 @@ internal class MethodResolver(private val classes: Set<ClassDef>, private val si
 
     // These functions do not require the constructor values, so they can be static.
     companion object {
-        fun resolveMethod(classNode: ClassDef, signature: MethodSignature): MethodSignatureScanResult? {
+        fun resolveFromProxy(classProxy: ClassProxy, signature: MethodSignature): SignatureResolverResult? {
-            for (method in classNode.methods) {
+            for (method in classProxy.immutableClass.methods) {
-                val (r, sr) = cmp(method, signature)
+                val (r, sr) = compareSignatureToMethod(signature, method)
                 if (!r || sr == null) continue
-                return MethodSignatureScanResult(
+                return SignatureResolverResult(
-                    method,
+                    classProxy,
-                    PatternScanData(0, 0) // opcode list is always ignored.
+                    method.name,
+                    null
                 )
             }
             return null
         }
 
-        private fun cmp(method: Method, signature: MethodSignature): Pair<Boolean, MethodResolverScanResult?> {
+        private fun compareSignatureToMethod(
+            signature: MethodSignature,
+            method: Method
+        ): Pair<Boolean, PatternScanData?> {
             // TODO: compare as generic object if not primitive
             signature.returnType?.let { _ ->
                 if (signature.returnType != method.returnType) {
-                    return@cmp false to null
+                    return@compareSignatureToMethod false to null
                 }
             }
 
             signature.accessFlags?.let { _ ->
                 if (signature.accessFlags != method.accessFlags) {
-                    return@cmp false to null
+                    return@compareSignatureToMethod false to null
                 }
             }
 
             // TODO: compare as generic object if the parameter is not primitive
             signature.methodParameters?.let { _ ->
                 if (signature.methodParameters != method.parameters) {
-                    return@cmp false to null
+                    return@compareSignatureToMethod false to null
                 }
             }
 
             signature.opcodes?.let { _ ->
                 val result = method.implementation?.instructions?.scanFor(signature.opcodes)
-                return@cmp if (result != null && result.found) true to result else false to null
+                return@compareSignatureToMethod if (result != null && result.found) true to result else false to null
             }
 
-            return true to MethodResolverScanResult(true)
+            return true to PatternScanData(true)
         }
     }
 }
@@ -92,7 +104,7 @@ internal class MethodResolver(private val classes: Set<ClassDef>, private val si
 private operator fun ClassDef.component1() = this
 private operator fun ClassDef.component2() = this.methods
 
-private fun <T> MutableIterable<T>.scanFor(pattern: Array<Opcode>): MethodResolverScanResult {
+private fun <T> MutableIterable<T>.scanFor(pattern: Array<Opcode>): PatternScanData {
     // TODO: create var for count?
     for (i in 0 until this.count()) {
         var occurrence = 0
@@ -101,12 +113,12 @@ private fun <T> MutableIterable<T>.scanFor(pattern: Array<Opcode>): MethodResolv
             if (!n.shouldSkip() && n != pattern[occurrence]) break
             if (++occurrence >= pattern.size) {
                 val current = i + occurrence
-                return MethodResolverScanResult(true, current - pattern.size, current)
+                return PatternScanData(true, current - pattern.size, current)
             }
         }
     }
 
-    return MethodResolverScanResult(false)
+    return PatternScanData(false)
 }
 
 // TODO: extend Opcode type, not T (requires a cast to Opcode)

src/main/kotlin/app/revanced/patcher/signature/MethodSignatureScanResult.kt

@@ -1,23 +0,0 @@
-package app.revanced.patcher.signature
-
-import app.revanced.patcher.resolver.MethodResolver
-import org.jf.dexlib2.iface.ClassDef
-import org.jf.dexlib2.iface.Method
-import org.jf.dexlib2.immutable.reference.ImmutableTypeReference
-
-// TODO: IMPORTANT: we might have to use a class proxy as well here
-data class MethodSignatureScanResult(
-    val method: Method,
-    val scanData: PatternScanData
-) {
-    @Suppress("Unused") // TODO(Sculas): remove this when we have coverage for this method.
-    fun findParentMethod(signature: MethodSignature): MethodSignatureScanResult? {
-        // TODO: find a way to get the classNode out of method.definingClass
-        return MethodResolver.resolveMethod(ImmutableTypeReference(method.definingClass) as ClassDef, signature)
-    }
-}
-
-data class PatternScanData(
-    val startIndex: Int,
-    val endIndex: Int
-)

src/main/kotlin/app/revanced/patcher/signature/SignatureResolverResult.kt

@@ -0,0 +1,20 @@
+package app.revanced.patcher.signature
+
+import app.revanced.patcher.proxy.ClassProxy
+import app.revanced.patcher.resolver.SignatureResolver
+
+data class SignatureResolverResult(
+    val definingClassProxy: ClassProxy,
+    val resolvedMethodName: String,
+    val scanData: PatternScanResult?
+) {
+    @Suppress("Unused") // TODO(Sculas): remove this when we have coverage for this method.
+    fun findParentMethod(signature: MethodSignature): SignatureResolverResult? {
+        return SignatureResolver.resolveFromProxy(definingClassProxy, signature)
+    }
+}
+
+data class PatternScanResult(
+    val startIndex: Int,
+    val endIndex: Int
+)

src/test/kotlin/patcher/PatcherTest.kt

@@ -1,5 +1,6 @@
-package app.revanced.patcher
+package patcher
 
+import app.revanced.patcher.Patcher
 import app.revanced.patcher.cache.Cache
 import app.revanced.patcher.patch.Patch
 import app.revanced.patcher.patch.PatchResult
@@ -8,13 +9,10 @@ import app.revanced.patcher.patch.PatchResultSuccess
 import app.revanced.patcher.signature.MethodSignature
 import org.jf.dexlib2.AccessFlags
 import org.jf.dexlib2.Opcode
-import org.jf.dexlib2.builder.instruction.BuilderInstruction21c
+import org.jf.dexlib2.builder.instruction.BuilderInstruction35c
-import org.jf.dexlib2.iface.instruction.formats.Instruction21c
-import org.jf.dexlib2.iface.reference.FieldReference
 import org.jf.dexlib2.immutable.reference.ImmutableMethodReference
 import java.io.File
 
-
 fun main() {
     val signatures = arrayOf(
         MethodSignature(
@@ -33,40 +31,57 @@ fun main() {
 
     val patcher = Patcher(
         File("black.apk"),
-        File("folder/"),
+        File("./"),
         signatures
     )
 
     val mainMethodPatchViaClassProxy = object : Patch("main-method-patch-via-proxy") {
         override fun execute(cache: Cache): PatchResult {
             val proxy = cache.findClass { classDef ->
-                classDef.methods.any { method ->
+                classDef.type.contains("XAdRemover")
-                    method.name == "main"
+            } ?: return PatchResultError("Class 'XAdRemover' could not be found")
-                }
-            } ?: return PatchResultError("Class with method 'mainMethod' could not be found")
 
-            val mainMethodClass = proxy.resolve()
+            val xAdRemoverClass = proxy.resolve()
-            val mainMethod = mainMethodClass.methods.single { method -> method.name == "main" }
+            val hideReelMethod = xAdRemoverClass.methods.single { method -> method.name.contains("HideReel") }
 
-            val hideReelMethodRef = ImmutableMethodReference(
+            val readSettingsMethodRef = ImmutableMethodReference(
-                "Lfi/razerman/youtube/XAdRemover;",
+                "Lfi/razerman/youtube/XGlobals;",
-                "HideReel",
+                "ReadSettings",
-                listOf("Landroid/view/View;"),
+                emptyList(),
                 "V"
             )
 
-            val mainMethodInstructions = mainMethod.implementation!!.instructions
+            val instructions = hideReelMethod.implementation!!.instructions
-            val printStreamFieldRef = (mainMethodInstructions.first() as Instruction21c).reference as FieldReference
+
-            // TODO: not sure how to use the registers yet, find a way
+            val readSettingsInstruction = BuilderInstruction35c(
-            mainMethodInstructions.add(BuilderInstruction21c(Opcode.SGET_OBJECT, 0, printStreamFieldRef))
+                Opcode.INVOKE_STATIC,
+                0,
+                0,
+                0,
+                0,
+                0,
+                0,
+                readSettingsMethodRef
+            )
+
+            // TODO: figure out control flow
+            //  otherwise the we would still jump over to the original instruction at index 21 instead to our new one
+            instructions.add(
+                21,
+                readSettingsInstruction
+            )
             return PatchResultSuccess()
         }
     }
 
     val mainMethodPatchViaSignature = object : Patch("main-method-patch-via-signature") {
         override fun execute(cache: Cache): PatchResult {
-            cache.resolvedMethods["main-method"].method
+            val mainMethodMap = cache.resolvedMethods["main-method"]
-            return  PatchResultSuccess()
+            mainMethodMap.definingClassProxy.immutableClass.methods.single { method ->
+                method.name == mainMethodMap.resolvedMethodName
+            }
+
+            return PatchResultSuccess()
         }
     }
     patcher.addPatches(mainMethodPatchViaClassProxy, mainMethodPatchViaSignature)