De-duplicate OnDiskCorpus

Cargo.toml

@@ -122,6 +122,7 @@ uuid = { version = "1.10.0", features = ["serde", "v4"] }
 which = "6.0.3"
 windows = "0.58.0"
 z3 = "0.12.1"
+fs2 = "0.4.3"
 
 
 [workspace.lints.rust]

libafl/Cargo.toml

@@ -59,6 +59,7 @@ std = [
   "libafl_bolts/std",
   "typed-builder",
   "fastbloom",
+  "fs2"
 ]
 
 ## Tracks the Feedbacks and the Objectives that were interesting for a Testcase
@@ -143,6 +144,9 @@ unicode = ["libafl_bolts/alloc", "ahash/std", "serde/rc", "bitvec"]
 ## Enable multi-part input formats and mutators
 multipart_inputs = ["arrayvec", "rand_trait"]
 
+## Share objectives across nodes
+share_objectives = []
+
 #! ## LibAFL-Bolts Features
 
 ## Provide the `#[derive(SerdeAny)]` macro.
@@ -293,6 +297,7 @@ clap = { workspace = true, optional = true }
 num_enum = { workspace = true, optional = true }
 libipt = { workspace = true, optional = true }
 fastbloom = { version = "0.8.0", optional = true }
+fs2 = { workspace = true, optional = true }
 
 [lints]
 workspace = true

libafl/src/corpus/inmemory_ondisk.rs

@@ -12,8 +12,10 @@ use std::{
     io,
     io::Write,
     path::{Path, PathBuf},
+    string::ToString,
 };
 
+use fs2::FileExt;
 #[cfg(feature = "gzip")]
 use libafl_bolts::compress::GzipCompressor;
 use serde::{Deserialize, Serialize};
@@ -87,7 +89,7 @@ where
     fn add(&mut self, testcase: Testcase<I>) -> Result<CorpusId, Error> {
         let id = self.inner.add(testcase)?;
         let testcase = &mut self.get(id).unwrap().borrow_mut();
-        self.save_testcase(testcase, id)?;
+        self.save_testcase(testcase)?;
         *testcase.input_mut() = None;
         Ok(id)
     }
@@ -97,7 +99,7 @@ where
     fn add_disabled(&mut self, testcase: Testcase<I>) -> Result<CorpusId, Error> {
         let id = self.inner.add_disabled(testcase)?;
         let testcase = &mut self.get_from_all(id).unwrap().borrow_mut();
-        self.save_testcase(testcase, id)?;
+        self.save_testcase(testcase)?;
         *testcase.input_mut() = None;
         Ok(id)
     }
@@ -108,7 +110,7 @@ where
         let entry = self.inner.replace(id, testcase)?;
         self.remove_testcase(&entry)?;
         let testcase = &mut self.get(id).unwrap().borrow_mut();
-        self.save_testcase(testcase, id)?;
+        self.save_testcase(testcase)?;
         *testcase.input_mut() = None;
         Ok(entry)
     }
@@ -375,42 +377,50 @@ impl<I> InMemoryOnDiskCorpus<I> {
         }
     }
 
