Add PE checksum

src/LibObjectFile.Tests/PE/PEReaderTests.cs

@@ -56,7 +56,11 @@ public async Task TestPrinter(string name)
 
         // Write the PE back to a byte buffer
         var output = new MemoryStream();
-        peImage.Write(output, new PEImageWriterOptions() { EnableStackTrace = true });
+        peImage.Write(output, new PEImageWriterOptions()
+        {
+            EnableStackTrace = true,
+            EnableChecksum = peImage.OptionalHeader.CheckSum != 0 // Recalculate the checksum if it was present
+        });
         output.Position = 0;
         byte[] outputBuffer = output.ToArray();
 
@@ -74,13 +78,13 @@ public async Task TestPrinter(string name)
     public void TestCreatePE()
     {
         var pe = new PEFile();
-        
+
         // ***************************************************************************
         // Code section
         // ***************************************************************************
         var codeSection = pe.AddSection(PESectionName.Text, 0x1000);
         var streamCode = new PEStreamSectionData();
-        
+
         streamCode.Stream.Write([
             // SUB RSP, 0x28
             0x48, 0x83, 0xEC, 0x28,
@@ -93,7 +97,7 @@ public void TestCreatePE()
         ]);
 
         codeSection.Content.Add(streamCode);
-        
+
         // ***************************************************************************
         // Data section
         // ***************************************************************************
@@ -112,7 +116,7 @@ public void TestCreatePE()
         {
             peImportAddressTable
         };
-        
+
         var peImportLookupTable = new PEImportLookupTable()
         {
             exitProcessFunction
@@ -145,10 +149,15 @@ public void TestCreatePE()
         pe.Write(output, new() { EnableStackTrace = true });
         output.Position = 0;
 
-        var sourceFile = Path.Combine(AppContext.BaseDirectory, "PE", "generated_win64.exe");
+        var sourceFile = Path.Combine(AppContext.BaseDirectory, "PE", "RawNativeConsoleWin64_Generated.exe");
         File.WriteAllBytes(sourceFile, output.ToArray());
-    }
 
+        // Check the generated exe
+        var process = Process.Start(sourceFile);
+        process.WaitForExit();
+        Assert.AreEqual(156, process.ExitCode);
+    }
+
     [DataTestMethod]
     [DynamicData(nameof(GetWindowsExeAndDlls), DynamicDataSourceType.Method)]
     public async Task TestWindows(string sourceFile)
@@ -191,7 +200,11 @@ public async Task TestWindows(string sourceFile)
 
         // Write the PE back to a byte buffer
         var output = new MemoryStream();
-        peImage.Write(output, new PEImageWriterOptions() { EnableStackTrace = true });
+        peImage.Write(output, new PEImageWriterOptions()
+        {
+            EnableStackTrace = true,
+            //EnableChecksum = peImage.OptionalHeader.CheckSum != 0 // Recalculate the checksum if it was present, we cannot enable it because some DLLs have an invalid checksum
+        });
         output.Position = 0;
         var outputBuffer = output.ToArray();
 

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

@@ -25,7 +25,7 @@ DOS Stub
 COFF Header
     Machine                    = Amd64
     NumberOfSections           = 6
-    TimeDateStamp              = 1727110447
+    TimeDateStamp              = 1727802524
     PointerToSymbolTable       = 0x0
     NumberOfSymbols            = 0
     SizeOfOptionalHeader       = 240
@@ -53,7 +53,7 @@ Optional Header
     Win32VersionValue          = 0x0
     SizeOfImage                = 0x9000
     SizeOfHeaders              = 0x400
-    CheckSum                   = 0x0
+    CheckSum                   = 0x7C57
     Subsystem                  = WindowsCui
     DllCharacteristics         = HighEntropyVirtualAddressSpace, DynamicBase, NxCompatible, TerminalServerAware
     SizeOfStackReserve         = 0x100000
@@ -243,17 +243,17 @@ Sections
         [03] PEStreamSectionData              Position = 0x000022D0, Size = 0x00000010, RVA = 0x000034D0, VirtualSize = 0x00000010
 
         [04] PEDebugDirectory                 Position = 0x000022E0, Size = 0x00000070, RVA = 0x000034E0, VirtualSize = 0x00000070
