feat: 修改完测试

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,14 @@ 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>;
+
     /// 将size与节对齐值进行对齐，返回对齐后的值
     /// size: 需要对齐的值
     fn align_size_with_section_alignment(&self, size: u32) -> Result<u32, PEParseError> {
@@ -295,8 +308,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 +325,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 +353,19 @@ 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,19 +386,19 @@ 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
         };
@@ -638,3 +651,31 @@ pub trait MutablePE: ReadOnlyPE + DerefMut<Target = [u8]> + AsMut<[u8]> {
         Ok(())
     }
 }
+
+pub trait FilePE: ReadOnlyPE {}
+
+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,30 @@ 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> {
+        // 返回ImageBase + rva
+        let result = self.rva_to_va(rva)?;
+        Ok(result 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 +33,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 +41,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 +49,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> {
+        let result = self.rva_to_va(rva)?;
+        Ok(result 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 +80,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,7 +89,7 @@ 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 { 
@@ -95,14 +112,14 @@ 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 { 
@@ -125,7 +142,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")