Added state restorer testcase, fixed restorer (#240)

* added state restorer testcase

* fixed testcase

* no_std, clippy

* printing less often

fuzzers/baby_fuzzer/src/main.rs

@@ -1,6 +1,5 @@
 use std::path::PathBuf;
 
-use libafl::inputs::{BytesInput, HasTargetBytes};
 use libafl::{
     bolts::{current_nanos, rands::StdRand, tuples::tuple_list},
     corpus::{InMemoryCorpus, OnDiskCorpus, QueueCorpusScheduler},
@@ -9,6 +8,7 @@ use libafl::{
     feedbacks::{CrashFeedback, MapFeedbackState, MaxMapFeedback},
     fuzzer::{Fuzzer, StdFuzzer},
     generators::RandPrintablesGenerator,
+    inputs::{BytesInput, HasTargetBytes},
     mutators::scheduled::{havoc_mutations, StdScheduledMutator},
     observers::StdMapObserver,
     stages::mutational::StdMutationalStage,

libafl/src/bolts/llmp.rs

@@ -802,10 +802,20 @@ where
     /// Waits for this sender to be save to unmap.
     /// If a receiver is involved, this function should always be called.
     pub fn await_save_to_unmap_blocking(&self) {
+        #[cfg(feature = "std")]
+        let mut ctr = 0_u16;
         loop {
             if self.save_to_unmap() {
                 return;
             }
+            // We log that we're looping -> see when we're blocking.
+            #[cfg(feature = "std")]
+            {
+                ctr = ctr.wrapping_add(1);
+                if ctr == 0 {
+                    println!("Awaiting save_to_unmap_blocking");
+                }
+            }
         }
     }
 
@@ -819,6 +829,14 @@ where
         }
     }
 
+    /// For debug purposes: Mark save to unmap, even though it might not have been read by a receiver yet.
+    /// # Safety
+    /// If this method is called, the page may be unmapped before it is read by any receiver.
+    pub unsafe fn mark_save_to_unmap(&mut self) {
+        // No need to do this volatile, as we should be the same thread in this scenario.
+        (*self.out_maps.last_mut().unwrap().page_mut()).save_to_unmap = 1;
+    }
+
     /// Reattach to a vacant `out_map`.
     /// It is essential, that the receiver (or someone else) keeps a pointer to this map
     /// else reattach will get a new, empty page, from the OS, or fail.
@@ -2431,6 +2449,16 @@ where
         self.sender.save_to_unmap()
     }
 
