use crate::{fuzzer::context::Context, Node, Visitor};
use libafl::{
corpus::Corpus,
events::EventFirer,
executors::Executor,
stages::{Restartable, Stage},
state::{HasCorpus, HasCurrentTestcase},
Evaluator, HasMetadata,
};
use serde::Serialize;
use std::{cell::RefCell, marker::PhantomData, rc::Rc};
#[derive(Debug)]
pub struct GenerateStage<I> {
visitor: Rc<RefCell<Visitor>>,
phantom: PhantomData<I>,
}
impl<I> GenerateStage<I> {
pub fn new(visitor: Rc<RefCell<Visitor>>) -> Self {
Self {
visitor,
phantom: PhantomData,
}
}
}
impl<E, EM, Z, S, I> Stage<E, EM, S, Z> for GenerateStage<I>
where
I: Node + Serialize,
S: HasCurrentTestcase<I> + HasCorpus<I> + HasMetadata,
E: Executor<EM, I, S, Z>,
EM: EventFirer<I, S>,
Z: Evaluator<E, EM, I, S>,
{
fn perform(
&mut self,
fuzzer: &mut Z,
executor: &mut E,
state: &mut S,
manager: &mut EM,
) -> Result<(), libafl::Error> {
let mut metadata = state.metadata_mut::<Context>()?;
metadata.generated_input();
let Some(generated) = generate(&mut self.visitor.borrow_mut()) else {
metadata.default_input();
return Ok(());
};
metadata.add_mutation(crate::fuzzer::context::MutationMetadata::Generate);
fuzzer.evaluate_input(state, executor, manager, &generated)?;
let mut metadata = state.metadata_mut::<Context>()?;
metadata.default_input();
Ok(())
}
}
pub fn generate<I>(visitor: &mut Visitor) -> Option<I>
where
I: Node,
{
I::__autarkie_generate(visitor, &mut visitor.generate_depth(), 0, None)
}
impl<I, S> Restartable<S> for GenerateStage<I> {
fn should_restart(&mut self, state: &mut S) -> Result<bool, libafl::Error> {
Ok(true)
}
fn clear_progress(&mut self, state: &mut S) -> Result<(), libafl::Error> {
Ok(())
}
}