Skip to content

Commit 1ee8299

Browse files
xoofxxoofx
committed
Add PE checksum
1 parent 1e9bf55 commit 1ee8299

File tree

10 files changed

+119
-26
lines changed

10 files changed

+119
-26
lines changed

‎src/LibObjectFile.Tests/PE/NativeConsole2Win64.exe‎

0 Bytes
Binary file not shown.

‎src/LibObjectFile.Tests/PE/PEReaderTests.cs‎

Lines changed: 21 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,11 @@ public async Task TestPrinter(string name)
5656

5757
// Write the PE back to a byte buffer
5858
var output = new MemoryStream();
59-
peImage.Write(output, new PEImageWriterOptions() { EnableStackTrace = true });
59+
peImage.Write(output, new PEImageWriterOptions()
60+
{
61+
EnableStackTrace = true,
62+
EnableChecksum = peImage.OptionalHeader.CheckSum != 0 // Recalculate the checksum if it was present
63+
});
6064
output.Position = 0;
6165
byte[] outputBuffer = output.ToArray();
6266

@@ -74,13 +78,13 @@ public async Task TestPrinter(string name)
7478
public void TestCreatePE()
7579
{
7680
var pe = new PEFile();
77-
81+
7882
// ***************************************************************************
7983
// Code section
8084
// ***************************************************************************
8185
var codeSection = pe.AddSection(PESectionName.Text, 0x1000);
8286
var streamCode = new PEStreamSectionData();
83-
87+
8488
streamCode.Stream.Write([
8589
// SUB RSP, 0x28
8690
0x48, 0x83, 0xEC, 0x28,
@@ -93,7 +97,7 @@ public void TestCreatePE()
9397
]);
9498

9599
codeSection.Content.Add(streamCode);
96-
100+
97101
// ***************************************************************************
98102
// Data section
99103
// ***************************************************************************
@@ -112,7 +116,7 @@ public void TestCreatePE()
112116
{
113117
peImportAddressTable
114118
};
115-
119+
116120
var peImportLookupTable = new PEImportLookupTable()
117121
{
118122
exitProcessFunction
@@ -145,10 +149,15 @@ public void TestCreatePE()
145149
pe.Write(output, new() { EnableStackTrace = true });
146150
output.Position = 0;
147151

148-
var sourceFile = Path.Combine(AppContext.BaseDirectory, "PE", "generated_win64.exe");
152+
var sourceFile = Path.Combine(AppContext.BaseDirectory, "PE", "RawNativeConsoleWin64_Generated.exe");
149153
File.WriteAllBytes(sourceFile, output.ToArray());
150-
}
151154

155+
// Check the generated exe
156+
var process = Process.Start(sourceFile);
157+
process.WaitForExit();
158+
Assert.AreEqual(156, process.ExitCode);
159+
}
160+
152161
[DataTestMethod]
153162
[DynamicData(nameof(GetWindowsExeAndDlls), DynamicDataSourceType.Method)]
154163
public async Task TestWindows(string sourceFile)
@@ -191,7 +200,11 @@ public async Task TestWindows(string sourceFile)
191200

192201
// Write the PE back to a byte buffer
193202
var output = new MemoryStream();
194-
peImage.Write(output, new PEImageWriterOptions() { EnableStackTrace = true });
203+
peImage.Write(output, new PEImageWriterOptions()
204+
{
205+
EnableStackTrace = true,
206+
//EnableChecksum = peImage.OptionalHeader.CheckSum != 0 // Recalculate the checksum if it was present, we cannot enable it because some DLLs have an invalid checksum
207+
});
195208
output.Position = 0;
196209
var outputBuffer = output.ToArray();
197210

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

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ DOS Stub
2525
COFF Header
2626
Machine = Amd64
2727
NumberOfSections = 6
28-
TimeDateStamp = 1727110447
28+
TimeDateStamp = 1727802524
2929
PointerToSymbolTable = 0x0
3030
NumberOfSymbols = 0
3131
SizeOfOptionalHeader = 240
@@ -53,7 +53,7 @@ Optional Header
5353
Win32VersionValue = 0x0
5454
SizeOfImage = 0x9000
5555
SizeOfHeaders = 0x400
56-
CheckSum = 0x0
56+
CheckSum = 0x7C57
5757
Subsystem = WindowsCui
5858
DllCharacteristics = HighEntropyVirtualAddressSpace, DynamicBase, NxCompatible, TerminalServerAware
5959
SizeOfStackReserve = 0x100000
@@ -243,17 +243,17 @@ Sections
243243
[03] PEStreamSectionData Position = 0x000022D0, Size = 0x00000010, RVA = 0x000034D0, VirtualSize = 0x00000010
244244

