From dbca25d9cb35bd844de98dd528b34f5acea4a391 Mon Sep 17 00:00:00 2001 From: unspeaker Date: Fri, 21 Feb 2025 22:51:39 +0200 Subject: [PATCH] resolve 1 level of imports --- crates/vestal/src/main.rs | 107 +++++++++++++++++++------------------- 1 file changed, 54 insertions(+), 53 deletions(-) diff --git a/crates/vestal/src/main.rs b/crates/vestal/src/main.rs index 9abd4a0..108ea5a 100644 --- a/crates/vestal/src/main.rs +++ b/crates/vestal/src/main.rs @@ -26,11 +26,11 @@ struct Dll { pe: Arc, /// Bytes of `.text` section code: Arc<[u8]>, - /// DLLs that this one depends on - deps: BTreeMap, Arc>, - /// Calls to deps by source address. + /// Addresses of each imported method from each dependency + deps: BTreeMap, BTreeMap, u32>>, + /// Calls to dependencies by source address calls_by_source: BTreeMap>, - /// Calls to deps by target address. + /// Calls to dependencies by target address calls_by_target: BTreeMap>>, } @@ -53,15 +53,20 @@ struct Call { impl Rebuilder { fn new (path: &impl AsRef) -> Usually { let mut build = Self { paths: Default::default(), dlls: Default::default() }; - let dll = Dll::new(&mut build, &Arc::new(PathBuf::from(path.as_ref()))); + 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)?; + build.dlls.insert(dll.name.clone(), dll); Ok(build) } } impl Dll { - fn new (build: &mut Rebuilder, path: &Arc) -> Usually> { + fn new ( + path: &Arc, + verbose: bool + ) -> Usually> { println!("\n(load {BOLD}{path:?}{RESET})"); - build.paths.insert(path.clone()); let name = path.file_name().expect("no file name"); let name: Arc = name.to_str().map(Arc::from).expect("non-unicode filename"); let (bang, data) = crate::bang::slice_shebang(read(path.as_path())?.as_slice()); @@ -72,6 +77,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 (modules_count, methods_count) = Self::deps( + &pe, + &mut deps, + false + )?; let calls = Self::calls( &name, &pe, @@ -81,38 +92,43 @@ impl Dll { &mut calls_by_target, false )?; - let (modules_count, methods_count) = Self::imports(&pe)?; let dll = Arc::new(Self { name: name.clone(), path: path.clone(), bang, pe, - code: Arc::from(text), - deps: Default::default(), + code: Arc::from(text), + deps: deps.clone(), calls_by_source, calls_by_target, }); println!(" (deps-modules {modules_count})"); println!(" (deps-methods {methods_count})"); println!(" (call-sites {calls})"); - for (call, sites) in dll.calls_by_target.iter() { - println!(" (0x{call:08x}\n {:?})", sites.iter() - .map(|call|format!("0x{:08x}", call.offset)) - .collect::>()); + 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:#?}"); } - build.dlls.insert(name.clone(), dll.clone()); Ok(dll) } - fn imports (pe: &VecPE) -> Usually<(usize, usize)> { + fn deps ( + pe: &VecPE, + deps: &mut BTreeMap, BTreeMap, u32>>, + verbose: bool, + ) -> Usually<(usize, usize)> { let directory = ImportDirectory::parse(pe)?; - let mut total = 0; + let mut modules = 0; + let mut methods = 0; for descriptor in directory.descriptors { - let dep = descriptor.get_name(pe)?.as_str()?.to_lowercase(); + let module_name = descriptor.get_name(pe)?.as_str()?.to_lowercase(); let imp = descriptor.get_imports(pe)?; let iat = descriptor.get_first_thunk(pe)?; let ilt = descriptor.get_original_first_thunk(pe)?; let lut = descriptor.get_lookup_thunks(pe)?; - let mut imports = Vec::new(); let unwrap_thunk = |thunk: &Thunk, name|match thunk { Thunk::Thunk32(t) => panic!("32 bit {name}"), Thunk::Thunk64(t) => t.0 @@ -124,32 +140,28 @@ impl Dll { lut.iter().map(|thunk|format!("0x{:08x}", unwrap_thunk(thunk, "lookup thunk"))), ).enumerate() { let call_via = descriptor.first_thunk.0 + index as u32 * 8; - let name = match import { - ImportData::Ordinal(x) => { - //print!("\n (import-ordinal {BOLD}0x{:>08x}{RESET} IAT={} ILT={} LU={} 0x{:>04x})", - //call_via, thunk, orig, lookup, x); - format!("___VESTAL___ORD___{x}___") - }, - ImportData::ImportByName(name) => { - //print!("\n (import-by-name {BOLD}0x{:>08x}{RESET} IAT={} ILT={} LU={} {:?})", - //call_via, thunk, orig, lookup, name); - format!("{name}") - }, + let method = match import { + ImportData::Ordinal(x) => format!("___VESTAL_ORDINAL_{x}"), + ImportData::ImportByName(name) => format!("{name}"), }; - println!(" ({index:5} 0x{call_via:08x} {dep:>20} {name}"); - imports.push((thunk, orig, import)); - total += 1; - //if let Some(existing) = self.addr_to_import.get(&call_via) { - //panic!("addr space overlap at 0x{call_via:x}: {}::{} vs {}::{}", - //existing.0, - //existing.1, - //dep.to_string(), - //name); - //} - //self.addr_to_import.insert(call_via, (dep.to_string(), name)); + let module_name: Arc = module_name.clone().into(); + if !deps.contains_key(&module_name) { + deps.insert(module_name.clone(), Default::default()); + modules += 1; + } + let module = deps.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); + methods += 1; + if verbose { + println!(" ({index:5} 0x{call_via:08x} {module_name:>20} {method})"); + } } } - Ok((0, total)) + Ok((modules, methods)) } fn calls ( name: &Arc, @@ -193,17 +205,6 @@ impl Dll { 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::>().join(" "), - //instruction, call_target, - //); - //} return Ok(Some(Arc::new(Call { offset: offset, source: offset_rva,