feat: 修改偏移

src/error.rs

@@ -20,6 +20,9 @@ pub enum PEParseError {
     InvalidDataDirectoryIndex(u16),
     #[error("错误的导出函数序号: {0}")]
     ExportOrdinalNotFound(u16),
+
+    #[error("目标的缓冲区不够大")]
+    TargetBufferTooSmall
 }
 /// PE操作的错误
 #[derive(Error, Debug)]

src/file_pe.rs

@@ -10,7 +10,7 @@ use std::{
 };
 
 pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
-    /// Get a reference to a type at a given offset.
+    /// 获取一个给定偏移量的类型的引用
     fn get_ref<T>(&self, offset: usize) -> Result<&T, PEParseError> {
         let len = std::mem::size_of::<T>();
         if offset + len > self.len() {
@@ -20,7 +20,7 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
         Ok(unsafe { &*ptr })
     }
 
-    // 判断是64位还是32位
+    /// 判断是64位还是32位
     fn is_64_bit(&self) -> Result<bool, PEParseError> {
         // 先以32位加载可选头，通过可选头的Magic字段判断
         let optional_header: &ImageOptionalHeader32 = self.get_optional_header_32()?;
@@ -28,6 +28,11 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
         Ok(is_64_bit)
     }
 
+    /// 通过RVA获取偏移
+    /// 如果在文件PE中，获取的结果是FOA。
+    /// 如果是在内存PE中，获取的结果是VA。
+    fn rva_to_offset(&self, rva: u32) -> Result<usize, PEParseError>;
+
     /// 获取数据目录
     fn get_data_directories(&self) -> Result<&[ImageDataDirectory], PEParseError> {
         // 1. 获取数据目录偏移
@@ -109,6 +114,24 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
         Err(PEParseError::RvaToFoaError(foa))
     }
 
+    /// RVA转换为VA
+    fn rva_to_va(&self, rva: u32) -> Result<usize, PEParseError> {
+        let image_base = self.get_image_base()?;
+        Ok(image_base + rva as usize)
+    }
+
+    fn get_image_base(&self) -> Result<usize, PEParseError>;
+
+    /// 获取SizeOfImage大小
+    fn get_size_of_image(&self) -> Result<u32, PEParseError> {
+        let nt_header = self.get_nt_header()?;
+        let size_of_image = match nt_header {
+            ImageNTHeader::NTHeader32(nt_header) => nt_header.optional_header.size_of_image,
+            ImageNTHeader::NTHeader64(nt_header) => nt_header.optional_header.size_of_image,
+        };
+        Ok(size_of_image)
+    }
+
     /// 将size与节对齐值进行对齐，返回对齐后的值
     /// size: 需要对齐的值
     fn align_size_with_section_alignment(&self, size: u32) -> Result<u32, PEParseError> {
@@ -295,8 +318,8 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
 
     /// 以传入的RVA为开始偏移，从中解析导出表
     fn parse_export_table(&self, rva: u32) -> Result<&ImageExportDirectoryTable, PEParseError> {
-        let foa = self.rva_to_foa(rva)?;
+        let offset = self.rva_to_offset(rva)?;
-        let export_table = self.get_ref::<ImageExportDirectoryTable>(foa as usize)?;
+        let export_table = self.get_ref::<ImageExportDirectoryTable>(offset)?;
         Ok(export_table)
     }
 
@@ -312,22 +335,22 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
         let base = export_table.base; // 序号基数
         let number_of_fun = export_table.number_of_functions;
         let number_of_names = export_table.number_of_names;
-        // 名称指针表的FOA
+        // 名称指针表的偏移
-        let names_table_foa = self.rva_to_foa(export_table.address_of_names)?;
+        let names_table_offset = self.rva_to_offset(export_table.address_of_names)?;
 
         let get_function_rva = |ordinal: u32| -> Result<u32, PEParseError> {
-            let functions_table_foa = self.rva_to_foa(export_table.address_of_functions)?;
+            let functions_table_offset = self.rva_to_offset(export_table.address_of_functions)?;
             let function_item_foa =
-                functions_table_foa + ordinal * std::mem::size_of::<u32>() as u32;
+                functions_table_offset + ordinal as usize * std::mem::size_of::<u32>();
             Ok(self.get_ref::<u32>(function_item_foa as usize)?.clone())
         };
 
         let check_forwarder = |rva: u32| -> Result<Option<ExportFunEnum>, PEParseError> {
             if rva >= export_table_rva && rva < export_table_rva + export_table_size {
-                let forwarder_foa = self.rva_to_foa(rva)?;
+                let forwarder_offset = self.rva_to_offset(rva)?;
                 let forwarder = unsafe {
                     std::ffi::CStr::from_ptr(
-                        self.get_ref::<i8>(forwarder_foa as usize)? as *const i8
+                        self.get_ref::<i8>(forwarder_offset as usize)? as *const i8
                     )
                     .to_str()
                     .unwrap()
@@ -340,19 +363,22 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
         match proc_name {
             ProcName::Named(proc_name_str) => {
                 for index in 0..number_of_names {
-                    let name_item_foa = names_table_foa + index * std::mem::size_of::<u32>() as u32;
+                    let name_item_foa =
+                        names_table_offset + index as usize * std::mem::size_of::<u32>() as usize;
                     let name_rva = self.get_ref::<u32>(name_item_foa as usize)?.clone();
-                    let name_foa = self.rva_to_foa(name_rva)?;
+                    let name_offset = self.rva_to_offset(name_rva)?;
                     let name = unsafe {
-                        std::ffi::CStr::from_ptr(self.get_ref::<i8>(name_foa as usize)? as *const i8)
+                        std::ffi::CStr::from_ptr(
+                            self.get_ref::<i8>(name_offset as usize)? as *const i8
+                        )
                         .to_str()
                         .unwrap()
                     };
                     if name == proc_name_str {
-                        let ordinals_table_foa =
+                        let ordinals_table_offset =
-                            self.rva_to_foa(export_table.address_of_name_ordinals)?;
+                            self.rva_to_offset(export_table.address_of_name_ordinals)?;
                         let ordinal_item_foa =
-                            ordinals_table_foa + index * std::mem::size_of::<u16>() as u32;
+                            ordinals_table_offset + index as usize * std::mem::size_of::<u16>();
                         let ordinal = self.get_ref::<u16>(ordinal_item_foa as usize)?.clone();
                         let rva = get_function_rva(ordinal as u32)?;
                         return check_forwarder(rva);
@@ -373,24 +399,35 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
 
     /// 以传入的RVA为开始偏移，从中解析导入表
     fn parse_import_tables(&self, rva: u32) -> Result<&[ImageImportDescriptor], PEParseError> {
-        let foa = self.rva_to_foa(rva)?;
+        let offset = self.rva_to_offset(rva)?;
-        let mut current_foa = foa;
+        let mut current_offset = offset;
         let mut numbers = 0;
         loop {
-            let import_table = self.get_ref::<ImageImportDescriptor>(current_foa as usize)?;
+            let import_table = self.get_ref::<ImageImportDescriptor>(current_offset as usize)?;
             if import_table.original_first_thunk == 0 {
                 break;
             }
             numbers += 1;
-            current_foa += std::mem::size_of::<ImageImportDescriptor>() as u32;
+            current_offset += std::mem::size_of::<ImageImportDescriptor>();
         }
         let result_tables = unsafe {
-            let ptr = self.as_ptr().wrapping_offset(foa as isize) as *const ImageImportDescriptor;
+            let ptr =
+                self.as_ptr().wrapping_offset(offset as isize) as *const ImageImportDescriptor;
             let result = std::slice::from_raw_parts(ptr, numbers);
             result
         };
         Ok(result_tables)
     }
+
+    /// 获取映射的头部大小
+    fn get_size_of_header(&self) -> Result<u32, PEParseError> {
+        let nt_header = self.get_nt_header()?;
+        let size_of_header = match nt_header {
+            ImageNTHeader::NTHeader32(nt_header) => nt_header.optional_header.size_of_headers,
+            ImageNTHeader::NTHeader64(nt_header) => nt_header.optional_header.size_of_headers,
+        };
+        Ok(size_of_header)
+    }
 }
 /// 可修改的PE trait
 pub trait MutablePE: ReadOnlyPE + DerefMut<Target = [u8]> + AsMut<[u8]> {
@@ -638,3 +675,63 @@ pub trait MutablePE: ReadOnlyPE + DerefMut<Target = [u8]> + AsMut<[u8]> {
         Ok(())
     }
 }
+
+pub trait FilePE: ReadOnlyPE {
+    /// 将PE文件拷贝到指定地址
+    /// 仅拷贝节表数据 拷贝后的数据已经对齐
+    fn copy_file_pe_to_memory(&self, target: &mut [u8]) -> Result<(), PEParseError> {
+        // 如果长度不够PE中的size_of_image对齐到sction_alignment的值，就返回错误
+        let size_of_image = self.get_size_of_image()?;
+        let section_alignment = self.align_size_with_section_alignment(size_of_image)?;
+        if target.len() < section_alignment as usize {
+            return Err(PEParseError::TargetBufferTooSmall);
+        }
+        // 长度够了 现在开始拷贝
+        // 1. 拷贝头部[0..size_of_headers)
+        let size_of_header = self.get_size_of_header()?;
+        target[..size_of_header as usize]
+            .copy_from_slice(&self.as_ref()[..size_of_header as usize]);
+        // 2. 遍历节表
+        let sections = self.get_section_headers()?;
+        for item in sections.iter() {
+            // 如果节表的SizeOfRawData是0，那么这个节区是没有数据的，直接跳过
+            if item.size_of_raw_data == 0 {
+                continue;
+            }
+            // 拷贝到target的VirtualAddress处
+            let va = item.virtual_address as usize;
+            let copy_size = item.size_of_raw_data;
+            // copy_size = copy_size对齐到file_alignment的值
+            let copy_size = self.align_size_with_file_alignment(copy_size)? as usize;
+            let foa = item.pointer_to_raw_data as usize;
+            target[va..va + copy_size].copy_from_slice(&self.as_ref()[foa..foa + copy_size]);
+        }
+        Ok(())
+    }
+}
+
+impl<T> ReadOnlyPE for T
+where
+    T: FilePE,
+{
+    fn rva_to_offset(&self, rva: u32) -> Result<usize, PEParseError> {
+        let result = self.rva_to_foa(rva)?;
+        Ok(result as usize)
+    }
+
+    fn get_image_base(&self) -> Result<usize, PEParseError> {
+        let nt_header = self.get_nt_header()?;
+        let image_base = match nt_header {
+            ImageNTHeader::NTHeader32(nt_header) => nt_header.optional_header.image_base as usize,
+            ImageNTHeader::NTHeader64(nt_header) => nt_header.optional_header.image_base as usize,
+        };
+        Ok(image_base)
+    }
+}
+
+// 默认实现MutablePE
+// impl<T> MutablePE for T where T: ReadOnlyPE + DerefMut<Target = [u8]> + AsMut<[u8]> {}
+
+pub trait FileMutPE: FilePE + MutablePE {}
+
+impl<T> FileMutPE for T where T: FilePE + MutablePE {}

src/lib.rs

@@ -1,17 +1,17 @@
-pub mod error;
-pub mod header;
-pub mod pe;
-pub mod memory_pe;
 mod constants;
+pub mod error;
+pub mod file_pe;
+pub mod header;
+pub mod memory_pe;
 mod utils;
 #[cfg(feature = "memmap2_impl")]
 mod third_party {
-    use memmap2::{MmapMut, Mmap};
+    use crate::file_pe::{FilePE, MutablePE};
-    use crate::pe::{MutablePE, ReadOnlyPE};
+    use memmap2::{Mmap, MmapMut};
     // 为文件映射实现PE结构
-    impl ReadOnlyPE for Mmap {}
+    impl FilePE for Mmap {}
 
-    impl ReadOnlyPE for MmapMut {}
     // 为可变文件映射实现可操作PE结构
+    impl FilePE for MmapMut {}
     impl MutablePE for MmapMut {}
 }

src/memory_pe.rs

@@ -2,21 +2,28 @@ use std::ops::{Deref, DerefMut};
 
 use crate::{align_to, error::PEParseError};
 
-use super::pe::{MutablePE, ReadOnlyPE};
+use super::file_pe::{MutablePE, ReadOnlyPE};
 use crate::constants::*;
-pub struct MemoryPE {
+pub struct MemoryPEStruct {
     len: usize,
     buf: *const u8, // 缓冲区的指针 指向内存中的数据
 }
 
-pub struct MemoryPEMut {
+pub struct MemoryPEMutStruct {
     len: usize,
     buf: *mut u8, // 缓冲区的指针 指向内存中的数据
 }
-// DerefMut<Target = [u8]> + AsMut<[u8]>
-impl MutablePE for MemoryPEMut{}
 
-impl DerefMut for MemoryPEMut{
+impl ReadOnlyPE for MemoryPEMutStruct{
+    fn rva_to_offset(&self, rva: u32) -> Result<usize, PEParseError> {
+        Ok(rva as usize)
+    }
+    fn get_image_base(&self) -> Result<usize, PEParseError> {
+        Ok(self.buf as usize)
+    }
+}
+
+impl DerefMut for MemoryPEMutStruct{
     fn deref_mut(&mut self) -> &mut Self::Target {
         unsafe {
             std::slice::from_raw_parts_mut(self.buf, self.len)
@@ -24,7 +31,7 @@ impl DerefMut for MemoryPEMut{
     }
 }
 
-impl AsMut<[u8]> for MemoryPEMut{
+impl AsMut<[u8]> for MemoryPEMutStruct{
     fn as_mut(&mut self) -> &mut [u8] {
         unsafe {
             std::slice::from_raw_parts_mut(self.buf, self.len)
@@ -32,17 +39,7 @@ impl AsMut<[u8]> for MemoryPEMut{
     }
 }
 
-impl Deref for MemoryPEMut{
+impl AsRef<[u8]> for MemoryPEMutStruct{
-    type Target = [u8];
-    fn deref(&self) -> &Self::Target {
-        unsafe {
-            std::slice::from_raw_parts(self.buf, self.len)
-        }
-    }
-}
-impl ReadOnlyPE for MemoryPEMut{}
-
-impl AsRef<[u8]> for MemoryPEMut{
     fn as_ref(&self) -> &[u8] {
         unsafe {
             std::slice::from_raw_parts(self.buf, self.len)
@@ -50,9 +47,27 @@ impl AsRef<[u8]> for MemoryPEMut{
     }
 }
 
-impl ReadOnlyPE for MemoryPE{}
+impl Deref for MemoryPEMutStruct{
+    type Target = [u8];
+    fn deref(&self) -> &Self::Target {
+        unsafe {
+            std::slice::from_raw_parts(self.buf, self.len)
+        }
+    }
+}
 
-impl Deref for MemoryPE{
+impl ReadOnlyPE for MemoryPEStruct{
+    fn rva_to_offset(&self, rva: u32) -> Result<usize, PEParseError> {
+        // 内存中直接返回
+        Ok(rva as usize)
+    }
+
+    fn get_image_base(&self) -> Result<usize, PEParseError> {
+        Ok(self.buf as usize)
+    }
+}
+
+impl Deref for MemoryPEStruct{
     type Target = [u8];
 
     fn deref(&self) -> &Self::Target {
@@ -63,7 +78,7 @@ impl Deref for MemoryPE{
 }
 
 
-impl AsRef<[u8]> for MemoryPE{
+impl AsRef<[u8]> for MemoryPEStruct{
     fn as_ref(&self) -> &[u8] {
         unsafe {
             std::slice::from_raw_parts(self.buf, self.len)
@@ -72,10 +87,12 @@ impl AsRef<[u8]> for MemoryPE{
 }
 
 /// 需要实现一个从内存中解析PE的方法
-pub fn parse_pe_from_memory(buf: *const u8) -> Result<MemoryPE, PEParseError> {
+pub fn parse_pe_from_memory(buf: *const u8) -> Result<MemoryPEStruct, PEParseError> {
     // 1. 检查是否为PE文件
     // 取出头两个字节
-    let dos_magic = unsafe { *buf as u16 };
+    let dos_magic = unsafe { 
+        *(buf as *const u16)
+     };
     if dos_magic != 0x5A4D { // 小尾
         return Err(PEParseError::InvalidDOSMagic);
     }
@@ -93,17 +110,19 @@ pub fn parse_pe_from_memory(buf: *const u8) -> Result<MemoryPE, PEParseError> {
     let image_size = unsafe { *(buf.add(pe_offset as usize + IMAGE_SIZE_OFFSET) as *const u32) };
     // 计算对齐后的内存大小
     let aligned_image_size = align_to!(image_size, section_alignment);
-    Ok(MemoryPE {
+    Ok(MemoryPEStruct {
         len: aligned_image_size as usize,
         buf,
     })
 }
 
 /// 还需要一个Mut的版本
-pub fn parse_pe_from_memory_mut(buf: *mut u8) -> Result<MemoryPEMut, PEParseError> {
+pub fn parse_pe_from_memory_mut(buf: *mut u8) -> Result<MemoryPEMutStruct, PEParseError> {
     // 1. 检查是否为PE文件
     // 取出头两个字节
-    let dos_magic = unsafe { *buf as u16 };
+    let dos_magic = unsafe { 
+        *(buf as *const u16)
+     };
     if dos_magic != 0x5A4D { // 小尾
         return Err(PEParseError::InvalidDOSMagic);
     }
@@ -121,7 +140,7 @@ pub fn parse_pe_from_memory_mut(buf: *mut u8) -> Result<MemoryPEMut, PEParseErro
     let image_size = unsafe { *(buf.add(pe_offset as usize + IMAGE_SIZE_OFFSET) as *const u32) };
     // 计算对齐后的内存大小
     let aligned_image_size = align_to!(image_size, section_alignment);
-    Ok(MemoryPEMut {
+    Ok(MemoryPEMutStruct {
         len: aligned_image_size as usize,
         buf,
     })

tests/test_pe.rs

@@ -13,7 +13,7 @@ mod test {
         use core::ffi;
 
         // 1. 创建映射
-        use pe_parse::pe::ReadOnlyPE;
+        use pe_parse::file_pe::ReadOnlyPE;
         let file = std::fs::OpenOptions::new()
             .read(true)
             .open("tests/test_dll.dll")
@@ -63,7 +63,7 @@ mod test {
         let find_function_str = ProcName::Named("test_def_add".to_string());
         // let find_function_str = ProcName::Ordinal("div");
         // 1. 还是先映射
-        use pe_parse::pe::ReadOnlyPE;
+        use pe_parse::file_pe::ReadOnlyPE;
         let file = std::fs::OpenOptions::new()
             .read(true)
             .open("tests/test_dll.dll")