From e0efc0d9abc9a437a69098d51c8f7ec103a94351 Mon Sep 17 00:00:00 2001
From: Matt Nadareski <mnadareski@outlook.com>
Date: Fri, 2 Dec 2022 22:44:55 -0800
Subject: [PATCH] Cache PE overlay data

---
 BurnOutSharp.Wrappers/PortableExecutable.cs | 59 +++++++++++++++++++--
 BurnOutSharp.Wrappers/WrapperBase.cs        | 29 +++++++++-
 2 files changed, 82 insertions(+), 6 deletions(-)

diff --git a/BurnOutSharp.Wrappers/PortableExecutable.cs b/BurnOutSharp.Wrappers/PortableExecutable.cs
index 6c5079d3..8fe1a9ea 100644
--- a/BurnOutSharp.Wrappers/PortableExecutable.cs
+++ b/BurnOutSharp.Wrappers/PortableExecutable.cs
@@ -325,6 +325,7 @@ namespace BurnOutSharp.Wrappers
         /// <summary>
         /// Overlay data, if it exists
         /// </summary>
+        /// <see href="https://www.autoitscript.com/forum/topic/153277-pe-file-overlay-extraction/"/>
         public byte[] Overlay
         {
             get
@@ -335,8 +336,60 @@ namespace BurnOutSharp.Wrappers
                     if (_overlayData != null)
                         return _overlayData;
 
-                    // TODO: Implement from https://www.autoitscript.com/forum/topic/153277-pe-file-overlay-extraction/
-                    return null;
+                    // Get the end of the file, if possible
+                    int endOfFile = GetEndOfFile();
+                    if (endOfFile == -1)
+                        return null;
+
+                    // If we have certificate data, use that as the end
+                    if (_executable.OptionalHeader?.CertificateTable != null)
+                    {
+                        var certificateTable = _executable.OptionalHeader.CertificateTable;
+                        int certificateTableAddress = (int)certificateTable.VirtualAddress.ConvertVirtualAddress(_executable.SectionTable);
+                        if (certificateTableAddress != 0)
+                            endOfFile = certificateTableAddress;
+                    }
+
+                    // Search through all sections and find the furthest a section goes
+                    int endOfSectionData = -1;
+                    foreach (var section in _executable.SectionTable)
+                    {
+                        // If we have an invalid section address
+                        int sectionAddress = (int)section.VirtualAddress.ConvertVirtualAddress(_executable.SectionTable);
+                        if (sectionAddress == 0)
+                            continue;
+
+                        // If we have an invalid section size
+                        if (section.SizeOfRawData == 0 && section.VirtualSize == 0)
+                            continue;
+
+                        // Get the real section size
+                        int sectionSize;
+                        if (section.SizeOfRawData < section.VirtualSize)
+                            sectionSize = (int)section.VirtualSize;
+                        else
+                            sectionSize = (int)section.SizeOfRawData;
+
+                        // Compare and set the end of section data
+                        if (sectionAddress + sectionSize > endOfSectionData)
+                            endOfSectionData = sectionAddress + sectionSize;
+                    }
+
+                    // If we didn't find the end of section data
+                    if (endOfSectionData <= 0)
+                        return null;
+
+                    // If we're at the end of the file, cache an empty byte array
+                    if (endOfSectionData == endOfFile)
+                    {
+                        _overlayData = new byte[0];
+                        return _overlayData;
+                    }
+
+                    // Otherwise, cache and return the data
+                    int overlayLength = endOfFile - endOfSectionData;
+                    _overlayData = ReadFromDataSource(endOfSectionData, overlayLength);
+                    return _overlayData;
                 }
             }
         }
@@ -517,8 +570,6 @@ namespace BurnOutSharp.Wrappers
 
         #endregion
 
-        // TODO: Determine what extension properties are needed
-
         #endregion
 
         #region Instance Variables
diff --git a/BurnOutSharp.Wrappers/WrapperBase.cs b/BurnOutSharp.Wrappers/WrapperBase.cs
index afeb74dc..604ab480 100644
--- a/BurnOutSharp.Wrappers/WrapperBase.cs
+++ b/BurnOutSharp.Wrappers/WrapperBase.cs
@@ -50,7 +50,7 @@ namespace BurnOutSharp.Wrappers
                 // Stream data requires both a valid stream
                 case DataSource.Stream:
                     return _streamData != null && _streamData.CanRead && _streamData.CanSeek;
-                
+
                 // Everything else is invalid
                 case DataSource.UNKNOWN:
                 default:
@@ -101,7 +101,7 @@ namespace BurnOutSharp.Wrappers
             if (!SegmentValid(position, length))
                 return null;
 
-            // Read and retuen the data
+            // Read and return the data
             byte[] sectionData = null;
             switch (_dataSource)
             {
@@ -121,6 +121,31 @@ namespace BurnOutSharp.Wrappers
             return sectionData;
         }
 
+        /// <summary>
+        /// Get the ending offset of the source
+        /// </summary>
+        /// <returns>Value greater than 0 for a valid end of file, -1 on error</returns>
+        protected int GetEndOfFile()
+        {
+            // Validate the data souece
+            if (!DataSourceIsValid())
+                return -1;
+
+            // Return the effective endpoint
+            switch (_dataSource)
+            {
+                case DataSource.ByteArray:
+                    return _byteArrayData.Length - _byteArrayOffset;
+
+                case DataSource.Stream:
+                    return (int)_streamData.Length;
+
+                case DataSource.UNKNOWN:
+                default:
+                    return -1;
+            }
+        }
+
         #endregion
 
         #region Printing