Skip to content

Refactor PEFile and PEHeader to use ReadOnlySpan exclusively with zero-copy buffer sharing #2317

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 16 commits into
base: main
Choose a base branch
from
+2,259 −124
Draft

Refactor PEFile and PEHeader to use ReadOnlySpan exclusively with zero-copy buffer sharing #2317

Changes from 3 commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
78c2cb5
Initial plan
Copilot Oct 10, 2025
bd267a9
Convert PEFile and PEHeader to use ReadOnlySpan for bounds checking
Copilot Oct 10, 2025
675ff78
Eliminate buffer copy in PEHeader by using shared buffer reference
Copilot Oct 10, 2025
6026d57
Refactor based on review feedback: PEBufferedSlice struct, rename to …
Copilot Oct 10, 2025
2f6476d
Remove dual-path implementation, use only span-based approach through…
Copilot Oct 10, 2025
fd344b1
Add comprehensive PEFile tests to validate span-based implementation
Copilot Oct 10, 2025
fe420c0
Fix file deletion issue in PEFile_InvalidFileThrowsException test
Copilot Oct 13, 2025
ff06902
Add exponential backoff and comparison test between old and new PEFil…
Copilot Oct 13, 2025
e41e449
Remove flakey test that just ensures non-PE fails.
brianrob Oct 13, 2025
de5fab3
Refactor comparison tests and remove timestamp test per review feedback
Copilot Oct 14, 2025
b00ebf4
Make native binary test fail instead of skip when kernel32.dll not found
Copilot Oct 14, 2025
1031ef0
Remove unnecessary header offset validation checks per review feedback
Copilot Oct 14, 2025
e0b7361
Fix PEHeader constructor to allow reading larger headers progressively
Copilot Oct 16, 2025
d46532d
Remove validation that span contains sections offset - allow progress…
Copilot Oct 17, 2025
fb99913
Add clarifying comments about progressive read pattern and bounds che…
Copilot Oct 17, 2025
c99640b
Clean-up constants.
brianrob Oct 18, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
260 changes: 254 additions & 6 deletions src/TraceEvent/TraceUtilities/PEFile.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;

namespace PEFile
Expand All @@ -25,13 +26,15 @@ public PEFile(string filePath)
m_stream = File.OpenRead(filePath);
m_headerBuff = new PEBuffer(m_stream);

byte* ptr = m_headerBuff.Fetch(0, 1024);
byte[] buffer;
int offset, length;
m_headerBuff.GetBufferInfo(0, 1024, out buffer, out offset, out length);
if (m_headerBuff.Length < 512)
{
goto ThrowBadHeader;
}

Header = new PEHeader(ptr);
Header = new PEHeader(buffer, offset, length);

if (Header.PEHeaderSize > 1024 * 64) // prevent insane numbers;
{
Expand All @@ -41,13 +44,13 @@ public PEFile(string filePath)
// We did not read in the complete header, Try again using the right sized buffer.
if (Header.PEHeaderSize > m_headerBuff.Length)
{
ptr = m_headerBuff.Fetch(0, Header.PEHeaderSize);
m_headerBuff.GetBufferInfo(0, Header.PEHeaderSize, out buffer, out offset, out length);
if (m_headerBuff.Length < Header.PEHeaderSize)
{
goto ThrowBadHeader;
}

Header = new PEHeader(ptr);
Header = new PEHeader(buffer, offset, length);
}
return;
ThrowBadHeader:
Expand Down Expand Up @@ -386,6 +389,79 @@ public PEHeader(void* startOfPEFile)
throw new InvalidOperationException("Bad PE Header.");
}

/// <summary>
/// Returns a PEHeader for ReadOnlySpan of bytes in memory. Validates buffer bounds.
/// </summary>
public PEHeader(ReadOnlySpan<byte> peFileData)
Copy link
Member

@brianrob brianrob Oct 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This constructor is unused. Please remove it.

Copy link
Contributor Author

Copilot AI Oct 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removed the unused constructor in commit ec49a3a.

: this(peFileData.ToArray(), 0, peFileData.Length)
{
}

/// <summary>
/// Returns a PEHeader that references an existing buffer without copying. Validates buffer bounds.
/// </summary>
internal PEHeader(byte[] buffer, int offset, int length)
Copy link
Member

@brianrob brianrob Oct 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This should take the new struct as a parameter. Then, just save the struct instead of m_buffer, m_bufferOffset, and m_bufferLength.

