cleanup local job_id directory upon clean exit

src/agent/atexit/src/lib.rs

@@ -25,6 +25,13 @@ pub fn exit_process(code: i32) -> ! {
     ATEXIT.exit_process(code)
 }
 
+/// Runs the registered functions but does *not* terminate the process
+///
+/// This function is not called automatically (e.g. via `drop`).
+pub fn execute() {
+    ATEXIT.execute()
+}
+
 impl AtExit {
     fn new() -> Arc<Self> {
         let result = Arc::new(AtExit {
@@ -56,9 +63,13 @@ impl AtExit {
     }
 
     fn exit_process(&self, code: i32) -> ! {
+        self.execute();
+        std::process::exit(code);
+    }
+
+    fn execute(&self) {
         for function in self.functions.write().unwrap().iter_mut() {
             function();
         }
-        std::process::exit(code);
     }
 }

src/agent/onefuzz-agent/Cargo.toml

@@ -35,6 +35,7 @@ reqwest-retry = { path = "../reqwest-retry" }
 onefuzz-telemetry = { path = "../onefuzz-telemetry" }
 stacktrace-parser = { path = "../stacktrace-parser" }
 path-absolutize = "3.0.6"
+atexit = { path = "../atexit" }
 
 [dev-dependencies]
 tempfile = "3.2"

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

@@ -1,21 +1,20 @@
 use crate::tasks::config::CommonConfig;
 use crate::tasks::utils::parse_key_value;
-use anyhow::Result;
+use anyhow::{Context, Result};
+use backoff::{future::retry, Error as BackoffError, ExponentialBackoff};
 use clap::{App, Arg, ArgMatches};
 use onefuzz::jitter::delay_with_jitter;
 use onefuzz::{blob::BlobContainerUrl, monitor::DirectoryMonitor, syncdir::SyncedDir};
+use path_absolutize::Absolutize;
 use reqwest::Url;
+use std::task::Poll;
 use std::{
     collections::HashMap,
     path::{Path, PathBuf},
     time::Duration,
 };
 use uuid::Uuid;
 
-use backoff::{future::retry, Error as BackoffError, ExponentialBackoff};
-use path_absolutize::Absolutize;
-use std::task::Poll;
-
 pub const SETUP_DIR: &str = "setup_dir";
 pub const INPUTS_DIR: &str = "inputs_dir";
 pub const CRASHES_DIR: &str = "crashes_dir";
@@ -162,6 +161,30 @@ pub fn get_synced_dirs(
     Ok(dirs?)
 }
 
+fn recursive_remove(path: &Path) -> Result<()> {
+    if path.metadata()?.is_dir() {
+        for child in std::fs::read_dir(&path)? {
+            let path = child?.path();
+            // allow children remove to fail, as tempdirs are cleaning up at this
+            // point.  the follow on remove_dir_all will catch errors.
+            recursive_remove(&path).ok();
+        }
+        std::fs::remove_dir_all(&path)
+            .with_context(|| format!("removing directory {}", path.display()))?;
+    } else {
+        std::fs::remove_file(&path).with_context(|| format!("removing {}", path.display()))?;
+    }
+    Ok(())
+}
+
+pub fn register_cleanup(job_id: Uuid) -> Result<()> {
+    let path = std::env::current_dir()?.join(job_id.to_string());
+    atexit::register(move || {
+        recursive_remove(&path).expect("cleanup failed");
+    });
+    Ok(())
+}
+
 pub fn get_synced_dir(
     name: &str,
     job_id: Uuid,

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

@@ -3,9 +3,10 @@
 
 use crate::{
     local::common::{
-        build_common_config, 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,
+        build_common_config, get_cmd_arg, get_cmd_exe, get_hash_map, get_synced_dir,
+        register_cleanup, 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,
     },
     tasks::{
         analysis::generic::{run as run_analysis, Config},
@@ -61,6 +62,7 @@ pub fn build_analysis_config(
 
 pub async fn run(args: &clap::ArgMatches<'_>) -> Result<()> {
     let common = build_common_config(args, true)?;
+    register_cleanup(common.job_id)?;
     let config = build_analysis_config(args, None, common)?;
     run_analysis(config).await
 }

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

@@ -3,10 +3,10 @@
 
 use crate::{
     local::common::{
-        build_common_config, 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,
+        build_common_config, get_cmd_arg, get_cmd_env, get_cmd_exe, get_synced_dir,
+        register_cleanup, 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,
     },
     tasks::{
         config::CommonConfig,
@@ -72,6 +72,7 @@ pub fn build_report_config(
 
 pub async fn run(args: &clap::ArgMatches<'_>) -> Result<()> {
     let common = build_common_config(args, true)?;
+    register_cleanup(common.job_id)?;
     let config = build_report_config(args, None, common)?;
     ReportTask::new(config).managed_run().await
 }

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

@@ -4,7 +4,7 @@
 use crate::{
     local::common::{
         build_common_config, get_cmd_arg, get_cmd_env, get_cmd_exe, get_synced_dir,
-        get_synced_dirs, CmdType, CHECK_ASAN_LOG, CHECK_RETRY_COUNT, CRASHES_DIR,
+        get_synced_dirs, register_cleanup, 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,
     },
@@ -61,6 +61,7 @@ pub fn build_fuzz_config(args: &clap::ArgMatches<'_>, common: CommonConfig) -> R
 
 pub async fn run(args: &clap::ArgMatches<'_>) -> Result<()> {
     let common = build_common_config(args, true)?;
+    register_cleanup(common.job_id)?;
     let config = build_fuzz_config(args, common)?;
     GeneratorTask::new(config).run().await
 }

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

@@ -4,8 +4,8 @@
 use crate::{
     local::{
         common::{
-            build_common_config, wait_for_dir, DirectoryMonitorQueue, ANALYZER_EXE, COVERAGE_DIR,
-            UNIQUE_REPORTS_DIR,
+            build_common_config, register_cleanup, wait_for_dir, DirectoryMonitorQueue,
+            ANALYZER_EXE, COVERAGE_DIR, UNIQUE_REPORTS_DIR,
         },
         generic_analysis::{build_analysis_config, build_shared_args as build_analysis_args},
         libfuzzer_coverage::{build_coverage_config, build_shared_args as build_coverage_args},
@@ -20,7 +20,6 @@ use crate::{
 };
 use anyhow::Result;
 use clap::{App, SubCommand};
-
 use onefuzz::utils::try_wait_all_join_handles;
 use std::collections::HashSet;
 use tokio::task::spawn;
@@ -35,6 +34,8 @@ pub async fn run(args: &clap::ArgMatches<'_>) -> Result<()> {
         .as_file_path()
         .expect("invalid crash dir remote location");
 
+    register_cleanup(common.job_id)?;
+
     let fuzzer = LibFuzzerFuzzTask::new(fuzz_config)?;
     fuzzer.check_libfuzzer().await?;
     let mut task_handles = vec![];

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

@@ -4,8 +4,8 @@
 use crate::{
     local::common::{
         build_common_config, get_cmd_arg, get_cmd_env, get_cmd_exe, get_synced_dir,
-        get_synced_dirs, CmdType, CHECK_FUZZER_HELP, COVERAGE_DIR, INPUTS_DIR, READONLY_INPUTS,
-        TARGET_ENV, TARGET_EXE, TARGET_OPTIONS,
+        get_synced_dirs, register_cleanup, CmdType, CHECK_FUZZER_HELP, COVERAGE_DIR, INPUTS_DIR,
+        READONLY_INPUTS, TARGET_ENV, TARGET_EXE, TARGET_OPTIONS,
     },
     tasks::{
         config::CommonConfig,
@@ -56,8 +56,8 @@ pub fn build_coverage_config(
 
 pub async fn run(args: &clap::ArgMatches<'_>) -> Result<()> {
     let common = build_common_config(args, true)?;
+    register_cleanup(common.job_id)?;
     let config = build_coverage_config(args, false, None, common)?;
-
     let mut task = CoverageTask::new(config);
     task.managed_run().await
 }

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

@@ -3,9 +3,10 @@
 
 use crate::{
     local::common::{
-        build_common_config, get_cmd_arg, get_cmd_env, get_cmd_exe, get_synced_dir, CmdType,
-        CHECK_FUZZER_HELP, CHECK_RETRY_COUNT, CRASHES_DIR, DISABLE_CHECK_QUEUE, NO_REPRO_DIR,
-        REPORTS_DIR, TARGET_ENV, TARGET_EXE, TARGET_OPTIONS, TARGET_TIMEOUT, UNIQUE_REPORTS_DIR,
+        build_common_config, get_cmd_arg, get_cmd_env, get_cmd_exe, get_synced_dir,
+        register_cleanup, CmdType, CHECK_FUZZER_HELP, CHECK_RETRY_COUNT, CRASHES_DIR,
+        DISABLE_CHECK_QUEUE, NO_REPRO_DIR, REPORTS_DIR, TARGET_ENV, TARGET_EXE, TARGET_OPTIONS,
+        TARGET_TIMEOUT, UNIQUE_REPORTS_DIR,
     },
     tasks::{
         config::CommonConfig,
@@ -64,6 +65,7 @@ pub fn build_report_config(
 
 pub async fn run(args: &clap::ArgMatches<'_>) -> Result<()> {
     let common = build_common_config(args, true)?;
+    register_cleanup(common.job_id)?;
     let config = build_report_config(args, None, common)?;
     ReportTask::new(config).managed_run().await
 }

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

@@ -3,9 +3,9 @@
 
 use crate::{
     local::common::{
-        build_common_config, get_cmd_arg, get_cmd_env, get_cmd_exe, get_synced_dir, CmdType,
-        CHECK_FUZZER_HELP, CRASHES_DIR, INPUTS_DIR, TARGET_ENV, TARGET_EXE, TARGET_OPTIONS,
-        TARGET_WORKERS,
+        build_common_config, get_cmd_arg, get_cmd_env, get_cmd_exe, get_synced_dir,
+        register_cleanup, CmdType, CHECK_FUZZER_HELP, CRASHES_DIR, INPUTS_DIR, TARGET_ENV,
+        TARGET_EXE, TARGET_OPTIONS, TARGET_WORKERS,
     },
     tasks::{
         config::CommonConfig,
@@ -51,7 +51,9 @@ pub fn build_fuzz_config(args: &clap::ArgMatches<'_>, common: CommonConfig) -> R
 
 pub async fn run(args: &clap::ArgMatches<'_>) -> Result<()> {
     let common = build_common_config(args, true)?;
+    register_cleanup(common.job_id)?;
     let config = build_fuzz_config(args, common)?;
+
     LibFuzzerFuzzTask::new(config)?.run().await
 }
 

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

@@ -4,8 +4,9 @@
 use crate::{
     local::common::{
         build_common_config, get_cmd_arg, get_cmd_env, get_cmd_exe, get_synced_dir,
-        get_synced_dirs, CmdType, ANALYSIS_INPUTS, ANALYSIS_UNIQUE_INPUTS, CHECK_FUZZER_HELP,
-        INPUTS_DIR, PRESERVE_EXISTING_OUTPUTS, TARGET_ENV, TARGET_EXE, TARGET_OPTIONS,
+        get_synced_dirs, register_cleanup, CmdType, ANALYSIS_INPUTS, ANALYSIS_UNIQUE_INPUTS,
+        CHECK_FUZZER_HELP, INPUTS_DIR, PRESERVE_EXISTING_OUTPUTS, TARGET_ENV, TARGET_EXE,
+        TARGET_OPTIONS,
     },
     tasks::{
         config::CommonConfig,
@@ -46,6 +47,7 @@ pub fn build_merge_config(
 
 pub async fn run(args: &clap::ArgMatches<'_>) -> Result<()> {
     let common = build_common_config(args, true)?;
+    register_cleanup(common.job_id)?;
     let config = build_merge_config(args, None, common)?;
     spawn(std::sync::Arc::new(config)).await
 }

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

@@ -3,7 +3,7 @@
 
 use crate::{
     local::{
-        common::{build_common_config, DirectoryMonitorQueue},
+        common::{build_common_config, register_cleanup, DirectoryMonitorQueue},
         generic_crash_report::{build_report_config, build_shared_args as build_crash_args},
         generic_generator::{build_fuzz_config, build_shared_args as build_fuzz_args},
     },
@@ -18,7 +18,9 @@ use uuid::Uuid;
 
 pub async fn run(args: &clap::ArgMatches<'_>) -> Result<()> {
     let common = build_common_config(args, true)?;
+    register_cleanup(common.job_id)?;
     let fuzz_config = build_fuzz_config(args, common.clone())?;
+
     let crash_dir = fuzz_config
         .crashes
         .url

src/agent/onefuzz-agent/src/main.rs

@@ -40,7 +40,11 @@ fn main() -> Result<()> {
     let matches = app.get_matches();
 
     let mut rt = tokio::runtime::Runtime::new()?;
-    rt.block_on(run(matches))
+    rt.block_on(run(matches))?;
+
+    // only call atexit if everything exits cleanly
+    atexit::execute();
+    Ok(())
 }
 
 async fn run(args: ArgMatches<'_>) -> Result<()> {