simplify PEDecoder using PEReader

Author: Max Charlamb

src/native/managed/cdacreader/src/Decoder/PEDecoder.cs

@@ -1,82 +1,87 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System;
 using System.IO;
-using Microsoft.Diagnostics.DataContractReader.Decoder.PETypes;
-using Microsoft.Diagnostics.DataContractReader.Legacy;
+using System.Reflection.PortableExecutable;
+using System.Text;
 
 namespace Microsoft.Diagnostics.DataContractReader.Decoder;
-internal sealed class PEDecoder
+
+internal sealed class PEDecoder : IDisposable
 {
-    private readonly ICLRDataTarget _dataTarget;
-    private readonly ulong _baseAddress;
-    private uint _peSigOffset;
-    private ushort _optHeaderMagic;
+    private struct IMAGE_EXPORT_DIRECTORY
+    {
+        public uint Characteristics;
+        public uint TimeDateStamp;
+        public ushort MajorVersion;
+        public ushort MinorVersion;
+        public uint Name;
+        public uint Base;
+        public uint NumberOfFunctions;
+        public uint NumberOfNames;
+        public uint AddressOfFunctions;
+        public uint AddressOfNames;
+        public uint AddressOfNameOrdinals;
+
+        public IMAGE_EXPORT_DIRECTORY(BinaryReader reader)
+        {
+            Characteristics = reader.ReadUInt32();
+            TimeDateStamp = reader.ReadUInt32();
+            MajorVersion = reader.ReadUInt16();
+            MinorVersion = reader.ReadUInt16();
+            Name = reader.ReadUInt32();
+            Base = reader.ReadUInt32();
+            NumberOfFunctions = reader.ReadUInt32();
+            NumberOfNames = reader.ReadUInt32();
+            AddressOfFunctions = reader.ReadUInt32();
+            AddressOfNames = reader.ReadUInt32();
+            AddressOfNameOrdinals = reader.ReadUInt32();
+        }
+    }
+
+    private readonly Stream _stream;
     private IMAGE_EXPORT_DIRECTORY _exportDir;
+    private bool _disposed;
 
     public bool IsValid { get; init; }
 
-    public PEDecoder(ICLRDataTarget dataTarget, ulong baseAddress)
+    /// <summary>
+    /// Create PEDecoder with stream beginning at the base address of the module.
+    /// </summary>
+    public PEDecoder(Stream stream)
     {
-        _dataTarget = dataTarget;
-        _baseAddress = baseAddress;
+        _stream = stream;
 
         IsValid = Initialize();
     }
 
     private bool Initialize()
     {
-        using BinaryReader reader = new(new DataTargetStream(_dataTarget, _baseAddress));
+        using PEReader peReader = new(_stream, PEStreamOptions.LeaveOpen);
 
-        ushort dosMagic = reader.ReadUInt16();
-        if (dosMagic != 0x5A4D) // "MZ"
-            return false;
-
-        // PE Header offset is at 0x3C in DOS header
-        reader.BaseStream.Seek(0x3C, SeekOrigin.Begin);
-        _peSigOffset = reader.ReadUInt32();
-
-        // Read PE signature
-        reader.BaseStream.Seek(_peSigOffset, SeekOrigin.Begin);
-        uint peSig = reader.ReadUInt32();
-        if (peSig != 0x00004550) // "PE00"
-            return false;
-
-        // Seek to beginning of opt header and read magic
-        reader.BaseStream.Seek(_peSigOffset + 0x18, SeekOrigin.Begin);
-        _optHeaderMagic = reader.ReadUInt16();
-
-        // Seek back to beginning of opt header and parse
-        reader.BaseStream.Seek(_peSigOffset + 0x18, SeekOrigin.Begin);
-        uint rva;
-        switch (_optHeaderMagic)
+        PEHeader? peHeader = peReader.PEHeaders.PEHeader;
+        if (peHeader is null)
         {
-            case 0x10B: // PE32
-                IMAGE_OPTIONAL_HEADER32 optHeader32 = new(reader);
-                rva = optHeader32.DataDirectory[0].VirtualAddress;
-                break;
-            case 0x20B: // PE32+
-                IMAGE_OPTIONAL_HEADER64 optHeader64 = new(reader);
-                rva = optHeader64.DataDirectory[0].VirtualAddress;
-                break;
-            // unknown type, invalid
-            default:
-                return false;
+            // If PEHeader is null, then the file is not a valid PE file.
+            return false;
         }
+        DirectoryEntry exportsDirectory = peHeader.ExportTableDirectory;
 
-        // Seek to export directory and parse
-        reader.BaseStream.Seek(rva, SeekOrigin.Begin);
+        _stream.Seek(exportsDirectory.RelativeVirtualAddress, SeekOrigin.Begin);
+        using BinaryReader reader = new(_stream, Encoding.UTF8, leaveOpen: true);
         _exportDir = new IMAGE_EXPORT_DIRECTORY(reader);
 
         return true;
     }
 
-    public TargetPointer GetSymbolAddress(string symbol)
+    public bool TryGetRelativeSymbolAddress(string symbol, out ulong address)
     {
+        address = 0;
         if (!IsValid)
-            return TargetPointer.Null;
+            return false;
 
-        using BinaryReader reader = new(new DataTargetStream(_dataTarget, _baseAddress));
+        using BinaryReader reader = new(_stream, Encoding.UTF8, leaveOpen: true);
 
         for (int nameIndex = 0; nameIndex < _exportDir.NumberOfNames; nameIndex++)
         {
@@ -99,10 +104,31 @@ public TargetPointer GetSymbolAddress(string symbol)
                 reader.BaseStream.Seek(_exportDir.AddressOfFunctions + sizeof(uint) * ordinalForNamedExport, SeekOrigin.Begin);
                 uint symbolRVA = reader.ReadUInt32();
 
-                return new TargetPointer(_baseAddress + symbolRVA);
+                address = symbolRVA;
+                return true;
+            }
+        }
+
+        return false;
+    }
+
+    private void Dispose(bool disposing)
+    {
+        if (!_disposed)
+        {
+            if (disposing)
+            {
+                _stream.Dispose();
             }
+
+            _disposed = true;
         }
+    }
 
-        return TargetPointer.Null;
+    void IDisposable.Dispose()
+    {
+        // Do not change this code. Put cleanup code in 'Dispose(bool disposing)' method
+        Dispose(disposing: true);
+        GC.SuppressFinalize(this);
     }
 }

src/native/managed/cdacreader/src/Decoder/PETypes.cs

@@ -87,9 +87,12 @@ private static unsafe int CLRDataCreateInstance(Guid* pIID, IntPtr /*ICLRDataTar
             dataTarget.GetImageBase("coreclr.dll", &baseAddress);
         }
 
-
-        PEDecoder peDecoder = new(dataTarget, baseAddress);
-        ulong contractDescriptor = peDecoder.GetSymbolAddress("DotNetRuntimeContractDescriptor");
+        DataTargetStream dataTargetStream = new(dataTarget, baseAddress);
+        using PEDecoder peDecoder = new(dataTargetStream);
+        if (!peDecoder.TryGetRelativeSymbolAddress("DotNetRuntimeContractDescriptor", out ulong contractDescriptor))
+        {
+            return -1;
+        }
 
         if (!ContractDescriptorTarget.TryCreate(contractDescriptor, (address, buffer) =>
         {