diff --git a/src/agent/onefuzz-agent/src/tasks/coverage/libfuzzer_coverage.rs b/src/agent/onefuzz-agent/src/tasks/coverage/libfuzzer_coverage.rs index ac00f1102..ec95ccd39 100644 --- a/src/agent/onefuzz-agent/src/tasks/coverage/libfuzzer_coverage.rs +++ b/src/agent/onefuzz-agent/src/tasks/coverage/libfuzzer_coverage.rs @@ -43,7 +43,7 @@ use onefuzz::{fs::list_files, libfuzzer::LibFuzzer, syncdir::SyncedDir}; use onefuzz_telemetry::{Event::coverage_data, EventData}; use reqwest::Url; use serde::Deserialize; -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap}; use std::{ ffi::OsString, path::{Path, PathBuf}, @@ -197,7 +197,7 @@ pub struct CoverageProcessor { config: Arc, pub recorder: CoverageRecorder, pub total: TotalCoverage, - pub module_totals: HashMap, + pub module_totals: BTreeMap, heartbeat_client: Option, } @@ -206,7 +206,7 @@ impl CoverageProcessor { let heartbeat_client = config.common.init_heartbeat().await?; let total = TotalCoverage::new(config.coverage.path.join(TOTAL_COVERAGE)); let recorder = CoverageRecorder::new(config.clone()); - let module_totals = HashMap::default(); + let module_totals = BTreeMap::default(); Ok(Self { config, @@ -244,8 +244,9 @@ impl CoverageProcessor { Ok(()) } - async fn collect_by_module(&mut self, path: &Path) -> Result { - let files = list_files(&path).await?; + async fn collect_by_module(&mut self, path: &Path) -> Result<()> { + let mut files = list_files(&path).await?; + files.sort(); let mut sum = Vec::new(); for file in &files { @@ -264,14 +265,25 @@ impl CoverageProcessor { .await .with_context(|| format!("unable to write combined coverage file: {:?}", combined))?; - Ok(combined.into()) + Ok(()) } pub async fn test_input(&mut self, input: &Path) -> Result<()> { info!("processing input {:?}", input); let new_coverage = self.recorder.record(input).await?; - let combined = self.collect_by_module(&new_coverage).await?; - self.total.update(&combined).await?; + self.collect_by_module(&new_coverage).await?; + self.update_total().await?; + Ok(()) + } + + async fn update_total(&mut self) -> Result<()> { + let mut total = Vec::new(); + for module_total in self.module_totals.values() { + if let Some(mut module_data) = module_total.data().await? { + total.append(&mut module_data); + } + } + self.total.write(&total).await?; Ok(()) } diff --git a/src/agent/onefuzz-agent/src/tasks/coverage/total.rs b/src/agent/onefuzz-agent/src/tasks/coverage/total.rs index b38dc3f30..879ec2733 100644 --- a/src/agent/onefuzz-agent/src/tasks/coverage/total.rs +++ b/src/agent/onefuzz-agent/src/tasks/coverage/total.rs @@ -43,22 +43,29 @@ impl TotalCoverage { &self.path } + pub async fn write(&self, data: &[u8]) -> Result<()> { + fs::write(self.path(), data).await?; + Ok(()) + } + pub async fn update_bytes(&self, new_data: &[u8]) -> Result<()> { match self.data().await { Ok(Some(mut total_data)) => { + if total_data.len() < new_data.len() { + total_data.resize_with(new_data.len(), || 0); + } for (i, b) in new_data.iter().enumerate() { if *b > 0 { total_data[i] = 1; } } - - fs::write(self.path(), total_data).await?; + self.write(&total_data).await?; } Ok(None) => { // Base case: we don't yet have any total coverage. Promote the // new coverage to being our total coverage. info!("initializing total coverage map {}", self.path().display()); - fs::write(self.path(), new_data).await?; + self.write(new_data).await?; } Err(err) => { // Couldn't read total for some other reason, so this is a real error. @@ -69,11 +76,6 @@ impl TotalCoverage { Ok(()) } - pub async fn update(&self, new: impl AsRef) -> Result<()> { - let new_data = fs::read(new).await?; - self.update_bytes(&new_data).await - } - pub async fn info(&self) -> Result { let data = self .data()