diff --git a/crates/vestal/src/call_sites.rs b/crates/vestal/src/call_sites.rs index e649416..ea63a83 100644 --- a/crates/vestal/src/call_sites.rs +++ b/crates/vestal/src/call_sites.rs @@ -2,11 +2,12 @@ use crate::*; impl Module { /// Collect all calls that point to imports. - pub fn load_call_sites (self: Arc) -> Usually> { + pub fn load_call_sites (self: Arc, recurse: bool) -> Usually> { if self.verbose { println!(" {DIM}(load-call-sites){RESET}"); } - let mut decoder = Decoder::with_ip(64, self.code.as_ref(), 0, 0); + let mut decoder = Decoder::with_ip(64, self.code.as_ref(), self.code_base as u64, 0); + let mut targets: BTreeMap>> = Default::default(); while decoder.can_decode() { let position = decoder.position(); let instruction = decoder.decode(); @@ -14,6 +15,35 @@ impl Module { if CallSite::matches(&instruction) && !CallSite::skip(opcodes) { let offset = (position + self.code_start) as u32; let source = self.pe.offset_to_rva(Offset(offset))?.0; + let target = CallSite::target(source, opcodes).unwrap_or(0); + let imports = self.imports.read().unwrap(); + let label = imports.get(&target) + .map(|(module, method)|format!("{module}::{method}")) + .unwrap_or_else(||format!("{RED}unresolved{RESET}")); + println!(" (call {} {} {} {DIM}{:20}{RESET} {} {BOLD}{}{RESET})", + &self.name, + fmt_num(offset as usize), + fmt_num(source as usize), + fmt_bytes(opcodes), + fmt_num(target as usize), + label); + if !targets.contains_key(&target) { + targets.insert(target, vec![]); + } + let import = imports.get(&target).clone(); + targets.get_mut(&target).unwrap().push(Arc::new(CallSite { + caller: self.clone(), + source, + offset, + length: opcodes.len(), + target, + method: import.map(|x|x.1.clone()), + module: if let Some(import) = import { + self.dependencies.read().unwrap().get(&import.0).cloned() + } else { + None + } + })); } } Ok(self) @@ -70,8 +100,8 @@ impl CallSite { } false } - pub fn target (opcodes: &[u8], offset_rva: u32) -> Option { - let rip = offset_rva + opcodes.len() as u32; + pub fn target (address: u32, opcodes: &[u8]) -> Option { + let rip = address + opcodes.len() as u32; match opcodes[0] { 0xff => match opcodes[1] { 0x15 | 0x25 => return Some(rip + u32::from_le_bytes([ diff --git a/crates/vestal/src/main.rs b/crates/vestal/src/main.rs index 729ebd7..a7a5d63 100644 --- a/crates/vestal/src/main.rs +++ b/crates/vestal/src/main.rs @@ -62,11 +62,11 @@ pub struct CallSite { /// Length of link in opcodes pub length: usize, /// CallSite trampoline address - pub address: u32, + pub target: u32, /// Module that is being called - pub target: Arc, + pub module: Option>, /// Name of method that is being called - pub method: Arc, + pub method: Option>, } fn main () -> Usually<()> { use clap::{arg, command, value_parser, ArgAction, Command}; @@ -77,7 +77,8 @@ fn main () -> Usually<()> { .search(std::env::current_dir()?) .search(canonicalize(path.clone().parent().expect("invalid parent path"))?) .search("/home/user/Lab/Cosmo/wineprefix/drive_c/windows/system32") - .load(true)?; + .load(true)? + .resolve(true)?; } else { println!("Pass a path to a VST DLL"); } @@ -117,17 +118,6 @@ impl Module { self.search_paths.write().unwrap().insert(path.as_ref().to_path_buf().into()); self } - /// Load the dependency tree, starting from this module. - fn load (self: Arc, recurse: bool) -> Usually> { - let module = self - .load_exports()? - .load_imports(recurse)? - .load_call_sites()?; - if module.verbose { - println!("{:?}", &module); - } - Ok(module) - } /// Search for DLL by name in search paths. fn find (&self, name: &impl AsRef) -> Usually> { let name = name.as_ref(); @@ -142,6 +132,22 @@ impl Module { } Ok(None) } + /// Load the dependency tree reachable from this module. + fn load (self: Arc, recurse: bool) -> Usually> { + let module = self.load_exports()?.load_imports(recurse)?; + if module.verbose { + println!("{:?}", &module); + } + Ok(module) + } + /// Identify call sites and try to match them with their new locations. + fn resolve (self: Arc, recurse: bool) -> Usually> { + let module = self.load_call_sites(recurse)?; + if module.verbose { + println!("{:?}", &module); + } + Ok(module) + } } fn to_dll_name (path: &impl AsRef) -> Arc { diff --git a/crates/vestal/src/show.rs b/crates/vestal/src/show.rs index 0bbb2fc..d16e50a 100644 --- a/crates/vestal/src/show.rs +++ b/crates/vestal/src/show.rs @@ -46,7 +46,7 @@ impl Show { println!(" {}{}{}", link.map(|link|format!(" (offset {})", fmt_num(link.offset as usize))).unwrap_or(String::new()), link.map(|link|format!(" (source {})", fmt_num(link.source as usize))).unwrap_or(String::new()), - link.map(|link|format!(" (target {})", fmt_num(link.address as usize))).unwrap_or(String::new())); + link.map(|link|format!(" (target {})", fmt_num(link.target as usize))).unwrap_or(String::new())); } pub fn link_dasm (bytes: &[u8], rip: usize) -> Arc { let mut decoder = Decoder::with_ip(64, bytes, 0x1000 + rip as u64, DecoderOptions::NONE); @@ -132,10 +132,10 @@ impl Module { impl CallSite { pub fn show (&self) { let caller = self.caller.name.as_ref(); - let module = self.target.name.as_ref(); + let module = self.module.as_ref(); let method = self.method.as_ref(); let style = GREEN; - println!("╰--------> {caller} -> {style}{module}::{method}{RESET}"); + println!("╰--------> {caller} -> {style}{module:?}::{method:?}{RESET}"); //module.map(|x|x.as_ref()).unwrap_or(&""), //method.map(|x|x.as_ref()).unwrap_or(&"")); //} else {