Merge pull request #589 from Washi1337/issue/duplicated-exports

UnmanagedPEFileBuilder: Do not duplicate exports and relocations

src/AsmResolver.PE.File/OptionalHeader.cs

@@ -414,8 +414,10 @@ public void SetDataDirectory(DataDirectoryIndex index, DataDirectory directory)
         /// Maps a segment to a data directory.
         /// </summary>
         /// <param name="index">The index.</param>
-        /// <param name="contents">The contents of the data directory.</param>
-        public void SetDataDirectory(DataDirectoryIndex index, ISegment contents)
+        /// <param name="contents">
+        /// The contents of the data directory, or <c>null</c> to indicate an absence of the directory.
+        /// </param>
+        public void SetDataDirectory(DataDirectoryIndex index, ISegment? contents)
         {
             SetDataDirectory(index, DataDirectory.CreateForSegment(contents));
         }

src/AsmResolver.PE/Builder/UnmanagedPEFileBuilder.cs

@@ -188,29 +188,26 @@ protected override void AssignDataDirectories(BuilderContext context, PEFile out
 
         header.EnsureDataDirectoryCount(OptionalHeader.DefaultNumberOfRvasAndSizes);
 
+        // NOTE: We need to explicitly check if imports.count == 0 because the buffer may not be populated if
+        // IAT trampolining (and thus IAT rebuilding) is disabled. In such a case, we want to preserve the existing IAT
+        // rather than remove it.
         if (!context.ImportDirectory.IsEmpty)
         {
             header.SetDataDirectory(DataDirectoryIndex.ImportDirectory, context.ImportDirectory);
             header.SetDataDirectory(DataDirectoryIndex.IatDirectory, context.ImportDirectory.ImportAddressDirectory);
         }
+        else if (context.Image.Imports.Count == 0)
+        {
+            header.SetDataDirectory(DataDirectoryIndex.ImportDirectory, null);
+            header.SetDataDirectory(DataDirectoryIndex.IatDirectory, null);
+        }
 
-        if (!context.ExportDirectory.IsEmpty)
-            header.SetDataDirectory(DataDirectoryIndex.ExportDirectory, context.ExportDirectory);
-
-        if (!context.DebugDirectory.IsEmpty)
-            header.SetDataDirectory(DataDirectoryIndex.DebugDirectory, context.DebugDirectory);
-
-        if (!context.ResourceDirectory.IsEmpty)
-            header.SetDataDirectory(DataDirectoryIndex.ResourceDirectory, context.ResourceDirectory);
-
-        if (context.Image.DotNetDirectory is not null)
-            header.SetDataDirectory(DataDirectoryIndex.ClrDirectory, context.Image.DotNetDirectory);
-
-        if (context.Image.TlsDirectory is not null)
-            header.SetDataDirectory(DataDirectoryIndex.TlsDirectory, context.Image.TlsDirectory);
-
-        if (!context.RelocationsDirectory.IsEmpty)
-            header.SetDataDirectory(DataDirectoryIndex.BaseRelocationDirectory, context.RelocationsDirectory);
+        header.SetDataDirectory(DataDirectoryIndex.ExportDirectory, !context.ExportDirectory.IsEmpty ? context.ExportDirectory : null);
+        header.SetDataDirectory(DataDirectoryIndex.DebugDirectory, !context.DebugDirectory.IsEmpty ? context.DebugDirectory : null);
+        header.SetDataDirectory(DataDirectoryIndex.ResourceDirectory, !context.ResourceDirectory.IsEmpty ? context.ResourceDirectory : null);
+        header.SetDataDirectory(DataDirectoryIndex.ClrDirectory, context.Image.DotNetDirectory);
+        header.SetDataDirectory(DataDirectoryIndex.TlsDirectory, context.Image.TlsDirectory);
+        header.SetDataDirectory(DataDirectoryIndex.BaseRelocationDirectory, !context.RelocationsDirectory.IsEmpty ? context.RelocationsDirectory : null);
     }
 
     /// <summary>
