Update Windows generic coverage recording (#699)

- Reimplement Windows generic coverage using new coverage format
- Remove old format
- Update and unify examples

src/agent/coverage/examples/block_coverage.rs

@@ -1,46 +1,45 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-use std::path::PathBuf;
+use std::process::Command;
 
 use anyhow::Result;
 use structopt::StructOpt;
 
+#[derive(Debug, PartialEq, StructOpt)]
+struct Opt {
+    #[cfg(target_os = "linux")]
+    #[structopt(short, long)]
+    filter: Option<std::path::PathBuf>,
+
+    #[structopt(min_values = 1)]
+    cmd: Vec<String>,
+}
+
 #[cfg(target_os = "windows")]
 fn main() -> Result<()> {
-    use std::process::Command;
-
     env_logger::init();
 
-    let mut args = std::env::args().skip(1);
+    let opt = Opt::from_args();
-    let exe = args.next().unwrap();
+    log::info!("recording coverage for: {:?}", opt.cmd);
-    let args: Vec<_> = args.collect();
 
-    let mut cmd = Command::new(exe);
+    let mut cmd = Command::new(&opt.cmd[0]);
-    cmd.args(&args);
+    cmd.args(&opt.cmd[1..]);
 
     let coverage = coverage::block::windows::record(cmd)?;
-    let hit = coverage.count_blocks_hit();
-    let found = coverage.count_blocks();
-    let percent = 100.0 * (hit as f64) / (found as f64);
 
-    log::info!("block coverage = {}/{} ({:.2}%)", hit, found, percent);
+    for (module, cov) in coverage.iter() {
+        let total = cov.blocks.len();
+        let hit: u32 = cov.blocks.values().map(|b| b.count).sum();
+        let percent = 100.0 * (hit as f64) / (total as f64);
+        log::info!("module = {}, {} / {} ({:.2}%)", module, hit, total, percent);
+    }
 
     Ok(())
 }
 
-#[derive(Debug, PartialEq, StructOpt)]
-struct Opt {
-    #[structopt(short, long)]
-    filter: Option<PathBuf>,
-
-    cmd: Vec<String>,
-}
-
 #[cfg(target_os = "linux")]
 fn main() -> Result<()> {
-    use std::process::Command;
-
     use coverage::block::linux::Recorder;
     use coverage::code::{CmdFilter, CmdFilterDef};
 

src/agent/coverage/src/block/windows.rs

@@ -7,162 +7,261 @@ use std::time::{Duration, Instant};
 
 use anyhow::Result;
 use debugger::{
-    dbghelp::SymInfo,
     debugger::{BreakpointId, BreakpointType, DebugEventHandler, Debugger},
     target::Module,
 };
 
-use crate::{AppCoverageBlocks, ModuleCoverageBlocks};
+use crate::block::CommandBlockCov;
+use crate::cache::ModuleCache;
+use crate::code::ModulePath;
 
-pub fn record(cmd: Command) -> Result<AppCoverageBlocks> {
+pub fn record(cmd: Command) -> Result<CommandBlockCov> {
-    let mut handler = BlockCoverageHandler::new();
+    let mut cache = ModuleCache::default();
-
+    let recorder = Recorder::new(&mut cache);
-    let (mut dbg, _child) = Debugger::init(cmd, &mut handler)?;
+    let timeout = Duration::from_secs(5);
-    dbg.run(&mut handler)?;
+    let mut handler = RecorderEventHandler::new(recorder, timeout);
-
+    handler.run(cmd)?;
-    Ok(handler.coverage)
+    Ok(handler.recorder.into_coverage())
 }
 
-pub struct BlockCoverageHandler {
+#[derive(Debug)]
-    bp_to_block: BTreeMap<BreakpointId, (usize, usize)>,
+pub struct RecorderEventHandler<'a> {
-    coverage: AppCoverageBlocks,
+    recorder: Recorder<'a>,
     started: Instant,
-    max_duration: Duration,
     timed_out: bool,
+    timeout: Duration,
 }
 
-impl BlockCoverageHandler {
+impl<'a> RecorderEventHandler<'a> {
-    pub fn new() -> Self {
+    pub fn new(recorder: Recorder<'a>, timeout: Duration) -> Self {
-        let coverage = AppCoverageBlocks::new();
-        let bp_to_block = BTreeMap::default();
         let started = Instant::now();
-        let max_duration = Duration::from_secs(5);
         let timed_out = false;
 
         Self {
-            bp_to_block,
+            recorder,
-            coverage,
-            max_duration,
             started,
             timed_out,
+            timeout,
         }
     }
 
-    pub fn pc(&self, dbg: &mut Debugger) -> Result<(u64, Option<SymInfo>)> {
+    pub fn time_out(&self) -> bool {
-        use iced_x86::Register::RIP;
+        self.timed_out
-
-        let pc = dbg.read_register_u64(RIP)?;
-        let sym = dbg.get_symbol(pc).ok();
-
-        Ok((pc, sym))
     }
 
-    fn add_module(&mut self, dbg: &mut Debugger, module: &Module) {
+    pub fn timeout(&self) -> Duration {
-        let bitset = crate::pe::process_image(module.path(), false);
+        self.timeout
-        let bitset = match bitset {
-            Ok(bitset) => bitset,
-            Err(err) => {
-                // If we can't add the module, continue debugging.
-                // We don't expect to have symbols for every module.
-                log::warn!(
-                    "cannot record coverage for module = {}, err = {}",
-                    module.path().display(),
-                    err,
-                );
-                return;
-            }
-        };
-
-        let module_coverage = ModuleCoverageBlocks::new(module.path(), module.name(), bitset);
-
-        let m = self.coverage.add_module(module_coverage);
-        let module_coverage = &self.coverage.modules()[m];
-        for (b, block) in module_coverage.blocks().iter().enumerate() {
-            let bp =
-                dbg.register_breakpoint(module.name(), block.rva() as u64, BreakpointType::OneTime);
-            self.bp_to_block.insert(bp, (m, b));
     }
 
-        log::debug!(
+    pub fn run(&mut self, cmd: Command) -> Result<()> {
-            "inserted {} breakpoints for module {}",
+        let (mut dbg, _child) = Debugger::init(cmd, self)?;
-            module_coverage.blocks().len(),
+        dbg.run(self)?;
-            module.path().display(),
+        Ok(())
-        );
+    }
+
+    fn on_poll(&mut self, dbg: &mut Debugger) {
+        if !self.timed_out && self.started.elapsed() > self.timeout {
+            self.timed_out = true;
+            dbg.quit_debugging();
+        }
     }
 
     fn stop(&self, dbg: &mut Debugger) {
         dbg.quit_debugging();
     }
+}
 
-    fn handle_on_create_process(&mut self, dbg: &mut Debugger, module: &Module) -> Result<()> {
+#[derive(Debug)]
-        dbg.target().sym_initialize()?;
+pub struct Recorder<'a> {
+    breakpoints: Breakpoints,
 
-        log::info!(
+    // Reference to allow in-memory reuse across runs.
-            "exe loaded: {}, {} bytes",
+    cache: &'a mut ModuleCache,
-            module.path().display(),
-            module.image_size(),
-        );
 
-        self.add_module(dbg, module);
+    // Note: this could also be a reference to enable reuse across runs, to
+    // support implicit calculation of total coverage for a corpus. For now,
+    // assume callers will merge this into a separate struct when needed.
+    coverage: CommandBlockCov,
+}
 
-        Ok(())
+impl<'a> Recorder<'a> {
+    pub fn new(cache: &'a mut ModuleCache) -> Self {
+        let breakpoints = Breakpoints::default();
+        let coverage = CommandBlockCov::default();
+
+        Self {
+            breakpoints,
+            cache,
+            coverage,
+        }
     }
 
-    fn handle_on_load_dll(&mut self, dbg: &mut Debugger, module: &Module) {
+    pub fn coverage(&self) -> &CommandBlockCov {
-        log::info!(
+        &self.coverage
-            "dll loaded: {}, {} bytes",
-            module.path().display(),
-            module.image_size(),
-        );
-
-        self.add_module(dbg, module);
     }
 
-    fn handle_on_breakpoint(&mut self, dbg: &mut Debugger, bp: BreakpointId) -> Result<()> {
+    pub fn into_coverage(self) -> CommandBlockCov {
-        let (pc, _sym) = self.pc(dbg)?;
+        self.coverage
+    }
 
-        if let Some(&(m, b)) = self.bp_to_block.get(&bp) {
+    pub fn on_create_process(&mut self, dbg: &mut Debugger, module: &Module) -> Result<()> {
+        log::debug!("process created: {}", module.path().display());
+
+        if let Err(err) = dbg.target().sym_initialize() {
+            log::error!(
+                "unable to initialize symbol handler for new process {}: {:?}",
+                module.path().display(),
+                err,
+            );
+        }
+
+        self.insert_module(dbg, module)
+    }
+
+    pub fn on_load_dll(&mut self, dbg: &mut Debugger, module: &Module) -> Result<()> {
+        log::debug!("DLL loaded: {}", module.path().display());
+
+        self.insert_module(dbg, module)
+    }
+
+    pub fn on_breakpoint(&mut self, dbg: &mut Debugger, id: BreakpointId) -> Result<()> {
+        if let Some(breakpoint) = self.breakpoints.get(id) {
             if log::max_level() == log::Level::Trace {
-                let module = &self.coverage.modules()[m];
+                let name = breakpoint.module.name().to_string_lossy();
-                let block = &module.blocks()[b];
+                let offset = breakpoint.offset;
-                let name = module.name().display();
+                let pc = dbg.read_program_counter()?;
-                log::trace!("{:>16x}: {}+{:x}", pc, name, block.rva());
+
+                if let Ok(sym) = dbg.get_symbol(pc) {
+                    log::trace!(
+                        "{:>16x}: {}+{:x} ({}+{:x})",
+                        pc,
+                        name,
+                        offset,
+                        sym.symbol(),
+                        sym.displacement(),
+                    );
+                } else {
+                    log::trace!("{:>16x}: {}+{:x}", pc, name, offset);
+                }
             }
 
-            self.coverage.report_block_hit(m, b);
+            self.coverage
+                .increment(breakpoint.module, breakpoint.offset)?;
         } else {
-            log::error!("no block for breakpoint");
+            let pc = if let Ok(pc) = dbg.read_program_counter() {
+                format!("{:x}", pc)
+            } else {
+                "???".into()
+            };
+
+            log::error!("hit breakpoint without data, id = {}, pc = {}", id.0, pc);
         }
 
         Ok(())
     }
 
-    fn handle_on_poll(&mut self, dbg: &mut Debugger) {
+    fn insert_module(&mut self, dbg: &mut Debugger, module: &Module) -> Result<()> {
-        if !self.timed_out && self.started.elapsed() > self.max_duration {
+        let path = ModulePath::new(module.path().to_owned())?;
-            self.timed_out = true;
+
-            dbg.quit_debugging();
+        match self.cache.fetch(&path) {
+            Ok(Some(info)) => {
+                let new = self.coverage.insert(&path, info.blocks.iter().copied());
+
+                if !new {
+                    return Ok(());
                 }
+
+                self.breakpoints
+                    .set(dbg, module, info.blocks.iter().copied())?;
+
+                log::debug!("set {} breakpoints for module {}", info.blocks.len(), path);
+            }
+            Ok(None) => {
+                log::warn!("could not find module: {}", path);
+            }
+            Err(err) => {
+                log::warn!("could not disassemble module {}: {:?}", path, err);
+            }
+        }
+
+        Ok(())
     }
 }
 
-impl DebugEventHandler for BlockCoverageHandler {
+impl<'a> DebugEventHandler for RecorderEventHandler<'a> {
     fn on_create_process(&mut self, dbg: &mut Debugger, module: &Module) {
-        if self.handle_on_create_process(dbg, module).is_err() {
+        if self.recorder.on_create_process(dbg, module).is_err() {
             self.stop(dbg);
         }
     }
 
     fn on_load_dll(&mut self, dbg: &mut Debugger, module: &Module) {
-        self.handle_on_load_dll(dbg, module);
+        if self.recorder.on_load_dll(dbg, module).is_err() {
+            self.stop(dbg);
+        }
     }
 
     fn on_breakpoint(&mut self, dbg: &mut Debugger, bp: BreakpointId) {
-        if self.handle_on_breakpoint(dbg, bp).is_err() {
+        if self.recorder.on_breakpoint(dbg, bp).is_err() {
             self.stop(dbg);
         }
     }
 
     fn on_poll(&mut self, dbg: &mut Debugger) {
-        self.handle_on_poll(dbg);
+        self.on_poll(dbg);
     }
 }
+
+/// Relates opaque, runtime-generated breakpoint IDs to their corresponding
+/// location, via module and offset.
+#[derive(Clone, Debug, Default)]
+struct Breakpoints {
+    // Breakpoint-associated module paths, referenced by index to save space and
+    // avoid copying.
+    modules: Vec<ModulePath>,
+
+    // Map of breakpoint IDs to data which pick out an code location. For a
+    // value `(module, offset)`, `module` is an index into `self.modules`, and
+    // `offset` is a VA offset relative to the module base.
+    registered: BTreeMap<BreakpointId, (usize, u64)>,
+}
+
+impl Breakpoints {
+    pub fn get(&self, id: BreakpointId) -> Option<BreakpointData<'_>> {
+        let (module_index, offset) = self.registered.get(&id).copied()?;
+        let module = self.modules.get(module_index)?;
+        Some(BreakpointData { module, offset })
+    }
+
+    pub fn set(
+        &mut self,
+        dbg: &mut Debugger,
+        module: &Module,
+        offsets: impl Iterator<Item = u64>,
+    ) -> Result<()> {
+        // From the `target::Module`, create and save a `ModulePath`.
+        let module_path = ModulePath::new(module.path().to_owned())?;
+        let module_index = self.modules.len();
+        self.modules.push(module_path);
+
+        for offset in offsets {
+            // Register the breakpoint in the running target address space.
+            let id = dbg.register_breakpoint(module.name(), offset, BreakpointType::OneTime);
+
+            // Associate the opaque `BreakpointId` with the module and offset.
+            self.registered.insert(id, (module_index, offset));
+        }
+
+        log::debug!("{} total registered modules", self.modules.len());
+        log::debug!("{} total registered breakpoints", self.registered.len());
+
+        Ok(())
+    }
+}
+
+/// Code location data associated with an opaque breakpoint ID.
+#[derive(Clone, Copy, Debug)]
+pub struct BreakpointData<'a> {
+    pub module: &'a ModulePath,
+    pub offset: u64,
+}

src/agent/coverage/src/cache.rs

@@ -6,7 +6,6 @@ use serde::{Deserialize, Serialize};
 use std::collections::{BTreeSet, HashMap};
 
 use crate::code::{ModuleIndex, ModulePath};
-use crate::disasm::ModuleDisassembler;
 
 #[derive(Clone, Debug, Default, Deserialize, Serialize)]
 pub struct ModuleCache {
@@ -28,11 +27,19 @@ impl ModuleCache {
         Ok(self.cached.get(path))
     }
 
+    #[cfg(target_os = "linux")]
     pub fn insert(&mut self, path: &ModulePath) -> Result<()> {
         let entry = ModuleInfo::new_elf(path)?;
         self.cached.insert(path.clone(), entry);
         Ok(())
     }
+
+    #[cfg(target_os = "windows")]
+    pub fn insert(&mut self, path: &ModulePath) -> Result<()> {
+        let entry = ModuleInfo::new_pe(path)?;
+        self.cached.insert(path.clone(), entry);
+        Ok(())
+    }
 }
 
 #[derive(Clone, Debug, Deserialize, Serialize)]
@@ -48,10 +55,24 @@ impl ModuleInfo {
     #[cfg(target_os = "linux")]
     pub fn new_elf(path: &ModulePath) -> Result<Self> {
         let data = std::fs::read(path)?;
-        let module = ModuleIndex::parse_elf(path.clone(), &data)?;
+        let elf = goblin::elf::Elf::parse(&data)?;
-        let disasm = ModuleDisassembler::new(&module, &data)?;
+        let module = ModuleIndex::index_elf(path.clone(), &elf)?;
+        let disasm = crate::disasm::ModuleDisassembler::new(&module, &data)?;
         let blocks = disasm.find_blocks();
 
         Ok(Self { module, blocks })
     }
+
+    #[cfg(target_os = "windows")]
+    pub fn new_pe(path: &ModulePath) -> Result<Self> {
+        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 blocks = offsets.ones().map(|off| off as u64).collect();
+
+        Ok(Self { module, blocks })
+    }
 }

src/agent/coverage/src/code.rs

@@ -88,25 +88,30 @@ impl From<ModulePath> for PathBuf {
     }
 }
 
+/// Index over an executable module and its symbols.
 #[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
 pub struct ModuleIndex {
+    /// Absolute path to the module's backing file.
     pub path: ModulePath,
+
+    /// Preferred virtual address of the module's base image.
     pub base_va: u64,
+
+    /// Index over the module's symbols.
     pub symbols: SymbolIndex,
 }
 
 impl ModuleIndex {
+    /// Build a new index over a parsed ELF module.
     #[cfg(target_os = "linux")]
-    pub fn parse_elf(path: ModulePath, data: &[u8]) -> Result<Self> {
+    pub fn index_elf(path: ModulePath, elf: &goblin::elf::Elf) -> Result<Self> {
         use anyhow::format_err;
-        use goblin::elf::{self, program_header::PT_LOAD};
+        use goblin::elf::program_header::PT_LOAD;
-
-        let object = elf::Elf::parse(data)?;
 
         // Calculate the module base address as the lowest preferred VA of any loadable segment.
         //
         // https://refspecs.linuxbase.org/elf/gabi4+/ch5.pheader.html#base_address
-        let base_va = object
+        let base_va = elf
             .program_headers
             .iter()
             .filter(|h| h.p_type == PT_LOAD)
@@ -116,14 +121,14 @@ impl ModuleIndex {
 
         let mut symbols = SymbolIndex::default();
 
-        for sym in object.syms.iter() {
+        for sym in elf.syms.iter() {
             if sym.st_size == 0 {
                 log::debug!("skipping size 0 symbol: {:x?}", sym);
                 continue;
             }
 
             if sym.is_function() {
-                let name = match object.strtab.get(sym.st_name) {
+                let name = match elf.strtab.get(sym.st_name) {
                     None => {
                         log::error!("symbol not found in symbol string table: {:?}", sym);
                         continue;
@@ -153,7 +158,7 @@ impl ModuleIndex {
                 // A symbol is defined relative to some section, identified by `st_shndx`, an index
                 // into the section header table. We'll use the section header to compute the file
                 // offset of the symbol.
-                let section = object
+                let section = elf
                     .section_headers
                     .get(sym.st_shndx)
                     .cloned()
@@ -195,6 +200,21 @@ impl ModuleIndex {
             symbols,
         })
     }
+
+    /// Build a new index over a parsed PE module.
+    #[cfg(target_os = "windows")]
+    pub fn index_pe(path: ModulePath, pe: &goblin::pe::PE) -> Self {
+        let base_va = pe.image_base as u64;
+
+        // Not yet implemented.
+        let symbols = SymbolIndex::default();
+
+        Self {
+            path,
+            base_va,
+            symbols,
+        }
+    }
 }
 
 #[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]

src/agent/coverage/src/lib.rs

@@ -10,10 +10,7 @@ mod intel;
 pub mod pe;
 
 pub mod block;
-
-#[cfg(target_os = "linux")]
 pub mod cache;
-
 pub mod code;
 pub mod demangle;
 
@@ -22,157 +19,3 @@ pub mod disasm;
 
 pub mod filter;
 mod region;
-
-use std::{
-    ffi::OsString,
-    fs::File,
-    path::{Path, PathBuf},
-};
-
-use anyhow::Result;
-use fixedbitset::FixedBitSet;
-use serde::{Deserialize, Serialize};
-
-pub const COVERAGE_MAP: &str = "coverage-map";
-
-#[derive(Deserialize, Serialize, Debug, Clone)]
-pub struct Block {
-    rva: u32,
-    hit: bool,
-}
-
-impl Block {
-    pub fn new(rva: u32, hit: bool) -> Self {
-        Block { rva, hit }
-    }
-
-    pub fn rva(&self) -> u32 {
-        self.rva
-    }
-
-    pub fn hit(&self) -> bool {
-        self.hit
-    }
-
-    pub fn set_hit(&mut self) {
-        self.hit = true;
-    }
-}
-
-#[derive(Debug, Deserialize, Serialize, Clone)]
-pub struct ModuleCoverageBlocks {
-    module: OsString,
-    path: PathBuf,
-    blocks: Vec<Block>,
-}
-
-impl ModuleCoverageBlocks {
-    pub fn new(
-        path: impl Into<PathBuf>,
-        module: impl Into<OsString>,
-        rvas_bitset: FixedBitSet,
-    ) -> Self {
-        let blocks: Vec<_> = rvas_bitset
-            .ones()
-            .map(|rva| Block::new(rva as u32, false))
-            .collect();
-
-        ModuleCoverageBlocks {
-            path: path.into(),
-            module: module.into(),
-            blocks,
-        }
-    }
-
-    pub fn path(&self) -> &Path {
-        &self.path
-    }
-
-    pub fn name(&self) -> &Path {
-        Path::new(&self.module)
-    }
-
-    pub fn blocks(&self) -> &[Block] {
-        &self.blocks
-    }
-
-    pub fn set_block_hit(&mut self, block_index: usize) {
-        self.blocks[block_index].set_hit();
-    }
-
-    pub fn count_blocks_hit(&self) -> usize {
-        self.blocks.iter().filter(|b| b.hit).count()
-    }
-}
-
-#[derive(Debug, Deserialize, Serialize, Clone)]
-pub struct AppCoverageBlocks {
-    modules: Vec<ModuleCoverageBlocks>,
-}
-
-impl AppCoverageBlocks {
-    pub fn new() -> Self {
-        let modules = vec![];
-        Self { modules }
-    }
-
-    pub fn modules(&self) -> &[ModuleCoverageBlocks] {
-        &self.modules
-    }
-
-    pub fn add_module(&mut self, module: ModuleCoverageBlocks) -> usize {
-        let idx = self.modules.len();
-        self.modules.push(module);
-        idx
-    }
-
-    pub fn report_block_hit(&mut self, module_index: usize, block_index: usize) {
-        self.modules[module_index].set_block_hit(block_index);
-    }
-
-    pub fn save(&self, path: impl AsRef<Path>) -> Result<()> {
-        let cov_file = File::create(path)?;
-        bincode::serialize_into(&cov_file, self)?;
-        Ok(())
-    }
-
-    pub fn count_blocks(&self) -> usize {
-        self.modules().iter().map(|m| m.blocks().len()).sum()
-    }
-
-    pub fn count_blocks_hit(&self) -> usize {
-        self.modules().iter().map(|m| m.count_blocks_hit()).sum()
-    }
-}
-
-/// Statically analyze the specified images to discover the basic block
-/// entry points and write out the results in a file in `output_dir`.
-#[cfg(target_os = "windows")]
-pub fn run_init(output_dir: PathBuf, modules: Vec<PathBuf>, function: bool) -> Result<()> {
-    let mut result = AppCoverageBlocks::new();
-    for module in modules {
-        if module.is_file() {
-            let rvas_bitset = pe::process_image(&module, function)?;
-
-            let module_name = module.file_stem().unwrap(); // Unwrap guaranteed by `is_file` test above.
-            let module_rvas = ModuleCoverageBlocks::new(module.clone(), &module_name, rvas_bitset);
-            result.add_module(module_rvas);
-        } else {
-            anyhow::bail!("Cannot find file `{}`", module.as_path().display());
-        }
-    }
-
-    let output_file = output_dir.join(COVERAGE_MAP);
-    result.save(&output_file)?;
-
-    Ok(())
-}
-
-/// Load a coverage map created by `run_init`.
-pub fn load_coverage_map(output_dir: &Path) -> Result<Option<AppCoverageBlocks>> {
-    if let Ok(cov_file) = File::open(output_dir.join(COVERAGE_MAP)) {
-        Ok(Some(bincode::deserialize_from(cov_file)?))
-    } else {
-        Ok(None)
-    }
-}

src/agent/coverage/src/pe.rs

@@ -152,7 +152,7 @@ fn find_blocks(
     blocks: &mut FixedBitSet,
     address_map: &AddressMap,
     pe: &PE,
-    mmap: &Mmap,
+    data: &[u8],
     functions_only: bool,
 ) -> Result<()> {
     let file_alignment = pe
@@ -217,7 +217,7 @@ fn find_blocks(
 
                 intel::find_blocks(
                     bitness,
-                    &mmap[file_offset..file_offset + (code_len as usize)],
+                    &data[file_offset..file_offset + (code_len as usize)],
                     rva.0,
                     blocks,
                 );
@@ -228,12 +228,12 @@ fn find_blocks(
     Ok(())
 }
 
-pub fn process_image<P: AsRef<Path>>(image_path: P, functions_only: bool) -> Result<FixedBitSet> {
+pub fn process_module(
-    let file = File::open(image_path.as_ref())?;
+    pe_path: impl AsRef<Path>,
-    let mmap = unsafe { Mmap::map(&file)? };
+    data: &[u8],
-
+    pe: &PE,
-    let pe = PE::parse(mmap.as_ref())?;
+    functions_only: bool,
-
+) -> Result<FixedBitSet> {
     if let Some(DebugData {
         image_debug_directory: _,
         codeview_pdb70_debug_info: Some(cv),
@@ -248,13 +248,13 @@ pub fn process_image<P: AsRef<Path>>(image_path: P, functions_only: bool) -> Res
             anyhow::bail!(
                 "pdb `{}` doesn't match image `{}`",
                 pdb_path,
-                image_path.as_ref().display()
+                pe_path.as_ref().display()
             );
         }
 
         let address_map = pdb.address_map()?;
 
-        let mut blocks = FixedBitSet::with_capacity(mmap.len());
+        let mut blocks = FixedBitSet::with_capacity(data.len());
 
         let proc_sym_info = collect_proc_symbols(&mut pdb.global_symbols()?.iter())?;
         find_blocks(
@@ -262,7 +262,7 @@ pub fn process_image<P: AsRef<Path>>(image_path: P, functions_only: bool) -> Res
             &mut blocks,
             &address_map,
             &pe,
-            &mmap,
+            data,
             functions_only,
         )?;
 
@@ -277,7 +277,7 @@ pub fn process_image<P: AsRef<Path>>(image_path: P, functions_only: bool) -> Res
                     &mut blocks,
                     &address_map,
                     &pe,
-                    &mmap,
+                    data,
                     functions_only,
                 )?;
             }
@@ -288,3 +288,11 @@ pub fn process_image<P: AsRef<Path>>(image_path: P, functions_only: bool) -> Res
 
     anyhow::bail!("PE missing codeview pdb debug info")
 }
+
+pub fn process_image(path: impl AsRef<Path>, functions_only: bool) -> Result<FixedBitSet> {
+    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)
+}

src/agent/debugger/src/debugger.rs

@@ -678,6 +678,10 @@ impl Debugger {
         self.target.read_register_u64(reg)
     }
 
+    pub fn read_program_counter(&mut self) -> Result<u64> {
+        self.target.read_program_counter()
+    }
+
     pub fn read_flags_register(&mut self) -> Result<u32> {
         self.target.read_flags_register()
     }

src/agent/debugger/src/target.rs

@@ -595,6 +595,10 @@ impl Target {
         Ok(current_context.get_register_u64(reg))
     }
 
+    pub fn read_program_counter(&mut self) -> Result<u64> {
+        self.read_register_u64(iced_x86::Register::RIP)
+    }
+
     pub fn read_flags_register(&mut self) -> Result<u32> {
         let current_context = self.get_current_context()?;
         Ok(current_context.get_flags())

src/agent/input-tester/src/crash_detector.rs

@@ -14,12 +14,11 @@ use std::{
 };
 
 use anyhow::Result;
-use coverage::AppCoverageBlocks;
+use coverage::{block::windows::Recorder as BlockCoverageRecorder, cache::ModuleCache};
 use debugger::{
-    debugger::{BreakpointId, BreakpointType, DebugEventHandler, Debugger},
+    debugger::{BreakpointId, DebugEventHandler, Debugger},
     target::Module,
 };
-use fnv::FnvHashMap;
 use log::{debug, error, trace};
 use win_util::{
     pipe_handle::{pipe, PipeReaderNonBlocking},
@@ -135,18 +134,17 @@ struct CrashDetectorEventHandler<'a> {
     stderr_buffer: Vec<u8>,
     debugger_output: String,
     exceptions: Vec<Exception>,
-    id_to_block: Option<FnvHashMap<BreakpointId, (usize, usize)>>,
+    coverage: Option<BlockCoverageRecorder<'a>>,
-    coverage_map: Option<&'a mut AppCoverageBlocks>,
 }
 
 impl<'a> CrashDetectorEventHandler<'a> {
     pub fn new(
-        coverage_map: Option<&'a mut AppCoverageBlocks>,
         stdout: PipeReaderNonBlocking,
         stderr: PipeReaderNonBlocking,
         ignore_first_chance_exceptions: bool,
         start_time: Instant,
         max_duration: Duration,
+        coverage: Option<BlockCoverageRecorder<'a>>,
     ) -> Self {
         Self {
             start_time,
@@ -160,48 +158,11 @@ impl<'a> CrashDetectorEventHandler<'a> {
             stderr_buffer: vec![],
             debugger_output: String::new(),
             exceptions: vec![],
-            id_to_block: None,
+            coverage,
-            coverage_map,
         }
     }
 }
 
-/// Register the coverage breakpoints and build a mapping from `BreakpointId` to the pair
-/// (module index, block index).
-///
-/// The debugger passes a `BreakpointId` to the `on_breakpoint` callback - the same id returned
-/// when the breakpoint is registered.
-///
-/// This mapping is convenient so that we do not need to track module load addresses and later
-/// map between breakpoint address and module/rva pairs.
-fn prepare_coverage_breakpoints(
-    debugger: &mut Debugger,
-    coverage_map: &Option<&mut AppCoverageBlocks>,
-) -> Option<FnvHashMap<BreakpointId, (usize, usize)>> {
-    if let Some(coverage_map) = coverage_map {
-        let mut id_to_block = fnv::FnvHashMap::default();
-        for (m, module) in coverage_map.modules().iter().enumerate() {
-            let name = module.name();
-            for (i, block) in module.blocks().iter().enumerate() {
-                // For better performance, we can skip registering breakpoints that have
-                // been hit as we only care about new coverage.
-                if !block.hit() {
-                    let id = debugger.register_breakpoint(
-                        name,
-                        block.rva() as u64,
-                        BreakpointType::OneTime,
-                    );
-
-                    id_to_block.insert(id, (m, i));
-                }
-            }
-        }
-        Some(id_to_block)
-    } else {
-        None
-    }
-}
-
 fn is_vcpp_notification(exception: &EXCEPTION_DEBUG_INFO, target_process_handle: HANDLE) -> bool {
     if exception.ExceptionRecord.ExceptionCode == vcpp_debugger::EXCEPTION_VISUALCPP_DEBUGGER {
         match VcppDebuggerExceptionInfo::from_exception_record(
@@ -281,10 +242,6 @@ impl<'a> DebugEventHandler for CrashDetectorEventHandler<'a> {
         DBG_EXCEPTION_NOT_HANDLED
     }
 
-    fn on_create_process(&mut self, debugger: &mut Debugger, _module: &Module) {
-        prepare_coverage_breakpoints(debugger, &self.coverage_map);
-    }
-
     fn on_output_debug_string(&mut self, _debugger: &mut Debugger, message: String) {
         self.debugger_output.push_str(&message);
     }
@@ -313,26 +270,43 @@ impl<'a> DebugEventHandler for CrashDetectorEventHandler<'a> {
         }
     }
 
-    fn on_breakpoint(&mut self, _debugger: &mut Debugger, id: BreakpointId) {
+    fn on_create_process(&mut self, dbg: &mut Debugger, module: &Module) {
-        if let Some(id_to_block) = &self.id_to_block {
+        if let Some(coverage) = &mut self.coverage {
-            if let Some(&(mod_idx, block_idx)) = id_to_block.get(&id) {
+            if let Err(err) = coverage.on_create_process(dbg, module) {
-                if let Some(coverage_map) = &mut self.coverage_map {
+                error!("error recording coverage on create process: {:?}", err);
-                    coverage_map.report_block_hit(mod_idx, block_idx);
+                dbg.quit_debugging();
             }
         }
     }
+
+    fn on_load_dll(&mut self, dbg: &mut Debugger, module: &Module) {
+        if let Some(coverage) = &mut self.coverage {
+            if let Err(err) = coverage.on_load_dll(dbg, module) {
+                error!("error recording coverage on load DLL: {:?}", err);
+                dbg.quit_debugging();
+            }
+        }
+    }
+
+    fn on_breakpoint(&mut self, dbg: &mut Debugger, bp: BreakpointId) {
+        if let Some(coverage) = &mut self.coverage {
+            if let Err(err) = coverage.on_breakpoint(dbg, bp) {
+                error!("error recording coverage on breakpoint: {:?}", err);
+                dbg.quit_debugging();
+            }
+        }
     }
 }
 
 /// This function runs the application under a debugger to detect any crashes in
 /// the process or any children processes.
-pub fn test_process(
+pub fn test_process<'a>(
     app_path: impl AsRef<OsStr>,
     args: &[impl AsRef<OsStr>],
     env: &HashMap<String, String>,
     max_duration: Duration,
     ignore_first_chance_exceptions: bool,
-    coverage_map: Option<&mut AppCoverageBlocks>,
+    cache: Option<&'a mut ModuleCache>,
 ) -> Result<DebuggerResult> {
     debug!("Running: {}", logging::command_invocation(&app_path, args));
 
@@ -350,14 +324,15 @@ pub fn test_process(
         command.env(k, v);
     }
 
+    let recorder = cache.map(BlockCoverageRecorder::new);
     let start_time = Instant::now();
     let mut event_handler = CrashDetectorEventHandler::new(
-        coverage_map,
         stdout_reader,
         stderr_reader,
         ignore_first_chance_exceptions,
         start_time,
         max_duration,
+        recorder,
     );
     let (mut debugger, mut child) = Debugger::init(command, &mut event_handler)?;
     debugger.run(&mut event_handler)?;

src/agent/input-tester/src/tester.rs

@@ -20,7 +20,6 @@ use std::{
 };
 
 use anyhow::{Context, Result};
-use coverage::AppCoverageBlocks;
 use log::{error, info, trace, warn};
 use num_cpus;
 use rayon::{prelude::*, ThreadPoolBuilder};
@@ -53,19 +52,13 @@ const MAX_CRASH_SAMPLES: usize = 10;
 pub struct InputTestResult {
     pub debugger_result: DebuggerResult,
     pub input_path: PathBuf,
-    pub blocks_covered: Option<usize>,
 }
 
 impl InputTestResult {
-    pub fn new(
+    pub fn new(debugger_result: DebuggerResult, input_path: PathBuf) -> Self {
-        debugger_result: DebuggerResult,
-        input_path: PathBuf,
-        blocks_covered: Option<usize>,
-    ) -> Self {
         InputTestResult {
             debugger_result,
             input_path,
-            blocks_covered,
         }
     }
 }
@@ -78,7 +71,6 @@ pub struct Tester {
     ignore_first_chance_exceptions: bool,
     appverif_controller: Option<AppVerifierController>,
     bugs_found_dir: PathBuf,
-    coverage_map: Option<AppCoverageBlocks>,
 }
 
 impl Tester {
@@ -116,7 +108,6 @@ impl Tester {
             None
         };
 
-        let coverage_map = coverage::load_coverage_map(&output_dir)?;
         Ok(Arc::new(Tester {
             appverif_controller,
             driver,
@@ -125,7 +116,6 @@ impl Tester {
             max_run_s,
             ignore_first_chance_exceptions,
             bugs_found_dir,
-            coverage_map,
         }))
     }
 
@@ -133,20 +123,16 @@ impl Tester {
     pub fn test_application(&self, input_path: impl AsRef<Path>) -> Result<InputTestResult> {
         let app_args = args_with_input_file_applied(&self.driver_args, &input_path)?;
 
-        let mut coverage_map = self.coverage_map.clone();
-
         crash_detector::test_process(
             &self.driver,
             &app_args,
             &self.driver_env,
             Duration::from_secs(self.max_run_s),
             self.ignore_first_chance_exceptions,
-            coverage_map.as_mut(),
+            None,
         )
         .and_then(|result| {
-            let blocks_hit = coverage_map.map_or(None, |map| Some(map.count_blocks_hit()));
+            let result = InputTestResult::new(result, PathBuf::from(input_path.as_ref()));
-            let result =
-                InputTestResult::new(result, PathBuf::from(input_path.as_ref()), blocks_hit);
             log_input_test_result(&result);
             Ok(result)
         })