From ae5b66d50dc43e4c054a79be213f33baca6cdd25 Mon Sep 17 00:00:00 2001
From: hasherezade <hasherezade@gmail.com>
Date: Tue, 29 Oct 2024 07:35:45 -0700
Subject: [PATCH] [BUGFIX] Trim raw section size to the actual raw size (if the
 section start is correct. Round up virtual size to PAGE_SIZE

---
 libpeconv/include/peconv/pe_hdrs_helper.h | 14 ++++++++++++++
 libpeconv/src/pe_hdrs_helper.cpp          |  8 +-------
 libpeconv/src/pe_raw_to_virtual.cpp       | 13 +++++++++----
 3 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/libpeconv/include/peconv/pe_hdrs_helper.h b/libpeconv/include/peconv/pe_hdrs_helper.h
index 4dc8fa4f3..ee957d44e 100644
--- a/libpeconv/include/peconv/pe_hdrs_helper.h
+++ b/libpeconv/include/peconv/pe_hdrs_helper.h
@@ -18,6 +18,20 @@ namespace peconv {
     */
     const ULONGLONG MAX_HEADER_SIZE = PAGE_SIZE;
 
+    template <typename INT_TYPE>
+    INT_TYPE round_up_to_unit(const INT_TYPE size, const INT_TYPE unit)
+    {
+        if (unit == 0) {
+            return size;
+        }
+        INT_TYPE units_count = size / unit;
+        INT_TYPE rounded_size = units_count * unit;
+        if (rounded_size < size) {
+            rounded_size += unit;
+        }
+        return rounded_size;
+    }
+
     /**
     Fetch image size from headers.
     */
diff --git a/libpeconv/src/pe_hdrs_helper.cpp b/libpeconv/src/pe_hdrs_helper.cpp
index f7543feee..b81038d0c 100644
--- a/libpeconv/src/pe_hdrs_helper.cpp
+++ b/libpeconv/src/pe_hdrs_helper.cpp
@@ -553,13 +553,7 @@ DWORD peconv::get_virtual_sec_size(IN const BYTE* pe_hdr, IN const PIMAGE_SECTIO
     }
     //TODO: calculate real size, round up to Virtual Alignment
     DWORD alignment = peconv::get_sec_alignment((const PBYTE)pe_hdr, false);
-    DWORD vsize = sec_hdr->Misc.VirtualSize;
-    if (alignment != 0) {
-        DWORD units = vsize / alignment;
-        if ((vsize % alignment) > 0) units++;
-
-        vsize = units * alignment;
-    }
+    DWORD vsize = peconv::round_up_to_unit(sec_hdr->Misc.VirtualSize, alignment);
     DWORD image_size = peconv::get_image_size(pe_hdr);
     //if it is bigger than the image size, use the size from the headers
     if ((sec_hdr->VirtualAddress + vsize) > image_size) {
diff --git a/libpeconv/src/pe_raw_to_virtual.cpp b/libpeconv/src/pe_raw_to_virtual.cpp
index bdb7d1d24..c08b7ed09 100644
--- a/libpeconv/src/pe_raw_to_virtual.cpp
+++ b/libpeconv/src/pe_raw_to_virtual.cpp
@@ -64,10 +64,15 @@ bool sections_raw_to_virtual(IN const BYTE* payload, IN SIZE_T payloadSize, OUT
             std::cerr << "[-] Raw section size is out ouf bounds: " << std::hex << sec_size << std::endl;
             return false;
         }
+
         // validate source:
         if (!validate_ptr((const LPVOID)payload, payloadSize, section_raw_ptr, sec_size)) {
-            std::cerr << "[-] Section " << i << ":  out ouf bounds, skipping... " << std::endl;
-            continue;
+            if (next_sec->PointerToRawData > payloadSize) {
+                std::cerr << "[-] Section " << i << ":  out ouf bounds, skipping... " << std::endl;
+                continue;
+            }
+            // trim section
+            sec_size = payloadSize - (next_sec->PointerToRawData);
         }
         // validate destination:
         if (!peconv::validate_ptr(destBuffer, destBufferSize, section_mapped, sec_size)) {
@@ -110,7 +115,7 @@ BYTE* peconv::pe_raw_to_virtual(
     }
     DWORD payloadImageSize = 0;
 
-    bool is64 = is64bit(payload);
+    const bool is64 = is64bit(payload);
     if (is64) {
         IMAGE_NT_HEADERS64* payload_nt_hdr = (IMAGE_NT_HEADERS64*)nt_hdr;
         payloadImageSize = payload_nt_hdr->OptionalHeader.SizeOfImage;
@@ -119,9 +124,9 @@ BYTE* peconv::pe_raw_to_virtual(
         IMAGE_NT_HEADERS32* payload_nt_hdr = (IMAGE_NT_HEADERS32*)nt_hdr;
         payloadImageSize = payload_nt_hdr->OptionalHeader.SizeOfImage;
     }
+    payloadImageSize = peconv::round_up_to_unit(payloadImageSize, (DWORD)PAGE_SIZE);
 
     DWORD protect = executable ? PAGE_EXECUTE_READWRITE : PAGE_READWRITE;
-
     //first we will prepare the payload image in the local memory, so that it will be easier to edit it, apply relocations etc.
     //when it will be ready, we will copy it into the space reserved in the target process
     BYTE* localCopyAddress = alloc_pe_buffer(payloadImageSize, protect, desired_base);