From 8c4dd5b3a309077fa9a3827b4931fc28b0517809 Mon Sep 17 00:00:00 2001 From: oSumAtrIX Date: Fri, 26 Jul 2024 04:45:31 +0200 Subject: [PATCH] fix: Merge all extensions before initializing lookup maps --- .../kotlin/app/revanced/patcher/Patcher.kt | 16 ++- .../patcher/patch/BytecodePatchContext.kt | 21 +++- .../app/revanced/patcher/patch/Patch.kt | 119 +++++++++++------- 3 files changed, 102 insertions(+), 54 deletions(-) diff --git a/src/main/kotlin/app/revanced/patcher/Patcher.kt b/src/main/kotlin/app/revanced/patcher/Patcher.kt index 214f7b4..5c020ed 100644 --- a/src/main/kotlin/app/revanced/patcher/Patcher.kt +++ b/src/main/kotlin/app/revanced/patcher/Patcher.kt @@ -39,9 +39,6 @@ class Patcher(private val config: PatcherConfig) : Closeable { patch.addRecursively() } - fun Patch<*>.anyRecursively(predicate: (Patch<*>) -> Boolean): Boolean = - predicate(this) || dependencies.any { dependency -> dependency.anyRecursively(predicate) } - context.allPatches.let { allPatches -> // Check, if what kind of resource mode is required. config.resourceMode = if (allPatches.any { patch -> patch.anyRecursively { it is ResourcePatch } }) { @@ -99,6 +96,17 @@ class Patcher(private val config: PatcherConfig) : Closeable { context.resourceContext.decodeResources(config.resourceMode) } + logger.info("Merging extensions") + + context.executablePatches.forEachRecursively { patch -> + if (patch is BytecodePatch && patch.extension != null) { + context.bytecodeContext.merge(patch.extension) + } + } + + // Initialize lookup maps. + context.bytecodeContext.lookupMaps + logger.info("Executing patches") val executedPatches = LinkedHashMap, PatchResult>() @@ -146,7 +154,7 @@ class Patcher(private val config: PatcherConfig) : Closeable { } } - override fun close() = context.bytecodeContext.lookupMaps.close() + override fun close() = context.bytecodeContext.close() /** * Compile and save patched APK files. diff --git a/src/main/kotlin/app/revanced/patcher/patch/BytecodePatchContext.kt b/src/main/kotlin/app/revanced/patcher/patch/BytecodePatchContext.kt index b5671ea..56bd6d2 100644 --- a/src/main/kotlin/app/revanced/patcher/patch/BytecodePatchContext.kt +++ b/src/main/kotlin/app/revanced/patcher/patch/BytecodePatchContext.kt @@ -31,7 +31,9 @@ import java.util.logging.Logger * @param config The [PatcherConfig] used to create this context. */ @Suppress("MemberVisibilityCanBePrivate") -class BytecodePatchContext internal constructor(private val config: PatcherConfig) : PatchContext> { +class BytecodePatchContext internal constructor(private val config: PatcherConfig) : + PatchContext>, + Closeable { private val logger = Logger.getLogger(BytecodePatchContext::class.java.name) /** @@ -57,6 +59,13 @@ class BytecodePatchContext internal constructor(private val config: PatcherConfi */ internal val lookupMaps by lazy { LookupMaps(classes) } + /** + * A map for lookup by [merge]. + */ + internal val classesByType = mutableMapOf().apply { + classes.forEach { classDef -> put(classDef.type, classDef) } + } + /** * Merge an extension to [classes]. * @@ -66,11 +75,11 @@ class BytecodePatchContext internal constructor(private val config: PatcherConfi val extension = extensionInputStream.readAllBytes() RawDexIO.readRawDexFile(extension, 0, null).classes.forEach { classDef -> - val existingClass = lookupMaps.classesByType[classDef.type] ?: run { + val existingClass = classesByType[classDef.type] ?: run { logger.fine("Adding class \"$classDef\"") - lookupMaps.classesByType[classDef.type] = classDef classes += classDef + classesByType[classDef.type] = classDef return@forEach } @@ -254,6 +263,12 @@ class BytecodePatchContext internal constructor(private val config: PatcherConfi methodsByStrings.clear() } } + + override fun close() { + lookupMaps.close() + classesByType.clear() + classes.clear() + } } /** diff --git a/src/main/kotlin/app/revanced/patcher/patch/Patch.kt b/src/main/kotlin/app/revanced/patcher/patch/Patch.kt index 79b5d54..53a7736 100644 --- a/src/main/kotlin/app/revanced/patcher/patch/Patch.kt +++ b/src/main/kotlin/app/revanced/patcher/patch/Patch.kt @@ -85,6 +85,31 @@ sealed class Patch>( override fun toString() = name ?: "Patch" } +internal fun Patch<*>.anyRecursively( + visited: MutableSet> = mutableSetOf(), + predicate: (Patch<*>) -> Boolean, +): Boolean { + if (this in visited) return false + + if (predicate(this)) return true + + visited += this + + return dependencies.any { it.anyRecursively(visited, predicate) } +} + +internal fun Iterable>.forEachRecursively( + visited: MutableSet> = mutableSetOf(), + action: (Patch<*>) -> Unit, +): Unit = forEach { + if (it in visited) return@forEach + + visited += it + action(it) + + it.dependencies.forEachRecursively(visited, action) +} + /** * A bytecode patch. * @@ -127,7 +152,6 @@ class BytecodePatch internal constructor( finalizeBlock, ) { override fun execute(context: PatcherContext) = with(context.bytecodeContext) { - extension?.let(::merge) fingerprints.forEach { it.match(this) } execute(this) @@ -332,6 +356,16 @@ sealed class PatchBuilder>( internal abstract fun build(): Patch } +/** + * Builds a [Patch]. + * + * @param B The [PatchBuilder] to build the patch with. + * @param block The block to build the patch. + * + * @return The built [Patch]. + */ +private fun > B.buildPatch(block: B.() -> Unit = {}) = apply(block).build() + /** * A [BytecodePatchBuilder] builder. * @@ -379,9 +413,10 @@ class BytecodePatchBuilder internal constructor( * * @param extension The name of the extension resource. */ + @Suppress("NOTHING_TO_INLINE") inline fun extendWith(extension: String) = apply { this.extension = object {}.javaClass.classLoader.getResourceAsStream(extension) - ?: throw PatchException("Extension resource \"$extension\" not found") + ?: throw PatchException("Extension \"$extension\" not found") } override fun build() = BytecodePatch( @@ -398,6 +433,24 @@ class BytecodePatchBuilder internal constructor( ) } +/** + * Create a new [BytecodePatch]. + * + * @param name The name of the patch. + * If null, the patch is named "Patch" and will not be loaded by [PatchLoader]. + * @param description The description of the patch. + * @param use Weather or not the patch should be used. + * @param block The block to build the patch. + * + * @return The created [BytecodePatch]. + */ +fun bytecodePatch( + name: String? = null, + description: String? = null, + use: Boolean = true, + block: BytecodePatchBuilder.() -> Unit = {}, +) = BytecodePatchBuilder(name, description, use).buildPatch(block) as BytecodePatch + /** * A [RawResourcePatch] builder. * @@ -425,6 +478,23 @@ class RawResourcePatchBuilder internal constructor( ) } +/** + * Create a new [RawResourcePatch]. + * + * @param name The name of the patch. + * If null, the patch is named "Patch" and will not be loaded by [PatchLoader]. + * @param description The description of the patch. + * @param use Weather or not the patch should be used. + * @param block The block to build the patch. + * @return The created [RawResourcePatch]. + */ +fun rawResourcePatch( + name: String? = null, + description: String? = null, + use: Boolean = true, + block: RawResourcePatchBuilder.() -> Unit = {}, +) = RawResourcePatchBuilder(name, description, use).buildPatch(block) as RawResourcePatch + /** * A [ResourcePatch] builder. * @@ -452,51 +522,6 @@ class ResourcePatchBuilder internal constructor( ) } -/** - * Builds a [Patch]. - * - * @param B The [PatchBuilder] to build the patch with. - * @param block The block to build the patch. - * - * @return The built [Patch]. - */ -private fun > B.buildPatch(block: B.() -> Unit = {}) = apply(block).build() - -/** - * Create a new [BytecodePatch]. - * - * @param name The name of the patch. - * If null, the patch is named "Patch" and will not be loaded by [PatchLoader]. - * @param description The description of the patch. - * @param use Weather or not the patch should be used. - * @param block The block to build the patch. - * - * @return The created [BytecodePatch]. - */ -fun bytecodePatch( - name: String? = null, - description: String? = null, - use: Boolean = true, - block: BytecodePatchBuilder.() -> Unit = {}, -) = BytecodePatchBuilder(name, description, use).buildPatch(block) as BytecodePatch - -/** - * Create a new [RawResourcePatch]. - * - * @param name The name of the patch. - * If null, the patch is named "Patch" and will not be loaded by [PatchLoader]. - * @param description The description of the patch. - * @param use Weather or not the patch should be used. - * @param block The block to build the patch. - * @return The created [RawResourcePatch]. - */ -fun rawResourcePatch( - name: String? = null, - description: String? = null, - use: Boolean = true, - block: RawResourcePatchBuilder.() -> Unit = {}, -) = RawResourcePatchBuilder(name, description, use).buildPatch(block) as RawResourcePatch - /** * Create a new [ResourcePatch]. *