+    /// For debug purposes: mark the client as save to unmap, even though it might not have been read.
+    ///
+    /// # Safety
+    /// This should only be called in a debug scenario.
+    /// Calling this in other contexts may lead to a premature page unmap and result in a crash in another process,
+    /// or an unexpected read from an empty page in a receiving process.
+    pub unsafe fn mark_save_to_unmap(&mut self) {
+        self.sender.mark_save_to_unmap();
+    }
+
     /// Creates a new [`LlmpClient`]
     pub fn new(
         mut shmem_provider: SP,

libafl/src/bolts/staterestore.rs

@@ -46,7 +46,10 @@ where
 
     /// Create a [`StateRrestore`] from `env` variable name
     pub fn from_env(shmem_provider: &mut SP, env_name: &str) -> Result<Self, Error> {
-        Ok(Self::new(shmem_provider.existing_from_env(env_name)?))
+        Ok(Self {
+            shmem: shmem_provider.existing_from_env(env_name)?,
+            phantom: PhantomData,
+        })
     }
 
     /// Create a new [`StateRestorer`].
@@ -83,7 +86,17 @@ where
 
             // write the filename to shmem
             let filename_buf = postcard::to_allocvec(&filename)?;
+
             let len = filename_buf.len();
+            if len > self.shmem.len() {
+                return Err(Error::IllegalState(format!(
+                    "The state restorer map is too small to fit anything, even the filename! 
+                        It needs to be at least {} bytes. 
+                        The tmpfile was written to {:?}.",
+                    len,
+                    temp_dir().join(&filename)
+                )));
+            }
 
             /*println!(
                 "Storing {} bytes to tmpfile {} (larger than map of {} bytes)",
@@ -134,16 +147,19 @@ where
         }
     }
 
+    /// The content is either the name of the tmpfile, or the serialized bytes directly, if they fit on a single page.
     fn content(&self) -> &StateShMemContent {
         #[allow(clippy::cast_ptr_alignment)] // Beginning of the page will always be aligned
         let ptr = self.shmem.map().as_ptr() as *const StateShMemContent;
         unsafe { &*(ptr) }
     }
 
+    /// Returns true, if this [`StateRestorer`] has contents.
     pub fn has_content(&self) -> bool {
         self.content().buf_len > 0
     }
 
+    /// Restores the contents saved in this [`StateRestorer`], if any are availiable.
     pub fn restore<S>(&self) -> Result<Option<S>, Error>
     where
         S: DeserializeOwned,

libafl/src/events/llmp.rs

@@ -892,7 +892,7 @@ where
         // If we're restarting, deserialize the old state.
         let (state, mut mgr) = if let Some((state, mgr_description)) = staterestorer.restore()? {
             (
-                state,
+                Some(state),
                 LlmpRestartingEventManager::new(
                     LlmpEventManager::existing_client_from_description(
                         new_shmem_provider,
@@ -925,3 +925,111 @@ where
         Ok((state, mgr))
     }
 }
+
+#[cfg(test)]
+#[cfg(feature = "std")]
+mod tests {
+    use crate::{
+        bolts::{
+            llmp::{LlmpClient, LlmpSharedMap},
+            rands::StdRand,
+            shmem::{ShMemProvider, StdShMemProvider},
+            staterestore::StateRestorer,
+            tuples::tuple_list,
+        },
+        corpus::{Corpus, InMemoryCorpus, RandCorpusScheduler, Testcase},
+        events::{llmp::_ENV_FUZZER_SENDER, LlmpEventManager},
+        executors::{ExitKind, InProcessExecutor},
+        inputs::BytesInput,
+        mutators::BitFlipMutator,
+        stages::StdMutationalStage,
+        state::StdState,
+        Fuzzer, StdFuzzer,
+    };
+    use core::sync::atomic::{compiler_fence, Ordering};
+
+    #[test]
+    fn test_mgr_state_restore() {
+        let rand = StdRand::with_seed(0);
+
+        let mut corpus = InMemoryCorpus::<BytesInput>::new();
+        let testcase = Testcase::new(vec![0; 4]);
+        corpus.add(testcase).unwrap();
+
+        let solutions = InMemoryCorpus::<BytesInput>::new();
+
+        let mut state = StdState::new(rand, corpus, solutions, tuple_list!());
+
+        let mut shmem_provider = StdShMemProvider::new().unwrap();
+
+        let mut llmp_client = LlmpClient::new(
+            shmem_provider.clone(),
+            LlmpSharedMap::new(0, shmem_provider.new_map(1024).unwrap()),
+        )
+        .unwrap();
+
+        // A little hack for CI. Don't do that in a real-world scenario.
+        unsafe {
+            llmp_client.mark_save_to_unmap();
+        }
+
+        let mut llmp_mgr =
+            LlmpEventManager::<BytesInput, (), _, _>::new(llmp_client, "fuzzer".to_string())
+                .unwrap();
+
+        let scheduler = RandCorpusScheduler::new();
+
+        let mut fuzzer = StdFuzzer::new(scheduler, (), ());
+
+        let mut harness = |_buf: &BytesInput| ExitKind::Ok;
+        let mut executor = InProcessExecutor::new(
+            &mut harness,
+            tuple_list!(),
+            &mut fuzzer,
+            &mut state,
+            &mut llmp_mgr,
+        )
+        .unwrap();
+
+        let mutator = BitFlipMutator::new();
+        let mut stages = tuple_list!(StdMutationalStage::new(mutator));
+
+        // First, create a channel from the current fuzzer to the next to store state between restarts.
+        let mut staterestorer = StateRestorer::<StdShMemProvider>::new(
+            shmem_provider.new_map(256 * 1024 * 1024).unwrap(),
+        );
+
+        staterestorer.reset();
+        staterestorer
+            .save(&(&mut state, &llmp_mgr.describe().unwrap()))
+            .unwrap();
+        assert!(staterestorer.has_content());
+
+        // Store the information to a map.
+        staterestorer.write_to_env(_ENV_FUZZER_SENDER).unwrap();
+
+        compiler_fence(Ordering::SeqCst);
+
+        let sc_cpy = StateRestorer::from_env(&mut shmem_provider, _ENV_FUZZER_SENDER).unwrap();
+        assert!(sc_cpy.has_content());
+
+        let (mut state_clone, mgr_description) = staterestorer.restore().unwrap().unwrap();
+        let mut llmp_clone = LlmpEventManager::existing_client_from_description(
+            shmem_provider,
+            &mgr_description,
+            "fuzzer".to_string(),
+        )
+        .unwrap();
+
+        if false {
+            fuzzer
+                .fuzz_one(
+                    &mut stages,
+                    &mut executor,
+                    &mut state_clone,
+                    &mut llmp_clone,
+                )
+                .unwrap();
+        }
+    }
+}