resolve callsites

crates/vestal/src/main.rs

@@ -26,8 +26,10 @@ struct Dll {
     pe:    Arc<VecPE>,
     /// Bytes of `.text` section
     code:  Arc<[u8]>,
-    /// Addresses of each imported method from each dependency
+    /// Addresses of imported methods by library
-    deps:  BTreeMap<Arc<str>, BTreeMap<Arc<str>, u32>>,
+    deps_by_library: BTreeMap<Arc<str>, BTreeMap<Arc<str>, u32>>,
+    /// Imported methods by address
+    deps_by_address: BTreeMap<u32, (Arc<str>, Arc<str>)>,
     /// Calls to dependencies by source address
     calls_by_source: BTreeMap<u32, Arc<Call>>,
     /// Calls to dependencies by target address
@@ -55,7 +57,7 @@ impl Rebuilder {
         let mut build = Self { paths: Default::default(), dlls: Default::default() };
         let path: Arc<PathBuf> = Arc::from(PathBuf::from(path.as_ref()));
         build.paths.insert(path.clone());
-        let dll = Dll::new(&Arc::new(PathBuf::from(path.as_ref())), false)?;
+        let dll = Dll::new(&Arc::new(PathBuf::from(path.as_ref())), true)?;
         build.dlls.insert(dll.name.clone(), dll);
         Ok(build)
     }
@@ -77,10 +79,12 @@ impl Dll {
         let text = &data[start..start+size];
         let mut calls_by_source = Default::default();
         let mut calls_by_target = Default::default();
-        let mut deps = Default::default();
+        let mut deps_by_library  = Default::default();
+        let mut deps_by_address  = Default::default();
         let (modules_count, methods_count) = Self::deps(
             &pe,
-            &mut deps,
+            &mut deps_by_library,
+            &mut deps_by_address,
             false
         )?;
         let calls = Self::calls(
@@ -88,36 +92,27 @@ impl Dll {
             &pe, 
             start,
             text,
+            Some(&deps_by_address),
             &mut calls_by_source,
             &mut calls_by_target,
-            false
+            true
         )?;
-        let dll = Arc::new(Self {
+        Ok(Arc::new(Self {
             name:  name.clone(),
             path:  path.clone(),
             bang,
             pe,
             code: Arc::from(text),
-            deps: deps.clone(),
+            deps_by_library: deps_by_library.clone(),
+            deps_by_address: deps_by_address.clone(),
             calls_by_source,
             calls_by_target,
-        });
+        }))
-        println!("  (deps-modules {modules_count})");
-        println!("  (deps-methods {methods_count})");
-        println!("  (call-sites {calls})");
-        if verbose {
-            for (call, sites) in dll.calls_by_target.iter() {
-                println!("  (0x{call:08x}\n    {:?})", sites.iter()
-                    .map(|call|format!("0x{:08x}", call.offset))
-                    .collect::<Vec<_>>());
-            }
-            println!("{deps:#?}");
-        }
-        Ok(dll)
     }
     fn deps (
         pe:              &VecPE,
-        deps:    &mut BTreeMap<Arc<str>, BTreeMap<Arc<str>, u32>>,
+        deps_by_library: &mut BTreeMap<Arc<str>, BTreeMap<Arc<str>, u32>>,
+        deps_by_address: &mut BTreeMap<u32, (Arc<str>, Arc<str>)>,
         verbose:         bool,
     ) -> Usually<(usize, usize)> {
         let directory = ImportDirectory::parse(pe)?;
@@ -145,22 +140,37 @@ impl Dll {
                     ImportData::ImportByName(name) => format!("{name}"),
                 };
                 let module_name: Arc<str> = module_name.clone().into();
-                if !deps.contains_key(&module_name) {
+                if !deps_by_library.contains_key(&module_name) {
-                    deps.insert(module_name.clone(), Default::default());
+                    deps_by_library.insert(module_name.clone(), Default::default());
                     modules += 1;
                 }
-                let module = deps.get_mut(&module_name).unwrap();
+                let module = deps_by_library.get_mut(&module_name).unwrap();
                 let method: Arc<str> = method.clone().into();
                 if module.contains_key(&method) {
                     panic!("duplicate method {method} in {module_name}");
                 }
                 module.insert(method.clone(), call_via);
+                if deps_by_address.contains_key(&call_via) {
+                    panic!("duplicate address {call_via} from {module_name}");
+                }
+                deps_by_address.insert(call_via, (module_name.clone(), method.clone()));
                 methods += 1;
                 if verbose {
                     println!("  ({index:5} 0x{call_via:08x} {module_name:>20} {method})");
                 }
             }
         }
+        println!("  (deps-modules {modules})");
+        println!("  (deps-methods {methods})");
+        if verbose {
+            for (module, methods) in deps_by_library.iter() {
+                print!("    ({module}");
+                for (method, addr) in methods.iter() {
+                    print!("\n      (0x{addr:08x} {method})")
+                }
+                println!(")");
+            }
+        }
         Ok((modules, methods))
     }
     fn calls (
@@ -168,6 +178,7 @@ impl Dll {
         pe:    &VecPE,
         start: usize,
         data:  &[u8],
+        deps:  Option<&BTreeMap<u32, (Arc<str>, Arc<str>)>>,
         calls_by_source: &mut BTreeMap<u32, Arc<Call>>,
         calls_by_target: &mut BTreeMap<u32, Vec<Arc<Call>>>,
         verbose: bool,
@@ -175,7 +186,7 @@ impl Dll {
         let mut decoder = iced_x86::Decoder::with_ip(64, data, 0x1000, 0);
         let mut calls = 0;
         while decoder.can_decode() {
-            if let Some(call) = Self::call(name, pe, start, data, &mut decoder, verbose)? {
+            if let Some(call) = Self::call(name, pe, start, data, deps, &mut decoder, verbose)? {
                 calls += 1;
                 calls_by_source.insert(call.source, call.clone());
                 if !calls_by_target.contains_key(&call.target) {
@@ -184,13 +195,36 @@ impl Dll {
                 calls_by_target.get_mut(&call.target).unwrap().push(call);
             }
         }
+        println!("  (call-sites {calls})");
+        if verbose {
+            for (target, sites) in calls_by_target.iter() {
+                let (_, _, external) = Self::dep_name(deps, *target);
+                println!("    ({:>5}x call 0x{target:08x} {external})", sites.len());
+                    //.map(|site|format!("0x{:08x}", site.offset))
+                    //.collect::<Vec<_>>());
+                //println!("    (call 0x{target:08x} {external}\n      {:?})", sites.iter()
+                    //.map(|site|format!("0x{:08x}", site.offset))
+                    //.collect::<Vec<_>>());
+            }
+        }
         Ok(calls)
     }
+    fn dep_name (deps: Option<&BTreeMap<u32, (Arc<str>, Arc<str>)>>, target: u32)
+        -> (Option<Arc<str>>, Option<Arc<str>>, Arc<str>)
+    {
+        let module = deps.and_then(|deps|deps.get(&target)).map(|dep|dep.0.clone());
+        let method = deps.and_then(|deps|deps.get(&target)).map(|dep|dep.1.clone());
+        let external = format!("{}::{}",
+            module.as_ref().map(|x|x.as_ref()).unwrap_or("???"),
+            method.as_ref().map(|x|x.as_ref()).unwrap_or("???"));
+        (module, method, external.into())
+    }
     fn call (
         name:    &Arc<str>,
         pe:      &VecPE,
         start:   usize,
         data:    &[u8],
+        deps:    Option<&BTreeMap<u32, (Arc<str>, Arc<str>)>>,
         decoder: &mut iced_x86::Decoder,
         verbose: bool,
     ) -> Usually<Option<Arc<Call>>> {
@@ -200,18 +234,22 @@ impl Dll {
         if Self::matches(&instruction) && !Self::skip(opcodes) {
             let offset = (position + start) as u32;
             let offset_rva = pe.offset_to_rva(Offset(offset))?.0;
-            if let Some(call_target) = Self::target(opcodes, offset_rva) {
+            if let Some(target) = Self::target(opcodes, offset_rva) {
+                let (module, method, external) = Self::dep_name(deps, target);
                 if verbose {
-                    print!("  ({BOLD}0x{offset:08x}{RESET} 0x{offset_rva:08x}");
+                    let external = format!("{}::{}",
-                    println!(" {BOLD}{instruction:30}{RESET} 0x{call_target:x})");
+                        module.as_ref().map(|x|x.as_ref()).unwrap_or("???"),
+                        method.as_ref().map(|x|x.as_ref()).unwrap_or("???"));
+                    println!("  ({BOLD}0x{:08x}{RESET} 0x{:08x} {BOLD}{:30}{RESET} 0x{:x} {}",
+                        offset, offset_rva, instruction, target, external);
                 }
                 return Ok(Some(Arc::new(Call {
                     offset: offset,
                     source: offset_rva,
                     length: opcodes.len(),
-                    target: call_target,
+                    target,
-                    module: None,
+                    module,
-                    method: None,
+                    method,
                 })))
             }
         }