you get a pre_fork, and you get a post_fork, forks for everyone

Author: domenukk

libafl/src/bolts/llmp.rs

@@ -1707,7 +1707,7 @@ where
         // TODO: handle broker_ids properly/at all.
         let map_description = Self::b2b_thread_on(
             stream,
-            &self.shmem_provider,
+            &mut self.shmem_provider,
             self.llmp_clients.len() as ClientId,
             &self.llmp_out.out_maps.first().unwrap().shmem.description(),
         )?;
@@ -1858,11 +1858,13 @@ where
     #[allow(clippy::let_and_return)]
     fn b2b_thread_on(
         mut stream: TcpStream,
-        shmem_provider: &SP,
+        shmem_provider: &mut SP,
         b2b_client_id: ClientId,
         broker_map_description: &ShMemDescription,
     ) -> Result<ShMemDescription, Error> {
         let broker_map_description = *broker_map_description;
+
+        shmem_provider.pre_fork()?;
         let mut shmem_provider_clone = shmem_provider.clone();
 
         // A channel to get the new "client's" sharedmap id from
@@ -1963,6 +1965,8 @@ where
             }
         });
 
+        shmem_provider.post_fork(false)?;
+
         let ret = recv.recv().map_err(|_| {
             Error::Unknown("Error launching background thread for b2b communcation".to_string())
         });
@@ -1980,7 +1984,7 @@ where
         request: &TcpRequest,
         current_client_id: &mut u32,
         sender: &mut LlmpSender<SP>,
-        shmem_provider: &SP,
+        shmem_provider: &mut SP,
         broker_map_description: &ShMemDescription,
     ) {
         match request {
@@ -2060,6 +2064,7 @@ where
         let tcp_out_map_description = tcp_out_map.shmem.description();
         self.register_client(tcp_out_map);
 
+        self.shmem_provider.pre_fork()?;
         let mut shmem_provider_clone = self.shmem_provider.clone();
 
         let ret = thread::spawn(move || {
@@ -2116,7 +2121,7 @@ where
                             &req,
                             &mut current_client_id,
                             &mut tcp_incoming_sender,
-                            &shmem_provider_clone,
+                            &mut shmem_provider_clone,
                             &broker_map_description,
                         );
                     }

libafl/src/bolts/shmem.rs

@@ -196,13 +196,15 @@ pub trait ShMemProvider: Send + Clone + Default + Debug {
 
     /// This method should be called before a fork or a thread creation event, allowing the [`ShMemProvider`] to
     /// get ready for a potential reset of thread specific info, and for potential reconnects.
+    /// Make sure to call [`Self::post_fork()`] after threading!
     fn pre_fork(&mut self) -> Result<(), Error> {
         // do nothing
         Ok(())
     }
 
     /// This method should be called after a fork or after cloning/a thread creation event, allowing the [`ShMemProvider`] to
     /// reset thread specific info, and potentially reconnect.
+    /// Make sure to call [`Self::pre_fork()`] before threading!
     fn post_fork(&mut self, _is_child: bool) -> Result<(), Error> {
         // do nothing
         Ok(())
@@ -356,7 +358,7 @@ where
                 Ok(())
             }
             None => Err(Error::IllegalState(
-                "Unexpected `None` Pipe in RcShMemProvider!".to_string(),
+                "Unexpected `None` Pipe in RcShMemProvider! Missing post_fork()?".to_string(),
             )),
         }
     }
@@ -378,7 +380,7 @@ where
                 }
             }
             None => Err(Error::IllegalState(
-                "Unexpected `None` Pipe in RcShMemProvider!".to_string(),
+                "Unexpected `None` Pipe in RcShMemProvider! Missing post_fork()?".to_string(),
             )),
         }
     }

libafl/src/events/llmp.rs

@@ -589,9 +589,9 @@ where
         .launch()
 }
 
-/// Provides a builder which can be used to build a restarting manager, which is a combination of a
-/// restarter and runner, that can be used on systems both with and without `fork` support. The
-/// restarter will start a nre process each time the child crashes or timesout.
+/// Provides a `builder` which can be used to build a [`RestartingMgr`], which is a combination of a
+/// `restarter` and `runner`, that can be used on systems both with and without `fork` support. The
+/// `restarter` will start a new process each time the child crashes or times out.
 #[cfg(feature = "std")]
 #[allow(clippy::default_trait_access)]
 #[derive(Builder, Debug)]
