From 32f901d17a4f2be0218eba69577f3cf32eda5dfc Mon Sep 17 00:00:00 2001 From: unspeaker Date: Sun, 23 Feb 2025 22:03:15 +0200 Subject: [PATCH] retracking imports/dependencies --- crates/vestal/src/call_sites.rs | 3 ++ crates/vestal/src/exports.rs | 3 ++ crates/vestal/src/imports.rs | 71 ++++++++++++++++++++------------- crates/vestal/src/main.rs | 6 +-- crates/vestal/src/show.rs | 15 ++----- 5 files changed, 56 insertions(+), 42 deletions(-) diff --git a/crates/vestal/src/call_sites.rs b/crates/vestal/src/call_sites.rs index 02f08cb..e649416 100644 --- a/crates/vestal/src/call_sites.rs +++ b/crates/vestal/src/call_sites.rs @@ -3,6 +3,9 @@ use crate::*; impl Module { /// Collect all calls that point to imports. pub fn load_call_sites (self: Arc) -> Usually> { + if self.verbose { + println!(" {DIM}(load-call-sites){RESET}"); + } let mut decoder = Decoder::with_ip(64, self.code.as_ref(), 0, 0); while decoder.can_decode() { let position = decoder.position(); diff --git a/crates/vestal/src/exports.rs b/crates/vestal/src/exports.rs index 79d268b..89a6aac 100644 --- a/crates/vestal/src/exports.rs +++ b/crates/vestal/src/exports.rs @@ -2,6 +2,9 @@ use crate::*; impl Module { /// Collect all exported methods. pub fn load_exports (self: Arc) -> Usually> { + if self.verbose { + println!(" {DIM}(load-exports){RESET}"); + } let directory = ExportDirectory::parse(self.pe.as_ref())?; let export_map = directory.get_export_map(self.pe.as_ref())?; let exports = export_map.into_iter().map(|(k, v)|(k.into(), v)).collect(); diff --git a/crates/vestal/src/imports.rs b/crates/vestal/src/imports.rs index b3c8fd2..36b38f7 100644 --- a/crates/vestal/src/imports.rs +++ b/crates/vestal/src/imports.rs @@ -2,22 +2,40 @@ use crate::*; impl Module { /// Collect all imported modules and methods. pub fn load_imports (self: Arc, recurse: bool) -> Usually> { + if self.verbose { + println!(" {DIM}(load-imports){RESET}"); + } let pe = self.pe.clone(); let directory = ImportDirectory::parse(pe.as_ref())?; let mut new_modules = vec![]; for descriptor in directory.descriptors.iter() { - let (name, imports) = Self::parse_import(pe.as_ref(), descriptor)?; - let exists = self.namespace.read().unwrap().contains_key(&name); - if !exists { - Log::dep(false, &name, &imports); - let path = self.find(&name)?.expect("not found"); + let (module_name, imports) = Self::parse_import(pe.as_ref(), descriptor)?; + let create = !self.namespace.read().unwrap().contains_key(&module_name); + if create { + Log::dep(false, &module_name, &imports); + let path = self.find(&module_name)?.expect("not found"); let module = Arc::new(Self { namespace: self.namespace.clone(), search_paths: self.search_paths.clone(), ..Self::from_path(&path, self.verbose)? }); + assert_eq!(module_name, module.name); self.namespace.write().unwrap().insert(module.name.clone(), module.clone()); - new_modules.push(module); + } + let module = self.namespace.read().unwrap().get(&module_name).unwrap().clone(); + self.dependencies.write().unwrap().insert(module.name.clone(), module.clone()); + for (call_via, method_name) in imports { + self.load_import(call_via, &module_name, &method_name); + } + if create { + new_modules.push(module.clone()); + } + if self.verbose { + if create { + println!(" (dep {:?})", &module.name); + } else { + println!(" {DIM}(dep {:?}){RESET}", &module.name); + } } } if recurse { @@ -25,38 +43,37 @@ impl Module { module.load(recurse)?; } } - if self.verbose { - self.show_dependencies(); - } Ok(self) } + fn load_import ( + &self, call_via: u32, module_name: &Arc, method_name: &Arc + ) { + if self.imports.read().unwrap().contains_key(&call_via) { + panic!("duplicate address {call_via} from {module_name}"); + } + self.imports.write().unwrap().insert(call_via, ( + module_name.clone(), + method_name.clone(), + )); + } + /// Collect an imported dependency's descriptor thunks into a temporary [Vec]. fn parse_import (pe: &VecPE, descriptor: &ImageImportDescriptor) - -> Usually<(Arc, Vec<(usize, Arc)>)> + -> Usually<(Arc, Vec<(u32, Arc)>)> { - let mut imports = vec![]; - for (index, import) in descriptor.get_imports(pe)?.iter().enumerate() { - imports.push((index, match import { - ImportData::Ordinal(x) => format!("___VESTAL_ORDINAL_{x}"), - ImportData::ImportByName(name) => format!("{name}"), - }.into())); - } Ok(( descriptor.get_name(pe)?.as_str()?.to_lowercase().into(), - imports + descriptor.get_imports(pe)?.iter().enumerate().map(|(index, import)|( + descriptor.first_thunk.0 + index as u32 * 8, + match import { + ImportData::Ordinal(x) => format!("___VESTAL_ORDINAL_{x}"), + ImportData::ImportByName(name) => format!("{name}"), + }.into() + )).collect() )) } - fn load_import ( - &mut self, call_via: u32, module_name: &Arc, method_name: &Arc - ) { - if self.imports.contains_key(&call_via) { - panic!("duplicate address {call_via} from {module_name}"); - } - self.imports.insert(call_via, (module_name.clone(), method_name.clone())); - } - fn load_dependency ( &mut self, module_name: &Arc, diff --git a/crates/vestal/src/main.rs b/crates/vestal/src/main.rs index fb0139e..729ebd7 100644 --- a/crates/vestal/src/main.rs +++ b/crates/vestal/src/main.rs @@ -20,7 +20,7 @@ fn cli () -> clap::Command { /// A DLL that will be linked into the output. pub struct Module { /// Scope. - pub namespace: Arc, Arc>>>, + pub namespace: Arc, Arc>>>, /// Search paths for resolving DLL names. pub search_paths: Arc>>>, /// Canonical name like `xxx.dll` (always lowercase) @@ -40,9 +40,9 @@ pub struct Module { /// Size of `.text` section pub code_size: usize, /// Modules that this one depends on. - pub dependencies: BTreeMap, Arc>, + pub dependencies: RwLock, Arc>>, /// Call targets to methods from imported modules. - pub imports: BTreeMap, Arc)>, + pub imports: RwLock, Arc)>>, /// Addresses of exported methods by name pub exports: RwLock, ThunkData>>, /// Locations in `.text` section that need to be patched diff --git a/crates/vestal/src/show.rs b/crates/vestal/src/show.rs index ee09ad6..0bbb2fc 100644 --- a/crates/vestal/src/show.rs +++ b/crates/vestal/src/show.rs @@ -24,7 +24,7 @@ impl Log { println!("{DIM}(load {:?}){RESET}", path.as_ref()); } } - pub fn dep (show: bool, name: &impl AsRef, exports: &[(usize, Arc)]) { + pub fn dep (show: bool, name: &impl AsRef, exports: &[(u32, Arc)]) { if show { println!("(dep {:?} {})", name.as_ref(), exports.len()); } @@ -62,8 +62,8 @@ impl Show { impl std::fmt::Debug for Module { fn fmt (&self, f: &mut std::fmt::Formatter) -> std::result::Result<(), std::fmt::Error> { - write!(f, "(dll {BOLD}{UNDERLINE}{:15}{RESET} [0x{:>08x}] {} (img 0x{:>08x} -> mem 0x{:>08x}))", - &self.name, + write!(f, "(dll {} [0x{:>08x}] {} (img 0x{:>08x} -> mem 0x{:>08x}))", + format!("{BOLD}{UNDERLINE}{:15}{RESET}", self.name), self.code_size, format!("({:>5})", self.exports.read().unwrap().len()), self.code_start, @@ -72,15 +72,6 @@ impl std::fmt::Debug for Module { } impl Module { - pub fn show_dependencies (&self) { - for (name, module) in self.dependencies.iter() { - print!(" ({module:?}"); - for (method, addr) in module.exports.read().unwrap().iter() { - //print!("\n (0x{addr:08x} {method})") - } - println!(")"); - } - } pub fn show_instruction (&self, addr: usize) { let mut decoder = Decoder::with_ip(64, &self.code[addr..], 0x1000, DecoderOptions::NONE); let instruction = decoder.decode();