update call site collection

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<Self>) -> Usually<Arc<Self>> {
+    pub fn load_call_sites (self: Arc<Self>, recurse: bool) -> Usually<Arc<Self>> {
         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<u32, Vec<Arc<CallSite>>> = 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<u32> {
-        let rip = offset_rva + opcodes.len() as u32;
+    pub fn target (address: u32, opcodes: &[u8]) -> Option<u32> {
+        let rip = address + opcodes.len() as u32;
         match opcodes[0] {
             0xff => match opcodes[1] {
                 0x15 | 0x25 => return Some(rip + u32::from_le_bytes([

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<Module>,
+    pub module: Option<Arc<Module>>,
     /// Name of method that is being called
-    pub method: Arc<str>,
+    pub method: Option<Arc<str>>,
 }
 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<Self>, recurse: bool) -> Usually<Arc<Self>> {
-        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<str>) -> Usually<Option<PathBuf>> {
         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<Self>, recurse: bool) -> Usually<Arc<Self>> {
+        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<Self>, recurse: bool) -> Usually<Arc<Self>> {
+        let module = self.load_call_sites(recurse)?;
+        if module.verbose {
+            println!("{:?}", &module);
+        }
+        Ok(module)
+    }
 }
 
 fn to_dll_name (path: &impl AsRef<Path>) -> Arc<str> {

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<str> {
         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 {