From a092aed5385bd4efa4685a40c7e7302bcbf750c5 Mon Sep 17 00:00:00 2001
From: Fabian Freyer <fabian.freyer@physik.tu-berlin.de>
Date: Thu, 21 Sep 2023 01:03:56 +0200
Subject: [PATCH] libafl_frida: Allow setting path for DrCovRuntime (#1536)

---
 libafl_frida/src/drcov_rt.rs | 25 +++++++++++++++++++------
 1 file changed, 19 insertions(+), 6 deletions(-)

diff --git a/libafl_frida/src/drcov_rt.rs b/libafl_frida/src/drcov_rt.rs
index d5d0d63593..143cfdb40e 100644
--- a/libafl_frida/src/drcov_rt.rs
+++ b/libafl_frida/src/drcov_rt.rs
@@ -2,6 +2,7 @@
 use std::{
     collections::HashMap,
     hash::{BuildHasher, Hasher},
+    path::{Path, PathBuf},
     rc::Rc,
 };
 
@@ -25,6 +26,7 @@ pub struct DrCovRuntime {
     /// The memory ragnes of this target
     ranges: RangeMap<usize, (u16, String)>,
     stalked_addresses: HashMap<usize, usize>,
+    coverage_directory: PathBuf,
 }
 
 impl FridaRuntime for DrCovRuntime {
@@ -36,7 +38,7 @@ impl FridaRuntime for DrCovRuntime {
         _module_map: &Rc<ModuleMap>,
     ) {
         self.ranges = ranges.clone();
-        std::fs::create_dir_all("./coverage")
+        std::fs::create_dir_all(&self.coverage_directory)
             .expect("failed to create directory for coverage files");
     }
 
@@ -51,7 +53,9 @@ impl FridaRuntime for DrCovRuntime {
         let mut hasher = RandomState::with_seeds(0, 0, 0, 0).build_hasher();
         hasher.write(input.target_bytes().as_slice());
 
-        let filename = format!("./coverage/{:016x}.drcov", hasher.finish(),);
+        let filename = self
+            .coverage_directory
+            .join(format!("{:016x}.drcov", hasher.finish(),));
         DrCovWriter::new(&self.ranges).write(filename, &self.drcov_basic_blocks)?;
         self.drcov_basic_blocks.clear();
 
@@ -63,10 +67,14 @@ impl DrCovRuntime {
     /// Creates a new [`DrCovRuntime`]
     #[must_use]
     pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Create a new [`DrCovRuntime`] that writes coverage to the specified directory
+    pub fn with_path<P: AsRef<Path>>(path: P) -> Self {
         Self {
-            drcov_basic_blocks: vec![],
-            ranges: RangeMap::new(),
-            stalked_addresses: HashMap::new(),
+            coverage_directory: path.as_ref().into(),
+            ..Self::default()
         }
     }
 
@@ -88,6 +96,11 @@ impl DrCovRuntime {
 
 impl Default for DrCovRuntime {
     fn default() -> Self {
-        Self::new()
+        Self {
+            drcov_basic_blocks: vec![],
+            ranges: RangeMap::new(),
+            stalked_addresses: HashMap::new(),
+            coverage_directory: PathBuf::from("./coverage"),
+        }
     }
 }