Change centralized launcher signature (AFLplusplus#2094)

* poc * ai suggestion * rename this * aaaa * fmt * simplify * delete blob * ignore * fixup? * some progress on cow-ification * some more * clippy fixes, finalise tests * whoops, missed a spot * no std compat * api change: Named now requires alloc feature * doc fix * missed a spot * additional fixes * libfuzzer fixes * fix tutorial * fix * add * aa * fix tutorial * fix * Rename * fix * aa * fmt * aa * aa * another closure * clp * fix stuff --------- Co-authored-by: Addison Crump <[email protected]>
cube0x8 · Apr 30, 2024 · 2f7c19e · 2f7c19e
1 parent b231803
commit 2f7c19e
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 9 deletions.
diff --git a/fuzzers/libfuzzer_libpng_centralized/Makefile.toml b/fuzzers/libfuzzer_libpng_centralized/Makefile.toml
@@ -84,7 +84,7 @@ windows_alias = "unsupported"
 [tasks.run_unix]
 script_runner = "@shell"
 script='''
-./${FUZZER_NAME} --cores 0 --input ./corpus
+./${FUZZER_NAME} --cores 0-1 --input ./corpus
 '''
 dependencies = [ "fuzzer" ]
 
@@ -98,7 +98,7 @@ windows_alias = "unsupported"
 script_runner = "@shell"
 script='''
 rm -rf libafl_unix_shmem_server || true
-timeout 31s ./${FUZZER_NAME} --cores 0 --input ./corpus 2>/dev/null | tee fuzz_stdout.log || true
+timeout 31s ./${FUZZER_NAME} --cores 0-1 --input ./corpus 2>/dev/null | tee fuzz_stdout.log || true
 if grep -qa "corpus: 30" fuzz_stdout.log; then
     echo "Fuzzer is working"
 else

diff --git a/fuzzers/libfuzzer_libpng_centralized/src/lib.rs b/fuzzers/libfuzzer_libpng_centralized/src/lib.rs
@@ -8,7 +8,7 @@ use std::{env, net::SocketAddr, path::PathBuf};
 use clap::{self, Parser};
 use libafl::{
     corpus::{Corpus, InMemoryCorpus, OnDiskCorpus},
-    events::{launcher::CentralizedLauncher, EventConfig},
+    events::{centralized::CentralizedEventManager, launcher::CentralizedLauncher, EventConfig},
     executors::{inprocess::InProcessExecutor, ExitKind},
     feedback_or, feedback_or_fast,
     feedbacks::{CrashFeedback, MaxMapFeedback, TimeFeedback, TimeoutFeedback},
@@ -135,7 +135,9 @@ pub extern "C" fn libafl_main() {
 
     let monitor = MultiMonitor::new(|s| println!("{s}"));
 
-    let mut run_client = |state: Option<_>, mut mgr, _core_id: CoreId| {
+    let mut run_client = |state: Option<_>,
+                          mut mgr: CentralizedEventManager<_, _>,
+                          _core_id: CoreId| {
         // Create an observation channel using the coverage map
         let edges_observer =
             HitcountsMapObserver::new(unsafe { std_edges_map_observer("edges") }).track_indices();
@@ -241,16 +243,23 @@ pub extern "C" fn libafl_main() {
                 .unwrap_or_else(|_| panic!("Failed to load initial corpus at {:?}", &opt.input));
             println!("We imported {} inputs from disk.", state.corpus().count());
         }
-
-        fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;
+        if !mgr.is_main() {
+            fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;
+        } else {
+            let mut empty_stages = tuple_list!();
+            fuzzer.fuzz_loop(&mut empty_stages, &mut executor, &mut state, &mut mgr)?;
+        }
         Ok(())
     };
 
+    let mut main_run_client = run_client.clone(); // clone it just for borrow checker
+
     match CentralizedLauncher::builder()
         .shmem_provider(shmem_provider)
         .configuration(EventConfig::from_name("default"))
         .monitor(monitor)
         .run_client(&mut run_client)
+        .main_run_client(&mut main_run_client)
         .cores(&cores)
         .broker_port(broker_port)
         .remote_broker_addr(opt.remote_broker_addr)

diff --git a/libafl/src/events/launcher.rs b/libafl/src/events/launcher.rs
@@ -437,16 +437,22 @@ where
 }
 
 /// Provides a Launcher, which can be used to launch a fuzzing run on a specified list of cores with a single main and multiple secondary nodes
+/// This is for centralized, the 4th argument of the closure should mean if this is the main node.
 #[cfg(all(unix, feature = "std", feature = "fork"))]
 #[derive(TypedBuilder)]
 #[allow(clippy::type_complexity, missing_debug_implementations)]
-pub struct CentralizedLauncher<'a, CF, MT, S, SP>
+pub struct CentralizedLauncher<'a, CF, MF, MT, S, SP>
 where
     CF: FnOnce(
         Option<S>,
         CentralizedEventManager<LlmpRestartingEventManager<(), S, SP>, SP>, // No hooks for centralized EM
         CoreId,
     ) -> Result<(), Error>,
+    MF: FnOnce(
+        Option<S>,
+        CentralizedEventManager<LlmpRestartingEventManager<(), S, SP>, SP>, // No hooks for centralized EM
+        CoreId,
+    ) -> Result<(), Error>,
     S::Input: 'a,
     MT: Monitor,
     SP: ShMemProvider + 'static,
@@ -461,6 +467,9 @@ where
     /// The 'main' function to run for each client forked. This probably shouldn't return
     #[builder(default, setter(strip_option))]
     run_client: Option<CF>,
+    /// The 'main' function to run for the main evaluator noed
+    #[builder(default, setter(strip_option))]
+    main_run_client: Option<MF>,
     /// The broker port to use (or to attach to, in case [`Self::spawn_broker`] is `false`)
     #[builder(default = 1337_u16)]
     broker_port: u16,
@@ -506,13 +515,18 @@ where
 }
 
 #[cfg(all(unix, feature = "std", feature = "fork"))]
-impl<CF, MT, S, SP> Debug for CentralizedLauncher<'_, CF, MT, S, SP>
+impl<CF, MF, MT, S, SP> Debug for CentralizedLauncher<'_, CF, MF, MT, S, SP>
 where
     CF: FnOnce(
         Option<S>,
         CentralizedEventManager<LlmpRestartingEventManager<(), S, SP>, SP>,
         CoreId,
     ) -> Result<(), Error>,
+    MF: FnOnce(
+        Option<S>,
+        CentralizedEventManager<LlmpRestartingEventManager<(), S, SP>, SP>, // No hooks for centralized EM
+        CoreId,
+    ) -> Result<(), Error>,
     MT: Monitor + Clone,
     SP: ShMemProvider + 'static,
     S: State,
@@ -531,13 +545,18 @@ where
 }
 
 #[cfg(all(unix, feature = "std", feature = "fork"))]
-impl<'a, CF, MT, S, SP> CentralizedLauncher<'a, CF, MT, S, SP>
+impl<'a, CF, MF, MT, S, SP> CentralizedLauncher<'a, CF, MF, MT, S, SP>
 where
     CF: FnOnce(
         Option<S>,
         CentralizedEventManager<LlmpRestartingEventManager<(), S, SP>, SP>,
         CoreId,
     ) -> Result<(), Error>,
+    MF: FnOnce(
+        Option<S>,
+        CentralizedEventManager<LlmpRestartingEventManager<(), S, SP>, SP>, // No hooks for centralized EM
+        CoreId,
+    ) -> Result<(), Error>,
     MT: Monitor + Clone,
     S: State + HasExecutions,
     SP: ShMemProvider + 'static,
@@ -659,6 +678,9 @@ where
                             self.time_obs,
                         )?;
 
+                        if index == 1 {
+                            return (self.main_run_client.take().unwrap())(state, c_mgr, *bind_to);
+                        }
                         return (self.run_client.take().unwrap())(state, c_mgr, *bind_to);
                     }
                 };