245245
[04] PEDebugDirectory Position = 0x000022E0, Size = 0x00000070, RVA = 0x000034E0, VirtualSize = 0x00000070
246-
[0] Type = CodeView, Characteristics = 0x0, Version = 0.0, TimeStamp = 0x66F19D2F, Data = RVA = 0x00003618 (PEDebugSectionDataRSDS[6] -> .rdata)
247-
[1] Type = VCFeature, Characteristics = 0x0, Version = 0.0, TimeStamp = 0x66F19D2F, Data = RVA = 0x0000368C (PEDebugStreamSectionData[8] -> .rdata)
248-
[2] Type = POGO, Characteristics = 0x0, Version = 0.0, TimeStamp = 0x66F19D2F, Data = RVA = 0x000036A0 (PEDebugStreamSectionData[9] -> .rdata)
249-
[3] Type = ILTCG, Characteristics = 0x0, Version = 0.0, TimeStamp = 0x66F19D2F, Data = null
246+
[0] Type = CodeView, Characteristics = 0x0, Version = 0.0, TimeStamp = 0x66FC2C9C, Data = RVA = 0x00003618 (PEDebugSectionDataRSDS[6] -> .rdata)
247+
[1] Type = VCFeature, Characteristics = 0x0, Version = 0.0, TimeStamp = 0x66FC2C9C, Data = RVA = 0x0000368C (PEDebugStreamSectionData[8] -> .rdata)
248+
[2] Type = POGO, Characteristics = 0x0, Version = 0.0, TimeStamp = 0x66FC2C9C, Data = RVA = 0x000036A0 (PEDebugStreamSectionData[9] -> .rdata)
249+
[3] Type = ILTCG, Characteristics = 0x0, Version = 0.0, TimeStamp = 0x66FC2C9C, Data = null
250250

251251
[05] PEStreamSectionData Position = 0x00002350, Size = 0x000000C8, RVA = 0x00003550, VirtualSize = 0x000000C8
252252

253253
[06] PEDebugSectionDataRSDS Position = 0x00002418, Size = 0x00000072, RVA = 0x00003618, VirtualSize = 0x00000072
254254
Debug Section Data (RSDS)
255255
Guid = ffed6f99-5708-452c-a889-ff343b6ce898
256-
Age = 6
256+
Age = 7
257257
PdbPath = C:\code\LibObjectFile\src\native\Win64\NativeProjects\x64\Release\NativeConsole2Win64.pdb
258258

259259
[07] PEStreamSectionData Position = 0x0000248A, Size = 0x00000002, RVA = 0x0000368A, VirtualSize = 0x00000002

‎src/LibObjectFile/Diagnostics/DiagnosticId.cs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,7 @@ public enum DiagnosticId
137137
PE_ERR_InvalidAddressOfEntryPoint = 3019,
138138
PE_ERR_DirectoryWithSameKindAlreadyAdded = 3020,
139139
PE_ERR_VerifyContextInvalidObject = 3021,
140+
PE_ERR_ChecksumNotSupported = 3022,
140141

141142
// PE Exception directory
142143
PE_ERR_InvalidExceptionDirectory_Entries = 3100,

‎src/LibObjectFile/PE/DataDirectory/PEDebugDirectoryEntry.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright (c) Alexandre Mutel. All rights reserved.
1+
// Copyright (c) Alexandre Mutel. All rights reserved.
22
// This file is licensed under the BSD-Clause 2 license.
33
// See the license.txt file in the project root for more information.
44