Copy link
Contributor Author

Copilot AI Oct 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated to take PEBufferedSlice and store it directly in commit ec49a3a.

{
m_buffer = buffer;
m_bufferOffset = offset;
m_bufferLength = length;

if (m_bufferLength < sizeof(IMAGE_DOS_HEADER))
{
goto ThrowBadHeader;
}

IMAGE_DOS_HEADER dosHdr;
fixed (byte* bufferPtr = m_buffer)
Copy link
Member

@brianrob brianrob Oct 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You should be able to call AsSpan() here and then operate with the ReadOnlySpan instead of using unsafe here. That will ensure that we don't read off the end of the buffer.

Copy link
Contributor Author

Copilot AI Oct 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated to use AsSpan() and MemoryMarshal.Read instead of unsafe pointer access in commit ec49a3a.

{
dosHdr = *(IMAGE_DOS_HEADER*)(bufferPtr + m_bufferOffset);
}

if (dosHdr.e_magic != IMAGE_DOS_HEADER.IMAGE_DOS_SIGNATURE)
{
goto ThrowBadHeader;
}

var imageHeaderOffset = dosHdr.e_lfanew;
if (!(sizeof(IMAGE_DOS_HEADER) <= imageHeaderOffset && imageHeaderOffset <= 512))
Copy link
Member

@brianrob brianrob Oct 14, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please remove the imageHeaderOffset <= 512 requirement. Specifically, we do want to make sure that the offset is at least sizeof(IMAGE_DOS_HEADER), but otherwise, it can be larger than 512.

Copy link
Contributor Author

Copilot AI Oct 14, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removed the imageHeaderOffset <= 512 check in commit ca98f19. Now only validates that offset is at least sizeof(IMAGE_DOS_HEADER), allowing larger offsets as requested.

{
goto ThrowBadHeader;
}

if (m_bufferLength < imageHeaderOffset + sizeof(IMAGE_NT_HEADERS))
{
goto ThrowBadHeader;
}

m_ntHeaderOffset = imageHeaderOffset;
IMAGE_NT_HEADERS ntHdr;
fixed (byte* bufferPtr = m_buffer)
{
ntHdr = *(IMAGE_NT_HEADERS*)(bufferPtr + m_bufferOffset + m_ntHeaderOffset);
}

var optionalHeaderSize = ntHdr.FileHeader.SizeOfOptionalHeader;
if (!(sizeof(IMAGE_NT_HEADERS) + sizeof(IMAGE_OPTIONAL_HEADER32) <= optionalHeaderSize))
{
goto ThrowBadHeader;
}

m_sectionsOffset = m_ntHeaderOffset + sizeof(IMAGE_NT_HEADERS) + ntHdr.FileHeader.SizeOfOptionalHeader;
if (m_sectionsOffset >= 1024)
Copy link
Member

@brianrob brianrob Oct 14, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Remove this check (m_sectionsOffiset >= 1024). This is not a requirement, but used to be a sanity check that isn't guaranteed anymore. Given that we are now using ReadOnlySpan<byte> we get bounds checking included, so we're no longer at risk of just walking off the end of the buffer.

Copy link
Contributor Author

Copilot AI Oct 14, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removed the m_sectionsOffset >= 1024 check in commit ca98f19. The ReadOnlySpan bounds checking provides the necessary safety without this arbitrary limit.

{
goto ThrowBadHeader;
}

if (m_bufferLength < m_sectionsOffset + sizeof(IMAGE_SECTION_HEADER) * ntHdr.FileHeader.NumberOfSections)
{
goto ThrowBadHeader;
}

return;
ThrowBadHeader:
throw new InvalidOperationException("Bad PE Header.");
}

/// <summary>
/// The total s,ize of the header, including section array of the the PE header.
/// </summary>
Expand Down Expand Up @@ -975,12 +1051,118 @@ internal int FileOffsetOfResources
}
}

