use crate::{FieldLocation, Id, Node, Visitor};
use libafl::{corpus::CorpusId, inputs::ToTargetBytes, SerdeAny};
use libafl_bolts::current_time;
use libafl_bolts::AsSlice;
use serde::{Deserialize, Serialize};
use std::{
collections::{HashMap, HashSet, VecDeque},
io::ErrorKind,
path::{Path, PathBuf},
time::Duration,
u128,
};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum InputCause {
Default,
Generated,
}
#[derive(Debug, Clone, SerdeAny, Serialize, Deserialize)]
pub struct Context {
render: bool,
mutations: HashSet<MutationMetadata>,
out_dir: PathBuf,
pub type_input_map: HashMap<Id, Vec<PathBuf>>,
input_cause: InputCause,
}
impl Context {
pub fn register_input<I, TC>(
&mut self,
input: &I,
visitor: &mut Visitor,
converter: &mut TC,
is_solution: bool,
) where
TC: ToTargetBytes<I>,
I: Node,
{
let generated_fields = match &self.input_cause {
InputCause::Default => visitor.serialized(),
InputCause::Generated => {
input.__autarkie_serialized(visitor);
visitor.serialized()
}
};
let string_ty = String::__autarkie_id();
for field in generated_fields {
let (data, ty) = field;
if ty == string_ty {
visitor.register_string(crate::deserialize(&mut data.as_slice()));
}
let path = self.out_dir.join("chunks").join(ty.to_string());
match std::fs::create_dir(&path) {
Ok(_) => {}
Err(e) => {
if !matches!(e.kind(), ErrorKind::AlreadyExists) {
panic!("{:?}", e)
}
}
};
let hash = twox_hash::XxHash64::oneshot(0, &data);
let path = path.join(hash.to_string());
if !std::fs::exists(&path).unwrap() {
std::fs::write(&path, data).unwrap();
if let Some(e) = self.type_input_map.get_mut(&ty) {
e.push(path);
} else {
self.type_input_map.insert(ty, vec![path]);
}
}
}
let rendered = converter.to_target_bytes(&input);
let path = if is_solution {
self.out_dir.join("rendered_crashes")
} else {
self.out_dir.join("rendered_corpus")
};
let hash = twox_hash::XxHash64::oneshot(0, &rendered);
let path = path.join(hash.to_string());
if !std::fs::exists(&path).unwrap() {
std::fs::write(&path, rendered.as_slice()).unwrap();
}
self.input_cause = InputCause::Default;
}
pub fn generated_input(&mut self) {
self.input_cause = InputCause::Generated;
}
pub fn default_input(&mut self) {
self.input_cause = InputCause::Default;
}
pub fn add_existing_chunk(&mut self, path: PathBuf) {
let ty = path
.parent()
.unwrap()
.file_name()
.unwrap()
.to_str()
.unwrap()
.parse::<Id>()
.expect("corrupt chunk ID!");
if let Some(e) = self.type_input_map.get_mut(&ty) {
e.push(path);
} else {
self.type_input_map.insert(ty, vec![path]);
}
}
pub fn get_inputs_for_type(&self, t: &Id) -> Option<&Vec<PathBuf>> {
self.type_input_map.get(t)
}
}
impl Context {
pub fn new(out_dir: PathBuf, render: bool) -> Self {
let type_input_map = HashMap::default();
Self {
mutations: HashSet::new(),
input_cause: InputCause::Default,
out_dir,
type_input_map,
render,
}
}
pub fn add_mutation(&mut self, m: MutationMetadata) {
self.mutations.insert(m);
}
pub fn clear_mutations(&mut self) -> HashSet<MutationMetadata> {
let cloned = self.mutations.clone();
self.mutations = HashSet::new();
cloned
}
}
#[derive(
Debug,
Clone,
serde::Serialize,
serde::Deserialize,
SerdeAny,
PartialEq,
Eq,
Hash,
PartialOrd,
Ord,
)]
pub enum MutationMetadata {
SpliceFull,
SpliceSingle,
SpliceSubSplice,
SpliceAppend,
GenerateAppend,
RandomMutateSingle,
RandomMutateSubsplice,
RecursiveMinimization,
IterableMinimization,
IterablePop,
NoveltyMinimization,
Afl,
Generate,
Cmplog,
CmplogBytes,
I2S,
Random,
}