Support binding to a single core for ForkserverExecutor (#3236)

* support bind to core for ForkserverExecutor

* fix for non-fork platforms

* clippy

* Remove redundant env

Author: wtdcode

libafl/src/executors/command.rs

@@ -653,6 +653,17 @@ impl CommandExecutorBuilder {
             command.stderr(Stdio::piped());
         }
 
+        if let Some(core) = self.child_env_inner.core {
+            #[cfg(feature = "fork")]
+            command.bind(core);
+
+            #[cfg(not(feature = "fork"))]
+            return Err(Error::illegal_argument(format!(
+                "Your host doesn't support fork and thus libafl can not bind to core {:?} right after children get spawned",
+                core
+            )));
+        }
+
         let configurator = StdCommandConfigurator {
             debug_child: self.child_env_inner.debug_child,
             stdout_cap,

libafl/src/executors/forkserver.rs

@@ -22,6 +22,7 @@ use std::{
 use libafl_bolts::tuples::{Handle, Handled};
 use libafl_bolts::{
     AsSlice, AsSliceMut, InputLocation, StdTargetArgs, StdTargetArgsInner, Truncate,
+    core_affinity::CoreId,
     fs::{InputFile, get_unique_std_input_file},
     os::{dup2, pipes::Pipe},
     shmem::{ShMem, ShMemProvider, UnixShMem, UnixShMemProvider},
@@ -179,6 +180,8 @@ pub trait ConfigTarget {
     ) -> &mut Self;
     /// dup2 the specific fd, used for stdio
     fn setdup2(&mut self, old_fd: RawFd, new_fd: RawFd) -> &mut Self;
+    /// Bind children to a single core
+    fn bind(&mut self, core: CoreId) -> &mut Self;
 }
 
 impl ConfigTarget for Command {
@@ -281,6 +284,19 @@ impl ConfigTarget for Command {
         // This calls our non-shady function from above.
         unsafe { self.pre_exec(func) }
     }
+
+    fn bind(&mut self, core: CoreId) -> &mut Self {
+        let func = move || {
+            if let Err(e) = core.set_affinity_forced() {
+                return Err(io::Error::other(e));
+            }
+
+            Ok(())
+        };
+        // # Safety
+        // This calls our non-shady function from above.
+        unsafe { self.pre_exec(func) }
+    }
 }
 
 /// The [`Forkserver`] is communication channel with a child process that forks on request of the fuzzer.
@@ -365,6 +381,7 @@ impl Forkserver {
         stdout_memfd: Option<RawFd>,
         stderr_memfd: Option<RawFd>,
         cwd: Option<PathBuf>,
+        core: Option<CoreId>,
     ) -> Result<Self, Error> {
         let Some(coverage_map_size) = coverage_map_size else {
             return Err(Error::unknown(
@@ -421,6 +438,10 @@ impl Forkserver {
             command.stderr(Stdio::null());
         }
 
+        if let Some(core) = core {
+            command.bind(core);
+        }
+
         command.env(AFL_MAP_SIZE_ENV_VAR, format!("{coverage_map_size}"));
 
         // Persistent, deferred forkserver
@@ -1056,6 +1077,7 @@ where
                         .expect("only memory fd backend is allowed for forkserver executor")
                 }),
                 self.child_env_inner.current_directory.clone(),
+                self.child_env_inner.core,
             )?,
             None => {
                 return Err(Error::illegal_argument(

libafl/src/executors/mod.rs

@@ -16,9 +16,9 @@ pub use inprocess::InProcessExecutor;
 pub use inprocess_fork::InProcessForkExecutor;
 #[cfg(unix)]
 use libafl_bolts::os::unix_signals::Signal;
-#[cfg(feature = "std")]
-use libafl_bolts::tuples::Handle;
 use libafl_bolts::tuples::RefIndexable;
+#[cfg(feature = "std")]
+use libafl_bolts::{core_affinity::CoreId, tuples::Handle};
 use serde::{Deserialize, Serialize};
 pub use shadow::ShadowExecutor;
 pub use with_observers::WithObservers;
@@ -247,6 +247,8 @@ pub struct StdChildArgsInner {
     pub current_directory: Option<PathBuf>,
     /// Whether debug child by inheriting stdout/stderr
     pub debug_child: bool,
+    /// Core to bind for the children
+    pub core: Option<CoreId>,
 }
 
 #[cfg(feature = "std")]
@@ -258,6 +260,7 @@ impl Default for StdChildArgsInner {
             stdout_observer: None,
             current_directory: None,
             debug_child: false,
+            core: None,
         }
     }
 }
@@ -312,6 +315,13 @@ pub trait StdChildArgs: Sized {
         self.inner_mut().debug_child = debug_child;
         self
     }
+
+    #[must_use]
+    /// Set the core to bind for the children
+    fn core(mut self, core: CoreId) -> Self {
+        self.inner_mut().core = Some(core);
+        self
+    }
 }
 
 #[cfg(test)]