Fix PE writer, first version working in unit tests

Author: xoofx

src/LibObjectFile.Tests/PE/PEReaderTests.cs

@@ -23,7 +23,8 @@ public partial class PEReaderTests
     public async Task TestPrinter(string name)
     {
 
-        await using var stream = File.OpenRead(Path.Combine(AppContext.BaseDirectory, "PE", name));
+        var sourceFile = Path.Combine(AppContext.BaseDirectory, "PE", name);
+        await using var stream = File.OpenRead(sourceFile);
         var peImage = PEFile.Read(stream);
         var afterReadWriter = new StringWriter();
         peImage.Print(afterReadWriter);
@@ -45,6 +46,23 @@ public async Task TestPrinter(string name)
             TestContext.WriteLine("Error while verifying UpdateLayout");
             await Verifier.Verify(afterUpdateText).UseParameters(name).DisableRequireUniquePrefix();
         }
+
+        // Read in input as raw bytes
+        stream.Position = 0;
+        var inputBuffer = new byte[stream.Length];
+        stream.ReadExactly(inputBuffer);
+
+        // Write the PE back to a byte buffer
+        var output = new MemoryStream();
+        peImage.Write(output);
+        output.Position = 0;
+        var outputBuffer = output.ToArray();
+
+        //await Verifier.Verify(outputBuffer, sourceFile  sourceFile).
+        await File.WriteAllBytesAsync($"{sourceFile}.bak", outputBuffer);
+
+        // Compare the input and output buffer
+        CollectionAssert.AreEqual(inputBuffer, outputBuffer);
     }
 
     [TestMethod]

src/LibObjectFile/PE/DataDirectory/PEBaseRelocationBlock.cs

@@ -152,6 +152,26 @@ public override unsafe void Read(PEImageReader reader)
         Debug.Assert(Size == block.SizeOfBlock);
     }
 
+    public override unsafe void Write(PEImageWriter writer)
+    {
+        var block = new ImageBaseRelocation
+        {
+            PageRVA = SectionLink.RVA(),
+            SizeOfBlock = (uint)Size
+        };
+
+        writer.Write(block);
+
+        var span = CollectionsMarshal.AsSpan(Relocations);
+        var spanBytes = MemoryMarshal.AsBytes(span);
+        writer.Write(spanBytes);
+
+        if ((Relocations.Count & 1) != 0)
+        {
+            writer.WriteZero((int)sizeof(PEBaseRelocation));
+        }
+    }
+
     protected override bool PrintMembers(StringBuilder builder)
     {
         if (base.PrintMembers(builder))

src/LibObjectFile/PE/DataDirectory/PEClrMetadata.cs

@@ -6,22 +6,28 @@
 
 namespace LibObjectFile.PE;
 
+/// <summary>
+/// Represents the CLR metadata directory in a PE file.
+/// </summary>
 public sealed class PEClrMetadata : PEDataDirectory
 {
+    /// <summary>
+    /// Initializes a new instance of the <see cref="PEClrMetadata"/> class.
+    /// </summary>
     public PEClrMetadata() : base(PEDataDirectoryKind.ClrMetadata)
     {
     }
 
-    protected override uint ComputeHeaderSize(PELayoutContext context)
-    {
-        return 0;
-    }
+    /// <inheritdoc/>
+    protected override uint ComputeHeaderSize(PELayoutContext context) => 0;
 
+    /// <inheritdoc/>
     public override void Read(PEImageReader reader)
     {
     }
 
+    /// <inheritdoc/>
     public override void Write(PEImageWriter writer)
     {
     }
-}
+}

src/LibObjectFile/PE/DataDirectory/PECompositeSectionData.cs

@@ -3,6 +3,7 @@
 // See the license.txt file in the project root for more information.
 
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.Linq;
 using LibObjectFile.Collections;
 using LibObjectFile.Utils;
@@ -86,7 +87,24 @@ internal void WriteHeaderAndContent(PEImageWriter writer)
 
         foreach (var table in Content)
         {
+            var position = writer.Position;
+            var alignment = table.GetRequiredPositionAlignment(writer.PEFile);
+            if (alignment > 1)
+            {
+                var zeroSize = AlignHelper.AlignUp(position, alignment) - (uint)position;
+                writer.WriteZero((int)zeroSize);
+            }
+
+            Debug.Assert(table.Position == writer.Position);
+
             table.Write(writer);
+
+            alignment = table.GetRequiredSizeAlignment(writer.PEFile);
+            if (alignment > 1)
+            {
+                var zeroSize = AlignHelper.AlignUp((uint)table.Size, alignment) - table.Size;
+                writer.WriteZero((int)zeroSize);
+            }
         }
     }
 

