fixed lots of issues with inner class reconstruction, particularly for inner class trees

also fixed lots of issues with reading jars that aren't Minecraft. =P

6 files changed, 156 insertions, 159 deletions

diff --git a/src/cuchaz/enigma/Deobfuscator.java b/src/cuchaz/enigma/Deobfuscator.java
index b7440a72..0b7808da 100644
--- a/src/cuchaz/enigma/Deobfuscator.java
+++ b/src/cuchaz/enigma/Deobfuscator.java
@@ -204,33 +204,36 @@ public class Deobfuscator {
                 }
         }
         
-        public CompilationUnit getSourceTree(String obfClassName) {
-                // is this class deobfuscated?
+        public CompilationUnit getSourceTree(String className) {
+                
+                // we don't know if this class name is obfuscated or deobfuscated
                 // we need to tell the decompiler the deobfuscated name so it doesn't get freaked out
-                // the decompiler only sees the deobfuscated class, so we need to load it by the deobfuscated name
-                String lookupClassName = obfClassName;
-                ClassMapping classMapping = m_mappings.getClassByObf(obfClassName);
-                if (classMapping != null && classMapping.getDeobfName() != null) {
-                        lookupClassName = classMapping.getDeobfName();
-                }
+                // the decompiler only sees classes after deobfuscation, so we need to load it by the deobfuscated name if there is one
                 
-                // is this class even in the jar?
-                if (!m_jarIndex.containsObfClass(new ClassEntry(obfClassName))) {
-                        return null;
+                // first, assume class name is deobf
+                String deobfClassName = className;
+                
+                // if it wasn't actually deobf, then we can find a mapping for it and get the deobf name
+                ClassMapping classMapping = m_mappings.getClassByObf(className);
+                if (classMapping != null && classMapping.getDeobfName() != null) {
+                        deobfClassName = classMapping.getDeobfName();
                 }
                 
                 // set the type loader
-                m_settings.setTypeLoader(new TranslatingTypeLoader(
+                TranslatingTypeLoader loader = new TranslatingTypeLoader(
                         m_jar,
                         m_jarIndex,
                         getTranslator(TranslationDirection.Obfuscating),
                         getTranslator(TranslationDirection.Deobfuscating)
-                ));
+                ); 
+                m_settings.setTypeLoader(loader);
 
                 // see if procyon can find the type
-                TypeReference type = new MetadataSystem(m_settings.getTypeLoader()).lookupType(lookupClassName);
+                TypeReference type = new MetadataSystem(loader).lookupType(deobfClassName);
                 if (type == null) {
-                        throw new Error("Unable to find type: " + lookupClassName + " (obf name: " + obfClassName + ")");
+                        throw new Error(String.format("Unable to find type: %s (deobf: %s)\nTried class names: %s",
+                                className, deobfClassName, loader.getClassNamesToTry(deobfClassName)
+                        ));
                 }
                 TypeDefinition resolvedType = type.resolve();
                 
diff --git a/src/cuchaz/enigma/TranslatingTypeLoader.java b/src/cuchaz/enigma/TranslatingTypeLoader.java
index 26d5e7a7..94ad6ebe 100644
--- a/src/cuchaz/enigma/TranslatingTypeLoader.java
+++ b/src/cuchaz/enigma/TranslatingTypeLoader.java
@@ -13,6 +13,7 @@ package cuchaz.enigma;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.util.List;
 import java.util.Map;
 import java.util.jar.JarEntry;
 import java.util.jar.JarFile;
@@ -24,6 +25,7 @@ import javassist.CtClass;
 import javassist.NotFoundException;
 import javassist.bytecode.Descriptor;
 
+import com.beust.jcommander.internal.Lists;
 import com.google.common.collect.Maps;
 import com.strobel.assembler.metadata.Buffer;
 import com.strobel.assembler.metadata.ClasspathTypeLoader;
@@ -65,19 +67,20 @@ public class TranslatingTypeLoader implements ITypeLoader {
         }
         
         @Override
-        public boolean tryLoadType(String deobfClassName, Buffer out) {
+        public boolean tryLoadType(String className, Buffer out) {
+                
                 // check the cache
                 byte[] data;
-                if (m_cache.containsKey(deobfClassName)) {
-                        data = m_cache.get(deobfClassName);
+                if (m_cache.containsKey(className)) {
+                        data = m_cache.get(className);
                 } else {
-                        data = loadType(deobfClassName);
-                        m_cache.put(deobfClassName, data);
+                        data = loadType(className);
+                        m_cache.put(className, data);
                 }
                 
                 if (data == null) {
                         // chain to default type loader
-                        return m_defaultTypeLoader.tryLoadType(deobfClassName, out);
+                        return m_defaultTypeLoader.tryLoadType(className, out);
                 }
                 
                 // send the class to the decompiler
@@ -88,6 +91,7 @@ public class TranslatingTypeLoader implements ITypeLoader {
         }
         
         public CtClass loadClass(String deobfClassName) {
+                
                 byte[] data = loadType(deobfClassName);
                 if (data == null) {
                         return null;
@@ -104,40 +108,35 @@ public class TranslatingTypeLoader implements ITypeLoader {
                 }
         }
         
-        private byte[] loadType(String deobfClassName) {
-                
-                ClassEntry deobfClassEntry = new ClassEntry(deobfClassName);
-                ClassEntry obfClassEntry = m_obfuscatingTranslator.translateEntry(deobfClassEntry);
-                
-                // is this an inner class referenced directly?
-                ClassEntry obfOuterClassEntry = m_jarIndex.getOuterClass(obfClassEntry);
-                if (obfOuterClassEntry != null) {
-                        // this class doesn't really exist. Reference it by outer$inner instead
-                        System.err.println(String.format("WARNING: class %s referenced by bare inner name instead of via outer class %s", deobfClassName, obfOuterClassEntry));
-                        return null;
-                }
-                
-                /* DEBUG
-                if( !Arrays.asList( "java", "org", "io" ).contains( deobfClassName.split( "/" )[0] ) ) {
-                        System.out.println( String.format( "Looking for %s (%s)", deobfClassEntry.getName(), obfClassEntry.getName() ) );
+        private byte[] loadType(String className) {
+                
+                // NOTE: don't know if class name is obf or deobf
+                ClassEntry classEntry = new ClassEntry(className);
+                ClassEntry obfClassEntry = m_obfuscatingTranslator.translateEntry(classEntry);
+                
+                // is this an inner class referenced directly? (ie trying to load b instead of a$b)
+                if (!obfClassEntry.isInnerClass()) {
+                        List<ClassEntry> classChain = m_jarIndex.getObfClassChain(obfClassEntry);
+                        if (classChain.size() > 1) {
+                                System.err.println(String.format("WARNING: no class %s after inner class reconstruction. Try %s",
+                                        className, obfClassEntry.buildClassEntry(classChain)
+                                ));
+                                return null;
+                        }
                 }
-                */
                 
-                // get the jar entry
-                String classFileName;
-                if (obfClassEntry.isInnerClass()) {
-                        // use just the inner class name for inner classes
-                        classFileName = obfClassEntry.getInnerClassName();
-                } else if (obfClassEntry.getPackageName().equals(Constants.NonePackage)) {
-                        // use the outer class simple name for classes in the none package
-                        classFileName = obfClassEntry.getSimpleName();
-                } else {
-                        // otherwise, just use the class name (ie for classes in packages)
-                        classFileName = obfClassEntry.getName();
+                // is this a class we should even know about?
+                if (!m_jarIndex.containsObfClass(obfClassEntry)) {
+                        return null;
                 }
                 
-                JarEntry entry = m_jar.getJarEntry(classFileName + ".class");
-                if (entry == null) {
+                // DEBUG
+                //System.out.println(String.format("Looking for %s (obf: %s)", classEntry.getName(), obfClassEntry.getName()));
+
+                // find the class in the jar
+                String classInJarName = findClassInJar(obfClassEntry);
+                if (classInJarName == null) {
+                        // couldn't find it
                         return null;
                 }
                 
@@ -145,7 +144,7 @@ public class TranslatingTypeLoader implements ITypeLoader {
                         // read the class file into a buffer
                         ByteArrayOutputStream data = new ByteArrayOutputStream();
                         byte[] buf = new byte[1024 * 1024]; // 1 KiB
-                        InputStream in = m_jar.getInputStream(entry);
+                        InputStream in = m_jar.getInputStream(m_jar.getJarEntry(classInJarName + ".class"));
                         while (true) {
                                 int bytesRead = in.read(buf);
                                 if (bytesRead <= 0) {
@@ -158,15 +157,15 @@ public class TranslatingTypeLoader implements ITypeLoader {
                         buf = data.toByteArray();
                         
                         // load the javassist handle to the raw class
-                        String javaClassFileName = Descriptor.toJavaName(classFileName);
                         ClassPool classPool = new ClassPool();
-                        classPool.insertClassPath(new ByteArrayClassPath(javaClassFileName, buf));
-                        CtClass c = classPool.get(javaClassFileName);
+                        String classInJarJavaName = Descriptor.toJavaName(classInJarName);
+                        classPool.insertClassPath(new ByteArrayClassPath(classInJarJavaName, buf));
+                        CtClass c = classPool.get(classInJarJavaName);
                         
                         c = transformClass(c);
                         
                         // sanity checking
-                        assertClassName(c, deobfClassEntry);
+                        assertClassName(c, classEntry);
                         
                         // DEBUG
                         //Util.writeClass( c );
@@ -178,6 +177,38 @@ public class TranslatingTypeLoader implements ITypeLoader {
                 }
         }
         
+        private String findClassInJar(ClassEntry obfClassEntry) {
+
+                // try to find the class in the jar
+                for (String className : getClassNamesToTry(obfClassEntry)) {
+                        JarEntry jarEntry = m_jar.getJarEntry(className + ".class");
+                        if (jarEntry != null) {
+                                return className;
+                        }
+                }
+                
+                // didn't find it  ;_;
+                return null;
+        }
+        
+        public List<String> getClassNamesToTry(String className) {
+                return getClassNamesToTry(m_obfuscatingTranslator.translateEntry(new ClassEntry(className)));
+        }
+        
+        public List<String> getClassNamesToTry(ClassEntry obfClassEntry) {
+                List<String> classNamesToTry = Lists.newArrayList();
+                classNamesToTry.add(obfClassEntry.getName());
+                if (obfClassEntry.getPackageName().equals(Constants.NonePackage)) {
+                        // taking off the none package, if any
+                        classNamesToTry.add(obfClassEntry.getSimpleName());
+                }
+                if (obfClassEntry.isInnerClass()) {
+                        // try just the inner class name
+                        classNamesToTry.add(obfClassEntry.getInnerClassName());
+                }
+                return classNamesToTry;
+        }
+
         public CtClass transformClass(CtClass c)
         throws IOException, NotFoundException, CannotCompileException {
                 
diff --git a/src/cuchaz/enigma/analysis/JarIndex.java b/src/cuchaz/enigma/analysis/JarIndex.java
index 1afcb76c..e0a8bf52 100644
--- a/src/cuchaz/enigma/analysis/JarIndex.java
+++ b/src/cuchaz/enigma/analysis/JarIndex.java
@@ -12,6 +12,7 @@ package cuchaz.enigma.analysis;
 
 import java.lang.reflect.Modifier;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -156,11 +157,8 @@ public class JarIndex {
                         
                         // step 6: update other indices with inner class info
                         Map<String,String> renames = Maps.newHashMap();
-                        for (Map.Entry<ClassEntry,ClassEntry> mapEntry : m_innerClassesByOuter.entries()) {
-                                ClassEntry outerClassEntry = mapEntry.getKey();
-                                ClassEntry innerClassEntry = mapEntry.getValue();
-                                outerClassEntry = EntryFactory.getChainedOuterClassName(this, outerClassEntry);
-                                String newName = outerClassEntry.getName() + "$" + innerClassEntry.getSimpleName();
+                        for (ClassEntry innerClassEntry : m_innerClassesByOuter.values()) {
+                                String newName = innerClassEntry.buildClassEntry(getObfClassChain(innerClassEntry)).getName();
                                 if (!innerClassEntry.getName().equals(newName)) {
                                         // DEBUG
                                         //System.out.println("REPLACE: " + innerClassEntry.getName() + " WITH " + newName);
@@ -780,4 +778,25 @@ public class JarIndex {
         public MethodEntry getBridgedMethod(MethodEntry bridgeMethodEntry) {
                 return m_bridgedMethods.get(bridgeMethodEntry);
         }
+
+        public List<ClassEntry> getObfClassChain(ClassEntry obfClassEntry) {
+                
+                // build class chain in inner-to-outer order
+                List<ClassEntry> obfClassChain = Lists.newArrayList(obfClassEntry);
+                ClassEntry checkClassEntry = obfClassEntry;
+                while (true) {
+                        ClassEntry obfOuterClassEntry = getOuterClass(checkClassEntry);
+                        if (obfOuterClassEntry != null) {
+                                obfClassChain.add(obfOuterClassEntry);
+                                checkClassEntry = obfOuterClassEntry;
+                        } else {
+                                break;
+                        }
+                }
+                
+                // switch to outer-to-inner order
+                Collections.reverse(obfClassChain);
+                
+                return obfClassChain;
+        }
 }
diff --git a/src/cuchaz/enigma/bytecode/InnerClassWriter.java b/src/cuchaz/enigma/bytecode/InnerClassWriter.java
index 3bebd176..dd21a780 100644
--- a/src/cuchaz/enigma/bytecode/InnerClassWriter.java
+++ b/src/cuchaz/enigma/bytecode/InnerClassWriter.java
@@ -11,7 +11,6 @@
 package cuchaz.enigma.bytecode;
 
 import java.util.Collection;
-import java.util.Collections;
 import java.util.List;
 
 import javassist.CtClass;
@@ -36,35 +35,14 @@ public class InnerClassWriter {
         
         public void write(CtClass c) {
                 
-                // build the class chain (inner to outer)
                 ClassEntry obfClassEntry = EntryFactory.getClassEntry(c);
-                List<ClassEntry> obfClassChain = Lists.newArrayList();
-                ClassEntry checkClassEntry = obfClassEntry;
-                while (checkClassEntry != null) {
-                        obfClassChain.add(checkClassEntry);
-                        checkClassEntry = m_index.getOuterClass(checkClassEntry);
-                }
-                
-                // change order: outer to inner
-                Collections.reverse(obfClassChain);
-                
-                // does this class have any inner classes?
-                Collection<ClassEntry> obfInnerClassEntries = m_index.getInnerClasses(obfClassEntry);
+                List<ClassEntry> obfClassChain = m_index.getObfClassChain(obfClassEntry);
                 
                 boolean isInnerClass = obfClassChain.size() > 1;
                 if (isInnerClass) {
                         
                         // it's an inner class, rename it to the fully qualified name
-                        StringBuilder buf = new StringBuilder();
-                        for (ClassEntry obfChainClassEntry : obfClassChain) {
-                                if (buf.length() == 0) {
-                                        buf.append(obfChainClassEntry.getName());
-                                } else {
-                                        buf.append("$");
-                                        buf.append(obfChainClassEntry.getSimpleName());
-                                }
-                        }
-                        c.setName(buf.toString());
+                        c.setName(obfClassEntry.buildClassEntry(obfClassChain).getName());
                         
                         BehaviorEntry caller = m_index.getAnonymousClassCaller(obfClassEntry);
                         if (caller != null) {
@@ -78,6 +56,9 @@ public class InnerClassWriter {
                         }
                 }
                 
+                // does this class have any inner classes?
+                Collection<ClassEntry> obfInnerClassEntries = m_index.getInnerClasses(obfClassEntry);
+                
                 if (isInnerClass || !obfInnerClassEntries.isEmpty()) {
                         
                         // create an inner class attribute
@@ -86,7 +67,11 @@ public class InnerClassWriter {
                         
                         // write the ancestry, but not the outermost class
                         for (int i=1; i<obfClassChain.size(); i++) {
-                                writeInnerClass(attr, obfClassChain, obfClassChain.get(i));
+                                ClassEntry obfInnerClassEntry = obfClassChain.get(i);
+                                writeInnerClass(attr, obfClassChain, obfInnerClassEntry);
+                                
+                                // update references to use the fully qualified inner class name
+                                c.replaceClassName(obfInnerClassEntry.getName(), obfInnerClassEntry.buildClassEntry(obfClassChain).getName());
                         }
                         
                         // write the inner classes
@@ -96,34 +81,34 @@ public class InnerClassWriter {
                                 List<ClassEntry> extendedObfClassChain = Lists.newArrayList(obfClassChain);
                                 extendedObfClassChain.add(obfInnerClassEntry);
                                 
-                                String fullyQualifiedInnerClassName = writeInnerClass(attr, extendedObfClassChain, obfInnerClassEntry);
+                                writeInnerClass(attr, extendedObfClassChain, obfInnerClassEntry);
                                 
-                                // make sure we only reference the fully qualified inner class name
-                                c.replaceClassName(obfInnerClassEntry.getName(), fullyQualifiedInnerClassName);
+                                // update references to use the fully qualified inner class name
+                                c.replaceClassName(obfInnerClassEntry.getName(), obfInnerClassEntry.buildClassEntry(extendedObfClassChain).getName());
                         }
                 }
         }
         
-        private String writeInnerClass(InnerClassesAttribute attr, List<ClassEntry> obfClassChain, ClassEntry obfClassEntry) {
+        private void writeInnerClass(InnerClassesAttribute attr, List<ClassEntry> obfClassChain, ClassEntry obfClassEntry) {
                 
                 // get the new inner class name
-                String obfInnerClassName = getFullyQualifiedName(obfClassChain, obfClassEntry);
-                String obfParentClassName = getFullyQualifiedParentName(obfClassChain, obfClassEntry);
+                ClassEntry obfInnerClassEntry = obfClassEntry.buildClassEntry(obfClassChain);
+                ClassEntry obfOuterClassEntry = obfInnerClassEntry.getOuterClassEntry();
                 
                 // here's what the JVM spec says about the InnerClasses attribute
                 // append(inner, parent, 0 if anonymous else simple name, flags);
                 
                 // update the attribute with this inner class
                 ConstPool constPool = attr.getConstPool();
-                int innerClassIndex = constPool.addClassInfo(obfInnerClassName);
-                int parentClassIndex = constPool.addClassInfo(obfParentClassName);
-                int innerClassSimpleNameIndex = 0;
+                int innerClassIndex = constPool.addClassInfo(obfInnerClassEntry.getName());
+                int parentClassIndex = constPool.addClassInfo(obfOuterClassEntry.getName());
+                int innerClassNameIndex = 0;
                 int accessFlags = 0;
                 if (!m_index.isAnonymousClass(obfClassEntry)) {
-                        innerClassSimpleNameIndex = constPool.addUtf8Info(obfClassEntry.getSimpleName());
+                        innerClassNameIndex = constPool.addUtf8Info(obfInnerClassEntry.getInnerClassName());
                 }
                 
-                attr.append(innerClassIndex, parentClassIndex, innerClassSimpleNameIndex, accessFlags);
+                attr.append(innerClassIndex, parentClassIndex, innerClassNameIndex, accessFlags);
                 
                 /* DEBUG 
                 System.out.println(String.format("\tOBF: %s -> ATTR: %s,%s,%s (replace %s with %s)",
@@ -135,32 +120,5 @@ public class InnerClassWriter {
                         obfClassEntry.getName()
                 ));
                 */
-                
-                return obfInnerClassName;
-        }
-
-        private String getFullyQualifiedParentName(List<ClassEntry> classChain, ClassEntry classEntry) {
-                assert(classChain.size() > 1);
-                assert(classChain.contains(classEntry));
-                StringBuilder buf = new StringBuilder();
-                for (int i=0; classChain.get(i) != classEntry; i++) {
-                        ClassEntry chainEntry = classChain.get(i);
-                        if (buf.length() == 0) {
-                                buf.append(chainEntry.getName());
-                        } else {
-                                buf.append("$");
-                                buf.append(chainEntry.getSimpleName());
-                        }
-                }
-                return buf.toString();
-        }
-        
-        private String getFullyQualifiedName(List<ClassEntry> classChain, ClassEntry classEntry) {
-                boolean isInner = classChain.size() > 1;
-                if (isInner) {
-                        return getFullyQualifiedParentName(classChain, classEntry) + "$" + classEntry.getSimpleName();
-                } else {
-                        return classEntry.getName();
-                }
         }
 }
diff --git a/src/cuchaz/enigma/mapping/ClassEntry.java b/src/cuchaz/enigma/mapping/ClassEntry.java
index 69f171a5..69e66bc0 100644
--- a/src/cuchaz/enigma/mapping/ClassEntry.java
+++ b/src/cuchaz/enigma/mapping/ClassEntry.java
@@ -11,6 +11,7 @@
 package cuchaz.enigma.mapping;
 
 import java.io.Serializable;
+import java.util.List;
 
 public class ClassEntry implements Entry, Serializable {
         
@@ -114,13 +115,28 @@ public class ClassEntry implements Entry, Serializable {
         }
         
         public String getSimpleName() {
-                if (isInnerClass()) {
-                        return getInnerClassName();
-                }
                 int pos = m_name.lastIndexOf('/');
                 if (pos > 0) {
                         return m_name.substring(pos + 1);
                 }
                 return m_name;
         }
+        
+        public ClassEntry buildClassEntry(List<ClassEntry> classChain) {
+                assert(classChain.contains(this));
+                StringBuilder buf = new StringBuilder();
+                for (ClassEntry chainEntry : classChain) {
+                        if (buf.length() == 0) {
+                                buf.append(chainEntry.getName());
+                        } else {
+                                buf.append("$");
+                                buf.append(chainEntry.isInnerClass() ? chainEntry.getInnerClassName() : chainEntry.getSimpleName());
+                        }
+                        
+                        if (chainEntry == this) {
+                                break;
+                        }
+                }
+                return new ClassEntry(buf.toString());
+        }
 }
diff --git a/src/cuchaz/enigma/mapping/EntryFactory.java b/src/cuchaz/enigma/mapping/EntryFactory.java
index bbdfa739..f4d62c8e 100644
--- a/src/cuchaz/enigma/mapping/EntryFactory.java
+++ b/src/cuchaz/enigma/mapping/EntryFactory.java
@@ -1,7 +1,5 @@
 package cuchaz.enigma.mapping;
 
-import java.util.List;
-
 import javassist.CtBehavior;
 import javassist.CtClass;
 import javassist.CtConstructor;
@@ -13,7 +11,6 @@ import javassist.expr.FieldAccess;
 import javassist.expr.MethodCall;
 import javassist.expr.NewExpr;
 
-import com.beust.jcommander.internal.Lists;
 import com.strobel.assembler.metadata.MethodDefinition;
 
 import cuchaz.enigma.analysis.JarIndex;
@@ -25,35 +22,8 @@ public class EntryFactory {
         }
         
         public static ClassEntry getObfClassEntry(JarIndex jarIndex, ClassMapping classMapping) {
-                return getChainedOuterClassName(jarIndex, new ClassEntry(classMapping.getObfFullName()));
-        }
-        
-        public static ClassEntry getChainedOuterClassName(JarIndex jarIndex, ClassEntry obfClassEntry) {
-                
-                // lookup the chain of outer classes
-                List<ClassEntry> obfClassChain = Lists.newArrayList(obfClassEntry);
-                ClassEntry checkClassEntry = obfClassEntry;
-                while (true) {
-                        ClassEntry obfOuterClassEntry = jarIndex.getOuterClass(checkClassEntry);
-                        if (obfOuterClassEntry != null) {
-                                obfClassChain.add(obfOuterClassEntry);
-                                checkClassEntry = obfOuterClassEntry;
-                        } else {
-                                break;
-                        }
-                }
-                
-                // build the chained class name
-                StringBuilder buf = new StringBuilder();
-                for (int i=obfClassChain.size()-1; i>=0; i--) {
-                        if (buf.length() == 0) {
-                                buf.append(obfClassChain.get(i).getName());
-                        } else {
-                                buf.append("$");
-                                buf.append(obfClassChain.get(i).getSimpleName());
-                        }
-                }
-                return new ClassEntry(buf.toString());
+                ClassEntry obfClassEntry = new ClassEntry(classMapping.getObfFullName());
+                return obfClassEntry.buildClassEntry(jarIndex.getObfClassChain(obfClassEntry));
         }
         
         public static ClassEntry getDeobfClassEntry(ClassMapping classMapping) {