collect calls by target

crates/vestal/src/bang.rs

@@ -13,3 +13,16 @@ pub fn slice_shebang (buffer: &[u8]) -> (Arc<[u8]>, Arc<[u8]>) {
         (vec![].into(), buffer.to_vec().into())
     }
 }
+
+impl Vestal {
+    pub fn load_bang_data (&mut self, path: &Arc<PathBuf>) -> Usually<Arc<[u8]>> {
+        let (bang, data) = crate::bang::slice_shebang(read(path.as_path())?.as_slice());
+        self.path_to_bang.insert(path.clone(), bang.clone());
+        if bang.len() > 0 {
+            println!("  (bang {path:?} {:x})", bang.len())
+        }
+        self.path_to_data.insert(path.clone(), data.clone());
+        println!("  (buffer {:p} 0x{:08x} {path:?})", data.as_ptr(), data.len());
+        Ok(data)
+    }
+}

crates/vestal/src/load.rs

@@ -87,16 +87,6 @@ impl Vestal {
     pub fn load_exports (&mut self, path: &PathBuf, dll: &VecPE) -> Usually<()> {
         Ok(())
     }
-    pub fn load_bang_data (&mut self, path: &Arc<PathBuf>) -> Usually<Arc<[u8]>> {
-        let (bang, data) = crate::bang::slice_shebang(read(path.as_path())?.as_slice());
-        self.path_to_bang.insert(path.clone(), bang.clone());
-        if bang.len() > 0 {
-            println!("  (bang {path:?} {:x})", bang.len())
-        }
-        self.path_to_data.insert(path.clone(), data.clone());
-        println!("  (buffer {:p} 0x{:08x} {path:?})", data.as_ptr(), data.len());
-        Ok(data)
-    }
     pub fn load_pe (&mut self, path: &Arc<PathBuf>, data: &Arc<[u8]>) -> Usually<Arc<VecPE>> {
         let pe = Arc::new(VecPE::from_disk_data(data));
         self.path_to_pe.insert(path.clone(), pe.clone());

crates/vestal/src/main.rs

@@ -5,8 +5,223 @@ mod show;
 mod link;
 mod bang;
 pub(crate) use self::util::*;
-use clap::{arg, command, value_parser, ArgAction, Command};
+
+#[derive(Debug)]
+struct Rebuilder {
+    /// Paths visited
+    paths: BTreeSet<Arc<PathBuf>>,
+    /// All DLLs in scope
+    dlls:  BTreeMap<Arc<str>, Arc<Dll>>,
+}
+
+#[derive(Debug)]
+struct Dll {
+    /// Canonical name like `xxx.dll` (always lower case!)
+    name:  Arc<str>,
+    /// Path to DLL on host filesystem.
+    path:  Arc<PathBuf>,
+    /// Bytes of `#!`-instruction
+    bang:  Arc<[u8]>,
+    /// Parsed portable executable
+    pe:    Arc<VecPE>,
+    /// Bytes of `.text` section
+    code:  Arc<[u8]>,
+    /// DLLs that this one depends on
+    deps:  BTreeMap<Arc<str>, Arc<Self>>,
+    /// Calls to deps by source address.
+    calls_by_source: BTreeMap<u32, Arc<Call>>,
+    /// Calls to deps by target address.
+    calls_by_target: BTreeMap<u32, Vec<Arc<Call>>>,
+}
+
+#[derive(Debug)]
+struct Call {
+    /// Address on disk
+    offset: u32,
+    /// Address in memory
+    source: u32,
+    /// Length of call in opcodes
+    length: usize,
+    /// Call trampoline address
+    target: u32,
+    /// Library being called
+    module: Option<Arc<str>>,
+    /// Method being called
+    method: Option<Arc<str>>,
+}
+
+impl Rebuilder {
+    fn new (path: &impl AsRef<Path>) -> Usually<Self> {
+        let mut build = Self { paths: Default::default(), dlls: Default::default() };
+        let dll = Dll::new(&mut build, &Arc::new(PathBuf::from(path.as_ref())));
+        Ok(build)
+    }
+}
+
+impl Dll {
+    fn new (build: &mut Rebuilder, path: &Arc<PathBuf>) -> Usually<Arc<Self>> {
+        println!("\n(load {BOLD}{path:?}{RESET})");
+        build.paths.insert(path.clone());
+        let name = path.file_name().expect("no file name");
+        let name: Arc<str> = name.to_str().map(Arc::from).expect("non-unicode filename");
+        let (bang, data) = crate::bang::slice_shebang(read(path.as_path())?.as_slice());
+        let pe = Arc::new(VecPE::from_disk_data(data.clone()));
+        let code = pe.get_section_by_name(".text")?;
+        let start = code.pointer_to_raw_data.0 as usize;
+        let size = code.size_of_raw_data as usize;
+        let text = &data[start..start+size];
+        let mut calls_by_source = Default::default();
+        let mut calls_by_target = Default::default();
+        let _ = Self::calls(
+            &name,
+            &pe, 
+            start,
+            text,
+            &mut calls_by_source,
+            &mut calls_by_target,
+            false
+        )?;
+        let dll = Arc::new(Self {
+            name:  name.clone(),
+            path:  path.clone(),
+            bang,
+            pe,
+            code:  Arc::from(text),
+            deps:  Default::default(),
+            calls_by_source,
+            calls_by_target,
+        });
+        build.dlls.insert(name.clone(), dll.clone());
+        Ok(dll)
+    }
+    fn calls (
+        name:  &Arc<str>,
+        pe:    &VecPE,
+        start: usize,
+        data:  &[u8],
+        calls_by_source: &mut BTreeMap<u32, Arc<Call>>,
+        calls_by_target: &mut BTreeMap<u32, Vec<Arc<Call>>>,
+        verbose: bool,
+    ) -> Usually<usize> {
+        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)? {
+                calls += 1;
+                calls_by_source.insert(call.source, call.clone());
+                if !calls_by_target.contains_key(&call.target) {
+                    calls_by_target.insert(call.target, Default::default());
+                }
+                calls_by_target.get_mut(&call.target).unwrap().push(call);
+            }
+        }
+        println!("  (call-sites {calls})");
+        for (call, sites) in calls_by_target.iter() {
+            println!("  (0x{call:08x}\n    {:?})", sites.iter()
+                .map(|call|format!("0x{:08x}", call.offset))
+                .collect::<Vec<_>>());
+        }
+        Ok(calls)
+    }
+    fn call (
+        name:    &Arc<str>,
+        pe:      &VecPE,
+        start:   usize,
+        data:    &[u8],
+        decoder: &mut iced_x86::Decoder,
+        verbose: bool,
+    ) -> Usually<Option<Arc<Call>>> {
+        let position    = decoder.position();
+        let instruction = decoder.decode();
+        let opcodes     = &data[position..position+instruction.len()];
+        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 verbose {
+                    print!("  ({BOLD}0x{offset:08x}{RESET} 0x{offset_rva:08x}");
+                    println!(" {BOLD}{instruction:30}{RESET} 0x{call_target:x})");
+                }
+                //let unknown  = (String::from("unknown"), String::from("unknown"));
+                //let target   = self.addr_to_import.get(call_target).unwrap_or(&unknown).0;
+                //let method   = self.addr_to_import.get(call_target).unwrap_or(&unknown).1;
+                //let external = format!("{}::{}", target, method);
+                //if verbose {
+                    //println!("  ({BOLD}{external}{RESET}\n   Offset(0x{:08x}) RVA(R=0x{:08x})\n   {:25} {:40} 0x{:08x}",
+                        //offset, offset_rva,
+                        //opcodes.iter().map(|x|format!("{x:>02x}")).collect::<Vec<_>>().join(" "),
+                        //instruction, call_target,
+                    //);
+                //}
+                return Ok(Some(Arc::new(Call {
+                    offset: offset,
+                    source: offset_rva,
+                    length: opcodes.len(),
+                    target: call_target,
+                    module: None,
+                    method: None,
+                })))
+            }
+        }
+        Ok(None)
+    }
+    fn matches (instruction: &iced_x86::Instruction) -> bool {
+        instruction.op0_kind() == iced_x86::OpKind::Memory && (
+            instruction.flow_control() == iced_x86::FlowControl::IndirectBranch ||
+            instruction.flow_control() == iced_x86::FlowControl::IndirectCall
+        )
+    }
+    fn skip (opcodes: &[u8]) -> bool {
+        match opcodes[0] {
+            0x41 | 0x42 | 0x43 | 0x49 => match opcodes[1] {
+                0xff => return true,
+                _ => {}
+            },
+            0x48 => match opcodes[2] {
+                0x20 | 0x60 | 0x62 | 0xa0 | 0xa2 => return true,
+                _ => {}
+            },
+            0xff => match opcodes[1] {
+                0x10 | 0x12 | 0x13 |
+                0x50 | 0x51 | 0x52 | 0x53 | 0x54 | 0x55 | 0x56 | 0x57 |
+                0x60 | 0x90 | 0x92 | 0x93 | 0x94 | 0x97 => return true,
+                _ => {}
+            },
+            _ => {}
+        }
+        false
+    }
+    fn target (opcodes: &[u8], offset_rva: u32) -> Option<u32> {
+        match opcodes[0] {
+            0xff => match opcodes[1] {
+                0x15 | 0x25 => return Some(offset_rva + opcodes.len() as u32 + u32::from_le_bytes([
+                    opcodes[2],
+                    opcodes[3],
+                    opcodes[4],
+                    opcodes[5]
+                ])),
+                _ => {}
+            },
+            0x48 => match opcodes[1] {
+                0xff => match opcodes[2] {
+                    0x15 | 0x25 => return Some(offset_rva + opcodes.len() as u32 + u32::from_le_bytes([
+                        opcodes[3],
+                        opcodes[4],
+                        opcodes[5],
+                        opcodes[6]
+                    ])),
+                    _ => {}
+                },
+                _ => {}
+            }
+            _ => {}
+        }
+        None
+    }
+}
+
 fn main () -> Usually<()> {
+    use clap::{arg, command, value_parser, ArgAction, Command};
     let matches = command!()
         .arg(arg!([path] "Path to VST DLL").value_parser(value_parser!(PathBuf)))
         .arg(arg!(-i --inspect "Show info, don't run").required(false))
@@ -20,7 +235,8 @@ fn main () -> Usually<()> {
             println!("(search {path:?})")
         }
         if let Some(path) = vestal.resolve(path.to_str().expect("path must be unicode"))? {
-            vestal.enter(&path)?;
+            let rebuilder = Rebuilder::new(&path)?;
+            //vestal.enter(&path)?;
         } else {
             panic!("Could not find: {path:?}")
         }
@@ -29,6 +245,9 @@ fn main () -> Usually<()> {
     }
     Ok(())
 }
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
 #[derive(Default, Debug)]
 struct Vestal {
     search_paths:   Vec<PathBuf>,
@@ -53,11 +272,11 @@ impl Vestal {
             self.show_calls(dll_path, dll_path.as_ref() == path)?;
             if dll_path.as_ref() == path {
                 println!("{:?}", dll_data.hex_dump());
-                let text_section = dll.get_section_by_name(".text")?;
+                let text = dll.get_section_by_name(".text")?;
-                println!("\n{:#?}", text_section);
+                println!("\n{:#?}", text);
-                let start = text_section.pointer_to_raw_data.0 as usize;
+                let start = text.pointer_to_raw_data.0 as usize;
-                let size  = text_section.size_of_raw_data as usize;
+                let size  = text.size_of_raw_data as usize;
-                let vsize = text_section.virtual_size     as usize;
+                let vsize = text.virtual_size     as usize;
                 println!("\n{:?}", &dll_data[start..start+size].hex_dump());
                 println!("\n{:?}", &dll_data[start..start+vsize].hex_dump());
                 let elf = crate::link::create_elf(&dll_data[start..start+vsize]);