diff --git a/src/agent/coverage/src/block/windows.rs b/src/agent/coverage/src/block/windows.rs index e4c8779be2..519fa4f266 100644 --- a/src/agent/coverage/src/block/windows.rs +++ b/src/agent/coverage/src/block/windows.rs @@ -171,7 +171,7 @@ impl<'c> Recorder<'c> { return Ok(()); } - match self.cache.fetch(&path) { + match self.cache.fetch(&path, dbg.target().process_handle()) { Ok(Some(info)) => { let new = self.coverage.insert(&path, info.blocks.iter().copied()); diff --git a/src/agent/coverage/src/cache.rs b/src/agent/coverage/src/cache.rs index 3b9eca565b..8c8085ca7e 100644 --- a/src/agent/coverage/src/cache.rs +++ b/src/agent/coverage/src/cache.rs @@ -6,6 +6,9 @@ use std::collections::{BTreeSet, HashMap}; use anyhow::Result; use serde::{Deserialize, Serialize}; +#[cfg(target_os = "windows")] +use winapi::um::winnt::HANDLE; + use crate::code::{ModuleIndex, ModulePath}; #[derive(Clone, Debug, Default, Deserialize, Serialize)] @@ -20,6 +23,7 @@ impl ModuleCache { Self { cached } } + #[cfg(target_os = "linux")] pub fn fetch(&mut self, path: &ModulePath) -> Result> { if !self.cached.contains_key(path) { self.insert(path)?; @@ -28,6 +32,19 @@ impl ModuleCache { Ok(self.cached.get(path)) } + #[cfg(target_os = "windows")] + pub fn fetch( + &mut self, + path: &ModulePath, + handle: impl Into>, + ) -> Result> { + if !self.cached.contains_key(path) { + self.insert(path, handle)?; + } + + Ok(self.cached.get(path)) + } + #[cfg(target_os = "linux")] pub fn insert(&mut self, path: &ModulePath) -> Result<()> { let entry = ModuleInfo::new_elf(path)?; @@ -36,8 +53,8 @@ impl ModuleCache { } #[cfg(target_os = "windows")] - pub fn insert(&mut self, path: &ModulePath) -> Result<()> { - let entry = ModuleInfo::new_pe(path)?; + pub fn insert(&mut self, path: &ModulePath, handle: impl Into>) -> Result<()> { + let entry = ModuleInfo::new_pe(path, handle)?; self.cached.insert(path.clone(), entry); Ok(()) } @@ -65,13 +82,13 @@ impl ModuleInfo { } #[cfg(target_os = "windows")] - pub fn new_pe(path: &ModulePath) -> Result { + pub fn new_pe(path: &ModulePath, handle: impl Into>) -> Result { let file = std::fs::File::open(path)?; let data = unsafe { memmap2::Mmap::map(&file)? }; let pe = goblin::pe::PE::parse(&data)?; let module = ModuleIndex::index_pe(path.clone(), &pe); - let offsets = crate::pe::process_module(path, &data, &pe, false)?; + let offsets = crate::pe::process_module(path, &data, &pe, false, handle.into())?; let blocks = offsets.ones().map(|off| off as u64).collect(); Ok(Self { module, blocks }) diff --git a/src/agent/coverage/src/pe.rs b/src/agent/coverage/src/pe.rs index 657e418231..06c4b16be3 100644 --- a/src/agent/coverage/src/pe.rs +++ b/src/agent/coverage/src/pe.rs @@ -233,13 +233,14 @@ pub fn process_module( data: &[u8], pe: &PE, functions_only: bool, + target_handle: Option, ) -> Result { if let Some(DebugData { image_debug_directory: _, codeview_pdb70_debug_info: Some(cv), }) = pe.debug_data { - let pdb_path = find_pdb_path(pe_path.as_ref(), &cv) + let pdb_path = find_pdb_path(pe_path.as_ref(), &cv, target_handle.into()) .with_context(|| format!("searching for PDB for PE: {}", pe_path.as_ref().display()))?; log::info!("found PDB: {}", pdb_path.display()); @@ -302,7 +303,11 @@ fn process_pdb(data: &[u8], pe: &PE, functions_only: bool, pdb_path: &Path) -> R // See: https://docs.microsoft.com/en-us/windows/win32/api/dbghelp/nf-dbghelp-syminitializew const PSEUDO_HANDLE: HANDLE = -2i64 as _; -fn find_pdb_path(pe_path: &Path, cv: &CodeviewPDB70DebugInfo) -> Result { +fn find_pdb_path( + pe_path: &Path, + cv: &CodeviewPDB70DebugInfo, + target_handle: Option, +) -> Result { let cv_filename = CStr::from_bytes_with_nul(cv.filename)?.to_str()?; // This field is named `filename`, but it may be an absolute path. @@ -314,11 +319,22 @@ fn find_pdb_path(pe_path: &Path, cv: &CodeviewPDB70DebugInfo) -> Result return Ok(cv_filename.to_owned()); } - let handle = PSEUDO_HANDLE; + // If we have one, use the the process handle for an existing debug + let handle = target_handle.unwrap_or(PSEUDO_HANDLE); let dbghelp = debugger::dbghelp::lock()?; - dbghelp.sym_initialize(handle)?; - let _cleanup = DbgHelpCleanupGuard::new(&dbghelp, handle); + + // If a target handle was provided, we assume the caller initialized the + // dbghelp symbol handler, and will clean up after itself. + // + // Otherwise, initialize a symbol handler with our own pseudo-path, and use + // a drop guard to ensure we don't leak resources. + let _cleanup = if target_handle.is_some() { + None + } else { + dbghelp.sym_initialize(handle)?; + Some(DbgHelpCleanupGuard::new(&dbghelp, handle)) + }; // Enable signature and age checking. let options = dbghelp.sym_get_options(); @@ -371,10 +387,14 @@ impl<'d> Drop for DbgHelpCleanupGuard<'d> { } } -pub fn process_image(path: impl AsRef, functions_only: bool) -> Result { +pub fn process_image( + path: impl AsRef, + functions_only: bool, + handle: Option, +) -> Result { let file = File::open(path.as_ref())?; let mmap = unsafe { Mmap::map(&file)? }; let pe = PE::parse(&mmap)?; - process_module(path, &mmap, &pe, functions_only) + process_module(path, &mmap, &pe, functions_only, handle) }