Use Iterators as Generator (#1101)

Also, remove seemingly-unused generate_dummy method to enable this instance.

Also, add an adapter that creates an Iterator from a Generator.

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>

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<I, S>
 where
@@ -28,9 +26,65 @@ where
 {
     /// Generate a new input
     fn generate(&mut self, state: &mut S) -> Result<I, Error>;
+}
 
-    /// 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<T, I, S> Generator<I, S> for T
+where
+    T: Iterator<Item = I>,
+    I: Input,
+{
+    fn generate(&mut self, _state: &mut S) -> Result<I, Error> {
+        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<I, S>,
+{
+    gen: G,
+    state: &'a mut S,
+    phantom: PhantomData<I>,
+}
+
+impl<'a, I, S, G> GeneratorIter<'a, I, S, G>
+where
+    I: Input,
+    G: Generator<I, S>,
+{
+    /// 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<I, S>,
+{
+    type Item = I;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        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<S> RandBytesGenerator<S>
@@ -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<S> RandPrintablesGenerator<S>
@@ -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<Vec<u8>> {
-                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<BytesInput, Error> {
             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

libafl/src/generators/nautilus.rs

@@ -71,10 +71,6 @@ impl<'a, S> Generator<NautilusInput, S> 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> {