From 41a130bce006d95b5211ad511c986b2a7a48ab25 Mon Sep 17 00:00:00 2001 From: unspeaker Date: Fri, 21 Feb 2025 23:28:32 +0200 Subject: [PATCH] resolve callsites --- crates/vestal/src/main.rs | 106 ++++++++++++++++++++++++++------------ 1 file changed, 72 insertions(+), 34 deletions(-) diff --git a/crates/vestal/src/main.rs b/crates/vestal/src/main.rs index 108ea5a..b14e145 100644 --- a/crates/vestal/src/main.rs +++ b/crates/vestal/src/main.rs @@ -26,8 +26,10 @@ struct Dll { pe: Arc, /// Bytes of `.text` section code: Arc<[u8]>, - /// Addresses of each imported method from each dependency - deps: BTreeMap, BTreeMap, u32>>, + /// Addresses of imported methods by library + deps_by_library: BTreeMap, BTreeMap, u32>>, + /// Imported methods by address + deps_by_address: BTreeMap, Arc)>, /// Calls to dependencies by source address calls_by_source: BTreeMap>, /// 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 = 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,37 +92,28 @@ 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::>()); - } - println!("{deps:#?}"); - } - Ok(dll) + })) } fn deps ( - pe: &VecPE, - deps: &mut BTreeMap, BTreeMap, u32>>, - verbose: bool, + pe: &VecPE, + deps_by_library: &mut BTreeMap, BTreeMap, u32>>, + deps_by_address: &mut BTreeMap, Arc)>, + verbose: bool, ) -> Usually<(usize, usize)> { let directory = ImportDirectory::parse(pe)?; let mut modules = 0; @@ -145,22 +140,37 @@ impl Dll { ImportData::ImportByName(name) => format!("{name}"), }; let module_name: Arc = module_name.clone().into(); - if !deps.contains_key(&module_name) { - deps.insert(module_name.clone(), Default::default()); + if !deps_by_library.contains_key(&module_name) { + 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 = 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, Arc)>>, calls_by_source: &mut BTreeMap>, calls_by_target: &mut BTreeMap>>, 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::>()); + //println!(" (call 0x{target:08x} {external}\n {:?})", sites.iter() + //.map(|site|format!("0x{:08x}", site.offset)) + //.collect::>()); + } + } Ok(calls) } + fn dep_name (deps: Option<&BTreeMap, Arc)>>, target: u32) + -> (Option>, Option>, Arc) + { + 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, pe: &VecPE, start: usize, data: &[u8], + deps: Option<&BTreeMap, Arc)>>, decoder: &mut iced_x86::Decoder, verbose: bool, ) -> Usually>> { @@ -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}"); - println!(" {BOLD}{instruction:30}{RESET} 0x{call_target:x})"); + let external = format!("{}::{}", + 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, - module: None, - method: None, + target, + module, + method, }))) } }