Bridge Method Fixes (#111)

* Detect synthetic bridges not marked as bridges, and add back flags to produced bytecode

* Remove debug check

* Remove more test code

* Remove unneeded change to `TranslationClassVisitor`

6 files changed, 182 insertions, 21 deletions

diff --git a/src/main/java/cuchaz/enigma/CompiledSourceTypeLoader.java b/src/main/java/cuchaz/enigma/CompiledSourceTypeLoader.java
index b1a8cd58..c746abed 100644
--- a/src/main/java/cuchaz/enigma/CompiledSourceTypeLoader.java
+++ b/src/main/java/cuchaz/enigma/CompiledSourceTypeLoader.java
@@ -15,6 +15,7 @@ import com.google.common.collect.Lists;
 import com.strobel.assembler.metadata.Buffer;
 import com.strobel.assembler.metadata.ITypeLoader;
 import cuchaz.enigma.translation.representation.entry.ClassEntry;
+import org.objectweb.asm.ClassVisitor;
 import org.objectweb.asm.ClassWriter;
 import org.objectweb.asm.Opcodes;
 import org.objectweb.asm.tree.AbstractInsnNode;
@@ -23,18 +24,25 @@ import org.objectweb.asm.tree.MethodInsnNode;
 import org.objectweb.asm.tree.MethodNode;
 
 import java.util.Collection;
+import java.util.LinkedList;
 import java.util.List;
+import java.util.function.Function;
 
 public class CompiledSourceTypeLoader extends CachingTypeLoader {
         //Store one instance as the classpath shouldn't change during load
         private static final ITypeLoader CLASSPATH_TYPE_LOADER = new CachingClasspathTypeLoader();
 
         private final CompiledSource compiledSource;
+        private final LinkedList<Function<ClassVisitor, ClassVisitor>> visitors = new LinkedList<>();
 
         public CompiledSourceTypeLoader(CompiledSource compiledSource) {
                 this.compiledSource = compiledSource;
         }
 
+        public void addVisitor(Function<ClassVisitor, ClassVisitor> visitor) {
+                this.visitors.addFirst(visitor);
+        }
+
         @Override
         protected byte[] doLoad(String className) {
                 byte[] data = loadType(className);
@@ -66,7 +74,13 @@ public class CompiledSourceTypeLoader extends CachingTypeLoader {
                 removeRedundantClassCalls(node);
 
                 ClassWriter writer = new ClassWriter(0);
-                node.accept(writer);
+
+                ClassVisitor visitor = writer;
+                for (Function<ClassVisitor, ClassVisitor> visitorFunction : this.visitors) {
+                        visitor = visitorFunction.apply(visitor);
+                }
+
+                node.accept(visitor);
 
                 // we have a transformed class!
                 return writer.toByteArray();
diff --git a/src/main/java/cuchaz/enigma/Deobfuscator.java b/src/main/java/cuchaz/enigma/Deobfuscator.java
index 843c7614..47cd05ce 100644
--- a/src/main/java/cuchaz/enigma/Deobfuscator.java
+++ b/src/main/java/cuchaz/enigma/Deobfuscator.java
@@ -24,6 +24,7 @@ import cuchaz.enigma.analysis.IndexTreeBuilder;
 import cuchaz.enigma.analysis.ParsedJar;
 import cuchaz.enigma.analysis.index.JarIndex;
 import cuchaz.enigma.api.EnigmaPlugin;
+import cuchaz.enigma.bytecode.translators.SourceFixVisitor;
 import cuchaz.enigma.bytecode.translators.TranslationClassVisitor;
 import cuchaz.enigma.translation.Translatable;
 import cuchaz.enigma.translation.Translator;
@@ -81,7 +82,10 @@ public class Deobfuscator {
 
                 listener.accept("Preparing...");
 
-                this.obfSourceProvider = new SourceProvider(SourceProvider.createSettings(), new CompiledSourceTypeLoader(parsedJar));
+                CompiledSourceTypeLoader typeLoader = new CompiledSourceTypeLoader(parsedJar);
+                typeLoader.addVisitor(visitor -> new SourceFixVisitor(Opcodes.ASM5, visitor, jarIndex));
+
+                this.obfSourceProvider = new SourceProvider(SourceProvider.createSettings(), typeLoader);
 
                 // init mappings
                 mapper = new EntryRemapper(jarIndex);
@@ -230,16 +234,19 @@ public class Deobfuscator {
                 }
 
                 //create a common instance outside the loop as mappings shouldn't be changing while this is happening
+                CompiledSourceTypeLoader typeLoader = new CompiledSourceTypeLoader(translatedClasses::get);
+                typeLoader.addVisitor(visitor -> new SourceFixVisitor(Opcodes.ASM5, visitor, jarIndex));
+
                 //synchronized to make sure the parallelStream doesn't CME with the cache
-                ITypeLoader typeLoader = new SynchronizedTypeLoader(new CompiledSourceTypeLoader(translatedClasses::get));
+                ITypeLoader synchronizedTypeLoader = new SynchronizedTypeLoader(typeLoader);
 
-                MetadataSystem metadataSystem = new Deobfuscator.NoRetryMetadataSystem(typeLoader);
+                MetadataSystem metadataSystem = new Deobfuscator.NoRetryMetadataSystem(synchronizedTypeLoader);
 
                 //ensures methods are loaded on classload and prevents race conditions
                 metadataSystem.setEagerMethodLoadingEnabled(true);
 
                 DecompilerSettings settings = SourceProvider.createSettings();
-                SourceProvider sourceProvider = new SourceProvider(settings, typeLoader, metadataSystem);
+                SourceProvider sourceProvider = new SourceProvider(settings, synchronizedTypeLoader, metadataSystem);
 
                 AtomicInteger count = new AtomicInteger();
 
diff --git a/src/main/java/cuchaz/enigma/analysis/index/BridgeMethodIndex.java b/src/main/java/cuchaz/enigma/analysis/index/BridgeMethodIndex.java
index 8f6bd462..de2ba1e5 100644
--- a/src/main/java/cuchaz/enigma/analysis/index/BridgeMethodIndex.java
+++ b/src/main/java/cuchaz/enigma/analysis/index/BridgeMethodIndex.java
@@ -1,22 +1,32 @@
 package cuchaz.enigma.analysis.index;
 
 import com.google.common.collect.Maps;
+import com.google.common.collect.Sets;
 import cuchaz.enigma.translation.mapping.EntryResolver;
 import cuchaz.enigma.translation.representation.AccessFlags;
+import cuchaz.enigma.translation.representation.MethodDescriptor;
+import cuchaz.enigma.translation.representation.TypeDescriptor;
+import cuchaz.enigma.translation.representation.entry.ClassEntry;
+import cuchaz.enigma.translation.representation.entry.MethodDefEntry;
 import cuchaz.enigma.translation.representation.entry.MethodEntry;
 
 import javax.annotation.Nullable;
 import java.util.Collection;
+import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 public class BridgeMethodIndex implements JarIndexer {
         private final EntryIndex entryIndex;
+        private final InheritanceIndex inheritanceIndex;
         private final ReferenceIndex referenceIndex;
 
-        private Map<MethodEntry, MethodEntry> accessedToBridge = Maps.newHashMap();
+        private final Set<MethodEntry> bridgeMethods = Sets.newHashSet();
+        private final Map<MethodEntry, MethodEntry> accessedToBridge = Maps.newHashMap();
 
-        public BridgeMethodIndex(EntryIndex entryIndex, ReferenceIndex referenceIndex) {
+        public BridgeMethodIndex(EntryIndex entryIndex, InheritanceIndex inheritanceIndex, ReferenceIndex referenceIndex) {
                 this.entryIndex = entryIndex;
+                this.inheritanceIndex = inheritanceIndex;
                 this.referenceIndex = referenceIndex;
         }
 
@@ -24,21 +34,26 @@ public class BridgeMethodIndex implements JarIndexer {
         public void processIndex(EntryResolver resolver) {
                 // look for access and bridged methods
                 for (MethodEntry methodEntry : entryIndex.getMethods()) {
-                        AccessFlags access = entryIndex.getMethodAccess(methodEntry);
+                        MethodDefEntry methodDefEntry = (MethodDefEntry) methodEntry;
+
+                        AccessFlags access = methodDefEntry.getAccess();
                         if (access == null || !access.isSynthetic()) {
                                 continue;
                         }
 
-                        indexSyntheticMethod(methodEntry, access);
+                        indexSyntheticMethod(methodDefEntry, access);
                 }
         }
 
-        private void indexSyntheticMethod(MethodEntry syntheticMethod, AccessFlags access) {
-                if (access.isBridge()) {
-                        MethodEntry accessedMethod = findAccessMethod(syntheticMethod);
-                        if (accessedMethod != null) {
-                                accessedToBridge.put(accessedMethod, syntheticMethod);
-                        }
+        private void indexSyntheticMethod(MethodDefEntry syntheticMethod, AccessFlags access) {
+                MethodEntry accessedMethod = findAccessMethod(syntheticMethod);
+                if (accessedMethod == null) {
+                        return;
+                }
+
+                if (access.isBridge() || isPotentialBridge(syntheticMethod, accessedMethod)) {
+                        bridgeMethods.add(syntheticMethod);
+                        accessedToBridge.put(accessedMethod, syntheticMethod);
                 }
         }
 
@@ -56,6 +71,60 @@ public class BridgeMethodIndex implements JarIndexer {
                 return referencedMethods.stream().findFirst().orElse(null);
         }
 
+        private boolean isPotentialBridge(MethodDefEntry bridgeMethod, MethodEntry accessedMethod) {
+                // Bridge methods only exist for inheritance purposes, if we're private, final, or static, we cannot be inherited
+                AccessFlags bridgeAccess = bridgeMethod.getAccess();
+                if (bridgeAccess.isPrivate() || bridgeAccess.isFinal() || bridgeAccess.isStatic()) {
+                        return false;
+                }
+
+                MethodDescriptor bridgeDesc = bridgeMethod.getDesc();
+                MethodDescriptor accessedDesc = accessedMethod.getDesc();
+                List<TypeDescriptor> bridgeArguments = bridgeDesc.getArgumentDescs();
+                List<TypeDescriptor> accessedArguments = accessedDesc.getArgumentDescs();
+
+                // A bridge method will always have the same number of arguments
+                if (bridgeArguments.size() != accessedArguments.size()) {
+                        return false;
+                }
+
+                // Check that all argument types are bridge-compatible
+                for (int i = 0; i < bridgeArguments.size(); i++) {
+                        if (!areTypesBridgeCompatible(bridgeArguments.get(i), accessedArguments.get(i))) {
+                                return false;
+                        }
+                }
+
+                // Check that the return type is bridge-compatible
+                return areTypesBridgeCompatible(bridgeDesc.getReturnDesc(), accessedDesc.getReturnDesc());
+        }
+
+        private boolean areTypesBridgeCompatible(TypeDescriptor bridgeDesc, TypeDescriptor accessedDesc) {
+                if (bridgeDesc.equals(accessedDesc)) {
+                        return true;
+                }
+
+                // Either the descs will be equal, or they are both types and different through a generic
+                if (bridgeDesc.isType() && accessedDesc.isType()) {
+                        ClassEntry bridgeType = bridgeDesc.getTypeEntry();
+                        ClassEntry accessedType = accessedDesc.getTypeEntry();
+
+                        // If the given types are completely unrelated to each other, this can't be bridge compatible
+                        InheritanceIndex.Relation relation = inheritanceIndex.computeClassRelation(accessedType, bridgeType);
+                        return relation != InheritanceIndex.Relation.UNRELATED;
+                }
+
+                return false;
+        }
+
+        public boolean isBridgeMethod(MethodEntry entry) {
+                return bridgeMethods.contains(entry);
+        }
+
+        public boolean isAccessedByBridge(MethodEntry entry) {
+                return accessedToBridge.containsKey(entry);
+        }
+
         @Nullable
         public MethodEntry getBridgeFromAccessed(MethodEntry entry) {
                 return accessedToBridge.get(entry);
diff --git a/src/main/java/cuchaz/enigma/analysis/index/InheritanceIndex.java b/src/main/java/cuchaz/enigma/analysis/index/InheritanceIndex.java
index 17bed54c..1b8d9a8d 100644
--- a/src/main/java/cuchaz/enigma/analysis/index/InheritanceIndex.java
+++ b/src/main/java/cuchaz/enigma/analysis/index/InheritanceIndex.java
@@ -22,13 +22,23 @@ import java.util.LinkedList;
 import java.util.Set;
 
 public class InheritanceIndex implements JarIndexer {
+        private final EntryIndex entryIndex;
+
         private Multimap<ClassEntry, ClassEntry> classParents = HashMultimap.create();
         private Multimap<ClassEntry, ClassEntry> classChildren = HashMultimap.create();
 
+        public InheritanceIndex(EntryIndex entryIndex) {
+                this.entryIndex = entryIndex;
+        }
+
         @Override
         public void indexClass(ClassDefEntry classEntry) {
+                if (classEntry.isJre()) {
+                        return;
+                }
+
                 ClassEntry superClass = classEntry.getSuperClass();
-                if (superClass != null) {
+                if (superClass != null && !superClass.getName().equals("java/lang/Object")) {
                         indexParent(classEntry, superClass);
                 }
 
@@ -38,9 +48,6 @@ public class InheritanceIndex implements JarIndexer {
         }
 
         private void indexParent(ClassEntry childEntry, ClassEntry parentEntry) {
-                if (childEntry.isJre() || parentEntry.isJre()) {
-                        return;
-                }
                 classParents.put(childEntry, parentEntry);
                 classChildren.put(parentEntry, childEntry);
         }
@@ -70,6 +77,21 @@ public class InheritanceIndex implements JarIndexer {
                 return ancestors;
         }
 
+        public Relation computeClassRelation(ClassEntry classEntry, ClassEntry potentialAncestor) {
+                if (potentialAncestor.getName().equals("java/lang/Object")) return Relation.RELATED;
+                if (!entryIndex.hasClass(classEntry)) return Relation.UNKNOWN;
+
+                for (ClassEntry ancestor : getAncestors(classEntry)) {
+                        if (potentialAncestor.equals(ancestor)) {
+                                return Relation.RELATED;
+                        } else if (!entryIndex.hasClass(ancestor)) {
+                                return Relation.UNKNOWN;
+                        }
+                }
+
+                return Relation.UNRELATED;
+        }
+
         public boolean isParent(ClassEntry classEntry) {
                 return classChildren.containsKey(classEntry);
         }
@@ -78,4 +100,10 @@ public class InheritanceIndex implements JarIndexer {
                 Collection<ClassEntry> parents = classParents.get(classEntry);
                 return parents != null && !parents.isEmpty();
         }
+
+        public enum Relation {
+                RELATED,
+                UNRELATED,
+                UNKNOWN
+        }
 }
diff --git a/src/main/java/cuchaz/enigma/analysis/index/JarIndex.java b/src/main/java/cuchaz/enigma/analysis/index/JarIndex.java
index 9b21cbae..cb58fced 100644
--- a/src/main/java/cuchaz/enigma/analysis/index/JarIndex.java
+++ b/src/main/java/cuchaz/enigma/analysis/index/JarIndex.java
@@ -46,9 +46,9 @@ public class JarIndex implements JarIndexer {
 
         public static JarIndex empty() {
                 EntryIndex entryIndex = new EntryIndex();
-                InheritanceIndex inheritanceIndex = new InheritanceIndex();
+                InheritanceIndex inheritanceIndex = new InheritanceIndex(entryIndex);
                 ReferenceIndex referenceIndex = new ReferenceIndex();
-                BridgeMethodIndex bridgeMethodIndex = new BridgeMethodIndex(entryIndex, referenceIndex);
+                BridgeMethodIndex bridgeMethodIndex = new BridgeMethodIndex(entryIndex, inheritanceIndex, referenceIndex);
                 return new JarIndex(entryIndex, inheritanceIndex, referenceIndex, bridgeMethodIndex);
         }
 
diff --git a/src/main/java/cuchaz/enigma/bytecode/translators/SourceFixVisitor.java b/src/main/java/cuchaz/enigma/bytecode/translators/SourceFixVisitor.java
new file mode 100644
index 00000000..99eef6ae
--- /dev/null
+++ b/src/main/java/cuchaz/enigma/bytecode/translators/SourceFixVisitor.java
@@ -0,0 +1,43 @@
+package cuchaz.enigma.bytecode.translators;
+
+import cuchaz.enigma.analysis.index.BridgeMethodIndex;
+import cuchaz.enigma.analysis.index.JarIndex;
+import cuchaz.enigma.translation.representation.entry.ClassDefEntry;
+import cuchaz.enigma.translation.representation.entry.MethodDefEntry;
+import cuchaz.enigma.translation.representation.entry.MethodEntry;
+import org.objectweb.asm.ClassVisitor;
+import org.objectweb.asm.MethodVisitor;
+import org.objectweb.asm.Opcodes;
+
+public class SourceFixVisitor extends ClassVisitor {
+        private final JarIndex index;
+        private ClassDefEntry ownerEntry;
+
+        public SourceFixVisitor(int api, ClassVisitor visitor, JarIndex index) {
+                super(api, visitor);
+                this.index = index;
+        }
+
+        @Override
+        public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {
+                ownerEntry = ClassDefEntry.parse(access, name, signature, superName, interfaces);
+                super.visit(version, access, name, signature, superName, interfaces);
+        }
+
+        @Override
+        public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
+                MethodDefEntry methodEntry = MethodDefEntry.parse(ownerEntry, access, name, descriptor, signature);
+
+                BridgeMethodIndex bridgeIndex = index.getBridgeMethodIndex();
+                if (bridgeIndex.isBridgeMethod(methodEntry)) {
+                        access |= Opcodes.ACC_BRIDGE;
+                } else {
+                        MethodEntry bridgeMethod = bridgeIndex.getBridgeFromAccessed(methodEntry);
+                        if (bridgeMethod != null) {
+                                name = bridgeMethod.getName();
+                        }
+                }
+
+                return super.visitMethod(access, name, descriptor, signature, exceptions);
+        }
+}