diff --git a/crates/vestal/src/call_sites.rs b/crates/vestal/src/call_sites.rs index 6bbf6a3..1f1e2cc 100644 --- a/crates/vestal/src/call_sites.rs +++ b/crates/vestal/src/call_sites.rs @@ -3,64 +3,147 @@ use crate::*; impl Module { /// Collect all calls that point to imports. 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(), self.code_base as u64, 0); - let mut targets: BTreeMap>> = Default::default(); + self.load_call_sites_slice(0, recurse.then_some(0))?; + Ok(self) + } + /// Collect all calls that point to imports, starting from an address. + fn load_call_sites_slice ( + self: &Arc, + start: usize, + recurse: Option, + ) -> 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()]; - if CallSite::matches(&instruction) && !CallSite::skip(opcodes) { - 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); - let import = self.imports.read().unwrap().get(&target).cloned(); - let call_site = Arc::new(CallSite { - caller: self.clone(), - source, - offset, - target, - length: opcodes.len(), - method: import.clone().map(|x|x.1), - module: if let Some(import) = import { - self.dependencies.read().unwrap().get(&import.0).cloned() - } else { - None - } - }); - call_site.show(Some(opcodes)); - if !targets.contains_key(&call_site.target) { - targets.insert(call_site.target, vec![]); + 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)?; } - targets.get_mut(&call_site.target).unwrap().push(call_site); } } - Ok(self) + //Log::call_sites(self.verbose && group, &self.name, &self.targets.read().unwrap(), recurse); + Ok(()) } - - fn parse_link (&self, link: &Arc) -> Option { - None - //Some(*self.deps_by_library - //.get(link.module.as_ref()?)? - //.get(link.method.as_ref()?)?) + /// Annotate a given instruction in the code section as a [CallSite]. + fn call_site (self: &Arc, position: usize, opcodes: &[u8]) -> Usually> { + 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); + let import = self.imports.read().unwrap().get(&target).cloned(); + let call_site = Arc::new(CallSite { + caller: self.clone(), + source, + offset, + target, + length: opcodes.len(), + method: import.clone().map(|x|x.1), + module: if let Some(import) = import { + self.dependencies.read().unwrap().get(&import.0).cloned() + } else { + None + } + }); + Ok(call_site) } - - pub fn resolve_forward (&self, rva: &RVA) -> Usually, Arc)>> { - let mut addr = (rva.0 - self.code_base) as usize; + /// Add a [CallSite] under the appropriate target + fn load_call_site (&self, call_site: &Arc) { + 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.clone()); + } + /// Follow the call site and resolve the next call site found in the dependency. + fn load_call_site_recurse (self: &Arc, call_site: &Arc, 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(()) + } + /// Follow a function call thunk. + fn load_call_site_recurse_export (self: &Arc, 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:?}"); + } + } + Ok(()) + } + /// Follow a function call site. + fn load_call_site_recurse_function (self: &Arc, 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, 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> { + 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, Arc)>> { + Ok(self.read_forward(rva)?.split_once(".").map(|(x, y)|(x.into(), y.into()))) } } diff --git a/crates/vestal/src/main.rs b/crates/vestal/src/main.rs index a7a5d63..f329a80 100644 --- a/crates/vestal/src/main.rs +++ b/crates/vestal/src/main.rs @@ -43,6 +43,8 @@ pub struct Module { pub dependencies: RwLock, Arc>>, /// Call targets to methods from imported modules. pub imports: RwLock, Arc)>>, + /// Mapping of call target to its invocations. + pub targets: RwLock>>>, /// Addresses of exported methods by name pub exports: RwLock, 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, diff --git a/crates/vestal/src/show.rs b/crates/vestal/src/show.rs index 8eefa75..6d56c23 100644 --- a/crates/vestal/src/show.rs +++ b/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, + ) { + 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, + targets: &BTreeMap>>, + depth: Option + ) { + 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),