diff --git a/src/agent/Cargo.lock b/src/agent/Cargo.lock index 18ec65440..bb97e427b 100644 --- a/src/agent/Cargo.lock +++ b/src/agent/Cargo.lock @@ -1702,6 +1702,7 @@ dependencies = [ "anyhow", "appinsights", "async-trait", + "atexit", "backoff", "clap", "env_logger", @@ -1713,6 +1714,7 @@ dependencies = [ "onefuzz", "onefuzz-telemetry", "path-absolutize", + "remove_dir_all 0.7.0", "reqwest", "reqwest-retry", "serde", @@ -2236,6 +2238,19 @@ dependencies = [ "winapi 0.3.9", ] +[[package]] +name = "remove_dir_all" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "882f368737489ea543bc5c340e6f3d34a28c39980bd9a979e47322b26f60ac40" +dependencies = [ + "libc", + "log", + "num_cpus", + "rayon", + "winapi 0.3.9", +] + [[package]] name = "reqwest" version = "0.10.10" @@ -2723,7 +2738,7 @@ dependencies = [ "libc", "rand 0.8.3", "redox_syscall", - "remove_dir_all", + "remove_dir_all 0.5.3", "winapi 0.3.9", ] diff --git a/src/agent/atexit/src/lib.rs b/src/agent/atexit/src/lib.rs index 678b611a5..e4e9477e7 100644 --- a/src/agent/atexit/src/lib.rs +++ b/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 { 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); } } diff --git a/src/agent/onefuzz-agent/Cargo.toml b/src/agent/onefuzz-agent/Cargo.toml index 1a423d1f2..1f8255169 100644 --- a/src/agent/onefuzz-agent/Cargo.toml +++ b/src/agent/onefuzz-agent/Cargo.toml @@ -35,6 +35,8 @@ reqwest-retry = { path = "../reqwest-retry" } onefuzz-telemetry = { path = "../onefuzz-telemetry" } stacktrace-parser = { path = "../stacktrace-parser" } path-absolutize = "3.0.6" +atexit = { path = "../atexit" } +remove_dir_all = "0.7" [dev-dependencies] tempfile = "3.2" diff --git a/src/agent/onefuzz-agent/src/local/common.rs b/src/agent/onefuzz-agent/src/local/common.rs index 221e67087..f1e371e72 100644 --- a/src/agent/onefuzz-agent/src/local/common.rs +++ b/src/agent/onefuzz-agent/src/local/common.rs @@ -1,10 +1,14 @@ use crate::tasks::config::CommonConfig; use crate::tasks::utils::parse_key_value; use anyhow::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 remove_dir_all::remove_dir_all; use reqwest::Url; +use std::task::Poll; use std::{ collections::HashMap, path::{Path, PathBuf}, @@ -12,10 +16,6 @@ use std::{ }; 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"; @@ -128,6 +128,12 @@ pub fn add_common_config(app: App<'static, 'static>) -> App<'static, 'static> { .takes_value(true) .required(false), ) + .arg( + Arg::with_name("keep_job_dir") + .long("keep_job_dir") + .required(false) + .help("keep the local directory created for running the task"), + ) } fn get_uuid(name: &str, args: &ArgMatches<'_>) -> Result { @@ -162,6 +168,14 @@ pub fn get_synced_dirs( Ok(dirs?) } +fn register_cleanup(job_id: Uuid) -> Result<()> { + let path = std::env::current_dir()?.join(job_id.to_string()); + atexit::register(move || { + remove_dir_all(&path).expect("cleanup failed"); + }); + Ok(()) +} + pub fn get_synced_dir( name: &str, job_id: Uuid, @@ -204,6 +218,10 @@ pub fn build_common_config(args: &ArgMatches<'_>, generate_task_id: bool) -> Res PathBuf::default() }; + if !args.is_present("keep_job_dir") { + register_cleanup(job_id)?; + } + let config = CommonConfig { job_id, task_id, diff --git a/src/agent/onefuzz-agent/src/local/libfuzzer.rs b/src/agent/onefuzz-agent/src/local/libfuzzer.rs index 74d5074af..cb3aaf48a 100644 --- a/src/agent/onefuzz-agent/src/local/libfuzzer.rs +++ b/src/agent/onefuzz-agent/src/local/libfuzzer.rs @@ -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; diff --git a/src/agent/onefuzz-agent/src/local/libfuzzer_coverage.rs b/src/agent/onefuzz-agent/src/local/libfuzzer_coverage.rs index 317b4e410..93136ec0f 100644 --- a/src/agent/onefuzz-agent/src/local/libfuzzer_coverage.rs +++ b/src/agent/onefuzz-agent/src/local/libfuzzer_coverage.rs @@ -57,7 +57,6 @@ pub fn build_coverage_config( pub async fn run(args: &clap::ArgMatches<'_>) -> Result<()> { let common = build_common_config(args, true)?; let config = build_coverage_config(args, false, None, common)?; - let mut task = CoverageTask::new(config); task.managed_run().await } diff --git a/src/agent/onefuzz-agent/src/main.rs b/src/agent/onefuzz-agent/src/main.rs index dfce3c327..56ca1b1b1 100644 --- a/src/agent/onefuzz-agent/src/main.rs +++ b/src/agent/onefuzz-agent/src/main.rs @@ -40,7 +40,9 @@ fn main() -> Result<()> { let matches = app.get_matches(); let mut rt = tokio::runtime::Runtime::new()?; - rt.block_on(run(matches)) + let result = rt.block_on(run(matches)); + atexit::execute(); + result } async fn run(args: ArgMatches<'_>) -> Result<()> {