Capture target output, allow Loader reuse (#2716)

src/agent/coverage/examples/coverage.rs

@@ -5,6 +5,8 @@ use anyhow::Result;
 use clap::Parser;
 use coverage::allowlist::{AllowList, TargetAllowList};
 use coverage::binary::BinaryCoverage;
+use coverage::record::CoverageRecorder;
+use debuggable_module::loader::Loader;
 
 #[derive(Parser, Debug)]
 struct Args {
@@ -45,14 +47,19 @@ fn main() -> Result<()> {
         allowlist.source_files = AllowList::load(path)?;
     }
 
-    let coverage = coverage::record::record(cmd, timeout, allowlist)?;
+    let loader = Loader::new();
+    let recorded = CoverageRecorder::new(cmd)
+        .allowlist(allowlist)
+        .loader(loader)
+        .timeout(timeout)
+        .record()?;
 
-    dump_modoff(coverage)?;
+    dump_modoff(&recorded.coverage)?;
 
     Ok(())
 }
 
-fn dump_modoff(coverage: BinaryCoverage) -> Result<()> {
+fn dump_modoff(coverage: &BinaryCoverage) -> Result<()> {
     for (module, coverage) in &coverage.modules {
         for (offset, count) in coverage.as_ref() {
             if count.reached() {

src/agent/coverage/src/lib.rs

@@ -14,4 +14,4 @@ mod timer;
 pub use allowlist::{AllowList, TargetAllowList};
 
 #[doc(inline)]
-pub use record::record;
+pub use record::{CoverageRecorder, Recorded};

src/agent/coverage/src/record.rs

@@ -1,14 +1,100 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+use std::process::{Command, Output, Stdio};
+use std::sync::Arc;
+use std::time::Duration;
+
+use anyhow::Result;
+use debuggable_module::loader::Loader;
+
+use crate::allowlist::TargetAllowList;
+use crate::binary::BinaryCoverage;
+
 #[cfg(target_os = "linux")]
 pub mod linux;
 
 #[cfg(target_os = "windows")]
 pub mod windows;
 
-#[cfg(target_os = "linux")]
-pub use crate::record::linux::record;
+pub struct CoverageRecorder {
+    allowlist: TargetAllowList,
+    cmd: Command,
+    loader: Arc<Loader>,
+    timeout: Duration,
+}
 
-#[cfg(target_os = "windows")]
-pub use crate::record::windows::record;
+impl CoverageRecorder {
+    pub fn new(mut cmd: Command) -> Self {
+        cmd.stdout(Stdio::piped());
+        cmd.stderr(Stdio::piped());
+
+        let allowlist = TargetAllowList::default();
+        let loader = Arc::new(Loader::new());
+        let timeout = Duration::from_secs(5);
+
+        Self {
+            allowlist,
+            cmd,
+            loader,
+            timeout,
+        }
+    }
+
+    pub fn allowlist(mut self, allowlist: TargetAllowList) -> Self {
+        self.allowlist = allowlist;
+        self
+    }
+
+    pub fn loader(mut self, loader: impl Into<Arc<Loader>>) -> Self {
+        self.loader = loader.into();
+        self
+    }
+
+    pub fn timeout(mut self, timeout: Duration) -> Self {
+        self.timeout = timeout;
+        self
+    }
+
+    #[cfg(target_os = "linux")]
+    pub fn record(self) -> Result<Recorded> {
+        use linux::debugger::Debugger;
+        use linux::LinuxRecorder;
+
+        let loader = self.loader.clone();
+
+        crate::timer::timed(self.timeout, move || {
+            let mut recorder = LinuxRecorder::new(&loader, self.allowlist);
+            let dbg = Debugger::new(&mut recorder);
+            let output = dbg.run(self.cmd)?;
+            let coverage = recorder.coverage;
+
+            Ok(Recorded { coverage, output })
+        })?
+    }
+
+    #[cfg(target_os = "windows")]
+    pub fn record(self) -> Result<Recorded> {
+        use debugger::Debugger;
+        use windows::WindowsRecorder;
+
+        let loader = self.loader.clone();
+
+        crate::timer::timed(self.timeout, move || {
+            let mut recorder = WindowsRecorder::new(&loader, self.allowlist);
+            let (mut dbg, child) = Debugger::init(self.cmd, &mut recorder)?;
+            dbg.run(&mut recorder)?;
+
+            let output = child.wait_with_output()?;
+            let coverage = recorder.coverage;
+
+            Ok(Recorded { coverage, output })
+        })?
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct Recorded {
+    pub coverage: BinaryCoverage,
+    pub output: Output,
+}

src/agent/coverage/src/record/linux.rs

@@ -2,8 +2,6 @@
 // Licensed under the MIT License.
 
 use std::collections::BTreeMap;
-use std::process::Command;
-use std::time::Duration;
 
 use anyhow::{bail, Result};
 use debuggable_module::linux::LinuxModule;
@@ -14,31 +12,14 @@ use debuggable_module::Address;
 use pete::Tracee;
 
 pub mod debugger;
-use debugger::{DebugEventHandler, Debugger, DebuggerContext, ModuleImage};
+use debugger::{DebugEventHandler, DebuggerContext, ModuleImage};
 
 use crate::allowlist::TargetAllowList;
 use crate::binary::{self, BinaryCoverage};
 
-pub fn record(
-    cmd: Command,
-    timeout: Duration,
-    allowlist: impl Into<Option<TargetAllowList>>,
-) -> Result<BinaryCoverage> {
-    let loader = Loader::new();
-    let allowlist = allowlist.into().unwrap_or_default();
-
-    crate::timer::timed(timeout, move || {
-        let mut recorder = LinuxRecorder::new(&loader, allowlist);
-        let dbg = Debugger::new(&mut recorder);
-        dbg.run(cmd)?;
-
-        Ok(recorder.coverage)
-    })?
-}
-
 pub struct LinuxRecorder<'data> {
     allowlist: TargetAllowList,
-    coverage: BinaryCoverage,
+    pub coverage: BinaryCoverage,
     loader: &'data Loader,
     modules: BTreeMap<FilePath, LinuxModule<'data>>,
 }

src/agent/coverage/src/record/linux/debugger.rs

@@ -2,7 +2,7 @@
 // Licensed under the MIT License.
 
 use std::collections::BTreeMap;
-use std::process::Command;
+use std::process::{Command, Output};
 
 use anyhow::{bail, format_err, Result};
 use debuggable_module::path::FilePath;
@@ -36,7 +36,7 @@ impl<'eh> Debugger<'eh> {
         }
     }
 
-    pub fn run(mut self, cmd: Command) -> Result<()> {
+    pub fn run(mut self, cmd: Command) -> Result<Output> {
         let mut child = self.context.tracer.spawn(cmd)?;
 
         if let Err(err) = self.wait_on_stops() {
@@ -46,7 +46,9 @@ impl<'eh> Debugger<'eh> {
             return Err(err);
         }
 
-        Ok(())
+        let output = child.wait_with_output()?;
+
+        Ok(output)
     }
 
     fn wait_on_stops(mut self) -> Result<()> {

src/agent/coverage/src/record/windows.rs

@@ -3,8 +3,6 @@
 
 use std::collections::BTreeMap;
 use std::path::Path;
-use std::process::Command;
-use std::time::Duration;
 
 use anyhow::{anyhow, Result};
 use debuggable_module::load_module::LoadModule;
@@ -17,27 +15,10 @@ use debugger::{BreakpointId, BreakpointType, DebugEventHandler, Debugger, Module
 use crate::allowlist::TargetAllowList;
 use crate::binary::{self, BinaryCoverage};
 
-pub fn record(
-    cmd: Command,
-    timeout: Duration,
-    allowlist: impl Into<Option<TargetAllowList>>,
-) -> Result<BinaryCoverage> {
-    let loader = Loader::new();
-    let allowlist = allowlist.into().unwrap_or_default();
-
-    crate::timer::timed(timeout, move || {
-        let mut recorder = WindowsRecorder::new(&loader, allowlist);
-        let (mut dbg, _child) = Debugger::init(cmd, &mut recorder)?;
-        dbg.run(&mut recorder)?;
-
-        Ok(recorder.coverage)
-    })?
-}
-
 pub struct WindowsRecorder<'data> {
     allowlist: TargetAllowList,
     breakpoints: Breakpoints,
-    coverage: BinaryCoverage,
+    pub coverage: BinaryCoverage,
     loader: &'data Loader,
     modules: BTreeMap<FilePath, WindowsModule<'data>>,
 }