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

* support bind to core for ForkserverExecutor

* fix for non-fork platforms

* clippy

* Remove redundant env
This commit is contained in:
lazymio 2025-05-17 19:55:44 +08:00 committed by GitHub
parent 4763ada075
commit ddd0930e67
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 45 additions and 2 deletions

View File

@ -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,

View File

@ -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(

View File

@ -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)]