Preserve BaseOfData when rewriting PE32 images

src/image.rs

@@ -3,6 +3,7 @@
 //! See <https://learn.microsoft.com/en-us/windows/win32/debug/pe-format> for more information.
 
 use alloc::{borrow::Cow, string::ToString, vec::Vec};
+use core::mem::size_of;
 
 use foldhash::fast::RandomState;
 use indexmap::IndexMap;
@@ -579,6 +580,14 @@ impl<'a> Image<'a> {
 
         new_image.extend_from_slice(coff_header.as_bytes());
         new_image.extend_from_slice(self.standard_header.as_bytes());
+        if matches!(windows_header, GenericWindowsHeader::WindowsHeader32(_)) {
+            let base_of_data_offset = self.coff_header_offset as usize
+                + size_of::<CoffHeader>()
+                + size_of::<StandardHeader>();
+            new_image.extend_from_slice(
+                &self.image[base_of_data_offset..base_of_data_offset + size_of::<u32>()],
+            );
+        }
         new_image.extend_from_slice(windows_header.as_bytes());
 
         for (_, data) in header_data_directory.iter() {

tests/tests.rs

@@ -2,6 +2,7 @@ use editpe::{constants::*, *};
 use std::sync::Once;
 
 static BINARY_PATH_SMALL: &str = "./tests/assets/smallbin.exe";
+static BINARY_PATH_SMALL_32: &str = "./tests/assets/smallbin-32.exe";
 static BINARY_PATH_LARGE: &str = "./tests/assets/smallbin-large.exe";
 static BINARY_PATH_UPX: &str = "./tests/assets/smallbin-large-upx.exe";
 static BINARY_PATH_WRAPPE: &str = "./tests/assets/smallbin-wrappe.exe";
@@ -167,6 +168,22 @@ fn set_resource_section() {
     assert_eq!(image, new_image, "original and rebuilt images equal");
 }
 
+#[test]
+fn set_pe32_resource_section() {
+    init_logger();
+
+    let data = std::fs::read(BINARY_PATH_SMALL_32).unwrap();
+    let mut image = Image::parse(&data[..]).unwrap();
+    let resource_directory = image.resource_directory().cloned().unwrap();
+    image.set_resource_directory(resource_directory).unwrap();
+
+    let new_data = image.data();
+    assert_eq!(data.len(), new_data.len(), "original and rebuilt PE32 data size equal");
+
+    let new_image = Image::parse(new_data).unwrap();
+    assert_eq!(image, new_image, "original and rebuilt PE32 images equal");
+}
+
 #[test]
 fn transfer_resource_section_small() {
     init_logger();