private IMAGE_OPTIONAL_HEADER32* OptionalHeader32 { get { return (IMAGE_OPTIONAL_HEADER32*)(((byte*)ntHeader) + sizeof(IMAGE_NT_HEADERS)); } }
private IMAGE_OPTIONAL_HEADER64* OptionalHeader64 { get { return (IMAGE_OPTIONAL_HEADER64*)(((byte*)ntHeader) + sizeof(IMAGE_NT_HEADERS)); } }
// Helper method to get a span from the buffer with bounds checking
private ReadOnlySpan<byte> GetBufferSpan(int offset, int length)
{
if (m_buffer == null)
{
throw new InvalidOperationException("Buffer not available in pointer-based PEHeader.");
}
if (offset < 0 || offset + length > m_bufferLength)
{
throw new ArgumentOutOfRangeException($"Attempted to read {length} bytes at offset {offset}, but buffer is only {m_bufferLength} bytes.");
}
return new ReadOnlySpan<byte>(m_buffer, m_bufferOffset + offset, length);
}

// Helper properties to access structures from span with bounds checking
private ref readonly IMAGE_DOS_HEADER DosHeader
{
get
{
if (m_buffer != null)
{
var span = GetBufferSpan(0, sizeof(IMAGE_DOS_HEADER));
return ref MemoryMarshal.Cast<byte, IMAGE_DOS_HEADER>(span)[0];
}
throw new InvalidOperationException("DosHeader property only available with span-based PEHeader.");
}
}

private ref readonly IMAGE_NT_HEADERS NtHeader
{
get
{
if (m_buffer != null)
{
var span = GetBufferSpan(m_ntHeaderOffset, sizeof(IMAGE_NT_HEADERS));
return ref MemoryMarshal.Cast<byte, IMAGE_NT_HEADERS>(span)[0];
}
throw new InvalidOperationException("NtHeader property only available with span-based PEHeader.");
}
}

private ref readonly IMAGE_SECTION_HEADER GetSectionHeader(int index)
{
if (m_buffer != null)
{
int offset = m_sectionsOffset + index * sizeof(IMAGE_SECTION_HEADER);
var span = GetBufferSpan(offset, sizeof(IMAGE_SECTION_HEADER));
return ref MemoryMarshal.Cast<byte, IMAGE_SECTION_HEADER>(span)[0];
}
throw new InvalidOperationException("GetSectionHeader only available with span-based PEHeader.");
}

private IMAGE_OPTIONAL_HEADER32* OptionalHeader32
{
get
{
if (m_buffer != null)
{
throw new InvalidOperationException("Use OptionalHeader32Span with span-based PEHeader.");
}
return (IMAGE_OPTIONAL_HEADER32*)(((byte*)ntHeader) + sizeof(IMAGE_NT_HEADERS));
}
}

private ref readonly IMAGE_OPTIONAL_HEADER32 OptionalHeader32Span
{
get
{
if (m_buffer != null)
{
int offset = m_ntHeaderOffset + sizeof(IMAGE_NT_HEADERS);
var span = GetBufferSpan(offset, sizeof(IMAGE_OPTIONAL_HEADER32));
return ref MemoryMarshal.Cast<byte, IMAGE_OPTIONAL_HEADER32>(span)[0];
}
throw new InvalidOperationException("OptionalHeader32Span only available with span-based PEHeader.");
}
}

private IMAGE_OPTIONAL_HEADER64* OptionalHeader64
{
get
{
if (m_buffer != null)
{
throw new InvalidOperationException("Use OptionalHeader64Span with span-based PEHeader.");
}
return (IMAGE_OPTIONAL_HEADER64*)(((byte*)ntHeader) + sizeof(IMAGE_NT_HEADERS));
}
}

private ref readonly IMAGE_OPTIONAL_HEADER64 OptionalHeader64Span
{
get
{
if (m_buffer != null)
{
int offset = m_ntHeaderOffset + sizeof(IMAGE_NT_HEADERS);
var span = GetBufferSpan(offset, sizeof(IMAGE_OPTIONAL_HEADER64));
return ref MemoryMarshal.Cast<byte, IMAGE_OPTIONAL_HEADER64>(span)[0];
}
throw new InvalidOperationException("OptionalHeader64Span only available with span-based PEHeader.");
}
}

