Fix UpdateLayout

Author: xoofx

src/LibObjectFile.Tests/PE/PEReaderTests.cs

@@ -7,8 +7,10 @@
 using System.Linq;
 using System.Runtime.InteropServices;
 using System.Threading.Tasks;
+using LibObjectFile.Diagnostics;
 using LibObjectFile.PE;
 using VerifyMSTest;
+using VerifyTests;
 
 namespace LibObjectFile.Tests.PE;
 
@@ -23,10 +25,28 @@ public partial class PEReaderTests
 
     public async Task TestPrinter(string name)
     {
-        var stream = File.OpenRead(Path.Combine(AppContext.BaseDirectory, "PE", name));
+
+        await using var stream = File.OpenRead(Path.Combine(AppContext.BaseDirectory, "PE", name));
         var peImage = PEFile.Read(stream);
-        var text = new StringWriter();
-        peImage.Print(text);
-        await Verifier.Verify(text).UseParameters(name);
+        var afterReadWriter = new StringWriter();
+        peImage.Print(afterReadWriter);
+
+        var afterReadText = afterReadWriter.ToString();
+
+        await Verifier.Verify(afterReadText).UseParameters(name);
+
+        // Update the layout
+        var diagnostics = new DiagnosticBag();
+        peImage.UpdateLayout(diagnostics);
+
+        var afterUpdateWriter = new StringWriter();
+        peImage.Print(afterUpdateWriter);
+        var afterUpdateText = afterUpdateWriter.ToString();
+
+        if (!string.Equals(afterReadText, afterUpdateText, StringComparison.Ordinal))
+        {
+            TestContext.WriteLine("Error while verifying UpdateLayout");
+            await Verifier.Verify(afterUpdateText).UseParameters(name).DisableRequireUniquePrefix();
+        }
     }
 }

src/LibObjectFile.Tests/Verified/PEReaderTests.TestPrinter_name=NativeConsole2Win64.exe.verified.txt

@@ -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.
 
@@ -325,7 +325,7 @@ public void Write(Stream stream)
         writer.Write();
     }
 
-    public override void UpdateLayout(ArVisitorContext visitorContext)
+    public override void UpdateLayout(ArVisitorContext context)
     {
 
     }

src/LibObjectFile.Tests/Verified/PEReaderTests.TestPrinter_name=NativeConsoleWin64.exe.verified.txt

@@ -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.
 
@@ -125,37 +125,12 @@ public enum DiagnosticId
     PE_ERR_InvalidSectionHeadersSize = 3007,
     PE_ERR_InvalidParent = 3008,
     PE_ERR_InvalidExtraData = 3009,
-    
-    // PE BaseRelocation
-    PE_ERR_BaseRelocationDirectoryInvalidEndOfStream = 3020,
-    PE_ERR_BaseRelocationDirectoryInvalidSection = 3021,
-    PE_ERR_BaseRelocationDirectoryInvalidSectionData = 3022,
-    PE_ERR_InvalidDataDirectorySection = 3023,
-    PE_ERR_BaseRelocationDirectoryInvalidVirtualAddress = 3024,
-
-    // PE Import
-    PE_ERR_ImportDirectoryInvalidEndOfStream = 3040,
-    PE_ERR_ImportLookupTableInvalidEndOfStream = 3041,
-    PE_ERR_ImportLookupTableInvalidHintNameTableRVA = 3042,
-    PE_ERR_ImportLookupTableInvalidParent = 3043,
-    PE_ERR_ImportDirectoryInvalidImportAddressTableRVA = 3044,
-    PE_ERR_ImportDirectoryInvalidImportLookupTableRVA = 3045,
-    PE_ERR_ImportAddressTableNotFound = 3046,
-    PE_ERR_InvalidInternalState = 3047,
-
-    // PE Export
-    PE_ERR_ExportAddressTableInvalidRVA = 3060,
-    PE_ERR_ExportDirectoryInvalidAddressOfNames = 3061,
-    PE_ERR_ExportDirectoryInvalidAddressOfFunctions = 3062,
-    PE_ERR_ExportDirectoryInvalidAddressOfNameOrdinals = 3063,
-    PE_ERR_ExportDirectoryInvalidName = 3064,
-    PE_ERR_ExportNameTableInvalidRVA = 3065,
-
-    // PE Resource directory
-    PE_ERR_InvalidResourceDirectory = 3080,
-    PE_ERR_InvalidResourceDirectoryEntry = 3081,
-    PE_ERR_InvalidResourceDirectoryEntryRVAOffsetToData = 3082,
-
+    PE_ERR_SectionSizeLargerThanVirtualSize = 3010,
+    PE_ERR_SectionRVALessThanPrevious = 3011,
+    PE_ERR_TooManySections = 3012,
+    PE_ERR_FileAlignmentNotPowerOfTwo = 3013,
+    PE_ERR_SectionAlignmentNotPowerOfTwo = 3014,
+    PE_ERR_SectionAlignmentLessThanFileAlignment = 315,
 
     // PE Exception directory
     PE_ERR_InvalidExceptionDirectory_Entries = 3100,
