sort of map of what needs to be linked

crates/vestal/src/call_sites.rs

@@ -3,16 +3,56 @@ use crate::*;
 impl Module {
     /// Collect all calls that point to imports.
     pub fn load_call_sites (self: Arc<Self>, recurse: bool) -> Usually<Arc<Self>> {
-        if self.verbose {
-            println!("  {DIM}(load-call-sites){RESET}");
+        self.load_call_sites_slice(0, recurse.then_some(0))?;
+        Ok(self)
     }
-        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();
+    /// Collect all calls that point to imports, starting from an address.
+    fn load_call_sites_slice (
+        self:    &Arc<Self>,
+        start:   usize, 
+        recurse: Option<usize>,
+    ) -> Usually<()> {
+        let group = false;
+        //if self.verbose {
+            //println!("  {DIM}(load-call-sites {} {} {} {RESET}",
+                //fmt_num(start),
+                //recurse.unwrap_or(0),
+                //self.name);
+        //}
+        let code = &self.code.as_ref()[start..];
+        //println!("{:x?}", &code[..64]);
+        let rip = self.code_base as u64 + start as u64;
+        let mut decoder = Decoder::with_ip(64, code, rip, 0);
         while decoder.can_decode() {
             let position    = decoder.position();
             let instruction = decoder.decode();
-            let opcodes     = &self.code[position..position+instruction.len()];
+            let opcodes     = &code[position..position+instruction.len()];
+            let address     = self.code_base + start as u32 + position as u32;
+            //println!("{:15} {}  {instruction} {}", self.name, fmt_num(address as usize), fmt_bytes(opcodes));
+            // Ascend on RET
+            if let Some(depth) = recurse {
+                if depth > 0 && opcodes == &[0xc3] {
+                    break
+                }
+            }
+            // Descend on CALL/JMP
             if CallSite::matches(&instruction) && !CallSite::skip(opcodes) {
+                let call_site = self.call_site(start + position, opcodes)?;
+                if !self.targets.read().unwrap().contains_key(&call_site.target) {
+                    Log::call_site(self.verbose && !group, &call_site, Some(opcodes), recurse);
+                }
+                self.load_call_site(&call_site);
+                if let Some(depth) = recurse {
+                    self.load_call_site_recurse(&call_site, depth)?;
+                }
+            }
+        }
+        //Log::call_sites(self.verbose && group, &self.name, &self.targets.read().unwrap(), recurse);
+        Ok(())
+    }
+    /// Annotate a given instruction in the code section as a [CallSite].
+    fn call_site (self: &Arc<Self>, position: usize, opcodes: &[u8]) -> Usually<Arc<CallSite>> {
+        let group  = false;
         let offset = (position + self.code_start) as u32;
         let source = self.pe.offset_to_rva(Offset((position + self.code_start) as u32))?.0;
         let target = CallSite::target(source, opcodes).unwrap_or(0);
@@ -30,37 +70,80 @@ impl Module {
                 None
             }
         });
-                call_site.show(Some(opcodes));
+        Ok(call_site)
+    }
+    /// Add a [CallSite] under the appropriate target
+    fn load_call_site (&self, call_site: &Arc<CallSite>) {
+        let mut targets = self.targets.write().unwrap();
         if !targets.contains_key(&call_site.target) {
             targets.insert(call_site.target, vec![]);
         }
-                targets.get_mut(&call_site.target).unwrap().push(call_site);
+        targets.get_mut(&call_site.target).unwrap().push(call_site.clone());
+    }
+    /// Follow the call site and resolve the next call site found in the dependency.
+    fn load_call_site_recurse (self: &Arc<Self>, call_site: &Arc<CallSite>, depth: usize) -> Usually<()> {
+        let CallSite { module, method, .. } = call_site.as_ref();
+        if let (Some(module), Some(method)) = (module, method) {
+            if let Some(method) = module.exports.read().unwrap().get(method) {
+                module.load_call_site_recurse_export(method, depth)?;
             }
         }
-        Ok(self)
+        Ok(())
     }