-    fn save_testcase(&self, testcase: &mut Testcase<I>, id: CorpusId) -> Result<(), Error>
+    fn save_testcase(&self, testcase: &mut Testcase<I>) -> Result<(), Error>
     where
         I: Input,
     {
-        let file_name_orig = testcase.filename_mut().take().unwrap_or_else(|| {
+        let file_name = testcase.filename_mut().take().unwrap_or_else(|| {
             // TODO walk entry metadata to ask for pieces of filename (e.g. :havoc in AFL)
-            testcase.input().as_ref().unwrap().generate_name(Some(id))
+            testcase.input().as_ref().unwrap().generate_name()
         });
 
-        // New testcase, we need to save it.
-        let mut file_name = file_name_orig.clone();
+        let mut ctr = String::new();
+        if self.locking {
+            let lockfile_name = format!(".{file_name}");
+            let lockfile_path = self.dir_path.join(lockfile_name);
+            let lockfile = try_create_new(&lockfile_path)?.unwrap_or(File::create(&lockfile_path)?);
 
-        let mut ctr = 2;
-        let file_name = if self.locking {
-            loop {
-                let lockfile_name = format!(".{file_name}.lafl_lock");
-                let lockfile_path = self.dir_path.join(lockfile_name);
+            lockfile.lock_exclusive()?;
 
-                if try_create_new(lockfile_path)?.is_some() {
-                    break file_name;
-                }
-
-                file_name = format!("{file_name_orig}-{ctr}");
-                ctr += 1;
+            // replace String herre too and try to reduce variables
+            ctr = fs::read_to_string(&lockfile_path)?;
+            if ctr.is_empty() {
+                ctr = String::from("1");
+            } else {
+                ctr = (ctr.parse::<u32>()? + 1).to_string();
             }
-        } else {
-            file_name
-        };
+
+            fs::write(lockfile_path, &ctr)?;
+        }
 
         if testcase.file_path().is_none() {
             *testcase.file_path_mut() = Some(self.dir_path.join(&file_name));
         }
         *testcase.filename_mut() = Some(file_name);
 
         if self.meta_format.is_some() {
-            let metafile_name = format!(".{}.metadata", testcase.filename().as_ref().unwrap());
+            let metafile_name;
+            if self.locking {
+                metafile_name = format!(
+                    ".{}_{}.metadata",
+                    testcase.filename().as_ref().unwrap(),
+                    ctr
+                );
+            } else {
+                metafile_name = format!(".{}.metadata", testcase.filename().as_ref().unwrap());
+            }
             let metafile_path = self.dir_path.join(&metafile_name);
             let mut tmpfile_path = metafile_path.clone();
             tmpfile_path.set_file_name(format!(".{metafile_name}.tmp",));
@@ -453,15 +463,26 @@ impl<I> InMemoryOnDiskCorpus<I> {
 
     fn remove_testcase(&self, testcase: &Testcase<I>) -> Result<(), Error> {
         if let Some(filename) = testcase.filename() {
+            if self.locking {
+                let lockfile_path = self.dir_path.join(format!(".{filename}"));
+                let lockfile = File::open(&lockfile_path)?;
+
+                lockfile.lock_exclusive()?;
+                let ctr = fs::read_to_string(&lockfile_path)?;
+
+                if ctr == "1" {
+                    lockfile.unlock()?;
+                    drop(fs::remove_file(lockfile_path));
+                } else {
+                    fs::write(lockfile_path, (ctr.parse::<u32>()? - 1).to_string())?;
+                    return Ok(());
+                }
+            }
+
             fs::remove_file(self.dir_path.join(filename))?;
             if self.meta_format.is_some() {
                 fs::remove_file(self.dir_path.join(format!(".{filename}.metadata")))?;
             }
-            // also try to remove the corresponding `.lafl_lock` file if it still exists
-            // (even though it shouldn't exist anymore, at this point in time)
-            drop(fs::remove_file(
-                self.dir_path.join(format!(".{filename}.lafl_lock")),
-            ));
         }
         Ok(())
     }

libafl/src/events/centralized.rs

@@ -285,15 +285,18 @@ where
         if !self.is_main {
             // secondary node
             let mut is_tc = false;
-            // Forward to main only if new tc or heartbeat
+            // Forward to main only if new tc, heartbeat, or optionally, a new objective
             let should_be_forwarded = match &mut event {
                 Event::NewTestcase { forward_id, .. } => {
                     *forward_id = Some(ClientId(self.inner.mgr_id().0 as u32));
                     is_tc = true;
                     true
                 }
-                Event::UpdateExecStats { .. } => true, // send it but this guy won't be handled. the only purpose is to keep this client alive else the broker thinks it is dead and will dc it
-                Event::Stop => true,
+                Event::UpdateExecStats { .. } | Event::Stop => true, // send UpdateExecStats but this guy won't be handled. the only purpose is to keep this client alive else the broker thinks it is dead and will dc it
+
+                #[cfg(feature = "share_objectives")]
+                Event::Objective { .. } => true,
+
                 _ => false,
             };
 
@@ -677,6 +680,12 @@ where
                     log::debug!("[{}] {} was discarded...)", process::id(), event_name);
                 }
             }
+
+            #[cfg(feature = "share_objectives")]
+            Event::Objective { .. } => {
+                log::debug!("Received new Objective");
+            }
+
             Event::Stop => {
                 state.request_stop();
             }

libafl/src/events/llmp/mgr.rs

@@ -469,6 +469,12 @@ where
                     }
                 }
             }
+
+            #[cfg(feature = "share_objectives")]
+            Event::Objective { .. } => {
+                log::debug!("Received new Objective");
+            }
+
             Event::CustomBuf { tag, buf } => {
                 for handler in &mut self.custom_buf_handlers {
                     if handler(state, &tag, &buf)? == CustomBufEventResult::Handled {

libafl/src/events/llmp/mod.rs

@@ -24,6 +24,11 @@ use crate::{
     state::{HasCorpus, HasExecutions, NopState, State, Stoppable, UsesState},
     Error, HasMetadata,
 };
+#[cfg(feature = "share_objectives")]
+use crate::{
+    corpus::Testcase,
+    state::{HasCurrentTestcase, HasSolutions},
+};
 
 /// The llmp event manager
 pub mod mgr;
@@ -284,6 +289,7 @@ where
     }
 
     // Handle arriving events in the client
+    #[cfg(not(feature = "share_objectives"))]
     fn handle_in_client<E, EM, Z>(
         &mut self,
         fuzzer: &mut Z,
@@ -340,7 +346,136 @@ where
         }
     }
 
+    #[cfg(feature = "share_objectives")]
+    fn handle_in_client<E, EM, Z>(
+        &mut self,
+        fuzzer: &mut Z,
+        executor: &mut E,
+        state: &mut S,
+        manager: &mut EM,
+        client_id: ClientId,
+        event: Event<DI>,
+    ) -> Result<(), Error>
+    where
+        E: Executor<EM, Z, State = S> + HasObservers,
+        EM: UsesState<State = S> + EventFirer,
+        S: HasSolutions + HasCurrentTestcase,
+        S::Corpus: Corpus<Input = S::Input>,
+        S::Solutions: Corpus<Input = S::Input>,
+        for<'a> E::Observers: Deserialize<'a>,
+        Z: ExecutionProcessor<EM, <S::Corpus as Corpus>::Input, E::Observers, S>
+            + EvaluatorObservers<E, EM, <S::Corpus as Corpus>::Input, S>,
+    {
+        match event {
+            Event::NewTestcase {
+                input, forward_id, ..
+            } => {
+                log::debug!("Received new Testcase to convert from {client_id:?} (forward {forward_id:?}, forward {forward_id:?})");
+
+                let Some(converter) = self.converter_back.as_mut() else {
+                    return Ok(());
+                };
+
+                let res = fuzzer.evaluate_input_with_observers(
+                    state,
+                    executor,
+                    manager,
+                    converter.convert(input)?,
+                    false,
+                )?;
+
+                if let Some(item) = res.1 {
+                    log::info!("Added received Testcase as item #{item}");
+                }
+                Ok(())
+            }
+            Event::Objective { input, .. } => {
+                log::debug!("Received new Objective");
+
+                let Some(converter) = self.converter_back.as_mut() else {
+                    return Ok(());
+                };
+
+                let converted_input = converter.convert(input)?;
+                let mut testcase = Testcase::from(converted_input);
+                testcase.set_parent_id_optional(*state.corpus().current());
+
+                if let Ok(mut tc) = state.current_testcase_mut() {
+                    tc.found_objective();
+                }
+
+                state.solutions_mut().add(testcase)?;
+                log::info!("Added received Objective to Corpus");
+
+                Ok(())
+            }
+            Event::CustomBuf { tag, buf } => {
+                for handler in &mut self.custom_buf_handlers {
+                    if handler(state, &tag, &buf)? == CustomBufEventResult::Handled {
+                        break;
+                    }
+                }
+                Ok(())
+            }
+            Event::Stop => Ok(()),
+            _ => Err(Error::unknown(format!(
+                "Received illegal message that message should not have arrived: {:?}.",
+                event.name()
+            ))),
+        }
+    }
+
+    /// Handle arriving events in the client
+    #[cfg(not(feature = "share_objectives"))]
+    pub fn process<E, EM, Z>(
+        &mut self,
+        fuzzer: &mut Z,
+        state: &mut S,
+        executor: &mut E,
+        manager: &mut EM,
+    ) -> Result<usize, Error>
+    where
+        E: Executor<EM, Z, State = S> + HasObservers,
+        EM: UsesState<State = S> + EventFirer,
+        S::Corpus: Corpus<Input = S::Input>,
+        for<'a> E::Observers: Deserialize<'a>,
+        Z: ExecutionProcessor<EM, <S::Corpus as Corpus>::Input, E::Observers, S>
+            + EvaluatorObservers<E, EM, <S::Corpus as Corpus>::Input, S>,
+    {
+        // TODO: Get around local event copy by moving handle_in_client
+        let self_id = self.llmp.sender().id();
+        let mut count = 0;
+        while let Some((client_id, tag, _flags, msg)) = self.llmp.recv_buf_with_flags()? {
+            assert!(
+                tag != _LLMP_TAG_EVENT_TO_BROKER,
+                "EVENT_TO_BROKER parcel should not have arrived in the client!"
+            );
+
+            if client_id == self_id {
+                continue;
+            }
+            #[cfg(not(feature = "llmp_compression"))]
+            let event_bytes = msg;
+            #[cfg(feature = "llmp_compression")]
+            let compressed;
+            #[cfg(feature = "llmp_compression")]
+            let event_bytes = if _flags & LLMP_FLAG_COMPRESSED == LLMP_FLAG_COMPRESSED {
+                compressed = self.compressor.decompress(msg)?;
+                &compressed
+            } else {
+                msg
+            };
+
+            let event: Event<DI> = postcard::from_bytes(event_bytes)?;
+            log::debug!("Processor received message {}", event.name_detailed());
+            self.handle_in_client(fuzzer, executor, state, manager, client_id, event)?;
+            count += 1;
+        }
+        Ok(count)
+    }
+
     /// Handle arriving events in the client
+    #[cfg(feature = "share_objectives")]
     pub fn process<E, EM, Z>(
         &mut self,
         fuzzer: &mut Z,
@@ -351,7 +486,9 @@ where
     where
         E: Executor<EM, Z, State = S> + HasObservers,
         EM: UsesState<State = S> + EventFirer,
+        S: HasSolutions + HasCurrentTestcase,
         S::Corpus: Corpus<Input = S::Input>,
+        S::Solutions: Corpus<Input = S::Input>,
         for<'a> E::Observers: Deserialize<'a>,
         Z: ExecutionProcessor<EM, <S::Corpus as Corpus>::Input, E::Observers, S>
             + EvaluatorObservers<E, EM, <S::Corpus as Corpus>::Input, S>,