fix: 修复一个错误解析32与64dos头的问题。修复了未处理转发函数的问题

src/constants.rs

@@ -1,6 +1,8 @@
 // DOS的Magic值
+#[allow(dead_code)]
 pub const DOS_MAGIC: u16 = 0x5A4D;
 // DOS头部大小
+#[allow(dead_code)]
 pub const DOS_HEADER_SIZE: usize = 0x40;
 // DOS头部中PE头部的偏移
 pub const DOS_HEADER_PE_OFFSET: usize = 0x3C;
@@ -12,4 +14,7 @@ pub const SECTION_ALIGNMENT_OFFSET: usize = 0x38;
 pub const IMAGE_SIZE_OFFSET: usize = 0x50;
 
 // 可选头Magic如果为32位的值
+#[allow(dead_code)]
 pub const OPTIONAL_MAGIC_32: u16 = 0x10B;
+// 可选头Magic如果为64位的值
+pub const OPTIONAL_MAGIC_64: u16 = 0x20B;

src/error.rs

@@ -16,8 +16,8 @@ pub enum PEParseError {
     RvaToFoaError(u32),
     #[error("RVA对应的是一个空节区数据!RVA: {0}")]
     RvaToEmptySection(u32),
-    #[error("错误的数据目录索引")]
+    #[error("错误的数据目录索: {0}")]
-    InvalidDataDirectoryIndex,
+    InvalidDataDirectoryIndex(u16),
     #[error("错误的导出函数序号: {0}")]
     ExportOrdinalNotFound(u16),
 }

src/header.rs

@@ -477,8 +477,16 @@ pub enum ExportAddressTable {
     // 转发函数字符串的RVA
     ForwarderStringRVA(u32),
 }
-
+/// 从导出表中查找的函数符号，有可能是名字，有可能是序号
 pub enum ProcName {
     Named(String),
     Ordinal(u16),
 }
+
+/// 导出函数
+pub enum ExportFunEnum {
+    // 导出函数的RVA
+    ExportFunctionRVA(u32),
+    // 转发函数字符串
+    ForwarderString(String),
+}

src/pe.rs

@@ -24,7 +24,7 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
     fn is_64_bit(&self) -> Result<bool, PEParseError> {
         // 先以32位加载可选头，通过可选头的Magic字段判断
         let optional_header: &ImageOptionalHeader32 = self.get_optional_header_32()?;
-        let is_64_bit = optional_header.magic == OPTIONAL_MAGIC_32;
+        let is_64_bit = optional_header.magic == OPTIONAL_MAGIC_64;
         Ok(is_64_bit)
     }
 
@@ -53,9 +53,9 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
         data_directory: ImageDirectoryEntry,
     ) -> Result<&ImageDataDirectory, PEParseError> {
         let data_directories = self.get_data_directories()?;
-        let data_directory = data_directories
+        let data_directory = data_directories.get(data_directory as usize).ok_or(
-            .get(data_directory as usize)
+            PEParseError::InvalidDataDirectoryIndex(data_directory as u16),
-            .ok_or(PEParseError::InvalidDataDirectoryIndex)?;
+        )?;
         Ok(data_directory)
     }
 
@@ -301,9 +301,13 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
     }
 
     /// 获取导出函数的RVA地址，类似于GetProcAddress
