diff --git a/blocks/DotNetUtils.cs b/blocks/DotNetUtils.cs index 0a44dd75..efe15591 100644 --- a/blocks/DotNetUtils.cs +++ b/blocks/DotNetUtils.cs @@ -23,6 +23,57 @@ using Mono.Cecil; using Mono.Cecil.Cil; namespace de4dot.blocks { + class TypeCache { + ModuleDefinition module; + Dictionary typeRefToDef = new Dictionary(); + + public TypeCache(ModuleDefinition module) { + this.module = module; + init(); + } + + void init() { + foreach (var type in module.GetTypes()) { + var key = new TypeReferenceKey(type); + typeRefToDef[key] = type; + } + } + + public TypeDefinition lookup(TypeReference typeReference) { + if (typeReference is TypeDefinition) + return (TypeDefinition)typeReference; + TypeDefinition typeDefinition; + typeRefToDef.TryGetValue(new TypeReferenceKey(typeReference), out typeDefinition); + return typeDefinition; + } + } + + public class TypeCaches { + Dictionary typeCaches = new Dictionary(); + + // Should be called when the whole module is reloaded or when a lot of types have been + // modified (eg. renamed) + public void invalidate(ModuleDefinition module) { + if (module == null) + return; + typeCaches.Remove(module); + } + + // Call this to invalidate all modules + public void invalidateAll() { + typeCaches.Clear(); + } + + public TypeDefinition lookup(ModuleDefinition module, TypeReference typeReference) { + if (typeReference is TypeDefinition) + return (TypeDefinition)typeReference; + TypeCache typeCache; + if (!typeCaches.TryGetValue(module, out typeCache)) + typeCaches[module] = typeCache = new TypeCache(module); + return typeCache.lookup(typeReference); + } + } + public class CallCounter { Dictionary calls = new Dictionary(); @@ -78,6 +129,8 @@ namespace de4dot.blocks { } public static class DotNetUtils { + public static readonly TypeCaches typeCaches = new TypeCaches(); + public static bool isLdcI4(Instruction instruction) { return isLdcI4(instruction.OpCode.Code); } @@ -309,11 +362,7 @@ namespace de4dot.blocks { return null; if (typeReference is TypeDefinition) return (TypeDefinition)typeReference; - foreach (var type in module.GetTypes()) { - if (MemberReferenceHelper.compareTypes(type, typeReference)) - return type; - } - return null; + return typeCaches.lookup(module, typeReference); } public static FieldDefinition getField(ModuleDefinition module, FieldReference field) { diff --git a/de4dot.code/AssemblyModule.cs b/de4dot.code/AssemblyModule.cs index beb5f8fa..b869a341 100644 --- a/de4dot.code/AssemblyModule.cs +++ b/de4dot.code/AssemblyModule.cs @@ -22,6 +22,7 @@ using System.IO; using System.Collections.Generic; using Mono.Cecil; using Mono.MyStuff; +using de4dot.blocks; namespace de4dot { class AssemblyModule { @@ -51,6 +52,7 @@ namespace de4dot { public ModuleDefinition reload(byte[] newModuleData, Dictionary dumpedMethods) { var assemblyResolver = AssemblyResolver.Instance; assemblyResolver.removeModule(module); + DotNetUtils.typeCaches.invalidate(module); this.dumpedMethods = dumpedMethods; var readerParameters = new ReaderParameters(ReadingMode.Deferred); diff --git a/de4dot.code/FilesDeobfuscator.cs b/de4dot.code/FilesDeobfuscator.cs index 56ec0af1..d78eb451 100644 --- a/de4dot.code/FilesDeobfuscator.cs +++ b/de4dot.code/FilesDeobfuscator.cs @@ -20,6 +20,8 @@ using System; using System.IO; using System.Collections.Generic; +using Mono.Cecil; +using de4dot.blocks; using de4dot.renamer; using de4dot.deobfuscators; using de4dot.AssemblyClient; @@ -70,9 +72,14 @@ namespace de4dot { deobfuscateAll(); } + void removeModule(ModuleDefinition module) { + AssemblyResolver.Instance.removeModule(module); + DotNetUtils.typeCaches.invalidate(module); + } + void detectObfuscators() { foreach (var file in loadAllFiles()) { - AssemblyResolver.Instance.removeModule(file.ModuleDefinition); + removeModule(file.ModuleDefinition); } } @@ -88,7 +95,7 @@ namespace de4dot { file.save(); - AssemblyResolver.Instance.removeModule(file.ModuleDefinition); + removeModule(file.ModuleDefinition); } catch (Exception ex) { Log.w("Could not deobfuscate {0}. Use -v to see stack trace", file.Filename); diff --git a/de4dot.code/renamer/DefinitionsRenamer.cs b/de4dot.code/renamer/DefinitionsRenamer.cs index 1bc8c060..7ed62198 100644 --- a/de4dot.code/renamer/DefinitionsRenamer.cs +++ b/de4dot.code/renamer/DefinitionsRenamer.cs @@ -124,6 +124,7 @@ namespace de4dot.renamer { renameMemberDefinitions(); renameMemberReferences(); renameResources(); + DotNetUtils.typeCaches.invalidateAll(); } void initAllTypes() {