Fix runtime vtable functions

samyycX · samyycX · commit 036b22e2aedf · 2025-11-13T21:05:01.000+08:00
diff --git a/s2binlib/src/module.rs b/s2binlib/src/module.rs
@@ -1,5 +1,3 @@
-use std::ffi::CString;
-
 #[derive(Debug)]
 pub struct ModuleInfo {
     pub base_address: usize,
@@ -9,31 +7,35 @@ pub struct ModuleInfo {
 #[derive(Debug)]
 pub enum ModuleError {
     NotFound,
-    InvalidName,
+    InvalidBase,
     SystemError(String),
 }
 
 impl std::fmt::Display for ModuleError {
     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
         match self {
             ModuleError::NotFound => write!(f, "Module not found"),
-            ModuleError::InvalidName => write!(f, "Invalid module name"),
+            ModuleError::InvalidBase => write!(f, "Invalid module base address"),
             ModuleError::SystemError(msg) => write!(f, "System error: {}", msg),
         }
     }
 }
 
 impl std::error::Error for ModuleError {}
 
-pub fn get_module_info(module_name: &str) -> Result<ModuleInfo, ModuleError> {
+pub fn get_module_info(module_base: u64) -> Result<ModuleInfo, ModuleError> {
+    if module_base == 0 || module_base > usize::MAX as u64 {
+        return Err(ModuleError::InvalidBase);
+    }
+
     #[cfg(target_os = "windows")]
     {
-        get_module_info_windows(module_name)
+        get_module_info_windows(module_base as usize)
     }
 
     #[cfg(target_os = "linux")]
     {
-        get_module_info_linux(module_name)
+        get_module_info_linux(module_base as usize)
     }
 
     #[cfg(not(any(target_os = "windows", target_os = "linux")))]
@@ -43,21 +45,18 @@ pub fn get_module_info(module_name: &str) -> Result<ModuleInfo, ModuleError> {
 }
 
 #[cfg(target_os = "windows")]
-fn get_module_info_windows(module_name: &str) -> Result<ModuleInfo, ModuleError> {
+fn get_module_info_windows(module_base: usize) -> Result<ModuleInfo, ModuleError> {
     use std::mem;
     use winapi::shared::minwindef::{FALSE, HMODULE};
-    use winapi::um::libloaderapi::GetModuleHandleA;
     use winapi::um::processthreadsapi::GetCurrentProcess;
     use winapi::um::psapi::{GetModuleInformation, MODULEINFO};
 
-    let c_name = CString::new(module_name).map_err(|_| ModuleError::InvalidName)?;
+    if module_base == 0 {
+        return Err(ModuleError::InvalidBase);
+    }
 
     unsafe {
-        let h_module: HMODULE = GetModuleHandleA(c_name.as_ptr());
-
-        if h_module.is_null() {
-            return Err(ModuleError::NotFound);
-        }
+        let h_module: HMODULE = module_base as HMODULE;
 
         let mut mod_info: MODULEINFO = mem::zeroed();
         let result = GetModuleInformation(
@@ -68,9 +67,14 @@ fn get_module_info_windows(module_name: &str) -> Result<ModuleInfo, ModuleError>
         );
 
         if result == FALSE {
-            return Err(ModuleError::SystemError(
-                "GetModuleInformation failed".to_string(),
-            ));
+            const ERROR_INVALID_HANDLE: i32 = 6;
+            let error = std::io::Error::last_os_error();
+            if error.raw_os_error() == Some(ERROR_INVALID_HANDLE) {
+                return Err(ModuleError::NotFound);
+            }
+            return Err(ModuleError::SystemError(format!(
+                "GetModuleInformation failed with error: {error}"
+            )));
         }
 
         Ok(ModuleInfo {
@@ -81,46 +85,97 @@ fn get_module_info_windows(module_name: &str) -> Result<ModuleInfo, ModuleError>
 }
 
 #[cfg(target_os = "linux")]
-fn get_module_info_linux(module_name: &str) -> Result<ModuleInfo, ModuleError> {
+fn get_module_info_linux(module_base: usize) -> Result<ModuleInfo, ModuleError> {
     use std::fs::File;
     use std::io::{BufRead, BufReader};
 
+    if module_base == 0 {
+        return Err(ModuleError::InvalidBase);
+    }
+
     let maps_file = File::open("/proc/self/maps")
         .map_err(|e| ModuleError::SystemError(format!("Cannot open /proc/self/maps: {}", e)))?;
 
     let reader = BufReader::new(maps_file);
     let mut base_address: Option<usize> = None;
     let mut end_address: Option<usize> = None;
+    let mut module_path: Option<String> = None;
+    let mut module_dev: Option<String> = None;
+    let mut module_inode: Option<String> = None;
 
     for line in reader.lines() {
         let line = line.map_err(|e| ModuleError::SystemError(format!("Read error: {}", e)))?;
 
-        if line.contains(module_name) {
-            let parts: Vec<&str> = line.split_whitespace().collect();
-            if parts.is_empty() {
-                continue;
+        let trimmed = line.trim();
+        if trimmed.is_empty() {
+            continue;
+        }
+
+        let parts: Vec<&str> = trimmed.split_whitespace().collect();
+        if parts.is_empty() {
+            continue;
+        }
+
+        let addr_range: Vec<&str> = parts[0].split('-').collect();
+        if addr_range.len() != 2 {
+            continue;
+        }
+
+        let start = usize::from_str_radix(addr_range[0], 16)
+            .map_err(|_| ModuleError::SystemError("Invalid address format".to_string()))?;
+        let end = usize::from_str_radix(addr_range[1], 16)
+            .map_err(|_| ModuleError::SystemError("Invalid address format".to_string()))?;
+
+        let current_dev = parts.get(3).map(|s| s.to_string());
+        let current_inode = parts.get(4).map(|s| s.to_string());
+        let current_path = if parts.len() > 5 {
+            Some(parts[5..].join(" "))
+        } else {
+            None
+        };
+
+        if base_address.is_none() {
+            if start == module_base {
+                base_address = Some(start);
+                end_address = Some(end);
+                module_path = current_path;
+                module_dev = current_dev;
+                module_inode = current_inode;
             }
+            continue;
+        }
+
+        let matches_path = match (&module_path, &current_path) {
+            (Some(expected), Some(actual)) => expected == actual,
+            (None, None) => true,
+            _ => false,
+        };
 
-            let addr_range: Vec<&str> = parts[0].split('-').collect();
-            if addr_range.len() != 2 {
-                continue;
+        let matches_identity = match (&module_dev, &module_inode, &current_dev, &current_inode) {
+            (Some(expected_dev), Some(expected_inode), Some(dev), Some(inode)) => {
+                expected_dev == dev && expected_inode == inode
             }
+            _ => false,
+        };
 
-            let start = usize::from_str_radix(addr_range[0], 16)
-                .map_err(|_| ModuleError::SystemError("Invalid address format".to_string()))?;
-            let end = usize::from_str_radix(addr_range[1], 16)
-                .map_err(|_| ModuleError::SystemError("Invalid address format".to_string()))?;
+        let contiguous = end_address.map_or(false, |current_end| start == current_end);
 
-            if base_address.is_none() {
-                base_address = Some(start);
+        if matches_path || matches_identity || contiguous {
+            match end_address {
+                Some(ref mut stored_end) if end > *stored_end => *stored_end = end,
+                None => end_address = Some(end),
+                _ => {}
             }
+            continue;
+        }
 
-            end_address = Some(end);
+        if start > module_base {
+            break;
         }
     }
 
     match (base_address, end_address) {
-        (Some(base), Some(end)) => Ok(ModuleInfo {
+        (Some(base), Some(end)) if end > base => Ok(ModuleInfo {
             base_address: base,
             size: end - base,
         }),
diff --git a/s2binlib/src/s2binlib.rs b/s2binlib/src/s2binlib.rs
@@ -26,7 +26,7 @@ use object::{Object, ObjectSection, ObjectSymbol, read::pe::ImageOptionalHeader}
 use crate::{
     VTableInfo, find_pattern_simd, is_executable,
     jit::JitTrampoline,
-    memory::{get_module_base_from_pointer, set_mem_access},
+    memory::{get_module_base_from_pointer, module_from_pointer, set_mem_access},
 };
 
 #[cfg(target_os = "windows")]
@@ -220,6 +220,24 @@ impl<'a> S2BinLib<'a> {
         self.manual_base_addresses.remove(lib_name);
     }
 
+    pub fn module_from_pointer(&self, ptr: u64) -> Result<u64> {
+        let module = module_from_pointer(ptr);
+
+        if module.is_none() {
+            bail!("Failed to get module from pointer.");
+        }
+
+        let (module_name, module_base) = module.unwrap();
+        // respect custom module base
+        for (manual_module_name, manual_module_base) in &self.manual_base_addresses  {
+            let lib_name = self.get_os_lib_name(&manual_module_name);
+            if module_name.contains(&lib_name) {
+                return Ok(*manual_module_base);
+            }
+        }
+        Ok(module_base)
+    }
+
     pub fn get_binary_path(&self, binary_name: &str) -> String {
         if self.os.to_lowercase() == "windows" {
             if let Some(path) = self.custom_binary_paths_windows.get(binary_name) {
diff --git a/s2binlib/src/view.rs b/s2binlib/src/view.rs
@@ -412,14 +412,9 @@ impl<'a> S2BinLib<'a> {
     }
 
     pub fn get_memory_view_from_ptr(&self, ptr: u64) -> Result<MemoryView<'_>> {
-        let result = module_from_pointer(ptr);
+        let module_base = self.module_from_pointer(ptr)?;
 
-        if result.is_none() {
-            bail!("Failed to get module from pointer.");
-        }
-
-        let (module_name, module_base) = result.unwrap();
-        let module_info = get_module_info(&module_name)?;
+        let module_info = get_module_info(module_base)?;
         let memory_view = unsafe {
             MemoryView::new(
                 module_info.base_address as *const u8,
diff --git a/s2binlib/src/vtable.rs b/s2binlib/src/vtable.rs
@@ -168,9 +168,13 @@ impl<'a> S2BinLib<'a> {
         let memory_view = self.get_memory_view_from_ptr(vtable_ptr)?;
 
         if !memory_view.contains(vtable_ptr - 8) {
-            bail!("Vtable pointer is not in the memory view.");
+            bail!("Vtable pointer {:X} is not in the memory view {:X}.", vtable_ptr - 8, memory_view.image_base());
         };
-        let rtti_ptr: u64 = vtable_ptr - 8;
+        let rtti_ptr = memory_view.read::<u64>(vtable_ptr - 8);
+        if rtti_ptr.is_none() {
+            bail!("Failed to read rtti pointer.");
+        }
+        let rtti_ptr = rtti_ptr.unwrap();
 
         #[cfg(target_os = "windows")]
         {
diff --git a/s2binlib_binding/src/c_bindings.rs b/s2binlib_binding/src/c_bindings.rs
@@ -2475,7 +2475,7 @@ pub extern "C" fn s2binlib_get_object_ptr_vtable_name(
 
                 0
             }
-            Err(_) => return_error!(-4, "Failed to get vtable info"),
+            Err(err) => return_error!(-4, format!("Failed to get vtable info: {:?}", err)),
         }
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -2475,7 +2475,7 @@ pub extern "C" fn s2binlib_get_object_ptr_vtable_name(`
`2475`	`2475`
`2476`	`2476`	`0`
`2477`	`2477`	`}`
`2478`		`- Err(_) => return_error!(-4, "Failed to get vtable info"),`
	`2478`	`+ Err(err) => return_error!(-4, format!("Failed to get vtable info: {:?}", err)),`
`2479`	`2479`	`}`
`2480`	`2480`	`}`
`2481`	`2481`	`}`