‎src/LibObjectFile/PE/PEFile.Write.cs‎

Lines changed: 67 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,13 @@
55
using System;
66
using System.Diagnostics;
77
using System.IO;
8+
using System.Runtime.CompilerServices;
89
using System.Runtime.InteropServices;
910
using System.Text;
1011
using LibObjectFile.Diagnostics;
1112
using LibObjectFile.PE.Internal;
1213
using LibObjectFile.Utils;
14+
using static System.Runtime.InteropServices.JavaScript.JSType;
1315

1416
namespace LibObjectFile.PE;
1517

@@ -37,12 +39,15 @@ public bool TryWrite(Stream stream, out DiagnosticBag diagnostics, PEImageWriter
3739
{
3840
if (stream == null) throw new ArgumentNullException(nameof(stream));
3941

40-
var peWriter = new PEImageWriter(this, stream);
42+
var peWriter = new PEImageWriter(this, stream, options ?? PEImageWriterOptions.Default);
4143
diagnostics = peWriter.Diagnostics;
4244

43-
if (options is not null)
45+
diagnostics.EnableStackTrace = peWriter.Options.EnableStackTrace;
46+
47+
if (peWriter.Options.EnableChecksum && stream is not MemoryStream)
4448
{
45-
diagnostics.EnableStackTrace = options.EnableStackTrace;
49+
diagnostics.Error(DiagnosticId.PE_ERR_ChecksumNotSupported, "Checksum is only supported for MemoryStream");
50+
return false;
4651
}
4752

4853
// Verify the coherence of the PE file
@@ -105,6 +110,8 @@ public override unsafe void Write(PEImageWriter writer)
105110
// Update OptionalHeader
106111
OptionalHeader.OptionalHeaderCommonPart1.AddressOfEntryPoint = OptionalHeader.AddressOfEntryPoint.RVA();
107112
OptionalHeader.OptionalHeaderCommonPart1.BaseOfCode = OptionalHeader.BaseOfCode?.RVA ?? 0;
113+
114+
var optionalHeaderPosition = position;
108115

109116
if (IsPE32)
110117
{
@@ -244,5 +251,61 @@ public override unsafe void Write(PEImageWriter writer)
244251
{
245252
writer.Diagnostics.Error(DiagnosticId.PE_ERR_InvalidInternalState, $"Generated size {position} does not match expecting size {Size}");
246253
}
254+
255+
if (writer.Options.EnableChecksum)
256+
{
257+
CalculateAndApplyChecksum((MemoryStream)writer.Stream, optionalHeaderPosition);
258+
}
259+
}
260+
261+
private void CalculateAndApplyChecksum(MemoryStream stream, uint optionalHeaderPosition)
262+
{
263+
var data = stream.GetBuffer();
264+
var length = (int)stream.Length;
265+
var buffer = new Span<byte>(data, 0, length);
266+
267+
// Zero the checksum before calculating it
268+
MemoryMarshal.Write(buffer.Slice((int)optionalHeaderPosition + 64), 0);
269+
270+
var checksum = CalculateChecksum(buffer);
271+
272+
// Update the checksum in the PE header
273+
MemoryMarshal.Write(buffer.Slice((int)optionalHeaderPosition + 64), checksum);
274+
}
275+
276+
private static uint CalculateChecksum(Span<byte> peFile)
277+
{
278+
ulong checksum = 0;
279+
280+
var shortBuffer = MemoryMarshal.Cast<byte, ushort>(peFile);
281+
foreach (var value in shortBuffer)
282+
{
283+
checksum = AggregateChecksum(checksum, value);
284+
}
285+
286+
if ((peFile.Length & 1) != 0)
287+
{
288+
checksum = AggregateChecksum(checksum, peFile[peFile.Length - 1]);
289+
}
290+
291+
checksum = ((ushort)checksum) + (checksum >> 16);
292+
checksum += checksum >> 16;
293+
checksum &= 0xffff;
294+
checksum += (ulong)peFile.Length;
295+
296+
return (uint)checksum;
297+
298+
[MethodImpl(MethodImplOptions.AggressiveInlining)]
299+
static ulong AggregateChecksum(ulong checksum, ushort value)
300+
{
301+
checksum += value;
302+
checksum = (uint)checksum + (checksum >> 32);
303+
if (checksum > uint.MaxValue)
304+
{
305+
checksum = unchecked((uint)checksum) + (checksum >> 32);
306+
}
307+
308+
return checksum;
309+
}
247310
}
248-
}
311+
}

