fix: Merge all extensions before initializing lookup maps

This commit is contained in:
oSumAtrIX 2024-07-26 04:45:31 +02:00
parent 736b3eebbf
commit 8c4dd5b3a3
3 changed files with 102 additions and 54 deletions

View file

@ -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<Patch<*>, 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.

View file

@ -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<Set<PatcherResult.PatchedDexFile>> {
class BytecodePatchContext internal constructor(private val config: PatcherConfig) :
PatchContext<Set<PatcherResult.PatchedDexFile>>,
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<String, ClassDef>().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()
}
}
/**

View file

@ -85,6 +85,31 @@ sealed class Patch<C : PatchContext<*>>(
override fun toString() = name ?: "Patch"
}
internal fun Patch<*>.anyRecursively(
visited: MutableSet<Patch<*>> = 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<Patch<*>>.forEachRecursively(
visited: MutableSet<Patch<*>> = 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<C : PatchContext<*>>(
internal abstract fun build(): Patch<C>
}
/**
* 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 : PatchBuilder<*>> 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 : PatchBuilder<*>> 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].
*