Wiring up the UI to the events

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

@@ -39,18 +39,20 @@ 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 args.subcommand() {
-            (RADAMSA, Some(sub)) => radamsa::run(sub).await,
+            (RADAMSA, Some(sub)) => radamsa::run(sub, event_sender).await,
             (LIBFUZZER, Some(sub)) => libfuzzer::run(sub, event_sender).await,
-            (LIBFUZZER_FUZZ, Some(sub)) => libfuzzer_fuzz::run(sub).await,
-            (LIBFUZZER_COVERAGE, Some(sub)) => libfuzzer_coverage::run(sub).await,
-            (LIBFUZZER_CRASH_REPORT, Some(sub)) => libfuzzer_crash_report::run(sub).await,
-            (LIBFUZZER_MERGE, Some(sub)) => libfuzzer_merge::run(sub).await,
-            (GENERIC_ANALYSIS, Some(sub)) => generic_analysis::run(sub).await,
-            (GENERIC_CRASH_REPORT, Some(sub)) => generic_crash_report::run(sub).await,
-            (GENERIC_GENERATOR, Some(sub)) => generic_generator::run(sub).await,
-            (GENERIC_TEST_INPUT, Some(sub)) => test_input::run(sub).await,
-            (LIBFUZZER_TEST_INPUT, Some(sub)) => libfuzzer_test_input::run(sub).await,
-            (LIBFUZZER_REGRESSION, Some(sub)) => libfuzzer_regression::run(sub).await,
+            (LIBFUZZER_FUZZ, Some(sub)) => libfuzzer_fuzz::run(sub, event_sender).await,
+            (LIBFUZZER_COVERAGE, Some(sub)) => libfuzzer_coverage::run(sub, event_sender).await,
+            (LIBFUZZER_CRASH_REPORT, Some(sub)) => {
+                libfuzzer_crash_report::run(sub, event_sender).await
+            }
+            (LIBFUZZER_MERGE, Some(sub)) => libfuzzer_merge::run(sub, event_sender).await,
+            (GENERIC_ANALYSIS, Some(sub)) => generic_analysis::run(sub, event_sender).await,
+            (GENERIC_CRASH_REPORT, Some(sub)) => generic_crash_report::run(sub, event_sender).await,
+            (GENERIC_GENERATOR, Some(sub)) => generic_generator::run(sub, event_sender).await,
+            (GENERIC_TEST_INPUT, Some(sub)) => test_input::run(sub, event_sender).await,
+            (LIBFUZZER_TEST_INPUT, Some(sub)) => libfuzzer_test_input::run(sub, event_sender).await,
+            (LIBFUZZER_REGRESSION, Some(sub)) => libfuzzer_regression::run(sub, event_sender).await,
             _ => {
                 anyhow::bail!("missing subcommand\nUSAGE: {}", args.usage());
             }

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

@@ -15,8 +15,7 @@ use std::{
     path::{Path, PathBuf},
     time::Duration,
 };
-use tokio::stream::StreamExt;
-use tokio::{sync::mpsc::UnboundedSender, task::JoinHandle, time::delay_for};
+use tokio::sync::mpsc::UnboundedSender;
 use uuid::Uuid;
 
 pub const SETUP_DIR: &str = "setup_dir";
@@ -69,6 +68,7 @@ pub enum CmdType {
 pub struct LocalContext {
     pub job_path: PathBuf,
     pub common_config: CommonConfig,
+    pub event_sender: Option<UnboundedSender<UiEvent>>,
 }
 
 pub fn get_hash_map(args: &clap::ArgMatches<'_>, name: &str) -> Result<HashMap<String, String>> {
@@ -212,7 +212,11 @@ pub fn get_synced_dir(
 // fuzzing tasks from generating random task id to using UUID::nil(). This
 // enables making the one-shot crash report generation, which isn't really a task,
 // consistent across multiple runs.
-pub fn build_local_context(args: &ArgMatches<'_>, generate_task_id: bool) -> Result<LocalContext> {
+pub fn build_local_context(
+    args: &ArgMatches<'_>,
+    generate_task_id: bool,
+    event_sender: Option<UnboundedSender<UiEvent>>,
+) -> Result<LocalContext> {
     let job_id = get_uuid("job_id", args).unwrap_or_else(|_| Uuid::nil());
     let task_id = get_uuid("task_id", args).unwrap_or_else(|_| {
         if generate_task_id {
@@ -250,6 +254,7 @@ pub fn build_local_context(args: &ArgMatches<'_>, generate_task_id: bool) -> Res
     Ok(LocalContext {
         job_path,
         common_config,
+        event_sender,
     })
 }
 
@@ -317,48 +322,31 @@ pub async fn wait_for_dir(path: impl AsRef<Path>) -> Result<()> {
     .await
 }
 
-pub fn spawn_file_count_monitor(
-    dir: PathBuf,
-    sender: UnboundedSender<UiEvent>,
-) -> JoinHandle<Result<()>> {
-    tokio::spawn(async move {
-        wait_for_dir(&dir).await?;
-
-        loop {
-            let mut rd = tokio::fs::read_dir(&dir).await?;
-            let mut count: usize = 0;
-
-            while let Some(Ok(entry)) = rd.next().await {
-                if entry.path().is_file() {
-                    count += 1;
-                }
-            }
+#[derive(Debug)]
+pub enum UiEvent {
+    MonitorDir(PathBuf),
+}
 
-            if sender
-                .send(UiEvent::FileCount {
-                    dir: dir.clone(),
-                    count,
-                })
-                .is_err()
-            {
-                return Ok(());
-            }
-            delay_for(Duration::from_secs(5)).await;
-        }
-    })
+pub trait SyncCountDirMonitor<T: Sized> {
+    fn monitor_count(self, event_sender: &Option<UnboundedSender<UiEvent>>) -> Result<T>;
 }
 
-pub fn monitor_file_urls(
-    urls: &[Option<impl AsRef<Path>>],
-    event_sender: UnboundedSender<UiEvent>,
-) -> Vec<JoinHandle<Result<()>>> {
-    urls.iter()
-        .filter_map(|x| x.as_ref())
-        .map(|path| spawn_file_count_monitor(path.as_ref().into(), event_sender.clone()))
-        .collect::<Vec<_>>()
+impl SyncCountDirMonitor<SyncedDir> for SyncedDir {
+    fn monitor_count(self, event_sender: &Option<UnboundedSender<UiEvent>>) -> Result<Self> {
+        if let (Some(event_sender), Some(p)) = (event_sender, self.url.as_file_path()) {
+            event_sender.send(UiEvent::MonitorDir(p))?;
+        }
+        Ok(self)
+    }
 }
 
-#[derive(Debug)]
-pub enum UiEvent {
-    FileCount { dir: PathBuf, count: usize },
+impl SyncCountDirMonitor<Option<SyncedDir>> for Option<SyncedDir> {
+    fn monitor_count(self, event_sender: &Option<UnboundedSender<UiEvent>>) -> Result<Self> {
+        if let Some(sd) = self {
+            let sd = sd.monitor_count(event_sender)?;
+            Ok(Some(sd))
+        } else {
+            Ok(self)
+        }
+    }
 }

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

@@ -4,8 +4,9 @@
 use crate::{
     local::common::{
         build_local_context, get_cmd_arg, get_cmd_exe, get_hash_map, get_synced_dir, CmdType,
-        ANALYSIS_DIR, ANALYZER_ENV, ANALYZER_EXE, ANALYZER_OPTIONS, CRASHES_DIR, NO_REPRO_DIR,
-        REPORTS_DIR, TARGET_ENV, TARGET_EXE, TARGET_OPTIONS, TOOLS_DIR, UNIQUE_REPORTS_DIR,
+        SyncCountDirMonitor, UiEvent, ANALYSIS_DIR, ANALYZER_ENV, ANALYZER_EXE, ANALYZER_OPTIONS,
+        CRASHES_DIR, NO_REPRO_DIR, REPORTS_DIR, TARGET_ENV, TARGET_EXE, TARGET_OPTIONS, TOOLS_DIR,
+        UNIQUE_REPORTS_DIR,
     },
     tasks::{
         analysis::generic::{run as run_analysis, Config},
@@ -15,30 +16,39 @@ use crate::{
 use anyhow::Result;
 use clap::{App, Arg, SubCommand};
 use storage_queue::QueueClient;
+use tokio::sync::mpsc::UnboundedSender;
 
 pub fn build_analysis_config(
     args: &clap::ArgMatches<'_>,
     input_queue: Option<QueueClient>,
     common: CommonConfig,
+    event_sender: Option<UnboundedSender<UiEvent>>,
 ) -> Result<Config> {
     let target_exe = get_cmd_exe(CmdType::Target, args)?.into();
     let target_options = get_cmd_arg(CmdType::Target, args);
 
     let analyzer_exe = value_t!(args, ANALYZER_EXE, String)?;
     let analyzer_options = args.values_of_lossy(ANALYZER_OPTIONS).unwrap_or_default();
     let analyzer_env = get_hash_map(args, ANALYZER_ENV)?;
-    let analysis = get_synced_dir(ANALYSIS_DIR, common.job_id, common.task_id, args)?;
+    let analysis = get_synced_dir(ANALYSIS_DIR, common.job_id, common.task_id, args)?
+        .monitor_count(&event_sender)?;
     let tools = get_synced_dir(TOOLS_DIR, common.job_id, common.task_id, args)?;
     let crashes = if input_queue.is_none() {
-        get_synced_dir(CRASHES_DIR, common.job_id, common.task_id, args).ok()
+        get_synced_dir(CRASHES_DIR, common.job_id, common.task_id, args)
+            .ok()
+            .monitor_count(&event_sender)?
     } else {
         None
     };
-
-    let reports = get_synced_dir(REPORTS_DIR, common.job_id, common.task_id, args).ok();
-    let no_repro = get_synced_dir(NO_REPRO_DIR, common.job_id, common.task_id, args).ok();
-    let unique_reports =
-        get_synced_dir(UNIQUE_REPORTS_DIR, common.job_id, common.task_id, args).ok();
+    let reports = get_synced_dir(REPORTS_DIR, common.job_id, common.task_id, args)
+        .ok()
+        .monitor_count(&event_sender)?;
+    let no_repro = get_synced_dir(NO_REPRO_DIR, common.job_id, common.task_id, args)
+        .ok()
+        .monitor_count(&event_sender)?;
+    let unique_reports = get_synced_dir(UNIQUE_REPORTS_DIR, common.job_id, common.task_id, args)
+        .ok()
+        .monitor_count(&event_sender)?;
 
     let config = Config {
         target_exe,
@@ -59,9 +69,12 @@ pub fn build_analysis_config(
     Ok(config)
 }
 
-pub async fn run(args: &clap::ArgMatches<'_>) -> Result<()> {
-    let context = build_local_context(args, true)?;
-    let config = build_analysis_config(args, None, context.common_config.clone())?;
+pub async fn run(
+    args: &clap::ArgMatches<'_>,
+    event_sender: Option<UnboundedSender<UiEvent>>,
+) -> Result<()> {
+    let context = build_local_context(args, true, event_sender.clone())?;
+    let config = build_analysis_config(args, None, context.common_config.clone(), event_sender)?;
     run_analysis(config).await
 }
 

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

@@ -4,9 +4,9 @@
 use crate::{
     local::common::{
         build_local_context, get_cmd_arg, get_cmd_env, get_cmd_exe, get_synced_dir, CmdType,
-        CHECK_ASAN_LOG, CHECK_RETRY_COUNT, CRASHES_DIR, DISABLE_CHECK_DEBUGGER,
-        DISABLE_CHECK_QUEUE, NO_REPRO_DIR, REPORTS_DIR, TARGET_ENV, TARGET_EXE, TARGET_OPTIONS,
-        TARGET_TIMEOUT, UNIQUE_REPORTS_DIR,
+        SyncCountDirMonitor, UiEvent, CHECK_ASAN_LOG, CHECK_RETRY_COUNT, CRASHES_DIR,
+        DISABLE_CHECK_DEBUGGER, DISABLE_CHECK_QUEUE, NO_REPRO_DIR, REPORTS_DIR, TARGET_ENV,
+        TARGET_EXE, TARGET_OPTIONS, TARGET_TIMEOUT, UNIQUE_REPORTS_DIR,
     },
     tasks::{
         config::CommonConfig,
@@ -16,11 +16,13 @@ use crate::{
 use anyhow::Result;
 use clap::{App, Arg, SubCommand};
 use storage_queue::QueueClient;
+use tokio::sync::mpsc::UnboundedSender;
 
 pub fn build_report_config(
     args: &clap::ArgMatches<'_>,
     input_queue: Option<QueueClient>,
     common: CommonConfig,
+    event_sender: Option<UnboundedSender<UiEvent>>,
 ) -> Result<Config> {
     let target_exe = get_cmd_exe(CmdType::Target, args)?.into();
     let target_env = get_cmd_env(CmdType::Target, args)?;
@@ -31,16 +33,22 @@ pub fn build_report_config(
         common.job_id,
         common.task_id,
         args,
-    )?);
-    let reports = get_synced_dir(REPORTS_DIR, common.job_id, common.task_id, args).ok();
-    let no_repro = get_synced_dir(NO_REPRO_DIR, common.job_id, common.task_id, args).ok();
+    )?)
+    .monitor_count(&event_sender)?;
+    let reports = get_synced_dir(REPORTS_DIR, common.job_id, common.task_id, args)
+        .ok()
+        .monitor_count(&event_sender)?;
+    let no_repro = get_synced_dir(NO_REPRO_DIR, common.job_id, common.task_id, args)
+        .ok()
+        .monitor_count(&event_sender)?;
 
     let unique_reports = Some(get_synced_dir(
         UNIQUE_REPORTS_DIR,
         common.job_id,
         common.task_id,
         args,
-    )?);
+    )?)
+    .monitor_count(&event_sender)?;
 
     let target_timeout = value_t!(args, TARGET_TIMEOUT, u64).ok();
 
@@ -70,9 +78,12 @@ pub fn build_report_config(
     Ok(config)
 }
 
-pub async fn run(args: &clap::ArgMatches<'_>) -> Result<()> {
-    let context = build_local_context(args, true)?;
-    let config = build_report_config(args, None, context.common_config.clone())?;
+pub async fn run(
+    args: &clap::ArgMatches<'_>,
+    event_sender: Option<UnboundedSender<UiEvent>>,
+) -> Result<()> {
+    let context = build_local_context(args, true, event_sender.clone())?;
+    let config = build_report_config(args, None, context.common_config.clone(), event_sender)?;
     ReportTask::new(config).managed_run().await
 }
 

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

@@ -4,9 +4,10 @@
 use crate::{
     local::common::{
         build_local_context, get_cmd_arg, get_cmd_env, get_cmd_exe, get_synced_dir,
-        get_synced_dirs, CmdType, CHECK_ASAN_LOG, CHECK_RETRY_COUNT, CRASHES_DIR,
-        DISABLE_CHECK_DEBUGGER, GENERATOR_ENV, GENERATOR_EXE, GENERATOR_OPTIONS, READONLY_INPUTS,
-        RENAME_OUTPUT, TARGET_ENV, TARGET_EXE, TARGET_OPTIONS, TARGET_TIMEOUT, TOOLS_DIR,
+        get_synced_dirs, CmdType, SyncCountDirMonitor, UiEvent, CHECK_ASAN_LOG, CHECK_RETRY_COUNT,
+        CRASHES_DIR, DISABLE_CHECK_DEBUGGER, GENERATOR_ENV, GENERATOR_EXE, GENERATOR_OPTIONS,
+        READONLY_INPUTS, RENAME_OUTPUT, TARGET_ENV, TARGET_EXE, TARGET_OPTIONS, TARGET_TIMEOUT,
+        TOOLS_DIR,
     },
     tasks::{
         config::CommonConfig,
@@ -15,25 +16,36 @@ use crate::{
 };
 use anyhow::Result;
 use clap::{App, Arg, SubCommand};
+use tokio::sync::mpsc::UnboundedSender;
 
-pub fn build_fuzz_config(args: &clap::ArgMatches<'_>, common: CommonConfig) -> Result<Config> {
-    let crashes = get_synced_dir(CRASHES_DIR, common.job_id, common.task_id, args)?;
+pub fn build_fuzz_config(
+    args: &clap::ArgMatches<'_>,
+    common: CommonConfig,
+    event_sender: Option<UnboundedSender<UiEvent>>,
+) -> Result<Config> {
+    let crashes = get_synced_dir(CRASHES_DIR, common.job_id, common.task_id, args)?
+        .monitor_count(&event_sender)?;
     let target_exe = get_cmd_exe(CmdType::Target, args)?.into();
     let target_options = get_cmd_arg(CmdType::Target, args);
     let target_env = get_cmd_env(CmdType::Target, args)?;
 
     let generator_exe = get_cmd_exe(CmdType::Generator, args)?;
     let generator_options = get_cmd_arg(CmdType::Generator, args);
     let generator_env = get_cmd_env(CmdType::Generator, args)?;
-    let readonly_inputs = get_synced_dirs(READONLY_INPUTS, common.job_id, common.task_id, args)?;
+    let readonly_inputs = 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 rename_output = args.is_present(RENAME_OUTPUT);
     let check_asan_log = args.is_present(CHECK_ASAN_LOG);
     let check_debugger = !args.is_present(DISABLE_CHECK_DEBUGGER);
     let check_retry_count = value_t!(args, CHECK_RETRY_COUNT, u64)?;
     let target_timeout = Some(value_t!(args, TARGET_TIMEOUT, u64)?);
 
-    let tools = get_synced_dir(TOOLS_DIR, common.job_id, common.task_id, args).ok();
+    let tools = get_synced_dir(TOOLS_DIR, common.job_id, common.task_id, args)
+        .ok()
+        .monitor_count(&event_sender)?;
 
     let ensemble_sync_delay = None;
 
@@ -59,9 +71,12 @@ pub fn build_fuzz_config(args: &clap::ArgMatches<'_>, common: CommonConfig) -> R
     Ok(config)
 }
 
-pub async fn run(args: &clap::ArgMatches<'_>) -> Result<()> {
-    let context = build_local_context(args, true)?;
-    let config = build_fuzz_config(args, context.common_config.clone())?;
+pub async fn run(
+    args: &clap::ArgMatches<'_>,
+    event_sender: Option<UnboundedSender<UiEvent>>,
+) -> Result<()> {
+    let context = build_local_context(args, true, event_sender.clone())?;
+    let config = build_fuzz_config(args, context.common_config.clone(), event_sender)?;
     GeneratorTask::new(config).run().await
 }