-    fn get_export_function_rva(&self, proc_name: ProcName) -> Result<Option<u32>, PEParseError> {
+    fn get_export_function_rva(
+        &self,
+        proc_name: ProcName,
+    ) -> Result<Option<ExportFunEnum>, PEParseError> {
         let export_table_dir = self.get_data_directory(ImageDirectoryEntry::Export)?;
         let export_table_rva = export_table_dir.virtual_address;
+        let export_table_size = export_table_dir.size;
         let export_table = self.parse_export_table(export_table_rva)?;
         let base = export_table.base; // 序号基数
         let number_of_fun = export_table.number_of_functions;
@@ -311,50 +315,57 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
         // 名称指针表的FOA
         let names_table_foa = self.rva_to_foa(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 function_item_foa =
+                functions_table_foa + ordinal * std::mem::size_of::<u32>() as 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 = unsafe {
+                    std::ffi::CStr::from_ptr(
+                        self.get_ref::<i8>(forwarder_foa as usize)? as *const i8
+                    )
+                    .to_str()
+                    .unwrap()
+                };
+                return Ok(Some(ExportFunEnum::ForwarderString(forwarder.to_string())));
+            }
+            Ok(Some(ExportFunEnum::ExportFunctionRVA(rva)))
+        };
+
         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_rva = self.get_ref::<u32>(name_item_foa as usize)?.clone();
                     let name_foa = self.rva_to_foa(name_rva)?;
                     let name = unsafe {
-                        std::ffi::CStr::from_ptr(
+                        std::ffi::CStr::from_ptr(self.get_ref::<i8>(name_foa as usize)? as *const i8)
-                            self.get_ref::<i8>(name_foa as usize)? as *const i8
-                        )
                             .to_str()
                             .unwrap()
                     };
-                    // 比较
                     if name == proc_name_str {
-                        // 找到了这个函数 那么需要找到这个函数的序号
                         let ordinals_table_foa =
                             self.rva_to_foa(export_table.address_of_name_ordinals)?;
                         let ordinal_item_foa =
                             ordinals_table_foa + index * std::mem::size_of::<u16>() as u32;
                         let ordinal = self.get_ref::<u16>(ordinal_item_foa as usize)?.clone();
-                        // 找到了序号 那么就可以找到函数地址
+                        let rva = get_function_rva(ordinal as u32)?;
-                        let functions_table_foa =
+                        return check_forwarder(rva);
-                            self.rva_to_foa(export_table.address_of_functions)?;
-                        let function_item_foa = functions_table_foa
-                            + ordinal as u32 * std::mem::size_of::<u32>() as u32;
-                        let rva = self.get_ref::<u32>(function_item_foa as usize)?.clone();
-                        return Ok(Some(rva));
                     }
                 }
             }
             ProcName::Ordinal(order) => {
-                // 通过序号查找
-                // 序号就在序号表的第order - base的下标位置
                 let find_ordinal = order as u32 - base;
                 if find_ordinal >= number_of_fun {
                     return Err(PEParseError::ExportOrdinalNotFound(order));
                 }
-                let functions_table_foa = self.rva_to_foa(export_table.address_of_functions)?;
+                let rva = get_function_rva(find_ordinal)?;
-                let function_item_foa =
+                return check_forwarder(rva);
-                    functions_table_foa + find_ordinal * std::mem::size_of::<u32>() as u32;
-                let rva = self.get_ref::<u32>(function_item_foa as usize)?.clone();
-                return Ok(Some(rva));
             }
         }
         Ok(None)
@@ -463,9 +474,9 @@ pub trait MutablePE: ReadOnlyPE + DerefMut<Target = [u8]> + AsMut<[u8]> {
         data_directory: ImageDirectoryEntry,
     ) -> Result<&mut ImageDataDirectory, PEParseError> {
         let data_directories = self.get_data_directories_mut()?;
-        let data_directory = data_directories
+        let data_directory = data_directories.get_mut(data_directory as usize).ok_or(
-            .get_mut(data_directory as usize)
+            PEParseError::InvalidDataDirectoryIndex(data_directory as u16),
-            .ok_or(PEParseError::InvalidDataDirectoryIndex)?;
+        )?;
         Ok(data_directory)
     }
 

tests/test_pe.rs

@@ -1,5 +1,4 @@
 mod test {
-    use super::*;
     use memmap2::Mmap;
     use pe_parse::header::ProcName;
     #[test]
@@ -61,8 +60,8 @@ mod test {
 
     #[test]
     fn test_find_export_function() {
-        // let find_function_str = ProcName::Named("test_def_add".to_string());
+        let find_function_str = ProcName::Named("test_def_add".to_string());
-        let find_function_str = ProcName::Ordinal(2);
+        // let find_function_str = ProcName::Ordinal("div");
         // 1. 还是先映射
         use pe_parse::pe::ReadOnlyPE;
         let file = std::fs::OpenOptions::new()
@@ -76,9 +75,14 @@ mod test {
                 println!("Can't find the function");
                 return;
             }
-            Some(rva) => {
+            Some(rva) => match rva {
-                println!("Find Function RVA: {:#x}", rva);
+                pe_parse::header::ExportFunEnum::ExportFunctionRVA(rva) => {
-            }
+                    println!("Find the function at: {:#x}", rva);
+                }
+                pe_parse::header::ExportFunEnum::ForwarderString(forward_str) => {
+                    println!("转发函数，转发到: {}", forward_str);
+                }
+            },
         }
     }
 }