src/LibObjectFile/PE/DataDirectory/PEDebugDirectory.cs

@@ -134,10 +134,11 @@ public override void Write(PEImageWriter writer)
         var entries = CollectionsMarshal.AsSpan(Entries);
         using var tempSpan = TempSpan<RawImageDebugDirectory>.Create(entries.Length, out var rawEntries);
 
-        RawImageDebugDirectory rawEntry = default;
         for (var i = 0; i < entries.Length; i++)
         {
             var entry = entries[i];
+            ref var rawEntry = ref rawEntries[i];
+
             rawEntry.Characteristics = entry.Characteristics;
             rawEntry.MajorVersion = entry.MajorVersion;
             rawEntry.MinorVersion = entry.MinorVersion;
@@ -147,15 +148,21 @@ public override void Write(PEImageWriter writer)
             if (entry.SectionData is not null)
             {
                 rawEntry.SizeOfData = (uint)entry.SectionData.Size;
-                rawEntry.AddressOfRawData = (uint)entry.SectionData.RVA;
-                rawEntry.PointerToRawData = 0;
+                rawEntry.AddressOfRawData = entry.SectionData.RVA;
+                rawEntry.PointerToRawData = (uint)entry.SectionData.Position;
             }
             else if (entry.ExtraData is not null)
             {
                 rawEntry.SizeOfData = (uint)entry.ExtraData.Size;
                 rawEntry.AddressOfRawData = 0;
                 rawEntry.PointerToRawData = (uint)entry.ExtraData.Position;
             }
+            else
+            {
+                rawEntry.SizeOfData = 0;
+                rawEntry.AddressOfRawData = 0;
+                rawEntry.PointerToRawData = 0;
+            }
 
             rawEntries[i] = rawEntry;
         }

src/LibObjectFile/PE/DataDirectory/PEDebugSectionDataRSDS.cs

@@ -5,6 +5,7 @@
 using System;
 using System.Buffers;
 using System.Diagnostics;
+using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
 using System.Text;
 using LibObjectFile.Collections;
@@ -18,6 +19,8 @@ namespace LibObjectFile.PE;
 [DebuggerDisplay("{ToString(),nq}")]
 public sealed class PEDebugSectionDataRSDS : PEDebugSectionData
 {
+    private const uint Signature = 0x53445352;
+
     /// <summary>
     /// Initializes a new instance of the <see cref="PEDebugSectionDataRSDS"/> class.
     /// </summary>
@@ -57,7 +60,7 @@ public override unsafe void Read(PEImageReader reader)
         }
 
         var signature = MemoryMarshal.Read<uint>(span);
-        if (signature != 0x53445352)
+        if (signature != Signature)
         {
             reader.Diagnostics.Error(DiagnosticId.PE_ERR_InvalidDebugDataRSDSSignature, $"Invalid signature for PEDebugDataRSDS");
             return;
@@ -74,8 +77,36 @@ public override unsafe void Read(PEImageReader reader)
         Guid = MemoryMarshal.Read<Guid>(span.Slice(4));
         Age = MemoryMarshal.Read<uint>(span.Slice(sizeof(uint) + sizeof(Guid)));
         PdbPath = System.Text.Encoding.UTF8.GetString(pdbPath.Slice(0, indexOfZero));
+
+        Debug.Assert(size == CalculateSize());
     }
 
