diff --git a/src/corpus/mod.rs b/src/corpus/mod.rs index b1df6b1912..2b13225a7e 100644 --- a/src/corpus/mod.rs +++ b/src/corpus/mod.rs @@ -4,6 +4,8 @@ pub use testcase::{Testcase, SimpleTestcase}; use crate::utils::Rand; use crate::AflError; +use std::path::PathBuf; + /// Corpus with all current testcases pub trait Corpus { /// Returns the number of elements @@ -24,7 +26,7 @@ pub trait Corpus { pub struct BaseCorpus<'a, RandT: Rand> { rand: &'a mut RandT, entries: Vec>, - dir_path: String, + dir_path: PathBuf, } impl Corpus for BaseCorpus<'_, RandT> { @@ -33,7 +35,13 @@ impl Corpus for BaseCorpus<'_, RandT> { self.entries.len() } - fn add(&mut self, entry: Box) { + fn add(&mut self, mut entry: Box) { + if entry.get_filename() == None { + // TODO walk entry metadatas to ask for pices of filename (e.g. :havoc in AFL) + let filename = &(String::from("id:") + &self.entries.len().to_string()); + let filename = self.dir_path.join(filename); + entry.set_filename(filename); + } self.entries.push(entry); } @@ -67,9 +75,9 @@ impl Corpus for BaseCorpus<'_, RandT> { } impl BaseCorpus<'_, RandT> { - pub fn new<'a>(rand: &'a mut RandT, dir_path: &str) -> BaseCorpus<'a, RandT> { + pub fn new<'a>(rand: &'a mut RandT, dir_path: PathBuf) -> BaseCorpus<'a, RandT> { BaseCorpus { - dir_path: dir_path.to_owned(), + dir_path: dir_path, entries: vec![], rand: rand, } @@ -118,7 +126,7 @@ impl Corpus for QueueCorpus<'_, RandT> { } impl QueueCorpus<'_, RandT> { - pub fn new<'a>(rand: &'a mut RandT, dir_path: &str) -> QueueCorpus<'a, RandT> { + pub fn new<'a>(rand: &'a mut RandT, dir_path: PathBuf) -> QueueCorpus<'a, RandT> { QueueCorpus { base: BaseCorpus::new(rand, dir_path), cycles: 0, @@ -144,16 +152,19 @@ mod tests { use crate::inputs::bytes::BytesInput; use crate::utils::Xoshiro256StarRand; + use std::path::PathBuf; + #[test] + fn test_queuecorpus() { let mut rand = Xoshiro256StarRand::new(); - let mut q = QueueCorpus::new(&mut rand, "fancy/path"); + let mut q = QueueCorpus::new(&mut rand, PathBuf::from("fancy/path")); let i = Box::new(BytesInput::new(vec![0; 4])); let mut t = Box::new(SimpleTestcase::new(i)); - t.set_filename("fancyfile".to_string()); + t.set_filename(PathBuf::from("fancyfile")); q.add(t); let filename = q.get().unwrap().get_filename().unwrap().to_owned(); assert_eq!(filename, q.get().unwrap().get_filename().unwrap().to_owned()); - assert_eq!(filename, "fancyfile".to_string()); + assert_eq!(filename, PathBuf::from("fancyfile")); } } diff --git a/src/corpus/testcase.rs b/src/corpus/testcase.rs index 42c3144485..86384f1b5d 100644 --- a/src/corpus/testcase.rs +++ b/src/corpus/testcase.rs @@ -2,6 +2,7 @@ use crate::inputs::Input; use crate::AflError; use hashbrown::HashMap; +use std::path::PathBuf; pub trait TestcaseMetadata {} @@ -13,9 +14,9 @@ pub trait Testcase { fn is_on_disk(&self) -> bool; - fn get_filename(&self) -> Option<& String>; + fn get_filename(&self) -> Option<& PathBuf>; - fn set_filename(&mut self, filename: String); + fn set_filename(&mut self, filename: PathBuf); fn get_metadatas(&mut self) -> &mut HashMap>; @@ -25,7 +26,7 @@ pub trait Testcase { pub struct SimpleTestcase { input: Option>, // is_on_disk: bool, // not needed, look at the Option - filename: Option, + filename: Option, metadatas: HashMap>, } @@ -43,11 +44,11 @@ impl Testcase for SimpleTestcase { !self.input.is_some() && self.filename.is_some() } - fn get_filename(&self) -> Option<& String> { + fn get_filename(&self) -> Option<& PathBuf> { self.filename.as_ref() } - fn set_filename(&mut self, filename: String) { + fn set_filename(&mut self, filename: PathBuf) { self.filename = Some(filename) }