From aa472eb9857145b53b49f843406a9764fbb7e5ce Mon Sep 17 00:00:00 2001 From: oSumAtrIX Date: Sun, 27 Oct 2024 16:00:30 +0100 Subject: [PATCH] fix: Merge extension only when patch executes (#315) --- .../kotlin/app/revanced/patcher/Patcher.kt | 12 ++--- .../patcher/patch/BytecodePatchContext.kt | 51 +++++++++---------- .../app/revanced/patcher/patch/Patch.kt | 8 ++- .../app/revanced/patcher/PatcherTest.kt | 2 +- 4 files changed, 36 insertions(+), 37 deletions(-) diff --git a/src/main/kotlin/app/revanced/patcher/Patcher.kt b/src/main/kotlin/app/revanced/patcher/Patcher.kt index 50717a4..e7d4915 100644 --- a/src/main/kotlin/app/revanced/patcher/Patcher.kt +++ b/src/main/kotlin/app/revanced/patcher/Patcher.kt @@ -91,19 +91,15 @@ class Patcher(private val config: PatcherConfig) : Closeable { }.also { executedPatches[this] = it } } - // Prevent from decoding the app manifest twice if it is not needed. + // Prevent decoding the app manifest twice if it is not needed. if (config.resourceMode != ResourcePatchContext.ResourceMode.NONE) { context.resourceContext.decodeResources(config.resourceMode) } - logger.info("Merging extensions") + logger.info("Initializing lookup maps") - with(context.bytecodeContext) { - context.executablePatches.mergeExtensions() - - // Initialize lookup maps. - lookupMaps - } + // Accessing the lazy lookup maps to initialize them. + context.bytecodeContext.lookupMaps logger.info("Executing patches") diff --git a/src/main/kotlin/app/revanced/patcher/patch/BytecodePatchContext.kt b/src/main/kotlin/app/revanced/patcher/patch/BytecodePatchContext.kt index e9e881a..a523511 100644 --- a/src/main/kotlin/app/revanced/patcher/patch/BytecodePatchContext.kt +++ b/src/main/kotlin/app/revanced/patcher/patch/BytecodePatchContext.kt @@ -59,42 +59,33 @@ class BytecodePatchContext internal constructor(private val config: PatcherConfi internal val lookupMaps by lazy { LookupMaps(classes) } /** - * Merge the extensions for this set of patches. + * Merge the extension of this patch. */ - internal fun Set>.mergeExtensions() { - // Lookup map to check if a class exists by its type quickly. - val classesByType = mutableMapOf().apply { - classes.forEach { classDef -> put(classDef.type, classDef) } - } + internal fun BytecodePatch.mergeExtension() { + extension?.use { extensionStream -> + RawDexIO.readRawDexFile(extensionStream, 0, null).classes.forEach { classDef -> + val existingClass = lookupMaps.classesByType[classDef.type] ?: run { + logger.fine("Adding class \"$classDef\"") - forEachRecursively { patch -> - if (patch !is BytecodePatch) return@forEachRecursively + classes += classDef + lookupMaps.classesByType[classDef.type] = classDef - patch.extension?.use { extensionStream -> - RawDexIO.readRawDexFile(extensionStream, 0, null).classes.forEach { classDef -> - val existingClass = classesByType[classDef.type] ?: run { - logger.fine("Adding class \"$classDef\"") + return@forEach + } - classes += classDef - classesByType[classDef.type] = classDef + logger.fine("Class \"$classDef\" exists already. Adding missing methods and fields.") - return@forEach + existingClass.merge(classDef, this@BytecodePatchContext).let { mergedClass -> + // If the class was merged, replace the original class with the merged class. + if (mergedClass === existingClass) { + return@let } - logger.fine("Class \"$classDef\" exists already. Adding missing methods and fields.") - - existingClass.merge(classDef, this@BytecodePatchContext).let { mergedClass -> - // If the class was merged, replace the original class with the merged class. - if (mergedClass === existingClass) { - return@let - } - - classes -= existingClass - classes += mergedClass - } + classes -= existingClass + classes += mergedClass } } - } + } ?: return logger.fine("No extension to merge") } /** @@ -185,6 +176,11 @@ class BytecodePatchContext internal constructor(private val config: PatcherConfi */ internal val methodsByStrings = MethodClassPairsLookupMap() + // Lookup map for fast checking if a class exists by its type. + val classesByType = mutableMapOf().apply { + classes.forEach { classDef -> put(classDef.type, classDef) } + } + init { classes.forEach { classDef -> classDef.methods.forEach { method -> @@ -231,6 +227,7 @@ class BytecodePatchContext internal constructor(private val config: PatcherConfi override fun close() { methodsByStrings.clear() + classesByType.clear() } } diff --git a/src/main/kotlin/app/revanced/patcher/patch/Patch.kt b/src/main/kotlin/app/revanced/patcher/patch/Patch.kt index 6397252..d56acd3 100644 --- a/src/main/kotlin/app/revanced/patcher/patch/Patch.kt +++ b/src/main/kotlin/app/revanced/patcher/patch/Patch.kt @@ -158,7 +158,13 @@ class BytecodePatch internal constructor( finalizeBlock, ) { override fun execute(context: PatcherContext) = with(context.bytecodeContext) { - fingerprints.forEach { it.match(this) } + with(context.bytecodeContext) { + mergeExtension() + } + + fingerprints.forEach { + it.match(this) + } execute(this) } diff --git a/src/test/kotlin/app/revanced/patcher/PatcherTest.kt b/src/test/kotlin/app/revanced/patcher/PatcherTest.kt index 7b3f2b5..2fd8c88 100644 --- a/src/test/kotlin/app/revanced/patcher/PatcherTest.kt +++ b/src/test/kotlin/app/revanced/patcher/PatcherTest.kt @@ -195,7 +195,7 @@ internal object PatcherTest { private operator fun Set>.invoke(): List { every { patcher.context.executablePatches } returns toMutableSet() every { patcher.context.bytecodeContext.lookupMaps } returns LookupMaps(patcher.context.bytecodeContext.classes) - every { with(patcher.context.bytecodeContext) { any>>().mergeExtensions() } } just runs + every { with(patcher.context.bytecodeContext) { any().mergeExtension() } } just runs return runBlocking { patcher().toList() } }