@@ -165,12 +140,6 @@ public enum DiagnosticId
     // PE Certificate directory
     PE_ERR_InvalidCertificateEntry = 3200,
 
-    // PE TLS directory
-    PE_ERR_InvalidTlsStartAddressOfRawData = 3300,
-    PE_ERR_InvalidTlsEndAddressOfRawData = 3301,
-    PE_ERR_InvalidTlsAddressOfIndex = 3302,
-    PE_ERR_InvalidTlsAddressOfCallBacks = 3303,
-
     // PE BoundImport directory
     PE_ERR_BoundImportDirectoryInvalidEndOfStream = 3400,
     PE_ERR_BoundImportDirectoryInvalidModuleName = 3401,
@@ -193,4 +162,34 @@ public enum DiagnosticId
     PE_ERR_InvalidDebugDataRSDSSignature = 3603,
     PE_ERR_InvalidDebugDataRSDSPdbPath = 3604,
     PE_ERR_DebugDirectoryExtraData = 3605,
+    
+    // PE BaseRelocation
+    PE_ERR_BaseRelocationDirectoryInvalidEndOfStream = 3700,
+    PE_ERR_BaseRelocationDirectoryInvalidSection = 3701,
+    PE_ERR_BaseRelocationDirectoryInvalidSectionData = 3702,
+    PE_ERR_InvalidDataDirectorySection = 3703,
+    PE_ERR_BaseRelocationDirectoryInvalidVirtualAddress = 3704,
+    
+    // PE Import
+    PE_ERR_ImportDirectoryInvalidEndOfStream = 3800,
+    PE_ERR_ImportLookupTableInvalidEndOfStream = 3801,
+    PE_ERR_ImportLookupTableInvalidHintNameTableRVA = 3802,
+    PE_ERR_ImportLookupTableInvalidParent = 3803,
+    PE_ERR_ImportDirectoryInvalidImportAddressTableRVA = 3804,
+    PE_ERR_ImportDirectoryInvalidImportLookupTableRVA = 3805,
+    PE_ERR_ImportAddressTableNotFound = 3806,
+    PE_ERR_InvalidInternalState = 3807,
+
+    // PE Export
+    PE_ERR_ExportAddressTableInvalidRVA = 3900,
+    PE_ERR_ExportDirectoryInvalidAddressOfNames = 3901,
+    PE_ERR_ExportDirectoryInvalidAddressOfFunctions = 3902,
+    PE_ERR_ExportDirectoryInvalidAddressOfNameOrdinals = 3903,
+    PE_ERR_ExportDirectoryInvalidName = 3904,
+    PE_ERR_ExportNameTableInvalidRVA = 3905,
+
+    // PE Resource directory
+    PE_ERR_InvalidResourceDirectory = 4000,
+    PE_ERR_InvalidResourceDirectoryEntry = 4001,
+    PE_ERR_InvalidResourceDirectoryEntryRVAOffsetToData = 4002,
 }

src/LibObjectFile.Tests/Verified/PEReaderTests.TestPrinter_name=NativeLibraryWin64.dll.verified.txt

@@ -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.
 
@@ -14,6 +14,11 @@ namespace LibObjectFile.PE;
 [DebuggerDisplay("{ToString(),nq}")]
 public readonly record struct PEBaseRelocation(PEBaseRelocationType Type, PESectionData? Container, RVO RVO) : IPELink<PESectionData>
 {
+    /// <summary>
+    /// Gets a value indicating whether the base relocation is zero padding.
+    /// </summary>
+    public bool IsZero => Type == PEBaseRelocationType.Absolute;
+
     /// <summary>
     /// Reads the address from the section data.
     /// </summary>
@@ -60,5 +65,5 @@ public ulong ReadAddress(PEFile file)
         }
     }
 
-    public override string ToString() => $"{Type} {this.ToDisplayTextWithRVA()}";
+    public override string ToString() => Type == PEBaseRelocationType.Absolute ? $"{Type} Zero Padding" : $"{Type} {this.ToDisplayTextWithRVA()}";
 }

src/LibObjectFile/Ar/ArArchiveFile.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.
 
@@ -53,7 +53,15 @@ internal unsafe uint CalculateSizeOf()
             return (uint)BlockBuffer.Length;
         }
 
