diff --git a/src/image.rs b/src/image.rs index c682a91..ea53427 100644 --- a/src/image.rs +++ b/src/image.rs @@ -564,15 +564,23 @@ impl<'a> Image<'a> { return Err(ImageWriteError::NotEnoughSpaceInHeader); } + let size_of_image = section_table + .iter() + .map(|section| u64::from(section.virtual_address) + u64::from(section.virtual_size)) + .max() + .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 { 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..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::()] + .try_into() + .unwrap(), + ) as usize; + let last_section_offset = pe_signature_offset + + size_of::() + + size_of::() + + image.coff_header().size_of_optional_header as usize + + (image.coff_header().number_of_sections as usize - 1) * size_of::(); + let virtual_size_offset = last_section_offset + size_of::(); + let virtual_address_offset = virtual_size_offset + size_of::(); + + data[virtual_size_offset..virtual_size_offset + size_of::()] + .copy_from_slice(&virtual_size.to_le_bytes()); + data[virtual_address_offset..virtual_address_offset + size_of::()] + .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(); @@ -234,6 +276,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!(