@@ -642,11 +642,12 @@ where
         )?;
 
         // We start ourself as child process to actually fuzz
-        let (sender, mut receiver, mut new_shmem_provider, core_id) = if std::env::var(
+        let (sender, mut receiver, new_shmem_provider, core_id) = if std::env::var(
             _ENV_FUZZER_SENDER,
         )
         .is_err()
         {
+            // We get here if we are on Unix, or we are a broker on Windows.
             let core_id = if mgr.is_broker() {
                 match self.kind {
                     ManagerKind::Broker | ManagerKind::Any => {
@@ -709,12 +710,19 @@ where
             loop {
                 dbg!("Spawning next client (id {})", ctr);
 
-                // On Unix, we fork (todo: measure if that is actually faster.)
+                // On Unix, we fork
                 #[cfg(unix)]
-                let child_status = match unsafe { fork() }? {
-                    ForkResult::Parent(handle) => handle.status(),
-                    ForkResult::Child => {
-                        break (sender, receiver, self.shmem_provider.clone(), core_id)
+                let child_status = {
+                    self.shmem_provider.pre_fork()?;
+                    match unsafe { fork() }? {
+                        ForkResult::Parent(handle) => {
+                            self.shmem_provider.post_fork(false)?;
+                            handle.status()
+                        }
+                        ForkResult::Child => {
+                            self.shmem_provider.post_fork(true)?;
+                            break (sender, receiver, self.shmem_provider.clone(), core_id);
+                        }
                     }
                 };
 
@@ -739,9 +747,9 @@ where
                 ctr = ctr.wrapping_add(1);
             }
         } else {
-            // We are the newly started fuzzing instance, first, connect to our own restore map.
+            // We are the newly started fuzzing instance (i.e. on Windows), first, connect to our own restore map.
+            // We get here *only on Windows*, if we were started by a restarting fuzzer.
             // A sender and a receiver for single communication
-            self.shmem_provider.post_fork(true)?;
             (
                 LlmpSender::on_existing_from_env(self.shmem_provider.clone(), _ENV_FUZZER_SENDER)?,
                 LlmpReceiver::on_existing_from_env(
@@ -753,8 +761,6 @@ where
             )
         };
 
-        new_shmem_provider.post_fork(false)?;
-
         if let Some(core_id) = core_id {
             core_affinity::set_for_current(core_id);
         }

libafl/src/utils.rs

@@ -658,6 +658,7 @@ where
         //spawn clients
         for (id, bind_to) in core_ids.iter().enumerate().take(num_cores) {
             if self.cores.iter().any(|&x| x == id) {
+                self.shmem_provider.pre_fork()?;
                 match unsafe { fork() }? {
                     ForkResult::Parent(child) => {
                         self.shmem_provider.post_fork(false)?;
@@ -698,27 +699,15 @@ where
         #[cfg(feature = "std")]
         println!("I am broker!!.");
 
-        if let Some(remote_broker_addr) = self.remote_broker_addr {
-            RestartingMgrBuilder::<I, S, SP, ST>::default()
-                .shmem_provider(self.shmem_provider.clone())
-                .stats(self.stats.clone())
-                .broker_port(self.broker_port)
-                .kind(ManagerKind::ConnectedBroker {
-                    connect_to: remote_broker_addr,
-                })
-                .build()
-                .unwrap()
-                .launch()?;
-        } else {
-            RestartingMgrBuilder::<I, S, SP, ST>::default()
-                .shmem_provider(self.shmem_provider.clone())
-                .stats(self.stats.clone())
-                .broker_port(self.broker_port)
-                .kind(ManagerKind::Broker)
-                .build()
-                .unwrap()
-                .launch()?;
-        }
+        RestartingMgrBuilder::<I, S, SP, ST>::default()
+            .shmem_provider(self.shmem_provider.clone())
+            .stats(self.stats.clone())
+            .broker_port(self.broker_port)
+            .kind(ManagerKind::Broker)
+            .remote_broker_addr(self.remote_broker_addr)
+            .build()
+            .unwrap()
+            .launch()?;
 
         //broker exited. kill all clients.
         for handle in &handles {