diff --git a/libafl/src/executors/command.rs b/libafl/src/executors/command.rs index b74946a0ca..22978264a5 100644 --- a/libafl/src/executors/command.rs +++ b/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, diff --git a/libafl/src/executors/forkserver.rs b/libafl/src/executors/forkserver.rs index cda34367a8..54eeb2bfe4 100644 --- a/libafl/src/executors/forkserver.rs +++ b/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, stderr_memfd: Option, cwd: Option, + core: Option, ) -> Result { 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( diff --git a/libafl/src/executors/mod.rs b/libafl/src/executors/mod.rs index eb076881cc..ee733005e9 100644 --- a/libafl/src/executors/mod.rs +++ b/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, /// Whether debug child by inheriting stdout/stderr pub debug_child: bool, + /// Core to bind for the children + pub core: Option, } #[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)]