add coverage to local commands (#1091)

src/agent/onefuzz-agent/src/local/cmd.rs

@@ -1,13 +1,13 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-#[cfg(any(target_os = "linux", target_os = "windows"))]
-use crate::local::libfuzzer_coverage;
 use crate::local::{
     common::add_common_config, generic_analysis, generic_crash_report, generic_generator,
     libfuzzer, libfuzzer_crash_report, libfuzzer_fuzz, libfuzzer_merge, libfuzzer_regression,
     libfuzzer_test_input, radamsa, test_input, tui::TerminalUi,
 };
+#[cfg(any(target_os = "linux", target_os = "windows"))]
+use crate::local::{coverage, libfuzzer_coverage};
 use anyhow::{Context, Result};
 use clap::{App, Arg, SubCommand};
 use crossterm::tty::IsTty;
@@ -21,6 +21,8 @@ use tokio::{select, time::timeout};
 #[strum(serialize_all = "kebab-case")]
 enum Commands {
     Radamsa,
+    #[cfg(any(target_os = "linux", target_os = "windows"))]
+    Coverage,
     LibfuzzerFuzz,
     LibfuzzerMerge,
     LibfuzzerCrashReport,
@@ -57,6 +59,8 @@ pub async fn run(args: clap::ArgMatches<'static>) -> Result<()> {
     let event_sender = terminal.as_ref().map(|t| t.task_events.clone());
     let command_run = tokio::spawn(async move {
         match command {
+            #[cfg(any(target_os = "linux", target_os = "windows"))]
+            Commands::Coverage => coverage::run(&sub_args, event_sender).await,
             Commands::Radamsa => radamsa::run(&sub_args, event_sender).await,
             Commands::LibfuzzerCrashReport => {
                 libfuzzer_crash_report::run(&sub_args, event_sender).await
@@ -115,6 +119,8 @@ pub fn args(name: &str) -> App<'static, 'static> {
 
     for subcommand in Commands::iter() {
         let app = match subcommand {
+            #[cfg(any(target_os = "linux", target_os = "windows"))]
+            Commands::Coverage => coverage::args(subcommand.into()),
             Commands::Radamsa => radamsa::args(subcommand.into()),
             Commands::LibfuzzerCrashReport => libfuzzer_crash_report::args(subcommand.into()),
             Commands::LibfuzzerFuzz => libfuzzer_fuzz::args(subcommand.into()),

src/agent/onefuzz-agent/src/local/common.rs

@@ -34,6 +34,9 @@ pub const CHECK_FUZZER_HELP: &str = "check_fuzzer_help";
 pub const DISABLE_CHECK_DEBUGGER: &str = "disable_check_debugger";
 pub const REGRESSION_REPORTS_DIR: &str = "regression_reports_dir";
 
+#[cfg(any(target_os = "linux", target_os = "windows"))]
+pub const COVERAGE_FILTER: &str = "coverage_filter";
+
 pub const TARGET_EXE: &str = "target_exe";
 pub const TARGET_ENV: &str = "target_env";
 pub const TARGET_OPTIONS: &str = "target_options";

src/agent/onefuzz-agent/src/local/coverage.rs

@@ -0,0 +1,134 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+use crate::{
+    local::common::{
+        build_local_context, get_cmd_arg, get_cmd_env, get_cmd_exe, get_synced_dir,
+        get_synced_dirs, CmdType, CHECK_FUZZER_HELP, COVERAGE_DIR, COVERAGE_FILTER, INPUTS_DIR,
+        READONLY_INPUTS, TARGET_ENV, TARGET_EXE, TARGET_OPTIONS, TARGET_TIMEOUT,
+    },
+    tasks::{
+        config::CommonConfig,
+        coverage::generic::{Config, CoverageTask},
+    },
+};
+use anyhow::Result;
+use clap::{App, Arg, SubCommand};
+use flume::Sender;
+use storage_queue::QueueClient;
+
+use super::common::{SyncCountDirMonitor, UiEvent};
+
+pub fn build_coverage_config(
+    args: &clap::ArgMatches<'_>,
+    local_job: bool,
+    input_queue: Option<QueueClient>,
+    common: CommonConfig,
+    event_sender: Option<Sender<UiEvent>>,
+) -> Result<Config> {
+    let target_exe = get_cmd_exe(CmdType::Target, args)?.into();
+    let target_env = get_cmd_env(CmdType::Target, args)?;
+    let target_options = get_cmd_arg(CmdType::Target, args);
+    let target_timeout = value_t!(args, TARGET_TIMEOUT, u64).ok();
+    let coverage_filter = value_t!(args, TARGET_TIMEOUT, String).ok();
+
+    let readonly_inputs = if local_job {
+        vec![
+            get_synced_dir(INPUTS_DIR, common.job_id, common.task_id, args)?
+                .monitor_count(&event_sender)?,
+        ]
+    } else {
+        get_synced_dirs(READONLY_INPUTS, common.job_id, common.task_id, args)?
+            .into_iter()
+            .map(|sd| sd.monitor_count(&event_sender))
+            .collect::<Result<Vec<_>>>()?
+    };
+
+    let coverage = get_synced_dir(COVERAGE_DIR, common.job_id, common.task_id, args)?
+        .monitor_count(&event_sender)?;
+
+    let config = Config {
+        target_exe,
+        target_env,
+        target_options,
+        target_timeout,
+        coverage_filter,
+        input_queue,
+        readonly_inputs,
+        coverage,
+        common,
+    };
+
+    Ok(config)
+}
+
+pub async fn run(args: &clap::ArgMatches<'_>, event_sender: Option<Sender<UiEvent>>) -> Result<()> {
+    let context = build_local_context(args, true, event_sender.clone())?;
+    let config = build_coverage_config(
+        args,
+        false,
+        None,
+        context.common_config.clone(),
+        event_sender,
+    )?;
+
+    let mut task = CoverageTask::new(config);
+    task.run().await
+}
+
+pub fn build_shared_args(local_job: bool) -> Vec<Arg<'static, 'static>> {
+    let mut args = vec![
+        Arg::with_name(TARGET_EXE)
+            .long(TARGET_EXE)
+            .takes_value(true)
+            .required(true),
+        Arg::with_name(TARGET_ENV)
+            .long(TARGET_ENV)
+            .takes_value(true)
+            .multiple(true),
+        Arg::with_name(TARGET_OPTIONS)
+            .long(TARGET_OPTIONS)
+            .takes_value(true)
+            .value_delimiter(" ")
+            .help("Use a quoted string with space separation to denote multiple arguments"),
+        Arg::with_name(TARGET_TIMEOUT)
+            .takes_value(true)
+            .long(TARGET_TIMEOUT),
+        Arg::with_name(COVERAGE_FILTER)
+            .takes_value(true)
+            .long(COVERAGE_FILTER),
+        Arg::with_name(TARGET_TIMEOUT)
+            .takes_value(true)
+            .long(TARGET_TIMEOUT),
+        Arg::with_name(COVERAGE_DIR)
+            .takes_value(true)
+            .required(!local_job)
+            .long(COVERAGE_DIR),
+        Arg::with_name(CHECK_FUZZER_HELP)
+            .takes_value(false)
+            .long(CHECK_FUZZER_HELP),
+    ];
+    if local_job {
+        args.push(
+            Arg::with_name(INPUTS_DIR)
+                .long(INPUTS_DIR)
+                .takes_value(true)
+                .required(true),
+        )
+    } else {
+        args.push(
+            Arg::with_name(READONLY_INPUTS)
+                .takes_value(true)
+                .required(true)
+                .long(READONLY_INPUTS)
+                .multiple(true),
+        )
+    }
+    args
+}
+
+pub fn args(name: &'static str) -> App<'static, 'static> {
+    SubCommand::with_name(name)
+        .about("execute a local-only coverage task")
+        .args(&build_shared_args(false))
+}

src/agent/onefuzz-agent/src/local/mod.rs

@@ -3,6 +3,8 @@
 
 pub mod cmd;
 pub mod common;
+#[cfg(any(target_os = "linux", target_os = "windows"))]
+pub mod coverage;
 pub mod generic_analysis;
 pub mod generic_crash_report;
 pub mod generic_generator;

src/agent/onefuzz-agent/src/tasks/coverage/generic.rs

@@ -188,6 +188,7 @@ impl<'a> TaskContext<'a> {
     }
 
     pub async fn record_input(&mut self, input: &Path) -> Result<()> {
+        debug!("recording coverage for {}", input.display());
         let attempts = MAX_COVERAGE_RECORDING_ATTEMPTS;
 
         for attempt in 1..=attempts {