+    public override unsafe void Write(PEImageWriter writer)
+    {
+        var size = (int)Size;
+        using var tempSpan = TempSpan<byte>.Create(size, out var span);
+
+        MemoryMarshal.Write(span, Signature);
+        span = span.Slice(sizeof(uint));
+        MemoryMarshal.Write(span, Guid);
+        span = span.Slice(sizeof(Guid));
+        MemoryMarshal.Write(span, Age);
+        span = span.Slice(sizeof(uint));
+        int written = Encoding.UTF8.GetBytes(PdbPath, span);
+        span[written] = 0;
+        span.Slice(written + 1);
+
+        writer.Write(tempSpan.AsBytes);
+    }
+
+    public override void UpdateLayout(PELayoutContext layoutContext)
+    {
+        Size = CalculateSize();
+    }
+
+    private unsafe uint CalculateSize()
+        => (uint)(sizeof(uint) + sizeof(Guid) + sizeof(uint) + Encoding.UTF8.GetByteCount(PdbPath) + 1);
+
     /// <inheritdoc />
     protected override bool PrintMembers(StringBuilder builder)
     {

src/LibObjectFile/PE/DataDirectory/PEDirectoryTable.cs

@@ -9,6 +9,7 @@
 using System.Diagnostics;
 using System.Reflection;
 using System.Runtime.CompilerServices;
+using LibObjectFile.Collections;
 using LibObjectFile.PE.Internal;
 
 namespace LibObjectFile.PE;
@@ -191,17 +192,22 @@ internal void Set(int index, PEObjectBase? directory)
 
     internal unsafe void Write(PEImageWriter writer, ref uint position)
     {
+        using var tempSpan = TempSpan<RawImageDataDirectory>.Create(stackalloc byte[16 * sizeof(RawImageDataDirectory)], Count, out var span);
+        span.Clear();
+
         for (int i = 0; i < Count; i++)
         {
-            RawImageDataDirectory rawDataDirectory = default;
             var entry = _entries[i];
             if (entry is not null)
             {
-                rawDataDirectory.RVA = entry is PEDataDirectory dataDirectory ? dataDirectory.RVA : (uint)entry.Position;
-                rawDataDirectory.Size = (uint)entry.Size;
+                ref var rawEntry = ref span[i];
+                rawEntry.RVA = entry is PEDataDirectory dataDirectory ? dataDirectory.RVA : (uint)entry.Position;
+                rawEntry.Size = (uint)entry.Size;
             }
         }
 
+        writer.Write(tempSpan);
+        
         position += (uint)(Count * sizeof(RawImageDataDirectory));
     }
 

src/LibObjectFile/PE/DataDirectory/PEExportAddressTable.cs

@@ -2,17 +2,10 @@
 // This file is licensed under the BSD-Clause 2 license.
 // See the license.txt file in the project root for more information.
 
-using System;
-using System.Buffers;
 using System.Collections.Generic;
-using System.ComponentModel;
-using System.Diagnostics;
 using System.Runtime.InteropServices;
-using System.Security.Cryptography;
-using System.Xml.Linq;
 using LibObjectFile.Collections;
 using LibObjectFile.Diagnostics;
-using static System.Collections.Specialized.BitVector32;
 
 namespace LibObjectFile.PE;
 
@@ -87,6 +80,14 @@ public override unsafe void Read(PEImageReader reader)
 
     public override void Write(PEImageWriter writer)
     {
-        throw new NotImplementedException();
+        using var tempSpan = TempSpan<RVA>.Create(Values.Count, out var spanRva);
+
+        for (int i = 0; i < Values.Count; i++)
+        {
+            var value = Values[i];
+            spanRva[i] = value.IsForwarderRVA ? value.ForwarderRVA.RVA() : value.ExportRVA.RVA();
+        }
+
+        writer.Write(tempSpan);
     }
 }

src/LibObjectFile/PE/DataDirectory/PEExportNameTable.cs

@@ -2,8 +2,6 @@
 // This file is licensed under the BSD-Clause 2 license.
 // See the license.txt file in the project root for more information.
 
-using System;
-using System.Buffers;
 using System.Collections.Generic;
 using System.Runtime.InteropServices;
 using LibObjectFile.Collections;
@@ -66,6 +64,14 @@ public override unsafe void Read(PEImageReader reader)
 
     public override void Write(PEImageWriter writer)
     {
-        throw new NotImplementedException();
+        using var tempSpan = TempSpan<RVA>.Create(Values.Count, out var spanRva);
+
+        for (int i = 0; i < Values.Count; i++)
+        {
+            var value = Values[i];
+            spanRva[i] = value.RVA();
+        }
+
+        writer.Write(tempSpan.AsBytes);
     }
 }

src/LibObjectFile/PE/DataDirectory/PEExportOrdinalTable.cs

@@ -1,4 +1,4 @@
-// Copyright (c) Alexandre Mutel. All rights reserved.
+// Copyright (c) Alexandre Mutel. All rights reserved.
 // This file is licensed under the BSD-Clause 2 license.
 // See the license.txt file in the project root for more information.
 
@@ -42,4 +42,10 @@ public override void Read(PEImageReader reader)
             return;
         }
     }
+
+    public override void Write(PEImageWriter writer)
+    {
+        var span = CollectionsMarshal.AsSpan(Values);
+        writer.Write(MemoryMarshal.AsBytes(span));
+    }
 }