diff --git a/src/main/kotlin/app/revanced/patcher/Matching.kt b/src/main/kotlin/app/revanced/patcher/Matching.kt index 84163c6..ef6fac1 100644 --- a/src/main/kotlin/app/revanced/patcher/Matching.kt +++ b/src/main/kotlin/app/revanced/patcher/Matching.kt @@ -4,8 +4,10 @@ package app.revanced.patcher import app.revanced.patcher.Matcher.MatchContext import app.revanced.patcher.dex.mutable.MutableMethod +import app.revanced.patcher.extensions.instructions import app.revanced.patcher.patch.BytecodePatchContext import com.android.tools.smali.dexlib2.HiddenApiRestriction +import com.android.tools.smali.dexlib2.Opcode import com.android.tools.smali.dexlib2.iface.* import com.android.tools.smali.dexlib2.iface.Annotation import com.android.tools.smali.dexlib2.iface.instruction.Instruction @@ -65,7 +67,7 @@ fun BytecodePatchContext.firstClassDefMutable(predicate: MatchPredicate)? = null -) = lookupMaps.classesByType[type]?.takeIf { +) = lookupMaps.classDefsByType[type]?.takeIf { predicate == null || with(predicate) { with(MatchContext()) { it.match() } } } @@ -81,18 +83,19 @@ fun BytecodePatchContext.firstClassDefMutable( type: String, predicate: (MatchPredicate)? = null ) = requireNotNull(firstClassDefMutableOrNull(type, predicate)) -fun BytecodePatchContext.firstMethodOrNull(predicate: MatchPredicate) = with(predicate) { +fun Iterable.firstMethodOrNull(predicate: MatchPredicate) = with(predicate) { with(MatchContext()) { - classDefs.asSequence().flatMap { it.methods.asSequence() }.firstOrNull { it.match() } + this@firstMethodOrNull.asSequence().flatMap { it.methods.asSequence() }.firstOrNull { it.match() } } } -fun BytecodePatchContext.firstMethod(predicate: MatchPredicate) = requireNotNull(firstMethodOrNull(predicate)) +fun Iterable.firstMethod(predicate: MatchPredicate) = requireNotNull(firstMethodOrNull(predicate)) -fun BytecodePatchContext.firstMethodMutableOrNull(predicate: MatchPredicate): MutableMethod? { +context(BytecodePatchContext) +fun Iterable.firstMethodMutableOrNull(predicate: MatchPredicate): MutableMethod? { with(predicate) { with(MatchContext()) { - classDefs.forEach { classDef -> + this@firstMethodMutableOrNull.forEach { classDef -> classDef.methods.firstOrNull { it.match() }?.let { method -> return classDef.mutable().methods.first { MethodUtil.methodSignaturesMatch(it, method) } } @@ -103,6 +106,20 @@ fun BytecodePatchContext.firstMethodMutableOrNull(predicate: MatchPredicate.firstMethodMutable(predicate: MatchPredicate) = + requireNotNull(firstMethodMutableOrNull(predicate)) + +fun BytecodePatchContext.firstMethodOrNull(predicate: MatchPredicate) = + classDefs.firstMethodOrNull(predicate) + +fun BytecodePatchContext.firstMethod(predicate: MatchPredicate) = + requireNotNull(firstMethodOrNull(predicate)) + + +fun BytecodePatchContext.firstMethodMutableOrNull(predicate: MatchPredicate) = + classDefs.firstMethodMutableOrNull(predicate) + fun BytecodePatchContext.firstMethodMutable(predicate: MatchPredicate) = requireNotNull(firstMethodMutableOrNull(predicate)) @@ -216,19 +233,21 @@ abstract class Matcher : MutableList by mutableListOf() { abstract operator fun invoke(haystack: Iterable): Boolean - class MatchContext internal constructor() : MutableMap by mutableMapOf() + class MatchContext internal constructor() : MutableMap by mutableMapOf() { + inline fun remember(key: String, defaultValue: () -> V) = + get(key) as? V ?: defaultValue().also { put(key, it) } + } } - -fun slidingWindowMatcher(builder: MutableList Boolean>.() -> Unit) = - SlidingWindowMatcher().apply(builder) +fun slidingWindowMatcher(build: SlidingWindowMatcher.() -> Unit) = + SlidingWindowMatcher().apply(build) context(MatchContext) -fun Iterable.matchSlidingWindow(key: String, builder: MutableList Boolean>.() -> Unit) = - (getOrPut(key) { slidingWindowMatcher(builder) } as Matcher Boolean>)(this) +fun Iterable.matchSlidingWindow(key: String, build: SlidingWindowMatcher.() -> Unit) = + remember(key) { slidingWindowMatcher(build) }(this) -fun Iterable.matchSlidingWindow(builder: MutableList Boolean>.() -> Unit) = - slidingWindowMatcher(builder)(this) +fun Iterable.matchSlidingWindow(build: SlidingWindowMatcher.() -> Unit) = + slidingWindowMatcher(build)(this) class SlidingWindowMatcher() : Matcher Boolean>() { override operator fun invoke(haystack: Iterable): Boolean { @@ -255,149 +274,124 @@ class SlidingWindowMatcher() : Matcher Boolean>() { } } -fun findStringsMatcher(builder: MutableList.() -> Unit) = - FindStringsMatcher().apply(builder) +fun findStringsMatcher(build: MutableList.() -> Unit) = + FindStringsMatcher().apply(build) class FindStringsMatcher() : Matcher() { - val matchedStrings = mutableMapOf() - var needles = toMutableSet() // Reduce O(n²) to O(log n) by removing from the set + private val _matchedStrings = mutableListOf>() + val matchedStrings: List> = _matchedStrings override fun invoke(haystack: Iterable): Boolean { - needles = toMutableSet() // Reset needles for each invocation - // (or do not use the set if set is too small for performance) + _matchedStrings.clear() + val remaining = indices.toMutableList() - val foundStrings = mutableMapOf() + haystack.forEachIndexed { hayIndex, instruction -> + val string = ((instruction as? ReferenceInstruction)?.reference as? StringReference)?.string + ?: return@forEachIndexed - haystack.forEachIndexed { index, instruction -> - if (instruction !is ReferenceInstruction) return@forEachIndexed - val reference = instruction.reference - if (reference !is StringReference) return@forEachIndexed - val string = reference.string + val index = remaining.firstOrNull { this[it] in string } ?: return@forEachIndexed - if (needles.removeIf { it in string }) { - foundStrings[string] = index - } + _matchedStrings += this[index] to hayIndex + remaining -= index } - return if (foundStrings.size == size) { - matchedStrings += foundStrings - - true - } else { - false - } + return remaining.isEmpty() } } -fun BytecodePatchContext.findStringIndices() { - val match = findStringsMatcher { - add("fullstring1") - add("fullstring2") - add("partialString") +fun BytecodePatchContext.a() { + val match = indexedMatcher { + first { opcode == Opcode.OR_INT_2ADDR } + after { opcode == Opcode.RETURN_VOID } + after(atLeast = 2, atMost = 5) { opcode == Opcode.MOVE_RESULT_OBJECT } + opcode(Opcode.RETURN_VOID) } - firstMethod("fullstring", "fullstring") { - implementation { - match(instructions) - } + val myMethod = firstMethod { + implementation { match(instructions) } } - match.matchedStrings.forEach { (key, value) -> - println("Found string '$key' at index $value") - } - - firstMethod { - implementation { - // Uncached usage - instructions.matchSlidingWindow { - - } || instructions.matchSlidingWindow("cached usage") { - - } - } - } + match._indices // Mapped in same order as defined } -fun BytecodePatchContext.anotherExample() { - val desiredStringIndices = listOf("fullstring1", "fullstring2", "partialString") - val matchedIndices = mutableMapOf() +fun IndexedMatcher.opcode(opcode: Opcode) { + after { this.opcode == opcode } +} - firstMethod("fullstring", "fullstring") { - val remaining = desiredStringIndices.toMutableSet() - val foundMap = mutableMapOf() +context(MatchContext) +fun Method.instructions(key: String, build: IndexedMatcher.() -> Unit) = + instructions.matchIndexed("instructions", build) - implementation { - instructions.withIndex().forEach { (index, instruction) -> - val string = (instruction as? ReferenceInstruction)?.reference - .let { it as? StringReference }?.string - ?: return@forEach +fun indexedMatcher(build: IndexedMatcher.() -> Unit) = + IndexedMatcher().apply(build) - val iterator = remaining.iterator() - while (iterator.hasNext()) { - val desired = iterator.next() - if (desired in string) { - foundMap[desired] = index - iterator.remove() +context(MatchContext) +fun Iterable.matchIndexed(key: String, build: IndexedMatcher.() -> Unit) = + remember>(key) { indexedMatcher(build) }(this) + +class IndexedMatcher() : Matcher Boolean>() { + private val _indices: MutableList = mutableListOf() + val indices: List = _indices + + private var lastMatchedIndex = -1 + private var currentIndex = -1 + + override fun invoke(haystack: Iterable): Boolean { + val hayList = haystack as? List ?: haystack.toList() + + _indices.clear() + + var firstNeedleIndex = 0 + + while (firstNeedleIndex <= hayList.lastIndex) { + lastMatchedIndex = -1 + + val tempIndices = mutableListOf() + + var matchedAll = true + var subIndex = firstNeedleIndex + + for (predicateIndex in _indices.indices) { + var predicateMatched = false + + while (subIndex <= hayList.lastIndex) { + currentIndex = subIndex + val element = hayList[subIndex] + if (this[predicateIndex](element)) { + tempIndices.add(subIndex) + lastMatchedIndex = subIndex + predicateMatched = true + subIndex++ + break } + subIndex++ } - if (remaining.isEmpty()) return@forEach - } - - if (remaining.isEmpty()) { - matchedIndices.putAll(foundMap) - true - } else { - false - } - } - } -} - -fun BytecodePatchContext.wrapperExample() { - fun Method.captureStrings( - desiredStringIndices: Set, - out: MutableMap - ): Boolean { - val remaining = desiredStringIndices.toMutableSet() - val foundMap = mutableMapOf() - - return implementation { - instructions.withIndex().forEach { (index, instruction) -> - val string = (instruction as? ReferenceInstruction)?.reference - .let { it as? StringReference }?.string - ?: return@forEach - - val iterator = remaining.iterator() - while (iterator.hasNext()) { - val desired = iterator.next() - if (desired in string) { - foundMap[desired] = index - iterator.remove() - } + if (!predicateMatched) { + // Restart from next possible first match + firstNeedleIndex = if (tempIndices.isNotEmpty()) tempIndices[0] + 1 else firstNeedleIndex + 1 + matchedAll = false + break } - - if (remaining.isEmpty()) return@forEach } - if (remaining.isEmpty()) { - out += foundMap - true - } else { - false + if (matchedAll) { + _indices.addAll(tempIndices) + return true } } + + return false } - val desiredStringIndices = setOf("fullstring1", "fullstring2", "partialString") - val matchedIndices = mutableMapOf() - - val method = firstMethod { - name == "desiredMethodName" && captureStrings(desiredStringIndices, matchedIndices) + fun first(predicate: T.() -> Boolean) = add { + if (lastMatchedIndex != -1) false + else predicate() } - for ((key, value) in matchedIndices) { - println("Found string '$key' at index $value in method '$method'") + fun after(atLeast: Int = 1, atMost: Int = 1, predicate: T.() -> Boolean) = add { + val distance = currentIndex - lastMatchedIndex + if (distance in atLeast..atMost) predicate() else false } } diff --git a/src/main/kotlin/app/revanced/patcher/patch/BytecodePatchContext.kt b/src/main/kotlin/app/revanced/patcher/patch/BytecodePatchContext.kt index 39c8b47..7fb5ee0 100644 --- a/src/main/kotlin/app/revanced/patcher/patch/BytecodePatchContext.kt +++ b/src/main/kotlin/app/revanced/patcher/patch/BytecodePatchContext.kt @@ -1,7 +1,6 @@ package app.revanced.patcher.patch import app.revanced.patcher.InternalApi -import app.revanced.patcher.Matcher import app.revanced.patcher.PatcherConfig import app.revanced.patcher.PatcherResult import app.revanced.patcher.dex.mutable.MutableClassDef @@ -67,7 +66,7 @@ class BytecodePatchContext internal constructor(private val config: PatcherConfi internal fun mergeExtension(bytecodePatch: BytecodePatch) { bytecodePatch.extensionInputStream?.get()?.use { extensionStream -> RawDexIO.readRawDexFile(extensionStream, 0, null).classes.forEach { classDef -> - val existingClass = lookupMaps.classesByType[classDef.type] ?: run { + val existingClass = lookupMaps.classDefsByType[classDef.type] ?: run { logger.fine { "Adding class \"$classDef\"" } classDefs += classDef @@ -85,7 +84,10 @@ class BytecodePatchContext internal constructor(private val config: PatcherConfi } classDefs -= existingClass + lookupMaps -= existingClass + classDefs += mergedClass + lookupMaps += mergedClass } } } ?: logger.fine("No extension to merge") @@ -97,8 +99,13 @@ class BytecodePatchContext internal constructor(private val config: PatcherConfi * * @return The mutable version of the [ClassDef]. */ - fun ClassDef.mutable(): MutableClassDef = - this as? MutableClassDef ?: also(classDefs::remove).toMutable().also(classDefs::add) + fun ClassDef.mutable(): MutableClassDef = this as? MutableClassDef ?: also { + classDefs -= this + lookupMaps -= this + }.toMutable().also { + classDefs += it + lookupMaps += it + } /** * Navigate a method. @@ -149,8 +156,6 @@ class BytecodePatchContext internal constructor(private val config: PatcherConfi return patchedDexFileResults } - internal val matchers =Map> - override fun close() { try { classDefs.clear() @@ -161,27 +166,48 @@ class BytecodePatchContext internal constructor(private val config: PatcherConfi } internal inner class LookupMaps { - private val _classesByType = mutableMapOf() - val classesByType: Map = _classesByType + private val _classDefsByType = mutableMapOf() + val classDefsByType: Map = _classDefsByType private val _methodsByStrings = mutableMapOf>() val methodsByStrings: Map> = _methodsByStrings + private val _methodsWithString = methodsByStrings.values.flatten().toMutableSet() + val methodsWithString: Set = _methodsWithString + init { classDefs.forEach(::plusAssign) } - operator fun plusAssign(classDef: ClassDef) { - classDef.methods.asSequence().forEach { method -> - method.instructionsOrNull?.asSequence() - ?.filterIsInstance() - ?.map { it.reference } - ?.filterIsInstance() - ?.map { it.string } - ?.forEach { string -> _methodsByStrings.getOrPut(string) { mutableListOf() } += method } - } + private fun ClassDef.forEachString(action: (Method, String) -> Unit) = methods.asSequence().forEach { method -> + method.instructionsOrNull?.asSequence() + ?.filterIsInstance() + ?.map { it.reference } + ?.filterIsInstance() + ?.map { it.string } + ?.forEach { string -> + action(method, string) + } + } - _classesByType[classDef.type] = classDef + operator fun plusAssign(classDef: ClassDef) { + _classDefsByType[classDef.type] = classDef + + classDef.forEachString { method, string -> + _methodsWithString += method + _methodsByStrings.getOrPut(string) { mutableListOf() } += method + } + } + + operator fun minusAssign(classDef: ClassDef) { + _classDefsByType -= classDef.type + + classDef.forEachString { method, string -> + _methodsWithString.remove(method) + + if (_methodsByStrings[string]?.also { it -= method }?.isEmpty() == true) + _methodsByStrings -= string + } } } } diff --git a/src/test/kotlin/app/revanced/patcher/PatcherTest.kt b/src/test/kotlin/app/revanced/patcher/PatcherTest.kt index f291614..a99e5b7 100644 --- a/src/test/kotlin/app/revanced/patcher/PatcherTest.kt +++ b/src/test/kotlin/app/revanced/patcher/PatcherTest.kt @@ -164,27 +164,25 @@ internal object PatcherTest { @Test fun `matches fingerprint`() { - every { patcher.context.bytecodeContext.classDefs } returns ProxyClassDefSet( - mutableListOf( - ImmutableClassDef( - "class", - 0, - null, - null, - null, - null, - null, - listOf( - ImmutableMethod( - "class", - "method", - emptyList(), - "V", - 0, - null, - null, - null, - ), + every { patcher.context.bytecodeContext.classDefs } returns mutableSetOf( + ImmutableClassDef( + "class", + 0, + null, + null, + null, + null, + null, + listOf( + ImmutableMethod( + "class", + "method", + emptyList(), + "V", + 0, + null, + null, + null, ), ), ),