From be14ffdc2d64bc7a317ca6b9bee3914237358caf Mon Sep 17 00:00:00 2001
From: Graham <gpe@openrs2.dev>
Date: Sat, 7 Mar 2020 01:11:34 +0000
Subject: [PATCH] Add custom Attribute for tracking original instruction
 indexes

---
 .../java/dev/openrs2/asm/NopClassVisitor.kt   |  6 ++
 .../dev/openrs2/asm/classpath/ClassPath.kt    | 12 ++++
 .../java/dev/openrs2/asm/classpath/Library.kt | 14 ++++-
 .../java/dev/openrs2/deob/Deobfuscator.kt     |  6 +-
 .../java/dev/openrs2/deob/OriginalPcTable.kt  | 37 +++++++++++++
 .../transform/OriginalPcRestoreTransformer.kt | 55 +++++++++++++++++++
 .../transform/OriginalPcSaveTransformer.kt    | 28 ++++++++++
 7 files changed, 156 insertions(+), 2 deletions(-)
 create mode 100644 asm/src/main/java/dev/openrs2/asm/NopClassVisitor.kt
 create mode 100644 deob/src/main/java/dev/openrs2/deob/OriginalPcTable.kt
 create mode 100644 deob/src/main/java/dev/openrs2/deob/transform/OriginalPcRestoreTransformer.kt
 create mode 100644 deob/src/main/java/dev/openrs2/deob/transform/OriginalPcSaveTransformer.kt

diff --git a/asm/src/main/java/dev/openrs2/asm/NopClassVisitor.kt b/asm/src/main/java/dev/openrs2/asm/NopClassVisitor.kt
new file mode 100644
index 0000000000..5285e79120
--- /dev/null
+++ b/asm/src/main/java/dev/openrs2/asm/NopClassVisitor.kt
@@ -0,0 +1,6 @@
+package dev.openrs2.asm
+
+import org.objectweb.asm.ClassVisitor
+import org.objectweb.asm.Opcodes
+
+object NopClassVisitor : ClassVisitor(Opcodes.ASM7)
diff --git a/asm/src/main/java/dev/openrs2/asm/classpath/ClassPath.kt b/asm/src/main/java/dev/openrs2/asm/classpath/ClassPath.kt
index 9a496535a0..1a7dc447f5 100644
--- a/asm/src/main/java/dev/openrs2/asm/classpath/ClassPath.kt
+++ b/asm/src/main/java/dev/openrs2/asm/classpath/ClassPath.kt
@@ -6,6 +6,7 @@ import dev.openrs2.asm.toBinaryClassName
 import dev.openrs2.common.collect.DisjointSet
 import dev.openrs2.common.collect.ForestDisjointSet
 import org.objectweb.asm.commons.Remapper