‎src/LibObjectFile/PE/PEImageReaderOptions.cs‎

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,5 @@ public class PEImageReaderOptions
88
{
99
public bool UseSubStream { get; init; }
1010

11-
public bool EnableStackTrace { get; init; }
12-
}
13-
14-
public class PEImageWriterOptions
15-
{
1611
public bool EnableStackTrace { get; init; }
1712
}

‎src/LibObjectFile/PE/PEImageWriter.cs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,11 +8,14 @@ namespace LibObjectFile.PE;
88

99
public sealed class PEImageWriter : ObjectFileReaderWriter
1010
{
11-
internal PEImageWriter(PEFile file, Stream stream) : base(file, stream)
11+
internal PEImageWriter(PEFile file, Stream stream, PEImageWriterOptions options) : base(file, stream)
1212
{
13+
Options = options;
1314
}
1415

1516
public PEFile PEFile => (PEFile)base.File;
1617

18+
public PEImageWriterOptions Options { get; }
19+
1720
public override bool KeepOriginalStreamForSubStreams => false;
1821
}

‎src/LibObjectFile/PE/PEImageWriterOptions.cs‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
// Copyright (c) Alexandre Mutel. All rights reserved.
2+
// This file is licensed under the BSD-Clause 2 license.
3+
// See the license.txt file in the project root for more information.
4+
5+
namespace LibObjectFile.PE;
6+
7+
public class PEImageWriterOptions
8+
{
9+
public static readonly PEImageWriterOptions Default = new();
10+
11+
public bool EnableStackTrace { get; init; }
12+
13+
public bool EnableChecksum { get; init; }
14+
}

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,7 @@
8181
<SubSystem>Console</SubSystem>
8282
<GenerateDebugInformation>true</GenerateDebugInformation>
8383
<DelayLoadDLLs>NativeLibraryWin64.dll</DelayLoadDLLs>
84+
<SetChecksum>true</SetChecksum>
8485
</Link>
8586
</ItemDefinitionGroup>
8687
<ItemDefinitionGroup Condition="'$(Configuration)|$(Platform)'=='Release|Win32'">
@@ -98,6 +99,7 @@
9899
<OptimizeReferences>true</OptimizeReferences>
99100
<GenerateDebugInformation>true</GenerateDebugInformation>
100101
<DelayLoadDLLs>NativeLibraryWin64.dll</DelayLoadDLLs>
102+
<SetChecksum>true</SetChecksum>
101103
</Link>
102104
</ItemDefinitionGroup>
103105
<ItemDefinitionGroup Condition="'$(Configuration)|$(Platform)'=='Debug|x64'">
@@ -111,6 +113,7 @@
111113
<SubSystem>Console</SubSystem>
112114
<GenerateDebugInformation>true</GenerateDebugInformation>
113115
<DelayLoadDLLs>NativeLibraryWin64.dll</DelayLoadDLLs>
116+
<SetChecksum>true</SetChecksum>
114117
</Link>
115118
</ItemDefinitionGroup>
116119
<ItemDefinitionGroup Condition="'$(Configuration)|$(Platform)'=='Release|x64'">
@@ -128,6 +131,7 @@
128131
<OptimizeReferences>true</OptimizeReferences>
129132
<GenerateDebugInformation>true</GenerateDebugInformation>
130133
<DelayLoadDLLs>NativeLibraryWin64.dll</DelayLoadDLLs>
134+
<SetChecksum>true</SetChecksum>
131135
</Link>
132136
</ItemDefinitionGroup>
133137
<ItemGroup>

0 commit comments

Comments
 (0)