-            [0] Type = CodeView, Characteristics = 0x0, Version = 0.0, TimeStamp = 0x66F19D2F, Data = RVA = 0x00003618 (PEDebugSectionDataRSDS[6] -> .rdata)
-            [1] Type = VCFeature, Characteristics = 0x0, Version = 0.0, TimeStamp = 0x66F19D2F, Data = RVA = 0x0000368C (PEDebugStreamSectionData[8] -> .rdata)
-            [2] Type = POGO, Characteristics = 0x0, Version = 0.0, TimeStamp = 0x66F19D2F, Data = RVA = 0x000036A0 (PEDebugStreamSectionData[9] -> .rdata)
-            [3] Type = ILTCG, Characteristics = 0x0, Version = 0.0, TimeStamp = 0x66F19D2F, Data = null
+            [0] Type = CodeView, Characteristics = 0x0, Version = 0.0, TimeStamp = 0x66FC2C9C, Data = RVA = 0x00003618 (PEDebugSectionDataRSDS[6] -> .rdata)
+            [1] Type = VCFeature, Characteristics = 0x0, Version = 0.0, TimeStamp = 0x66FC2C9C, Data = RVA = 0x0000368C (PEDebugStreamSectionData[8] -> .rdata)
+            [2] Type = POGO, Characteristics = 0x0, Version = 0.0, TimeStamp = 0x66FC2C9C, Data = RVA = 0x000036A0 (PEDebugStreamSectionData[9] -> .rdata)
+            [3] Type = ILTCG, Characteristics = 0x0, Version = 0.0, TimeStamp = 0x66FC2C9C, Data = null
 
         [05] PEStreamSectionData              Position = 0x00002350, Size = 0x000000C8, RVA = 0x00003550, VirtualSize = 0x000000C8
 
         [06] PEDebugSectionDataRSDS           Position = 0x00002418, Size = 0x00000072, RVA = 0x00003618, VirtualSize = 0x00000072
             Debug Section Data (RSDS)
                 Guid                       = ffed6f99-5708-452c-a889-ff343b6ce898
-                Age                        = 6
+                Age                        = 7
                 PdbPath                    = C:\code\LibObjectFile\src\native\Win64\NativeProjects\x64\Release\NativeConsole2Win64.pdb
 
         [07] PEStreamSectionData              Position = 0x0000248A, Size = 0x00000002, RVA = 0x0000368A, VirtualSize = 0x00000002

src/LibObjectFile/Diagnostics/DiagnosticId.cs

@@ -137,6 +137,7 @@ public enum DiagnosticId
     PE_ERR_InvalidAddressOfEntryPoint = 3019,
     PE_ERR_DirectoryWithSameKindAlreadyAdded = 3020,
     PE_ERR_VerifyContextInvalidObject = 3021,
+    PE_ERR_ChecksumNotSupported = 3022,
 
     // PE Exception directory
     PE_ERR_InvalidExceptionDirectory_Entries = 3100,

src/LibObjectFile/PE/DataDirectory/PEDebugDirectoryEntry.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.
 

src/LibObjectFile/PE/PEFile.Write.cs

@@ -5,11 +5,13 @@
 using System;
 using System.Diagnostics;
 using System.IO;
+using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
 using System.Text;
 using LibObjectFile.Diagnostics;
 using LibObjectFile.PE.Internal;
 using LibObjectFile.Utils;
+using static System.Runtime.InteropServices.JavaScript.JSType;
 
 namespace LibObjectFile.PE;
 
@@ -37,12 +39,15 @@ public bool TryWrite(Stream stream, out DiagnosticBag diagnostics, PEImageWriter
     {
         if (stream == null) throw new ArgumentNullException(nameof(stream));
 
-        var peWriter = new PEImageWriter(this, stream);
+        var peWriter = new PEImageWriter(this, stream, options ?? PEImageWriterOptions.Default);
         diagnostics = peWriter.Diagnostics;
 
-        if (options is not null)
+        diagnostics.EnableStackTrace = peWriter.Options.EnableStackTrace;
+
+        if (peWriter.Options.EnableChecksum && stream is not MemoryStream)
         {
-            diagnostics.EnableStackTrace = options.EnableStackTrace;
+            diagnostics.Error(DiagnosticId.PE_ERR_ChecksumNotSupported, "Checksum is only supported for MemoryStream");
+            return false;
         }
 
         // Verify the coherence of the PE file
@@ -105,6 +110,8 @@ public override unsafe void Write(PEImageWriter writer)
         // Update OptionalHeader
         OptionalHeader.OptionalHeaderCommonPart1.AddressOfEntryPoint = OptionalHeader.AddressOfEntryPoint.RVA();
         OptionalHeader.OptionalHeaderCommonPart1.BaseOfCode = OptionalHeader.BaseOfCode?.RVA ?? 0;
+
+        var optionalHeaderPosition = position;
 
         if (IsPE32)
         {
@@ -244,5 +251,61 @@ public override unsafe void Write(PEImageWriter writer)
         {
             writer.Diagnostics.Error(DiagnosticId.PE_ERR_InvalidInternalState, $"Generated size {position} does not match expecting size {Size}");
         }
+
+        if (writer.Options.EnableChecksum)
+        {
+            CalculateAndApplyChecksum((MemoryStream)writer.Stream, optionalHeaderPosition);
+        }
+    }
+
+    private void CalculateAndApplyChecksum(MemoryStream stream, uint optionalHeaderPosition)
+    {
+        var data = stream.GetBuffer();
+        var length = (int)stream.Length;
+        var buffer = new Span<byte>(data, 0, length);
+
+        // Zero the checksum before calculating it
+        MemoryMarshal.Write(buffer.Slice((int)optionalHeaderPosition + 64), 0);
+
+        var checksum = CalculateChecksum(buffer);
+
+        // Update the checksum in the PE header
+        MemoryMarshal.Write(buffer.Slice((int)optionalHeaderPosition + 64), checksum);
+    }
+
+    private static uint CalculateChecksum(Span<byte> peFile)
+    {
+        ulong checksum = 0;
+
+        var shortBuffer = MemoryMarshal.Cast<byte, ushort>(peFile);
+        foreach (var value in shortBuffer)
+        {
+            checksum = AggregateChecksum(checksum, value);
+        }
+
+        if ((peFile.Length & 1) != 0)
+        {
+            checksum = AggregateChecksum(checksum, peFile[peFile.Length - 1]);
+        }
+
+        checksum = ((ushort)checksum) + (checksum >> 16);
+        checksum += checksum >> 16;
+        checksum &= 0xffff;
+        checksum += (ulong)peFile.Length;
+
+        return (uint)checksum;
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        static ulong AggregateChecksum(ulong checksum, ushort value)
+        {
+            checksum += value;
+            checksum = (uint)checksum + (checksum >> 32);
+            if (checksum > uint.MaxValue)
+            {
+                checksum = unchecked((uint)checksum) + (checksum >> 32);
+            }
+
+            return checksum;
+        }
     }
-}
+}