-
-    fn parse_link (&self, link: &Arc<CallSite>) -> Option<u32> {
-        None
-        //Some(*self.deps_by_library
-            //.get(link.module.as_ref()?)?
-            //.get(link.method.as_ref()?)?)
+    /// Follow a function call thunk.
+    fn load_call_site_recurse_export (self: &Arc<Self>, method: &ThunkData, depth: usize) -> Usually<()> {
+        match method {
+            ThunkData::Function(rva) => {
+                self.load_call_site_recurse_function(rva, depth + 1)?;
+            },
+            ThunkData::ForwarderString(rva) => {
+                self.load_call_site_recurse_forward(rva, depth + 1)?;
+            },
+            x => {
+                unimplemented!("{x:?}");
             }
-
-    pub fn resolve_forward (&self, rva: &RVA) -> Usually<Option<(Arc<str>, Arc<str>)>> {
-        let mut addr = (rva.0 - self.code_base) as usize;
+        }
+        Ok(())
+    }
+    /// Follow a function call site.
+    fn load_call_site_recurse_function (self: &Arc<Self>, rva: &RVA, depth: usize) -> Usually<()> {
+        let index = (rva.0 - self.code_base) as usize;
+        let slice = &self.code[index..index+64];
+        self.load_call_sites_slice(index, Some(depth))?;
+        Ok(())
+    }
+    /// Follow a forwarded call site.
+    fn load_call_site_recurse_forward (self: &Arc<Self>, rva: &RVA, depth: usize) -> Usually<()> {
+        if let Some((mut module_name, method_name)) = self.resolve_forward(rva)? {
+            let mut name = module_name.to_lowercase();
+            if !name.ends_with(".dll") {
+                name = format!("{name}.dll");
+            }
+            if let Some(module) = self.dependencies.read().unwrap().get(name.as_str()) {
+                if let Some(method) = module.exports.read().unwrap().get(&method_name) {
+                    module.load_call_site_recurse_export(method, depth)?;
+                }
+            }
+            //panic!("{} {name}::{method_name}", self.name);
+        } else {
+            panic!("unresolved fwd {rva:x?} {}", self.read_forward(rva)?);
+        }
+        Ok(())
+    }
+    pub fn read_forward (&self, rva: &RVA) -> Usually<Arc<str>> {
+        let mut address = self.pe.rva_to_offset(*rva)?.0 as usize;
         let mut forward = vec![];
-        while let Some(c) = self.pe.as_slice().get(addr) {
+        while let Some(c) = self.pe.as_slice().get(address) {
             if *c == 0x00 {
                 break
             }
             forward.push(*c);
-            addr += 1;
+            address += 1;
         }
-        Ok(String::from_utf8(forward)?
-            .as_str()
-            .split_once(".")
-            .map(|(x, y)|(x.into(), y.into())))
+        Ok(String::from_utf8(forward)?.as_str().into())
+    }
+    pub fn resolve_forward (&self, rva: &RVA) -> Usually<Option<(Arc<str>, Arc<str>)>> {
+        Ok(self.read_forward(rva)?.split_once(".").map(|(x, y)|(x.into(), y.into())))
     }
 }
 

crates/vestal/src/main.rs

@@ -43,6 +43,8 @@ pub struct Module {
     pub dependencies: RwLock<BTreeMap<Arc<str>, Arc<Self>>>,
     /// Call targets to methods from imported modules.
     pub imports: RwLock<BTreeMap<u32, (Arc<str>, Arc<str>)>>,
+    /// Mapping of call target to its invocations.
+    pub targets: RwLock<BTreeMap<u32, Vec<Arc<CallSite>>>>,
     /// Addresses of exported methods by name
     pub exports: RwLock<BTreeMap<Arc<str>, ThunkData>>,
     /// Locations in `.text` section that need to be patched
@@ -106,6 +108,7 @@ impl Module {
             search_paths: Default::default(),
             dependencies: Default::default(),
             imports:      Default::default(),
+            targets:      Default::default(),
             exports:      Default::default(),
             call_sites:   Default::default(),
             pe,

crates/vestal/src/show.rs

@@ -39,6 +39,44 @@ impl Log {
             println!("(found {} {:?})", name.as_ref(), path.as_ref());
         }
     }
+    pub fn call_site (
+        show:    bool,
+        call:    &CallSite,
+        opcodes: Option<&[u8]>,
+        depth:   Option<usize>,
+    ) {
+        if show {
+            if let Some(depth) = depth {
+                for i in 0..depth {
+                    print!("    ");
+                }
+            }
+            call.show(opcodes);
+        }
+    }
+    pub fn call_sites (
+        show:    bool,
+        name:    &impl AsRef<str>,
+        targets: &BTreeMap<u32, Vec<Arc<CallSite>>>,
+        depth:   Option<usize>
+    ) {
+        for (target, call_sites) in targets.iter() {
+            if let Some(depth) = depth {
+                for i in 0..depth {
+                    print!("    ");
+                }
+            }
+            //println!("    (call {:15} {})", name.as_ref(), fmt_num(*target as usize));
+            for call_site in call_sites.iter() {
+                if let Some(depth) = depth {
+                    for i in 0..depth {
+                        print!("    ");
+                    }
+                }
+                call_site.show(None);
+            }
+        }
+    }
 }
 
 impl Show {
@@ -72,6 +110,24 @@ impl std::fmt::Debug for Module {
 }
 
 impl Module {
+    pub fn show_layout (&self) {
+        println!("\n█ = 64k");
+        for module in self.namespace.read().unwrap().values() {
+            print!("\n{} ", module.name);
+            let page_size = 256;
+            for i in 0..(module.code_start + module.code_size) / page_size {
+                if i % 64 == 0 {
+                    print!("\n");
+                }
+                if i * page_size < module.code_start {
+                    print!("░");
+                } else {
+                    print!("▒");
+                }
+            }
+            print!("\n");
+        }
+    }
     pub fn show_instruction (&self, addr: usize) {
         let mut decoder = Decoder::with_ip(64, &self.code[addr..], 0x1000, DecoderOptions::NONE);
         let instruction = decoder.decode();
@@ -132,9 +188,9 @@ impl Module {
 impl CallSite {
     pub fn show (&self, opcodes: Option<&[u8]>) {
         let label = self.caller.imports.read().unwrap().get(&self.target)
-            .map(|(module, method)|format!("{module}::{method}"))
+            .map(|(module, method)|format!("{GREEN}{module}::{method}{RESET}"))
             .unwrap_or_else(||format!("{RED}unresolved{RESET}"));
-        println!("    ╰-> (call {} {} {} {DIM}{:20}{RESET} {} {BOLD}{}{RESET})",
+        println!("    ╰-> (call {:15} {} {} {DIM}{:20}{RESET} {} {BOLD}{}{RESET})",
             &self.caller.name,
             fmt_num(self.offset as usize),
             fmt_num(self.source as usize),