BUGFIX: Check if .net directory exists before using it while preservi…

…ng tokens.
Washi1337 · Aug 19, 2022 · bf3b6cb · bf3b6cb
1 parent 505da7a
commit bf3b6cb
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 2 deletions.
diff --git a/src/AsmResolver.DotNet/Builder/Discovery/MemberDiscoverer.cs b/src/AsmResolver.DotNet/Builder/Discovery/MemberDiscoverer.cs
@@ -121,6 +121,9 @@ public static MemberDiscoveryResult DiscoverMembersInModule(ModuleDefinition mod
 
         private void CollectExistingMembers()
         {
+            if (_module.DotNetDirectory?.Metadata is null)
+                return;
+
             if ((_flags & MemberDiscoveryFlags.PreserveTypeOrder) != 0)
                 CollectMembersFromTable<TypeDefinition>(TableIndex.TypeDef);
             if ((_flags & MemberDiscoveryFlags.PreserveFieldOrder) != 0)
@@ -139,8 +142,8 @@ private void CollectMembersFromTable<TMember>(TableIndex tableIndex)
             where TMember: IMetadataMember, IModuleProvider
         {
             // Get original number of elements in the table.
-            int count = _module.DotNetDirectory!.Metadata
-                !.GetStream<TablesStream>()
+            int count = _module.DotNetDirectory!.Metadata!
+                .GetStream<TablesStream>()
                 .GetTable(tableIndex)
                 .Count;
 

diff --git a/src/AsmResolver.DotNet/Builder/DotNetDirectoryFactory.cs b/src/AsmResolver.DotNet/Builder/DotNetDirectoryFactory.cs
@@ -230,6 +230,9 @@ private void ImportBasicTablesIfSpecified(ModuleDefinition module, DotNetDirecto
 
         private void ImportTypeSpecsAndMemberRefsIfSpecified(ModuleDefinition module, DotNetDirectoryBuffer buffer)
         {
+            if (module.DotNetDirectory is null)
+                return;
+
             if ((MetadataBuilderFlags & MetadataBuilderFlags.PreserveTypeSpecificationIndices) != 0)
             {
                 ImportTables<TypeSpecification>(module, TableIndex.TypeSpec,

diff --git a/test/AsmResolver.DotNet.Tests/Builder/ManagedPEImageBuilderTest.cs b/test/AsmResolver.DotNet.Tests/Builder/ManagedPEImageBuilderTest.cs
@@ -1,5 +1,6 @@
 using System.IO;
 using System.Linq;
+using AsmResolver.DotNet.Builder;
 using AsmResolver.PE;
 using Xunit;
 
@@ -42,5 +43,41 @@ public void ExecutableImportDirectoryShouldContainMsCoreeCorDllMain()
                 relocation.Location.CanRead
                 && relocation.Location.CreateReader().ReadUInt32() == image.ImageBase + symbol.AddressTableEntry!.Rva);
         }
+
+        [Fact]
+        public void ConstructPEImageFromNewModuleWithNoPreservation()
+        {
+            var module = new ModuleDefinition("Module");
+            var result = module.ToPEImage();
+            var newModule = ModuleDefinition.FromImage(result);
+            Assert.Equal(module.Name, newModule.Name);
+        }
+
+        [Fact]
+        public void ConstructPEImageFromNewModuleWithPreservation()
+        {
+            var module = new ModuleDefinition("Module");
+            var result = module.ToPEImage(new ManagedPEImageBuilder(MetadataBuilderFlags.PreserveAll));
+            var newModule = ModuleDefinition.FromImage(result);
+            Assert.Equal(module.Name, newModule.Name);
+        }
+
+        [Fact]
+        public void ConstructPEImageFromExistingModuleWithNoPreservation()
+        {
+            var module = ModuleDefinition.FromBytes(Properties.Resources.HelloWorld);
+            var result = module.ToPEImage();
+            var newModule = ModuleDefinition.FromImage(result);
+            Assert.Equal(module.Name, newModule.Name);
+        }
+
+        [Fact]
+        public void ConstructPEImageFromExistingModuleWithPreservation()
+        {
+            var module = ModuleDefinition.FromBytes(Properties.Resources.HelloWorld);
+            var result = module.ToPEImage(new ManagedPEImageBuilder(MetadataBuilderFlags.PreserveAll));
+            var newModule = ModuleDefinition.FromImage(result);
+            Assert.Equal(module.Name, newModule.Name);
+        }
     }
 }