+import org.objectweb.asm.tree.AbstractInsnNode
 import org.objectweb.asm.tree.ClassNode
 
 class ClassPath(
@@ -15,6 +16,17 @@ class ClassPath(
 ) {
     private val cache = mutableMapOf<String, ClassMetadata?>()
 
+    /*
+     * XXX(gpe): this is a bit of a hack, as it makes the asm module contain
+     * some details that are only relevant in the deobfuscator. However, I
+     * can't think of a better way of storing this state at the moment - ASM
+     * doesn't have support for attaching arbitrary state to an
+     * AbstractInsnNode. We need to persist the state across all of our
+     * Transformers to avoid adding extraneous labels until the last possible
+     * moment, which would confuse some of our analyses if added earlier.
+     */
+    val originalPcs = mutableMapOf<AbstractInsnNode, Int>()
+
     val libraryClasses: List<ClassMetadata>
         get() {
             val classes = mutableListOf<ClassMetadata>()
diff --git a/asm/src/main/java/dev/openrs2/asm/classpath/Library.kt b/asm/src/main/java/dev/openrs2/asm/classpath/Library.kt
index 36e7eff6dd..a3218a984f 100644
--- a/asm/src/main/java/dev/openrs2/asm/classpath/Library.kt
+++ b/asm/src/main/java/dev/openrs2/asm/classpath/Library.kt
@@ -1,6 +1,7 @@
 package dev.openrs2.asm.classpath
 
 import com.github.michaelbull.logging.InlineLogger
+import dev.openrs2.asm.NopClassVisitor
 import dev.openrs2.asm.remap
 import dev.openrs2.common.crypto.Pkcs12KeyStore
 import dev.openrs2.common.io.DeterministicJarOutputStream
@@ -76,10 +77,21 @@ class Library constructor() : Iterable<ClassNode> {
             for (clazz in classes.values) {
                 val writer = ClassWriter(0)
 
-                clazz.accept(CheckClassAdapter(writer, true))
+                clazz.accept(writer)
 
                 jar.putNextEntry(JarEntry(clazz.name + CLASS_SUFFIX))
                 jar.write(writer.toByteArray())
+
+                /*
+                 * XXX(gpe): CheckClassAdapter breaks the Label offset
+                 * calculation in the OriginalPcTable's write method, so we do
+                 * a second pass without any attributes to check the class,
+                 * feeding the callbacks into a no-op visitor.
+                 */
+                for (method in clazz.methods) {
+                    method.attrs?.clear()
+                }
+                clazz.accept(CheckClassAdapter(NopClassVisitor, true))
             }
         }
     }
diff --git a/deob/src/main/java/dev/openrs2/deob/Deobfuscator.kt b/deob/src/main/java/dev/openrs2/deob/Deobfuscator.kt
index 4c31c1e9a1..fc5d6223a5 100644
--- a/deob/src/main/java/dev/openrs2/deob/Deobfuscator.kt
+++ b/deob/src/main/java/dev/openrs2/deob/Deobfuscator.kt
@@ -24,6 +24,8 @@ import dev.openrs2.deob.transform.InvokeSpecialTransformer
 import dev.openrs2.deob.transform.MethodOrderTransformer
 import dev.openrs2.deob.transform.OpaquePredicateTransformer
 import dev.openrs2.deob.transform.OriginalNameTransformer
+import dev.openrs2.deob.transform.OriginalPcRestoreTransformer
+import dev.openrs2.deob.transform.OriginalPcSaveTransformer
 import dev.openrs2.deob.transform.OverrideTransformer
 import dev.openrs2.deob.transform.RemapTransformer
 import dev.openrs2.deob.transform.ResetTransformer
@@ -151,6 +153,7 @@ class Deobfuscator(private val input: Path, private val output: Path) {
     companion object {
         private val logger = InlineLogger()
         private val TRANSFORMERS = listOf(
+            OriginalPcSaveTransformer(),
             OriginalNameTransformer(),
             *Bundler.TRANSFORMERS.toTypedArray(),
             ResourceTransformer(),
@@ -173,7 +176,8 @@ class Deobfuscator(private val input: Path, private val output: Path) {
             EmptyClassTransformer(),
             MethodOrderTransformer(),
             VisibilityTransformer(),
-            OverrideTransformer()
+            OverrideTransformer(),
+            OriginalPcRestoreTransformer()
         )
     }
 }
diff --git a/deob/src/main/java/dev/openrs2/deob/OriginalPcTable.kt b/deob/src/main/java/dev/openrs2/deob/OriginalPcTable.kt
new file mode 100644
index 0000000000..91c094d993
--- /dev/null
+++ b/deob/src/main/java/dev/openrs2/deob/OriginalPcTable.kt
@@ -0,0 +1,37 @@
+package dev.openrs2.deob
+
+import org.objectweb.asm.Attribute
+import org.objectweb.asm.ByteVector
+import org.objectweb.asm.ClassWriter
+import org.objectweb.asm.Label
+import org.objectweb.asm.tree.LabelNode
+
+class OriginalPcTable(private val pcs: Map<LabelNode, Int>) : Attribute("OriginalPcTable") {
+    override fun isCodeAttribute(): Boolean {
+        return true
+    }
+
+    override fun isUnknown(): Boolean {
+        return false
+    }
+
+    override fun getLabels(): Array<Label> {
+        return pcs.keys.map(LabelNode::getLabel).toTypedArray()
+    }
+
+    override fun write(
+        classWriter: ClassWriter,
+        code: ByteArray?,
+        codeLength: Int,
+        maxStack: Int,
+        maxLocals: Int
+    ): ByteVector {
+        val buf = ByteVector()
+        buf.putShort(pcs.size)
+        for ((label, pc) in pcs) {
+            buf.putShort(label.label.offset)
+            buf.putShort(pc)
+        }
+        return buf
+    }
+}
diff --git a/deob/src/main/java/dev/openrs2/deob/transform/OriginalPcRestoreTransformer.kt b/deob/src/main/java/dev/openrs2/deob/transform/OriginalPcRestoreTransformer.kt
new file mode 100644
index 0000000000..1ad235fabd
--- /dev/null
+++ b/deob/src/main/java/dev/openrs2/deob/transform/OriginalPcRestoreTransformer.kt
@@ -0,0 +1,55 @@
+package dev.openrs2.deob.transform
+
+import com.github.michaelbull.logging.InlineLogger
+import dev.openrs2.asm.classpath.ClassPath
+import dev.openrs2.asm.classpath.Library
+import dev.openrs2.asm.removeDeadCode
+import dev.openrs2.asm.transform.Transformer
+import dev.openrs2.deob.OriginalPcTable
+import org.objectweb.asm.tree.ClassNode
+import org.objectweb.asm.tree.LabelNode
+import org.objectweb.asm.tree.MethodNode
+
+class OriginalPcRestoreTransformer : Transformer() {
+    private var originalPcsRestored = 0
+
+    override fun preTransform(classPath: ClassPath) {
+        originalPcsRestored = 0
+    }
+
+    override fun transformCode(classPath: ClassPath, library: Library, clazz: ClassNode, method: MethodNode): Boolean {
+        method.removeDeadCode(clazz.name)
+
+        val pcs = mutableMapOf<LabelNode, Int>()
+
+        for (insn in method.instructions) {
+            if (insn.opcode == -1) {
+                continue
+            }
+
+            val originalPc = classPath.originalPcs[insn] ?: continue
+
+            val label = LabelNode()
+            method.instructions.insertBefore(insn, label)
+
+            pcs[label] = originalPc
+
+            originalPcsRestored++
+        }
+
+        if (method.attrs == null) {
+            method.attrs = mutableListOf()
+        }
+        method.attrs.add(OriginalPcTable(pcs))
+
+        return false
+    }
+
+    override fun postTransform(classPath: ClassPath) {
+        logger.info { "Restored $originalPcsRestored original instruction indexes" }
+    }
+
+    companion object {
+        private val logger = InlineLogger()
+    }
+}
diff --git a/deob/src/main/java/dev/openrs2/deob/transform/OriginalPcSaveTransformer.kt b/deob/src/main/java/dev/openrs2/deob/transform/OriginalPcSaveTransformer.kt
new file mode 100644
index 0000000000..1e181e7962
--- /dev/null
+++ b/deob/src/main/java/dev/openrs2/deob/transform/OriginalPcSaveTransformer.kt
@@ -0,0 +1,28 @@
+package dev.openrs2.deob.transform
+
+import com.github.michaelbull.logging.InlineLogger
+import dev.openrs2.asm.classpath.ClassPath
+import dev.openrs2.asm.classpath.Library
+import dev.openrs2.asm.transform.Transformer
+import org.objectweb.asm.tree.ClassNode
+import org.objectweb.asm.tree.MethodNode
+
+class OriginalPcSaveTransformer : Transformer() {
+    override fun transformCode(classPath: ClassPath, library: Library, clazz: ClassNode, method: MethodNode): Boolean {
+        for ((pc, insn) in method.instructions.withIndex()) {
+            if (insn.opcode != -1) {
+                classPath.originalPcs[insn] = pc
+            }
+        }
+
+        return false
+    }
+
+    override fun postTransform(classPath: ClassPath) {
+        logger.info { "Saved ${classPath.originalPcs.size} original instruction indexes" }
+    }
+
+    companion object {
+        private val logger = InlineLogger()
+    }
+}