-        return (uint)(Relocations.Count * sizeof(ushort));
+        var count = Relocations.Count;
+
+        // If we have an odd number of relocations, we need to add an extra 0x0
+        if (count > 0 && (count & 1) != 0)
+        {
+            count++;
+        }
+
+        return (uint)(count * sizeof(ushort));
     }
 
     internal void ReadAndBind(PEImageReader reader)
@@ -62,35 +70,33 @@ internal void ReadAndBind(PEImageReader reader)
 
         var relocSpan = MemoryMarshal.Cast<byte, RawImageBaseRelocation>(buffer.Span);
 
-        // Remove padding zeros at the end of the block
-        if (relocSpan.Length > 0 && relocSpan[^1].IsZero)
-        {
-            relocSpan = relocSpan.Slice(0, relocSpan.Length - 1);
-        }
-
         var section = SectionLink.Container!;
         var blockBaseAddress = SectionLink.RVA();
 
         // Iterate on all relocations
-        foreach (var relocation in relocSpan)
+        foreach (var rawReloc in relocSpan)
         {
-            if (relocation.IsZero)
+            PEBaseRelocation reloc;
+            if (rawReloc.IsZero)
             {
-                continue;
+                reloc = new PEBaseRelocation();
             }
+            else
+            {
+                var va = blockBaseAddress + rawReloc.OffsetInBlockPart;
 
-            var va = blockBaseAddress + relocation.OffsetInBlockPart;
+                // Find the section data containing the virtual address
+                if (!section.TryFindSectionData(va, out var sectionData))
+                {
+                    reader.Diagnostics.Error(DiagnosticId.PE_ERR_BaseRelocationDirectoryInvalidVirtualAddress, $"Unable to find the section data containing the virtual address 0x{va:X4}");
+                    continue;
+                }
 
-            // Find the section data containing the virtual address
-            if (!section.TryFindSectionData(va, out var sectionData))
-            {
-                reader.Diagnostics.Error(DiagnosticId.PE_ERR_BaseRelocationDirectoryInvalidVirtualAddress, $"Unable to find the section data containing the virtual address 0x{va:X4}");
-                continue;
-            }
+                var offsetInSectionData = va - sectionData.RVA;
+                reloc = new PEBaseRelocation(rawReloc.Type, sectionData, offsetInSectionData);
 
-            var offsetInSectionData = va - sectionData.RVA;
-            var newRelocation = new PEBaseRelocation(relocation.Type, sectionData, offsetInSectionData);
-            Relocations.Add(newRelocation);
+            }
+            Relocations.Add(reloc);
         }
 
         // Clear the buffer, as we don't need it anymore

src/LibObjectFile/Diagnostics/DiagnosticId.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.
 
@@ -19,12 +19,12 @@ public PEBaseRelocationDirectory() : base(PEDataDirectoryKind.BaseRelocation)
 
     public List<PEBaseRelocationBlock> Blocks { get; } = new();
 
-    protected override uint ComputeHeaderSize(PEVisitorContext context)
+    protected override unsafe uint ComputeHeaderSize(PEVisitorContext context)
     {
         var size = 0U;
         foreach (var block in Blocks)
         {
-            size += block.CalculateSizeOf();
+            size += (uint)(block.CalculateSizeOf() + sizeof(ImageBaseRelocation));
         }
 
         return size;

src/LibObjectFile/PE/DataDirectory/PEBaseRelocation.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.
 
@@ -43,22 +43,22 @@ public sealed override void UpdateLayout(PELayoutContext context)
 
         // A directory could have a content in addition to the header
         // So we update the VirtualAddress of each content and update the layout
-        var position = Position;
-        foreach (var table in Content)
+        var position = Position + headerSize;
+        foreach (var subData in Content)
         {
-            table.RVA = va;
+            subData.RVA = va;
 
             // Update layout will update virtual address
             if (!context.UpdateSizeOnly)
             {
-                table.Position = position;
+                subData.Position = position;
             }
 
-            table.UpdateLayout(context);
+            subData.UpdateLayout(context);
 
-            va += (uint)table.Size;
-            size += table.Size;
-            position += table.Size;
+            va += (uint)subData.Size;
+            size += subData.Size;
+            position += subData.Size;
         }
 
         Size = size;

src/LibObjectFile/PE/DataDirectory/PEBaseRelocationBlock.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.
 
@@ -146,6 +146,21 @@ internal void Set(PEDataDirectoryKind kind, PEDataDirectory? directory)
             _count++;
         }
     }
+
+    internal int CalculateNumberOfEntries()
+    {
+        int count = 0;
+        ReadOnlySpan<PEObjectBase?> span = _entries;
+        for(int i = 0; i < span.Length; i++) 
+        {
+            if (_entries[i] is not null)
+            {
+                count = i + 1;
+            }
+        }
+
+        return count;
+    }
 
     [InlineArray(15)]
     private struct InternalArray