diff --git a/src-tauri/src/pe_parse/pe.rs b/src-tauri/src/pe_parse/pe.rs index 470211a..5954505 100644 --- a/src-tauri/src/pe_parse/pe.rs +++ b/src-tauri/src/pe_parse/pe.rs @@ -1,12 +1,10 @@ -use memmap2::MmapMut; - use super::{ error::{MutablePEError, PEParseError}, header::*, }; use std::{ - mem::{self, offset_of}, - ops::{Deref, DerefMut, IndexMut}, + mem::{self}, + ops::{Deref, DerefMut}, }; pub trait ReadOnlyPE: Deref + Sized + AsRef<[u8]> { @@ -290,17 +288,24 @@ pub trait MutablePE: ReadOnlyPE + DerefMut + AsMut<[u8]> { // 检查一下把,万一调用者没有扩大文件大小 let sections = self.get_section_headers()?; let self_len = self.len(); - let last_section = sections.last().unwrap(); - let last_section_end = last_section.pointer_to_raw_data + last_section.size_of_raw_data; - if last_section_end as usize + aligned_file_size as usize > self_len { - return Err(MutablePEError::CannotExpandFileSize( - last_section_end as u32 + aligned_file_size, - self_len as u32, - )); + // 需要反向遍历,找到最后一个size_of_raw_data不为0的节区 + // 然后计算这个节区的PointerToRawData+(SizeOfRawData对齐到FileAlignment的值) + let mut last_section_end = 0; + for section in sections.iter().rev() { + if section.size_of_raw_data != 0 { + last_section_end = section.pointer_to_raw_data + + self.align_size_with_file_alignment(section.size_of_raw_data)?; + if last_section_end as usize + aligned_file_size as usize > self_len { + return Err(MutablePEError::CannotExpandFileSize( + last_section_end as u32 + aligned_file_size, + self_len as u32, + )); + } + break; + } } let mut nt_header = self.get_nt_header_mut()?; let op = nt_header.get_optional_header_mut(); - // 更新可选头的SizeOfHeaders和SizeOfImage match op { ImageOptionalHeaderMut::OptionalHeader32(op) => { @@ -310,9 +315,7 @@ pub trait MutablePE: ReadOnlyPE + DerefMut + AsMut<[u8]> { op.size_of_headers += aligned_file_size; } } - let mut_section = self.get_section_headers_mut()?; - let start_offset = mut_section.first().unwrap().pointer_to_raw_data; let end_offset = last_section_end; // 更新节区头的PointerToRawData @@ -379,10 +382,8 @@ pub trait MutablePE: ReadOnlyPE + DerefMut + AsMut<[u8]> { } new_section.virtual_size = section_alignment; // 如果section_size是0,那么virtual_size必须得是最小的section_alignment - // 新节区的virtual_address是上一个节区的virtual_address+(pointer_to_raw_data 对齐到 section_alignment的值) let last_section = self.get_section_headers()?.last().unwrap(); - // TODO: VA计算有问题 let mut last_section_size_of_raw_data_aligned = self.align_size_with_section_alignment(last_section.size_of_raw_data)?; if last_section_size_of_raw_data_aligned == 0 { @@ -390,19 +391,24 @@ pub trait MutablePE: ReadOnlyPE + DerefMut + AsMut<[u8]> { } new_section.virtual_address = last_section.virtual_address + last_section_size_of_raw_data_aligned; - // 新节区的size_of_raw_data是section_size对齐到file_alignment的值 - // 如果section_size是0,那么size_of_raw_data也是0 new_section.size_of_raw_data = file_alignment; - // 新节区的pointer_to_raw_data是上一个节区的pointer_to_raw_data+上一个节区的size_of_raw_data - new_section.pointer_to_raw_data = - last_section.pointer_to_raw_data + last_section.size_of_raw_data; + // 如果上一个节区的size_of_raw_data是0,那么上一个节区的pointer_to_raw_data是可以乱写的 + // 正确的做法应该是,找到上一个size_of_raw_data不为0的节区,然后计算 + let sections = self.get_section_headers()?; + // 反向遍历 + for section in sections.iter().rev() { + if section.size_of_raw_data != 0 { + new_section.pointer_to_raw_data = section.pointer_to_raw_data + + self.align_size_with_file_alignment(section.size_of_raw_data)?; + break; + } + } new_section.characteristics = section_characteristics; // 其他字段最好填充为0 new_section.number_of_linenumbers = 0; new_section.number_of_relocations = 0; new_section.pointer_to_linenumbers = 0; new_section.pointer_to_relocations = 0; - // 更新文件头的NumberOfSections let file_header = self.get_file_header_mut()?; file_header.number_of_sections += 1;