From 73c18f11f3011d0c92da07dc1c548ff7fe0098fa Mon Sep 17 00:00:00 2001
From: Zanie Blue <contact@zanie.dev>
Date: Wed, 5 Aug 2026 09:30:12 -0500
Subject: [PATCH 1/2] Align PE image size to section boundaries

---
 src/image.rs   | 11 +++++++++--
 tests/tests.rs |  8 ++++++++
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/src/image.rs b/src/image.rs
index c682a91..2bfd179 100644
--- a/src/image.rs
+++ b/src/image.rs
@@ -564,15 +564,22 @@ impl<'a> Image<'a> {
             return Err(ImageWriteError::NotEnoughSpaceInHeader);
         }
 
+        let size_of_image = section_table
+            .iter()
+            .map(|section| section.virtual_address + section.virtual_size)
+            .max()
+            .map(|size| aligned_to(size, windows_header.section_alignment()))
+            .unwrap_or_else(|| windows_header.size_of_image());
+
         match windows_header {
             GenericWindowsHeader::WindowsHeader32(ref mut header) => {
                 header.number_of_rva_and_sizes = header_data_directory.len() as u32;
-                header.size_of_image += new_section_data.len() as u32;
+                header.size_of_image = size_of_image;
                 header.check_sum = 0;
             }
             GenericWindowsHeader::WindowsHeader64(ref mut header) => {
                 header.number_of_rva_and_sizes = header_data_directory.len() as u32;
-                header.size_of_image += new_section_data.len() as u32;
+                header.size_of_image = size_of_image;
                 header.check_sum = 0;
             }
         }
diff --git a/tests/tests.rs b/tests/tests.rs
index 8e2a7af..ee66461 100644
--- a/tests/tests.rs
+++ b/tests/tests.rs
@@ -234,6 +234,14 @@ fn transfer_resource_section_large() {
     );
     image_small.set_resource_directory(target_resource_directory.clone()).unwrap();
 
+    let section_alignment = image_small.windows_header().section_alignment();
+    let size_of_image = image_small.windows_header().size_of_image();
+    assert_eq!(
+        size_of_image % section_alignment,
+        0,
+        "image size is aligned to section boundaries"
+    );
+
     let data_small_rebuilt = image_small.data();
 
     assert!(

From afb30b65706defa4f74c0faf753e4029b2312630 Mon Sep 17 00:00:00 2001
From: Zanie Blue <contact@zanie.dev>
Date: Wed, 5 Aug 2026 10:10:55 -0500
Subject: [PATCH 2/2] Handle overflow when calculating image size

---
 src/image.rs   |  5 +++--
 tests/tests.rs | 46 ++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 47 insertions(+), 4 deletions(-)

diff --git a/src/image.rs b/src/image.rs
index 2bfd179..ea53427 100644
--- a/src/image.rs
+++ b/src/image.rs
@@ -566,9 +566,10 @@ impl<'a> Image<'a> {
 
         let size_of_image = section_table
             .iter()
-            .map(|section| section.virtual_address + section.virtual_size)
+            .map(|section| u64::from(section.virtual_address) + u64::from(section.virtual_size))
             .max()
-            .map(|size| aligned_to(size, windows_header.section_alignment()))
+            .map(|size| aligned_to(size, u64::from(windows_header.section_alignment())))
+            .and_then(|size| u32::try_from(size).ok())
             .unwrap_or_else(|| windows_header.size_of_image());
 
         match windows_header {
diff --git a/tests/tests.rs b/tests/tests.rs
index ee66461..434c125 100644
--- a/tests/tests.rs
+++ b/tests/tests.rs
@@ -1,5 +1,5 @@
-use editpe::{constants::*, *};
-use std::sync::Once;
+use editpe::{constants::*, types::*, *};
+use std::{mem::size_of, sync::Once};
 
 static BINARY_PATH_SMALL: &str = "./tests/assets/smallbin.exe";
 static BINARY_PATH_LARGE: &str = "./tests/assets/smallbin-large.exe";
@@ -167,6 +167,48 @@ fn set_resource_section() {
     assert_eq!(image, new_image, "original and rebuilt images equal");
 }
 
+#[test]
+fn set_resource_section_with_overflowing_image_size() {
+    init_logger();
+
+    for (virtual_size, description) in [
+        (0x2000_u32, "virtual section extent"),
+        (0x0800_u32, "aligned image size"),
+    ] {
+        let mut data = std::fs::read(BINARY_PATH_SMALL).unwrap();
+        let image = Image::parse(&data[..]).unwrap();
+        let original_image_size = image.windows_header().size_of_image();
+
+        let pe_signature_offset = u32::from_le_bytes(
+            data[PE_PTR_OFFSET as usize..PE_PTR_OFFSET as usize + size_of::<u32>()]
+                .try_into()
+                .unwrap(),
+        ) as usize;
+        let last_section_offset = pe_signature_offset
+            + size_of::<u32>()
+            + size_of::<CoffHeader>()
+            + image.coff_header().size_of_optional_header as usize
+            + (image.coff_header().number_of_sections as usize - 1) * size_of::<SectionHeader>();
+        let virtual_size_offset = last_section_offset + size_of::<u64>();
+        let virtual_address_offset = virtual_size_offset + size_of::<u32>();
+
+        data[virtual_size_offset..virtual_size_offset + size_of::<u32>()]
+            .copy_from_slice(&virtual_size.to_le_bytes());
+        data[virtual_address_offset..virtual_address_offset + size_of::<u32>()]
+            .copy_from_slice(&0xffff_f000_u32.to_le_bytes());
+
+        let mut image = Image::parse(&data[..]).unwrap();
+        let resource_directory = image.resource_directory().cloned().unwrap();
+        image.set_resource_directory(resource_directory).unwrap();
+
+        assert_eq!(
+            image.windows_header().size_of_image(),
+            original_image_size,
+            "original image size is preserved when {description} overflows"
+        );
+    }
+}
+
 #[test]
 fn transfer_resource_section_small() {
     init_logger();