@@ -234,10 +231,6 @@ protected override void AssignDataDirectories(BuilderContext context, PEFile out
         if (TrampolineImports && !context.ImportDirectory.IsEmpty)
             contents.Add(context.ImportDirectory);
 
-        // Reconstructed exports.
-        if (!context.ExportDirectory.IsEmpty)
-            contents.Add(context.ExportDirectory);
-
         // Reconstructed debug directory.
         if (!context.DebugDirectory.IsEmpty)
         {
@@ -258,7 +251,7 @@ protected override void AssignDataDirectories(BuilderContext context, PEFile out
             context.ImportTrampolines.ApplyPatches(context.ClonedSections);
         }
 
-        // Code for newly added exports.
+        // Add newly added exports code.
         if (context.Image.Exports is { Entries: { Count: > 0 } entries })
         {
             for (int i = 0; i < entries.Count; i++)
@@ -350,6 +343,13 @@ void AddOrPatch(ISegment? segment, ISegment? originalSegment)
                 AddOrPatch(directory.CallbackFunctions, originalDirectory?.CallbackFunctions);
         }
 
+        // Add export directory.
+        if (image.Exports is { Entries.Count: > 0 })
+        {
+            if (!TryPatchDataDirectory(context, context.ExportDirectory, DataDirectoryIndex.ExportDirectory))
+                contents.Add(context.ExportDirectory, (uint) context.Platform.PointerSize);
+        }
+
         if (contents.Count == 0)
             return null;
 
@@ -386,10 +386,6 @@ void AddOrPatch(ISegment? newSegment, ISegment? originalSegment)
                 contents.Add(fixups[i].Tokens, (uint) context.Platform.PointerSize);
         }
 
-        // Add export directory.
-        if (image.Exports is { Entries.Count: > 0 })
-            contents.Add(context.ExportDirectory, (uint) context.Platform.PointerSize);
-
         if (image.TlsDirectory is { } directory)
         {
             // Add TLS index segment.

src/AsmResolver.PE/Relocations/Builder/RelocationBlock.cs

@@ -1,5 +1,4 @@
 using System.Collections.Generic;
-using System.Threading;
 using AsmResolver.IO;
 
 namespace AsmResolver.PE.Relocations.Builder
@@ -36,18 +35,27 @@ public IList<RelocationEntry> Entries
         }
 
         /// <inheritdoc />
-        public override uint GetPhysicalSize() => (uint) (Entries.Count + 1) * sizeof(ushort) + 2 * sizeof(uint);
+        public override uint GetPhysicalSize()
+        {
+            // Ensure the ending zero entry is included in the size.
+            int totalCount = Entries.Count + (Entries[Entries.Count - 1].IsEmpty ? 0 : 1);
+            return (uint) totalCount * sizeof(ushort) + 2 * sizeof(uint);
+        }
 
         /// <inheritdoc />
         public override void Write(BinaryStreamWriter writer)
         {
+            // Block header.
             writer.WriteUInt32(PageRva);
             writer.WriteUInt32(GetPhysicalSize());
 
+            // Write all entries in block.
             for (int i = 0; i < Entries.Count; i++)
                 Entries[i].Write(writer);
 
-            default(RelocationEntry).Write(writer);
+            // Ensure block ends with zero entry.
+            if (!Entries[Entries.Count - 1].IsEmpty)
+                default(RelocationEntry).Write(writer);
         }
 
     }

src/AsmResolver.PE/Relocations/Builder/RelocationEntry.cs

@@ -47,6 +47,11 @@ public RelocationEntry(RelocationType type, int offset)
         /// </summary>
         public int Offset => _value & 0xFFF;
 
+        /// <summary>
+        /// Gets a value indicating the entry is a zero entry.
+        /// </summary>
+        public bool IsEmpty => _value == 0;
+
         /// <inheritdoc />
         public uint GetPhysicalSize() => sizeof(ushort);
 

test/AsmResolver.PE.Tests/Builder/UnmanagedPEFileBuilderTest.cs

@@ -3,6 +3,7 @@
 using System.Linq;
 using AsmResolver.PE.Builder;
 using AsmResolver.PE.DotNet.Metadata;
+using AsmResolver.PE.Exports;
 using AsmResolver.PE.File;
 using AsmResolver.PE.Imports;
 using AsmResolver.Shims;
@@ -233,4 +234,19 @@ public void AddMetadataToMixedModeAssembly()
             expectedOutput
         );
     }
+
+    [Fact]
+    public void AddExportToExistingDirectory()
+    {
+        var image = PEImage.FromBytes(Properties.Resources.SimpleDll_Exports, TestReaderParameters);
+        image.Exports!.Entries.Add(new ExportedSymbol(new VirtualAddress(0x13371337), "MySymbol"));
+
+        var file = image.ToPEFile(new UnmanagedPEFileBuilder());
+        using var stream = new MemoryStream();
+        file.Write(stream);
+
+        var newImage = PEImage.FromBytes(stream.ToArray(), TestReaderParameters);
+        Assert.NotNull(newImage.Exports);
+        Assert.Equal(image.Exports.Entries.Select(x => x.Name), newImage.Exports.Entries.Select(x => x.Name));
+    }
 }