diff --git a/src/agent/onefuzz-agent/src/main.rs b/src/agent/onefuzz-agent/src/main.rs index 28e5ec867..aff91f680 100644 --- a/src/agent/onefuzz-agent/src/main.rs +++ b/src/agent/onefuzz-agent/src/main.rs @@ -48,10 +48,10 @@ fn main() -> Result<()> { async fn run(args: ArgMatches<'_>) -> Result<()> { match args.subcommand() { - (LICENSE_CMD, Some(_)) => return licenses(), - (DEBUG_CMD, Some(sub)) => return debug::cmd::run(sub).await, - (LOCAL_CMD, Some(sub)) => return local::cmd::run(sub).await, - (MANAGED_CMD, Some(sub)) => return managed::cmd::run(sub).await, + (LICENSE_CMD, Some(_)) => licenses(), + (DEBUG_CMD, Some(sub)) => debug::cmd::run(sub).await, + (LOCAL_CMD, Some(sub)) => local::cmd::run(sub).await, + (MANAGED_CMD, Some(sub)) => managed::cmd::run(sub).await, _ => { anyhow::bail!("missing subcommand\nUSAGE: {}", args.usage()); } diff --git a/src/agent/onefuzz-agent/src/tasks/analysis/generic.rs b/src/agent/onefuzz-agent/src/tasks/analysis/generic.rs index 9c8e2a61d..1926201e4 100644 --- a/src/agent/onefuzz-agent/src/tasks/analysis/generic.rs +++ b/src/agent/onefuzz-agent/src/tasks/analysis/generic.rs @@ -114,9 +114,7 @@ async fn poll_inputs(config: &Config, tmp_dir: OwnedDir) -> Result<()> { } pub async fn run_tool(input: impl AsRef, config: &Config) -> Result<()> { - let mut expand = Expand::new(); - - expand + let expand = Expand::new() .input_path(&input) .target_exe(&config.target_exe) .target_options(&config.target_options) diff --git a/src/agent/onefuzz-agent/src/tasks/fuzz/generator.rs b/src/agent/onefuzz-agent/src/tasks/fuzz/generator.rs index d0ff1084b..0a6c861ea 100644 --- a/src/agent/onefuzz-agent/src/tasks/fuzz/generator.rs +++ b/src/agent/onefuzz-agent/src/tasks/fuzz/generator.rs @@ -91,20 +91,18 @@ impl GeneratorTask { } async fn fuzzing_loop(&self, heartbeat_client: Option) -> Result<()> { - let mut tester = Tester::new( + let tester = Tester::new( &self.config.common.setup_dir, &self.config.target_exe, &self.config.target_options, &self.config.target_env, - ); - - tester - .check_asan_log(self.config.check_asan_log) - .check_debugger(self.config.check_debugger) - .check_retry_count(self.config.check_retry_count); - if let Some(timeout) = self.config.target_timeout { - tester.timeout(timeout); - } + ) + .check_asan_log(self.config.check_asan_log) + .check_debugger(self.config.check_debugger) + .check_retry_count(self.config.check_retry_count) + .set_optional(self.config.target_timeout, |tester, timeout| { + tester.timeout(timeout) + }); loop { for corpus_dir in &self.config.readonly_inputs { @@ -154,18 +152,16 @@ impl GeneratorTask { ) -> Result<()> { utils::reset_tmp_dir(&output_dir).await?; let (mut generator, generator_path) = { - let mut expand = Expand::new(); - expand + let expand = Expand::new() .generated_inputs(&output_dir) .input_corpus(&corpus_dir) .generator_exe(&self.config.generator_exe) .generator_options(&self.config.generator_options) .job_id(&self.config.common.job_id) - .task_id(&self.config.common.task_id); - - if let Some(tools) = &self.config.tools { - expand.tools_dir(&tools.path); - } + .task_id(&self.config.common.task_id) + .set_optional_ref(&self.config.tools, |expand, tools| { + expand.tools_dir(&tools.path) + }); let generator_path = expand.evaluate_value(&self.config.generator_exe)?; diff --git a/src/agent/onefuzz-agent/src/tasks/fuzz/supervisor.rs b/src/agent/onefuzz-agent/src/tasks/fuzz/supervisor.rs index 5b31221c3..49f6a0edb 100644 --- a/src/agent/onefuzz-agent/src/tasks/fuzz/supervisor.rs +++ b/src/agent/onefuzz-agent/src/tasks/fuzz/supervisor.rs @@ -173,8 +173,7 @@ async fn start_supervisor( inputs: &SyncedDir, reports_dir: PathBuf, ) -> Result { - let mut expand = Expand::new(); - expand + let expand = Expand::new() .supervisor_exe(&config.supervisor_exe) .supervisor_options(&config.supervisor_options) .runtime_dir(&runtime_dir) @@ -183,23 +182,17 @@ async fn start_supervisor( .reports_dir(&reports_dir) .setup_dir(&config.common.setup_dir) .job_id(&config.common.job_id) - .task_id(&config.common.task_id); - - if let Some(tools) = &config.tools { - expand.tools_dir(&tools.path); - } - - if let Some(target_exe) = &config.target_exe { - expand.target_exe(target_exe); - } - - if let Some(target_options) = &config.target_options { - expand.target_options(target_options); - } - - if let Some(input_marker) = &config.supervisor_input_marker { - expand.input_marker(input_marker); - } + .task_id(&config.common.task_id) + .set_optional_ref(&config.tools, |expand, tools| expand.tools_dir(&tools.path)) + .set_optional_ref(&config.target_exe, |expand, target_exe| { + expand.target_exe(target_exe) + }) + .set_optional_ref(&config.supervisor_input_marker, |expand, input_marker| { + expand.input_marker(input_marker) + }) + .set_optional_ref(&config.target_options, |expand, target_options| { + expand.target_options(target_options) + }); let supervisor_path = expand.evaluate_value(&config.supervisor_exe)?; let mut cmd = Command::new(supervisor_path); diff --git a/src/agent/onefuzz-agent/src/tasks/merge/generic.rs b/src/agent/onefuzz-agent/src/tasks/merge/generic.rs index 99afba6ba..1c9ca4c5c 100644 --- a/src/agent/onefuzz-agent/src/tasks/merge/generic.rs +++ b/src/agent/onefuzz-agent/src/tasks/merge/generic.rs @@ -129,9 +129,7 @@ async fn try_delete_blob(input_url: Url) -> Result<()> { } async fn merge(config: &Config, output_dir: impl AsRef) -> Result<()> { - let mut expand = Expand::new(); - - expand + let expand = Expand::new() .input_marker(&config.supervisor_input_marker) .input_corpus(&config.unique_inputs.path) .target_options(&config.target_options) diff --git a/src/agent/onefuzz-agent/src/tasks/report/generic.rs b/src/agent/onefuzz-agent/src/tasks/report/generic.rs index b6857d082..5847176b5 100644 --- a/src/agent/onefuzz-agent/src/tasks/report/generic.rs +++ b/src/agent/onefuzz-agent/src/tasks/report/generic.rs @@ -118,17 +118,15 @@ pub struct GenericReportProcessor<'a> { impl<'a> GenericReportProcessor<'a> { pub fn new(config: &'a Config, heartbeat_client: Option) -> Self { - let mut tester = Tester::new( + let tester = Tester::new( &config.common.setup_dir, &config.target_exe, &config.target_options, &config.target_env, - ); - - tester - .check_asan_log(config.check_asan_log) - .check_debugger(config.check_debugger) - .check_retry_count(config.check_retry_count); + ) + .check_asan_log(config.check_asan_log) + .check_debugger(config.check_debugger) + .check_retry_count(config.check_retry_count); Self { config, diff --git a/src/agent/onefuzz/src/asan.rs b/src/agent/onefuzz/src/asan.rs index 502074ec1..ee4b63e57 100644 --- a/src/agent/onefuzz/src/asan.rs +++ b/src/agent/onefuzz/src/asan.rs @@ -177,7 +177,7 @@ pub fn add_asan_log_env(env: &mut HashMap, as pub async fn check_asan_string(mut data: String) -> Result> { let asan = AsanLog::parse(data.clone()); if asan.is_some() { - return Ok(asan); + Ok(asan) } else { if data.len() > ASAN_LOG_TRUNCATE_SIZE { data.truncate(ASAN_LOG_TRUNCATE_SIZE); diff --git a/src/agent/onefuzz/src/expand.rs b/src/agent/onefuzz/src/expand.rs index 410ad4954..8c6562983 100644 --- a/src/agent/onefuzz/src/expand.rs +++ b/src/agent/onefuzz/src/expand.rs @@ -2,8 +2,8 @@ // Licensed under the MIT License. use anyhow::Result; -use std::collections::HashMap; use std::path::{Path, PathBuf}; +use std::{collections::HashMap, hash::Hash}; use strum::IntoEnumIterator; use strum_macros::EnumIter; use uuid::Uuid; @@ -117,137 +117,138 @@ impl<'a> Expand<'a> { } } - pub fn set_value(&mut self, name: PlaceHolder, value: ExpandedValue<'a>) -> &mut Self { - self.values.insert(name.get_string(), value); - self + pub fn set_value(self, name: PlaceHolder, value: ExpandedValue<'a>) -> Self { + let mut values = self.values; + values.insert(name.get_string(), value); + Self { values } } - pub fn generated_inputs(&mut self, arg: impl AsRef) -> &mut Self { + pub fn set_optional_ref<'l, T: 'l>( + self, + value: &'l Option, + setter: impl FnOnce(Self, &'l T) -> Self, + ) -> Self { + if let Some(value) = value { + setter(self, value) + } else { + self + } + } + + pub fn set_optional(self, value: Option, setter: impl FnOnce(Self, T) -> Self) -> Self { + if let Some(value) = value { + setter(self, value) + } else { + self + } + } + + pub fn generated_inputs(self, arg: impl AsRef) -> Self { let arg = arg.as_ref(); let path = String::from(arg.to_string_lossy()); - self.set_value(PlaceHolder::GeneratedInputs, ExpandedValue::Path(path)); - self + self.set_value(PlaceHolder::GeneratedInputs, ExpandedValue::Path(path)) } - pub fn crashes(&mut self, arg: impl AsRef) -> &mut Self { + pub fn crashes(self, arg: impl AsRef) -> Self { let arg = arg.as_ref(); let path = String::from(arg.to_string_lossy()); - self.set_value(PlaceHolder::Crashes, ExpandedValue::Path(path)); - self + self.set_value(PlaceHolder::Crashes, ExpandedValue::Path(path)) } - pub fn input_path(&mut self, arg: impl AsRef) -> &mut Self { + pub fn input_path(self, arg: impl AsRef) -> Self { let arg = arg.as_ref(); let path = String::from(arg.to_string_lossy()); - self.set_value(PlaceHolder::Input, ExpandedValue::Path(path)); - self + self.set_value(PlaceHolder::Input, ExpandedValue::Path(path)) } - pub fn input_marker(&mut self, arg: &str) -> &mut Self { - self.set_value(PlaceHolder::Input, ExpandedValue::Scalar(String::from(arg))); - self + pub fn input_marker(self, arg: &str) -> Self { + self.set_value(PlaceHolder::Input, ExpandedValue::Scalar(String::from(arg))) } - pub fn input_corpus(&mut self, arg: impl AsRef) -> &mut Self { + pub fn input_corpus(self, arg: impl AsRef) -> Self { let arg = arg.as_ref(); let path = String::from(arg.to_string_lossy()); - self.set_value(PlaceHolder::InputCorpus, ExpandedValue::Path(path)); - self + self.set_value(PlaceHolder::InputCorpus, ExpandedValue::Path(path)) } - pub fn generator_exe(&mut self, arg: impl AsRef) -> &mut Self { + pub fn generator_exe(self, arg: impl AsRef) -> Self { let arg = arg.as_ref(); let path = String::from(arg.to_string_lossy()); - self.set_value(PlaceHolder::GeneratorExe, ExpandedValue::Path(path)); - self + self.set_value(PlaceHolder::GeneratorExe, ExpandedValue::Path(path)) } - pub fn generator_options(&mut self, arg: &'a [String]) -> &mut Self { - self.set_value(PlaceHolder::GeneratorOptions, ExpandedValue::List(arg)); - self + pub fn generator_options(self, arg: &'a [String]) -> Self { + self.set_value(PlaceHolder::GeneratorOptions, ExpandedValue::List(arg)) } - pub fn target_exe(&mut self, arg: impl AsRef) -> &mut Self { + pub fn target_exe(self, arg: impl AsRef) -> Self { let arg = arg.as_ref(); let path = String::from(arg.to_string_lossy()); - self.set_value(PlaceHolder::TargetExe, ExpandedValue::Path(path)); - self + self.set_value(PlaceHolder::TargetExe, ExpandedValue::Path(path)) } - pub fn target_options(&mut self, arg: &'a [String]) -> &mut Self { - self.set_value(PlaceHolder::TargetOptions, ExpandedValue::List(arg)); - self + pub fn target_options(self, arg: &'a [String]) -> Self { + self.set_value(PlaceHolder::TargetOptions, ExpandedValue::List(arg)) } - pub fn analyzer_exe(&mut self, arg: impl AsRef) -> &mut Self { + pub fn analyzer_exe(self, arg: impl AsRef) -> Self { let arg = arg.as_ref(); let path = String::from(arg.to_string_lossy()); - self.set_value(PlaceHolder::AnalyzerExe, ExpandedValue::Path(path)); - self + self.set_value(PlaceHolder::AnalyzerExe, ExpandedValue::Path(path)) } - pub fn analyzer_options(&mut self, arg: &'a [String]) -> &mut Self { - self.set_value(PlaceHolder::AnalyzerOptions, ExpandedValue::List(arg)); - self + pub fn analyzer_options(self, arg: &'a [String]) -> Self { + self.set_value(PlaceHolder::AnalyzerOptions, ExpandedValue::List(arg)) } - pub fn supervisor_exe(&mut self, arg: impl AsRef) -> &mut Self { + pub fn supervisor_exe(self, arg: impl AsRef) -> Self { let arg = arg.as_ref(); let path = String::from(arg.to_string_lossy()); - self.set_value(PlaceHolder::SupervisorExe, ExpandedValue::Path(path)); - self + self.set_value(PlaceHolder::SupervisorExe, ExpandedValue::Path(path)) } - pub fn supervisor_options(&mut self, arg: &'a [String]) -> &mut Self { - self.set_value(PlaceHolder::SupervisorOptions, ExpandedValue::List(arg)); - self + pub fn supervisor_options(self, arg: &'a [String]) -> Self { + self.set_value(PlaceHolder::SupervisorOptions, ExpandedValue::List(arg)) } - pub fn output_dir(&mut self, arg: impl AsRef) -> &mut Self { + pub fn output_dir(self, arg: impl AsRef) -> Self { let arg = arg.as_ref(); let path = String::from(arg.to_string_lossy()); - self.set_value(PlaceHolder::OutputDir, ExpandedValue::Path(path)); - self + self.set_value(PlaceHolder::OutputDir, ExpandedValue::Path(path)) } - pub fn reports_dir(&mut self, arg: impl AsRef) -> &mut Self { + pub fn reports_dir(self, arg: impl AsRef) -> Self { let arg = arg.as_ref(); let path = String::from(arg.to_string_lossy()); - self.set_value(PlaceHolder::ReportsDir, ExpandedValue::Path(path)); - self + self.set_value(PlaceHolder::ReportsDir, ExpandedValue::Path(path)) } - pub fn tools_dir(&mut self, arg: impl AsRef) -> &mut Self { + pub fn tools_dir(self, arg: impl AsRef) -> Self { let arg = arg.as_ref(); let path = String::from(arg.to_string_lossy()); - self.set_value(PlaceHolder::ToolsDir, ExpandedValue::Path(path)); - self + self.set_value(PlaceHolder::ToolsDir, ExpandedValue::Path(path)) } - pub fn runtime_dir(&mut self, arg: impl AsRef) -> &mut Self { + pub fn runtime_dir(self, arg: impl AsRef) -> Self { let arg = arg.as_ref(); let path = String::from(arg.to_string_lossy()); - self.set_value(PlaceHolder::RuntimeDir, ExpandedValue::Path(path)); - self + self.set_value(PlaceHolder::RuntimeDir, ExpandedValue::Path(path)) } - pub fn setup_dir(&mut self, arg: impl AsRef) -> &mut Self { + pub fn setup_dir(self, arg: impl AsRef) -> Self { let arg = arg.as_ref(); let path = String::from(arg.to_string_lossy()); - self.set_value(PlaceHolder::SetupDir, ExpandedValue::Path(path)); - self + self.set_value(PlaceHolder::SetupDir, ExpandedValue::Path(path)) } - pub fn task_id(&mut self, arg: &Uuid) -> &mut Self { + pub fn task_id(self, arg: &Uuid) -> Self { let value = arg.to_hyphenated().to_string(); - self.set_value(PlaceHolder::TaskId, ExpandedValue::Scalar(value)); - self + self.set_value(PlaceHolder::TaskId, ExpandedValue::Scalar(value)) } - pub fn job_id(&mut self, arg: &Uuid) -> &mut Self { + pub fn job_id(self, arg: &Uuid) -> Self { let value = arg.to_hyphenated().to_string(); - self.set_value(PlaceHolder::JobId, ExpandedValue::Scalar(value)); - self + self.set_value(PlaceHolder::JobId, ExpandedValue::Scalar(value)) } fn replace_value( diff --git a/src/agent/onefuzz/src/input_tester.rs b/src/agent/onefuzz/src/input_tester.rs index 8ea0b8047..7a207f303 100644 --- a/src/agent/onefuzz/src/input_tester.rs +++ b/src/agent/onefuzz/src/input_tester.rs @@ -61,29 +61,47 @@ impl<'a> Tester<'a> { } } - pub fn timeout(&mut self, value: u64) -> &mut Self { - self.timeout = Duration::from_secs(value); - self + pub fn timeout(self, value: u64) -> Self { + Self { + timeout: Duration::from_secs(value), + ..self + } } - pub fn check_asan_log(&mut self, value: bool) -> &mut Self { - self.check_asan_log = value; - self + pub fn check_asan_log(self, value: bool) -> Self { + Self { + check_asan_log: value, + ..self + } } - pub fn check_asan_stderr(&mut self, value: bool) -> &mut Self { - self.check_asan_stderr = value; - self + pub fn check_asan_stderr(self, value: bool) -> Self { + Self { + check_asan_stderr: value, + ..self + } } - pub fn check_debugger(&mut self, value: bool) -> &mut Self { - self.check_debugger = value; - self + pub fn check_debugger(self, value: bool) -> Self { + Self { + check_debugger: value, + ..self + } } - pub fn check_retry_count(&mut self, value: u64) -> &mut Self { - self.check_retry_count = value; - self + pub fn check_retry_count(self, value: u64) -> Self { + Self { + check_retry_count: value, + ..self + } + } + + pub fn set_optional(self, value: Option, setter: impl FnOnce(Self, T) -> Self) -> Self { + if let Some(value) = value { + setter(self, value) + } else { + self + } } #[cfg(target_os = "windows")] @@ -211,13 +229,11 @@ impl<'a> Tester<'a> { }; let (argv, env) = { - let mut expand = Expand::new(); - expand + let expand = Expand::new() .input_path(input_file) .target_exe(&self.exe_path) - .target_options(&self.arguments); - - expand.setup_dir(&self.setup_dir); + .target_options(&self.arguments) + .setup_dir(&self.setup_dir); let argv = expand.evaluate(&self.arguments)?; let mut env: HashMap = HashMap::new(); diff --git a/src/agent/onefuzz/src/libfuzzer.rs b/src/agent/onefuzz/src/libfuzzer.rs index 289a52225..fc0ac1c27 100644 --- a/src/agent/onefuzz/src/libfuzzer.rs +++ b/src/agent/onefuzz/src/libfuzzer.rs @@ -56,8 +56,7 @@ impl<'a> LibFuzzer<'a> { .stderr(Stdio::piped()) .arg("-help=1"); - let mut expand = Expand::new(); - expand + let expand = Expand::new() .target_exe(&self.exe) .target_options(&self.options) .setup_dir(&self.setup_dir); @@ -92,8 +91,7 @@ impl<'a> LibFuzzer<'a> { let corpus_dir = corpus_dir.as_ref(); let fault_dir = fault_dir.as_ref(); - let mut expand = Expand::new(); - expand + let expand = Expand::new() .target_exe(&self.exe) .target_options(&self.options) .input_corpus(&corpus_dir) @@ -162,11 +160,10 @@ impl<'a> LibFuzzer<'a> { let mut options = self.options.to_owned(); options.push("{input}".to_string()); - let mut tester = Tester::new(&self.setup_dir, &self.exe, &options, &self.env); - tester.check_asan_stderr(true).check_retry_count(retry); - if let Some(timeout) = timeout { - tester.timeout(timeout); - } + let tester = Tester::new(&self.setup_dir, &self.exe, &options, &self.env) + .check_asan_stderr(true) + .check_retry_count(retry) + .set_optional(timeout, |tester, timeout| tester.timeout(timeout)); tester.test_input(test_input.as_ref()).await } @@ -175,8 +172,7 @@ impl<'a> LibFuzzer<'a> { corpus_dir: impl AsRef, corpus_dirs: &[impl AsRef], ) -> Result { - let mut expand = Expand::new(); - expand + let expand = Expand::new() .target_exe(&self.exe) .target_options(&self.options) .input_corpus(&corpus_dir)