feat: parse and use Kotlin SourceDebugExtension for rename classes and packages

jadx-cli/build.gradle.kts

@@ -18,6 +18,7 @@ dependencies {
 	runtimeOnly(project(":jadx-plugins:jadx-smali-input"))
 	runtimeOnly(project(":jadx-plugins:jadx-rename-mappings"))
 	runtimeOnly(project(":jadx-plugins:jadx-kotlin-metadata"))
+	runtimeOnly(project(":jadx-plugins:jadx-kotlin-source-debug-extension"))
 	runtimeOnly(project(":jadx-plugins:jadx-script:jadx-script-plugin"))
 	runtimeOnly(project(":jadx-plugins:jadx-xapk-input"))
 	runtimeOnly(project(":jadx-plugins:jadx-aab-input"))

jadx-plugins/jadx-kotlin-source-debug-extension/build.gradle.kts

@@ -0,0 +1,13 @@
+plugins {
+	id("jadx-library")
+	id("jadx-kotlin")
+}
+
+dependencies {
+	api(project(":jadx-core"))
+
+	testImplementation(project.project(":jadx-core").sourceSets.getByName("test").output)
+	testImplementation("org.apache.commons:commons-lang3:3.17.0")
+
+	testRuntimeOnly(project(":jadx-plugins:jadx-smali-input"))
+}

jadx-plugins/jadx-kotlin-source-debug-extension/src/main/kotlin/jadx/plugins/kotlin/smap/KotlinSmapOptions.kt

@@ -0,0 +1,24 @@
+package jadx.plugins.kotlin.smap
+
+import jadx.api.plugins.options.impl.BasePluginOptionsBuilder
+import jadx.plugins.kotlin.smap.KotlinSmapPlugin.Companion.PLUGIN_ID
+
+class KotlinSmapOptions : BasePluginOptionsBuilder() {
+	var isClassAliasSourceDbg: Boolean = true
+		private set
+
+	override fun registerOptions() {
+		boolOption(CLASS_ALIAS_SOURCE_DBG_OPT)
+			.description("rename class alias from SourceDebugExtension")
+			.defaultValue(false)
+			.setter { isClassAliasSourceDbg = it }
+	}
+
+	fun isClassSourceDbg(): Boolean {
+		return isClassAliasSourceDbg
+	}
+
+	companion object {
+		const val CLASS_ALIAS_SOURCE_DBG_OPT = "$PLUGIN_ID.class-alias-source-dbg"
+	}
+}

jadx-plugins/jadx-kotlin-source-debug-extension/src/main/kotlin/jadx/plugins/kotlin/smap/KotlinSmapPlugin.kt

@@ -0,0 +1,27 @@
+package jadx.plugins.kotlin.smap
+
+import jadx.api.plugins.JadxPlugin
+import jadx.api.plugins.JadxPluginContext
+import jadx.api.plugins.JadxPluginInfo
+import jadx.plugins.kotlin.smap.pass.KotlinSourceDebugExtensionPass
+
+class KotlinSmapPlugin : JadxPlugin {
+
+	private val options = KotlinSmapOptions()
+
+	override fun getPluginInfo(): JadxPluginInfo {
+		return JadxPluginInfo(PLUGIN_ID, "Kotlin SMAP", "Use kotlin.SourceDebugExtension annotation for rename class alias")
+	}
+
+	override fun init(context: JadxPluginContext) {
+		context.registerOptions(options)
+
+		if (options.isClassSourceDbg()) {
+			context.addPass(KotlinSourceDebugExtensionPass(options))
+		}
+	}
+
+	companion object {
+		const val PLUGIN_ID = "kotlin-smap"
+	}
+}

jadx-plugins/jadx-kotlin-source-debug-extension/src/main/kotlin/jadx/plugins/kotlin/smap/model/ClassAliasRename.kt

@@ -0,0 +1,6 @@
+package jadx.plugins.kotlin.smap.model
+
+data class ClassAliasRename(
+	val pkg: String,
+	val name: String,
+)

jadx-plugins/jadx-kotlin-source-debug-extension/src/main/kotlin/jadx/plugins/kotlin/smap/model/Constants.kt

@@ -0,0 +1,5 @@
+package jadx.plugins.kotlin.smap.model
+
+object Constants {
+	const val KOTLIN_SOURCE_DEBUG_EXTENSION = "Lkotlin/jvm/internal/SourceDebugExtension;"
+}

jadx-plugins/jadx-kotlin-source-debug-extension/src/main/kotlin/jadx/plugins/kotlin/smap/model/SMAP.kt

@@ -0,0 +1,78 @@
+/*
+ * Copyright 2010-2024 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+package jadx.plugins.kotlin.smap.model
+
+import kotlin.math.max
+
+const val KOTLIN_STRATA_NAME = "Kotlin"
+const val KOTLIN_DEBUG_STRATA_NAME = "KotlinDebug"
+
+/**
+ * Represents SMAP as a structure that is contained in `SourceDebugExtension` attribute of a class.
+ * This structure is immutable, we can only query for a result.
+ */
+class SMAP(val fileMappings: List<FileMapping>) {
+	// assuming disjoint line mappings (otherwise binary search can't be used anyway)
+	private val intervals = fileMappings.flatMap { it.lineMappings }.sortedBy { it.dest }
+
+	fun findRange(lineNumber: Int): RangeMapping? {
+		val index = intervals.binarySearch { if (lineNumber in it) 0 else it.dest - lineNumber }
+		return if (index < 0) null else intervals[index]
+	}
+
+	companion object {
+		const val FILE_SECTION = "*F"
+		const val LINE_SECTION = "*L"
+		const val STRATA_SECTION = "*S"
+		const val END = "*E"
+	}
+}
+
+class FileMapping(val name: String, val path: String) {
+	val lineMappings = arrayListOf<RangeMapping>()
+
+	fun toSourceInfo(): SourceInfo =
+		SourceInfo(
+			name,
+			path,
+			lineMappings.fold(0) { result, mapping -> max(result, mapping.source + mapping.range - 1) },
+		)
+
+	fun mapNewLineNumber(source: Int, currentIndex: Int, callSite: SourcePosition?): Int {
+		// Save some space in the SMAP by reusing (or extending if it's the last one) the existing range.
+		// TODO some *other* range may already cover `source`; probably too slow to check them all though.
+		//   Maybe keep the list ordered by `source` and use binary search to locate the closest range on the left?
+		val mapping = lineMappings.lastOrNull()?.takeIf { it.canReuseFor(source, currentIndex, callSite) }
+			?: lineMappings.firstOrNull()?.takeIf { it.canReuseFor(source, currentIndex, callSite) }
+			?: mapNewInterval(source, currentIndex + 1, 1, callSite)
+		mapping.range = max(mapping.range, source - mapping.source + 1)
+		return mapping.mapSourceToDest(source)
+	}
+
+	private fun RangeMapping.canReuseFor(newSource: Int, globalMaxDest: Int, newCallSite: SourcePosition?): Boolean =
+		callSite == newCallSite && (newSource - source) in 0 until range + (if (globalMaxDest in this) 10 else 0)
+
+	fun mapNewInterval(source: Int, dest: Int, range: Int, callSite: SourcePosition? = null): RangeMapping =
+		RangeMapping(source, dest, range, callSite, parent = this).also { lineMappings.add(it) }
+}
+
+data class RangeMapping(val source: Int, val dest: Int, var range: Int, val callSite: SourcePosition?, val parent: FileMapping) {
+	operator fun contains(destLine: Int): Boolean =
+		dest <= destLine && destLine < dest + range
+
+	fun hasMappingForSource(sourceLine: Int): Boolean =
+		source <= sourceLine && sourceLine < source + range
+
+	fun mapDestToSource(destLine: Int): SourcePosition =
+		SourcePosition(source + (destLine - dest), parent.name, parent.path)
+
+	fun mapSourceToDest(sourceLine: Int): Int =
+		dest + (sourceLine - source)
+}
+
+val RangeMapping.toRange: IntRange
+	get() = dest until dest + range
+
+data class SourcePosition(val line: Int, val file: String, val path: String)

jadx-plugins/jadx-kotlin-source-debug-extension/src/main/kotlin/jadx/plugins/kotlin/smap/model/SourceInfo.kt

@@ -0,0 +1,7 @@
+package jadx.plugins.kotlin.smap.model
+
+data class SourceInfo(
+	val sourceFileName: String?,
+	val pathOrCleanFQN: String,
+	val linesInFile: Int,
+)

jadx-plugins/jadx-kotlin-source-debug-extension/src/main/kotlin/jadx/plugins/kotlin/smap/pass/KotlinSourceDebugExtensionPass.kt

@@ -0,0 +1,39 @@
+package jadx.plugins.kotlin.smap.pass
+
+import jadx.api.plugins.pass.JadxPassInfo
+import jadx.api.plugins.pass.impl.OrderedJadxPassInfo
+import jadx.api.plugins.pass.types.JadxPreparePass
+import jadx.core.dex.attributes.AFlag
+import jadx.core.dex.nodes.RootNode
+import jadx.plugins.kotlin.smap.KotlinSmapOptions
+import jadx.plugins.kotlin.smap.utils.KotlinSmapUtils
+
+class KotlinSourceDebugExtensionPass(
+	private val options: KotlinSmapOptions,
+) : JadxPreparePass {
+
+	override fun getInfo(): JadxPassInfo {
+		return OrderedJadxPassInfo(
+			"SourceDebugExtensionPrepare",
+			"Use kotlin.jvm.internal.SourceDebugExtension annotation to rename class & package",
+		)
+			.before("RenameVisitor")
+	}
+
+	override fun init(root: RootNode) {
+		if (options.isClassAliasSourceDbg) {
+			for (cls in root.classes) {
+				if (cls.contains(AFlag.DONT_RENAME)) {
+					continue
+				}
+
+				// rename class & package
+				val kotlinCls = KotlinSmapUtils.getClassAlias(cls)
+				if (kotlinCls != null) {
+					cls.rename(kotlinCls.name)
+					cls.packageNode.rename(kotlinCls.pkg)
+				}
+			}
+		}
+	}
+}

jadx-plugins/jadx-kotlin-source-debug-extension/src/main/kotlin/jadx/plugins/kotlin/smap/utils/Extensions.kt

@@ -0,0 +1,24 @@
+@file:Suppress("UNCHECKED_CAST")
+
+package jadx.plugins.kotlin.smap.utils
+
+import jadx.api.plugins.input.data.annotations.EncodedType
+import jadx.api.plugins.input.data.annotations.EncodedValue
+import jadx.api.plugins.input.data.annotations.IAnnotation
+import jadx.core.dex.nodes.ClassNode
+import jadx.plugins.kotlin.smap.model.Constants
+import jadx.plugins.kotlin.smap.model.SMAP
+
+fun ClassNode.getSourceDebugExtension(): SMAP? {
+	val annotation: IAnnotation? = getAnnotation(Constants.KOTLIN_SOURCE_DEBUG_EXTENSION)
+	return annotation?.run {
+		val smapParser = SMAPParser.parseOrNull(getParamsAsList("value")?.get(0)?.value.toString())
+		return smapParser
+	}
+}
+
+private fun IAnnotation.getParamsAsList(paramName: String): List<EncodedValue>? {
+	val encodedValue = values[paramName]
+		?.takeIf { it.type == EncodedType.ENCODED_ARRAY && it.value is List<*> }
+	return encodedValue?.value?.let { it as List<EncodedValue> }
+}

jadx-plugins/jadx-kotlin-source-debug-extension/src/main/kotlin/jadx/plugins/kotlin/smap/utils/KotlinSmapUtils.kt

@@ -0,0 +1,72 @@
+package jadx.plugins.kotlin.smap.utils
+
+import jadx.core.deobf.NameMapper
+import jadx.core.dex.attributes.nodes.RenameReasonAttr
+import jadx.core.dex.nodes.ClassNode
+import jadx.core.utils.Utils
+import jadx.plugins.kotlin.smap.model.ClassAliasRename
+import jadx.plugins.kotlin.smap.model.SMAP
+import org.slf4j.Logger
+import org.slf4j.LoggerFactory
+import kotlin.jvm.java
+
+object KotlinSmapUtils {
+
+	val LOG: Logger = LoggerFactory.getLogger(KotlinSmapUtils::class.java)
+
+	@JvmStatic
+	fun getClassAlias(cls: ClassNode): ClassAliasRename? {
+		val annotation = cls.getSourceDebugExtension() ?: return null
+		return getClassAlias(cls, annotation)
+	}
+
+	private fun getClassAlias(cls: ClassNode, annotation: SMAP): ClassAliasRename? {
+		val firstValue = annotation.fileMappings[0].path.replace("/", ".")
+		try {
+			val clsName = firstValue.trim()
+				.takeUnless(String::isEmpty)
+				?.let(Utils::cleanObjectName)
+				?: return null
+
+			val alias = splitAndCheckClsName(cls, clsName)
+			if (alias != null) {
+				RenameReasonAttr.forNode(cls).append("from SourceDebugExtension")
+				return alias
+			}
+		} catch (e: Exception) {
+			LOG.error("Failed to parse SourceDebugExtension", e)
+		}
+		return null
+	}
+
+	// Don't use ClassInfo facility to not pollute class into cache
+	private fun splitAndCheckClsName(originCls: ClassNode, fullClsName: String): ClassAliasRename? {
+		if (!NameMapper.isValidFullIdentifier(fullClsName)) {
+			return null
+		}
+		val pkg: String
+		val name: String
+		val dot = fullClsName.lastIndexOf('.')
+		if (dot == -1) {
+			pkg = ""
+			name = fullClsName
+		} else {
+			pkg = fullClsName.substring(0, dot)
+			name = fullClsName.substring(dot + 1)
+		}
+		val originClsInfo = originCls.classInfo
+		val originName = originClsInfo.shortName
+		if (originName == name || name.contains("$") ||
+			!NameMapper.isValidIdentifier(name) || pkg.startsWith("java.")
+		) {
+			return null
+		}
+		val newClsNode = originCls.root().resolveClass(fullClsName)
+		return if (newClsNode != null) {
+			// class with alias name already exist
+			null
+		} else {
+			ClassAliasRename(pkg, name)
+		}
+	}
+}