Refs #overwrite-fix, temp commit

src/PePatch.hpp

@@ -6,6 +6,7 @@
 namespace PeEpIntercept {
     class PePatch : public PeFile {
     protected:
+        SectionHeader new_section_header;
     public:
         std::vector<char> Assemble(const std::string &assembly);
         explicit PePatch(std::string &path);

src/PePatchX64.cpp

@@ -1,5 +1,6 @@
 #include <cstring>
 #include "PePatchX64.hpp"
+#include "PeStructs.hpp"
 
 namespace PeEpIntercept {
     PePatchX64::PePatchX64(std::string &path) : PePatch(path) {
@@ -25,6 +26,16 @@ namespace PeEpIntercept {
             uint32_t next_section = first_section + section_index;
             auto hdr = *(SectionHeaderPtr) &raw_buffer[next_section];
             section_headers.push_back(hdr);
+
+            // Make a copy of the first n bytes of each section
+            auto raw_data_offset = hdr.PointerToRawData;
+            auto raw_data_size = hdr.SizeOfRawData;
+            uint32_t bytes = (raw_data_size < 100) ? raw_data_size : 100;
+            auto start = file_buffer.begin() + raw_data_offset;
+            auto end = start + bytes;
+
+            std::vector<char> section_bytes;
+            section_bytes.assign(start, end);
         }
 
         auto rest_of_data = first_section;
@@ -46,9 +57,24 @@ namespace PeEpIntercept {
         new_section.Misc.VirtualSize = Align(
                 aligned_size,
                 optional_header.SectionAlignment);
-        new_section.PointerToRawData = Align(
+        auto new_pointer_to_raw_data = Align(
                 last_section.SizeOfRawData + last_section.PointerToRawData,
                 optional_header.FileAlignment);
+
+        // Sanity check that the new section's raw data does not overwrite
+        for (auto header : section_headers) {
+            auto existing_pointer = header.PointerToRawData;
+
+            if (new_pointer_to_raw_data == existing_pointer) {
+                // Remove reinterpret
+                std::string section_name = reinterpret_cast<char *>(header.Name);
+                throw std::runtime_error(
+                        "Cannot create new section. Section, \""
+                        + section_name + "\" already has that starting offset.");
+            }
+        }
+
+        new_section.PointerToRawData = new_pointer_to_raw_data;
         new_section.VirtualAddress = Align(
                 last_section.Misc.VirtualSize + last_section.VirtualAddress,
                 optional_header.SectionAlignment);
@@ -69,7 +95,7 @@ namespace PeEpIntercept {
         optional_header.AddressOfEntryPoint = new_section.VirtualAddress;
         optional_header.SizeOfImage =
                 new_section.VirtualAddress + new_section.Misc.VirtualSize;
-        section_headers.push_back(new_section);
+        new_section_header = new_section;
     }
 
     void PePatchX64::SaveFile(std::string new_path, std::vector<char> code_buffer) {
@@ -101,16 +127,33 @@ namespace PeEpIntercept {
             address += sizeof(section_header);
         }
 
-        auto new_section = &section_headers.back();
-        uint32_t code_position = new_section->PointerToRawData;
-        file_input.seekg(code_position);
-
-        // Padding might be required otherwise the loader will fail
-        // when loading the executable
-        while (code_buffer.size() < new_section->SizeOfRawData) {
-            code_buffer.push_back(0);
+        uint32_t code_position = new_section_header.PointerToRawData;
+
+        // The file may have data appended to it.
+        // This is outside the PE image size
+        if (code_position < file_buffer.size()) {
+            SectionHeader last_section = section_headers.back();
+            auto code_start = last_section.PointerToRawData;
+            auto code_size = last_section.SizeOfRawData;
+            auto offset = code_start + code_size;
+            file_input.seekg(offset);
+
+            while(file_input.tellp() != file_input.eof()) {
+                file_input.put(0);
+//            file_input.seekg(new_section_header.PointerToRawData + new_section_header.SizeOfRawData);
+//            char *raw_data = file_buffer.data();
+//            file_input.write(&raw_data[offset], file_buffer.size() - offset);
+            }
         }
 
-        file_input.write(code_buffer.data(), code_buffer.size());
+//        file_input.seekg(code_position);
+//
+//        // Padding might be required otherwise the loader will fail
+//        // when loading the executable
+//        while (code_buffer.size() < new_section_header.SizeOfRawData) {
+//            code_buffer.push_back(0);
+//        }
+//
+//        file_input.write(code_buffer.data(), code_buffer.size());
     }
 }