diff --git a/libafl/src/generators/gramatron.rs b/libafl/src/generators/gramatron.rs index 0f7c21e1a8..07ac077c8a 100644 --- a/libafl/src/generators/gramatron.rs +++ b/libafl/src/generators/gramatron.rs @@ -51,10 +51,6 @@ where self.append_generated_terminals(&mut input, state); Ok(input) } - - fn generate_dummy(&self, _state: &mut S) -> GramatronInput { - GramatronInput::new(vec![]) - } } impl<'a, S> GramatronGenerator<'a, S> diff --git a/libafl/src/generators/mod.rs b/libafl/src/generators/mod.rs index 6a7dbe1c13..6e22a8c22f 100644 --- a/libafl/src/generators/mod.rs +++ b/libafl/src/generators/mod.rs @@ -1,13 +1,14 @@ //! Generators may generate bytes or, in general, data, for inputs. +use alloc::string::String; use alloc::vec::Vec; -use core::{cmp::min, marker::PhantomData}; +use core::marker::PhantomData; use crate::{ bolts::rands::Rand, inputs::{bytes::BytesInput, Input}, state::HasRand, - Error, + Error, ErrorBacktrace, }; pub mod gramatron; @@ -18,9 +19,6 @@ pub mod nautilus; #[cfg(feature = "nautilus")] pub use nautilus::*; -/// The maximum size of dummy bytes generated by _dummy generator methods -const DUMMY_BYTES_MAX: usize = 64; - /// Generators can generate ranges of bytes. pub trait Generator where @@ -28,9 +26,65 @@ where { /// Generate a new input fn generate(&mut self, state: &mut S) -> Result; +} - /// Generate a new dummy input - fn generate_dummy(&self, state: &mut S) -> I; +/// Iterators may be used as generators. +/// +/// `generate` throws a [`Error::Empty`] if an input is requested but +/// [`Iterator::next`] returns `None`. +impl Generator for T +where + T: Iterator, + I: Input, +{ + fn generate(&mut self, _state: &mut S) -> Result { + match self.next() { + Some(i) => Ok(i), + None => Err(Error::Empty( + String::from("No more items in iterator when generating inputs"), + ErrorBacktrace::new(), + )), + } + } +} + +/// An [`Iterator`] built from a [`Generator`]. +#[derive(Debug)] +pub struct GeneratorIter<'a, I, S, G> +where + I: Input, + G: Generator, +{ + gen: G, + state: &'a mut S, + phantom: PhantomData, +} + +impl<'a, I, S, G> GeneratorIter<'a, I, S, G> +where + I: Input, + G: Generator, +{ + /// Create a new [`GeneratorIter`] + pub fn new(gen: G, state: &'a mut S) -> Self { + Self { + gen, + state, + phantom: PhantomData::default(), + } + } +} + +impl<'a, I, S, G> Iterator for GeneratorIter<'a, I, S, G> +where + I: Input, + G: Generator, +{ + type Item = I; + + fn next(&mut self) -> Option { + self.gen.generate(self.state).ok() + } } #[derive(Clone, Debug)] @@ -57,12 +111,6 @@ where .collect(); Ok(BytesInput::new(random_bytes)) } - - /// Generates up to `DUMMY_BYTES_MAX` non-random dummy bytes (0) - fn generate_dummy(&self, _state: &mut S) -> BytesInput { - let size = min(self.max_size, DUMMY_BYTES_MAX); - BytesInput::new(vec![0; size]) - } } impl RandBytesGenerator @@ -104,12 +152,6 @@ where .collect(); Ok(BytesInput::new(random_bytes)) } - - /// Generates up to `DUMMY_BYTES_MAX` non-random dummy bytes (0) - fn generate_dummy(&self, _state: &mut S) -> BytesInput { - let size = min(self.max_size, DUMMY_BYTES_MAX); - BytesInput::new(vec![0_u8; size]) - } } impl RandPrintablesGenerator @@ -163,16 +205,6 @@ pub mod pybind { .unwrap(); Ok(BytesInput::new(bytes)) } - - fn generate_dummy(&self, state: &mut PythonStdState) -> BytesInput { - let bytes = Python::with_gil(|py| -> PyResult> { - self.inner - .call_method1(py, "generate_dummy", (PythonStdStateWrapper::wrap(state),))? - .extract(py) - }) - .unwrap(); - BytesInput::new(bytes) - } } #[pyclass(unsendable, name = "RandBytesGenerator")] @@ -249,20 +281,6 @@ pub mod pybind { wrapper: PythonGeneratorWrapper, } - macro_rules! unwrap_me { - ($wrapper:expr, $name:ident, $body:block) => { - crate::unwrap_me_body!($wrapper, $name, $body, PythonGeneratorWrapper, - { RandBytes, RandPrintables }, - { - Python(py_wrapper) => { - let $name = py_wrapper; - $body - } - } - ) - }; - } - macro_rules! unwrap_me_mut { ($wrapper:expr, $name:ident, $body:block) => { crate::unwrap_me_mut_body!($wrapper, $name, $body, PythonGeneratorWrapper, @@ -314,10 +332,6 @@ pub mod pybind { fn generate(&mut self, state: &mut PythonStdState) -> Result { unwrap_me_mut!(self.wrapper, g, { g.generate(state) }) } - - fn generate_dummy(&self, state: &mut PythonStdState) -> BytesInput { - unwrap_me!(self.wrapper, g, { g.generate_dummy(state) }) - } } /// Register the classes to the python module diff --git a/libafl/src/generators/nautilus.rs b/libafl/src/generators/nautilus.rs index 65da432c4b..fb842000d5 100644 --- a/libafl/src/generators/nautilus.rs +++ b/libafl/src/generators/nautilus.rs @@ -71,10 +71,6 @@ impl<'a, S> Generator for NautilusGenerator<'a> { self.generate_from_nonterminal(&mut input, nonterm, len); Ok(input) } - - fn generate_dummy(&self, _state: &mut S) -> NautilusInput { - NautilusInput::empty() - } } impl<'a> NautilusGenerator<'a> {