src/LibObjectFile/PE/PEImageReaderOptions.cs

@@ -8,10 +8,5 @@ public class PEImageReaderOptions
 {
     public bool UseSubStream { get; init; }
 
-    public bool EnableStackTrace { get; init; }
-}
-
-public class PEImageWriterOptions
-{
     public bool EnableStackTrace { get; init; }
 }

src/LibObjectFile/PE/PEImageWriter.cs

@@ -8,11 +8,14 @@ namespace LibObjectFile.PE;
 
 public sealed class PEImageWriter : ObjectFileReaderWriter
 {
-    internal PEImageWriter(PEFile file, Stream stream) : base(file, stream)
+    internal PEImageWriter(PEFile file, Stream stream, PEImageWriterOptions options) : base(file, stream)
     {
+        Options = options;
     }
 
     public PEFile PEFile => (PEFile)base.File;
 
+    public PEImageWriterOptions Options { get; }
+
     public override bool KeepOriginalStreamForSubStreams => false;
 }

src/LibObjectFile/PE/PEImageWriterOptions.cs

@@ -0,0 +1,14 @@
+// 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.
+
+namespace LibObjectFile.PE;
+
+public class PEImageWriterOptions
+{
+    public static readonly PEImageWriterOptions Default = new();
+
+    public bool EnableStackTrace { get; init; }
+
+    public bool EnableChecksum { get; init; }
+}

src/native/Win64/NativeProjects/NativeConsole2Win64/NativeConsole2Win64.vcxproj

@@ -81,6 +81,7 @@
       <SubSystem>Console</SubSystem>
       <GenerateDebugInformation>true</GenerateDebugInformation>
       <DelayLoadDLLs>NativeLibraryWin64.dll</DelayLoadDLLs>
+      <SetChecksum>true</SetChecksum>
     </Link>
   </ItemDefinitionGroup>
   <ItemDefinitionGroup Condition="'$(Configuration)|$(Platform)'=='Release|Win32'">
@@ -98,6 +99,7 @@
       <OptimizeReferences>true</OptimizeReferences>
       <GenerateDebugInformation>true</GenerateDebugInformation>
       <DelayLoadDLLs>NativeLibraryWin64.dll</DelayLoadDLLs>
+      <SetChecksum>true</SetChecksum>
     </Link>
   </ItemDefinitionGroup>
   <ItemDefinitionGroup Condition="'$(Configuration)|$(Platform)'=='Debug|x64'">
@@ -111,6 +113,7 @@
       <SubSystem>Console</SubSystem>
       <GenerateDebugInformation>true</GenerateDebugInformation>
       <DelayLoadDLLs>NativeLibraryWin64.dll</DelayLoadDLLs>
+      <SetChecksum>true</SetChecksum>
     </Link>
   </ItemDefinitionGroup>
   <ItemDefinitionGroup Condition="'$(Configuration)|$(Platform)'=='Release|x64'">
@@ -128,6 +131,7 @@
       <OptimizeReferences>true</OptimizeReferences>
       <GenerateDebugInformation>true</GenerateDebugInformation>
       <DelayLoadDLLs>NativeLibraryWin64.dll</DelayLoadDLLs>
+      <SetChecksum>true</SetChecksum>
     </Link>
   </ItemDefinitionGroup>
   <ItemGroup>