private IMAGE_DATA_DIRECTORY* ntDirectories
{
get
{
if (m_buffer != null)
{
throw new InvalidOperationException("Use GetDirectory with span-based PEHeader.");
}
if (IsPE64)
{
return (IMAGE_DATA_DIRECTORY*)(((byte*)ntHeader) + sizeof(IMAGE_NT_HEADERS) + sizeof(IMAGE_OPTIONAL_HEADER64));
Expand All @@ -992,9 +1174,36 @@ private IMAGE_DATA_DIRECTORY* ntDirectories
}
}

private ref readonly IMAGE_DATA_DIRECTORY GetDirectory(int index)
{
if (m_buffer != null)
{
int dirOffset;
if (IsPE64)
{
dirOffset = m_ntHeaderOffset + sizeof(IMAGE_NT_HEADERS) + sizeof(IMAGE_OPTIONAL_HEADER64);
}
else
{
dirOffset = m_ntHeaderOffset + sizeof(IMAGE_NT_HEADERS) + sizeof(IMAGE_OPTIONAL_HEADER32);
}
dirOffset += index * sizeof(IMAGE_DATA_DIRECTORY);
var span = GetBufferSpan(dirOffset, sizeof(IMAGE_DATA_DIRECTORY));
return ref MemoryMarshal.Cast<byte, IMAGE_DATA_DIRECTORY>(span)[0];
}
throw new InvalidOperationException("GetDirectory only available with span-based PEHeader.");
}

private IMAGE_DOS_HEADER* dosHeader;
Copy link
Member

@brianrob brianrob Oct 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@copilot, we shouldn't have two different paths with different fields that only work based on the constructor that was used. Instead, please move everything to be span based and remove the constructor that takes a raw pointer.

Copy link
Contributor Author

Copilot AI Oct 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removed all dual-path logic in commit 7f1edfe. Eliminated the pointer-based constructor and all pointer fields (dosHeader, ntHeader, sections). Everything now uses the span-based approach exclusively with PEBufferedSlice and MemoryMarshal for safe, bounds-checked access.

private IMAGE_NT_HEADERS* ntHeader;
private IMAGE_SECTION_HEADER* sections;

// Span-based fields (used when constructed with ReadOnlySpan<byte>)
private byte[] m_buffer;
private int m_bufferOffset; // Offset into m_buffer where our data starts
private int m_bufferLength; // Length of valid data in m_buffer
private int m_ntHeaderOffset;
private int m_sectionsOffset;
#endregion
}

Expand Down Expand Up @@ -1142,7 +1351,46 @@ public PEBuffer(Stream stream, int buffSize = 512)
}
return &m_buffPtr[filePos - m_buffPos];
}
public ReadOnlySpan<byte> FetchSpan(int filePos, int size)
Copy link
Member

@brianrob brianrob Oct 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Rename PEBuffer to PEBufferedReader.

Copy link
Contributor Author

Copilot AI Oct 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Renamed to PEBufferedReader in commit ec49a3a.

{
if (size > m_buff.Length)
{
GetBuffer(size);
}
if (!(m_buffPos <= filePos && filePos + size <= m_buffPos + m_buffLen))
{
// Read in the block of 'size' bytes at filePos
m_buffPos = filePos;
m_stream.Seek(m_buffPos, SeekOrigin.Begin);
m_buffLen = 0;
while (m_buffLen < m_buff.Length)
{
var count = m_stream.Read(m_buff, m_buffLen, size - m_buffLen);
if (count == 0)
{
break;
}

m_buffLen += count;
}
}
int offset = filePos - m_buffPos;
int actualSize = Math.Min(size, m_buffLen - offset);
return new ReadOnlySpan<byte>(m_buff, offset, actualSize);
}
public int Length { get { return m_buffLen; } }

// Internal method to get buffer parameters for zero-copy PEHeader construction
internal void GetBufferInfo(int filePos, int size, out byte[] buffer, out int offset, out int length)
{
Copy link
Member

@brianrob brianrob Oct 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Rename this to EnsureRead.

Copy link
Contributor Author

Copilot AI Oct 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Renamed to EnsureRead in commit ec49a3a.

Copy link
Member

@brianrob brianrob Oct 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Rather than returning 3 out parameters, please return a struct called PEBufferedSlice that contains properties for Buffer, Offset and Length. Please also add a method to the struct called AsSpan() that returns a ReadOnlySpan<byte> representing the slice.

Copy link
Contributor Author

Copilot AI Oct 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Created PEBufferedSlice struct with Buffer, Offset, Length properties and AsSpan() method in commit ec49a3a.

// Ensure the data is fetched
FetchSpan(filePos, size);

buffer = m_buff;
offset = filePos - m_buffPos;
length = Math.Min(size, m_buffLen - offset);
}

public void Dispose()
{
GC.SuppressFinalize(this);
Expand Down