From cb2c6f3630d66a79c604ec156db366c0332681ea Mon Sep 17 00:00:00 2001
From: "381848900@qq.com" <asahichyan33@gmail.com>
Date: Wed, 18 Dec 2024 14:20:41 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BF=AE=E5=A4=8D=E4=B8=80=E4=B8=AAtodo?=
 =?UTF-8?q?=E9=97=AE=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/pe.rs        | 10 ++++------
 tests/test_pe.rs | 15 ++++++++++++---
 2 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/src/pe.rs b/src/pe.rs
index 638715d..670c8d7 100644
--- a/src/pe.rs
+++ b/src/pe.rs
@@ -301,7 +301,7 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
 
 
     /// 获取导出函数的RVA地址，类似于GetProcAddress
-    fn get_export_function_rva(&self, proc_name: ProcName) -> Result<u32, PEParseError>{
+    fn get_export_function_rva(&self, proc_name: ProcName) -> Result<Option<u32>, PEParseError>{
         let export_table_dir = self.get_data_directory(ImageDirectoryEntry::Export)?;
         let export_table_rva = export_table_dir.virtual_address;
         let export_table = self.parse_export_table(export_table_rva)?;
@@ -335,7 +335,7 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
                         let functions_table_foa = 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(rva);
+                        return Ok(Some(rva));
                     }
                 }
             }
@@ -349,13 +349,11 @@ pub trait ReadOnlyPE: Deref<Target = [u8]> + Sized + AsRef<[u8]> {
                 let functions_table_foa = self.rva_to_foa(export_table.address_of_functions)?;
                 let function_item_foa = 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(rva);
+                return Ok(Some(rva));
             }
         }
-
-        todo!()
+        Ok(None)
     }
-
 }
 /// 可修改的PE trait
 pub trait MutablePE: ReadOnlyPE + DerefMut<Target = [u8]> + AsMut<[u8]> {
diff --git a/tests/test_pe.rs b/tests/test_pe.rs
index ba539c7..4b95c01 100644
--- a/tests/test_pe.rs
+++ b/tests/test_pe.rs
@@ -60,8 +60,9 @@ mod test {
     }
 
     #[test]
-    fn test_find_export_function(){
-        let find_function_str = ProcName::Named("test_def_add".to_string());
+    fn test_find_export_function() {
+        // let find_function_str = ProcName::Named("test_def_add".to_string());
+        let find_function_str = ProcName::Ordinal(2);
         // 1. 还是先映射
         use pe_parse::pe::ReadOnlyPE;
         let file = std::fs::OpenOptions::new()
@@ -70,6 +71,14 @@ mod test {
             .unwrap();
         let mmap = unsafe { Mmap::map(&file).unwrap() };
         let result_rva = mmap.get_export_function_rva(find_function_str).unwrap();
-        println!("Find Function RVA: {:#x}", result_rva);
+        match result_rva {
+            None => {
+                println!("Can't find the function");
+                return;
+            }
+            Some(rva) => {
+                println!("Find Function RVA: {:#x}", rva);
+            }
+        }
     }
 }