Skip to content

Commit

Permalink
Add memfd shmem backend (#2647)
Browse files Browse the repository at this point in the history
  • Loading branch information
bernhl authored Nov 6, 2024
1 parent 36a24ab commit cfe1240
Show file tree
Hide file tree
Showing 2 changed files with 179 additions and 0 deletions.
1 change: 1 addition & 0 deletions libafl_bolts/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -149,6 +149,7 @@ miniz_oxide = { version = "0.8.0", optional = true }
hostname = { version = "0.4.0", optional = true } # Is there really no gethostname in the stdlib?
rand_core = { version = "0.6.4", optional = true }
nix = { version = "0.29.0", optional = true, default-features = false, features = [
"fs",
"signal",
"socket",
"poll",
Expand Down
178 changes: 178 additions & 0 deletions libafl_bolts/src/shmem.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1340,6 +1340,184 @@ pub mod unix_shmem {
}
}
}

/// Module containing `memfd` shared memory support, usable on Linux and Android.
#[cfg(all(unix, feature = "std", not(target_vendor = "apple")))]
pub mod memfd {
use alloc::string::ToString;
use core::{
ops::{Deref, DerefMut},
ptr, slice,
};
use std::{ffi::CString, os::fd::IntoRawFd};

use libc::{
c_void, close, fstat, ftruncate, mmap, munmap, MAP_SHARED, PROT_READ, PROT_WRITE,
};
use nix::sys::memfd::{memfd_create, MemFdCreateFlag};

use crate::{
shmem::{ShMem, ShMemId, ShMemProvider},
Error,
};

/// An memfd based impl for linux/android
#[cfg(unix)]
#[derive(Clone, Debug)]
pub struct MemfdShMem {
id: ShMemId,
map: *mut u8,
map_size: usize,
}

impl MemfdShMem {
/// Create a new shared memory mapping, using shmget/shmat
pub fn new(map_size: usize) -> Result<Self, Error> {
unsafe {
let c_str = CString::new("libAFL").unwrap();
let Ok(fd) = memfd_create(&c_str, MemFdCreateFlag::empty()) else {
return Err(Error::last_os_error("Failed to create memfd".to_string()));
};
let fd = fd.into_raw_fd();

#[allow(clippy::cast_possible_wrap)]
if ftruncate(fd, map_size as i64) == -1 {
close(fd);
return Err(Error::last_os_error(format!(
"Failed to ftruncate memfd to {map_size}"
)));
}
let map = mmap(
ptr::null_mut(),
map_size,
PROT_READ | PROT_WRITE,
MAP_SHARED,
fd,
0,
);
if map == usize::MAX as *mut c_void {
close(fd);
return Err(Error::unknown(
"Failed to map the memfd mapping".to_string(),
));
}
Ok(Self {
id: ShMemId::from_int(fd),
map: map as *mut u8,
map_size,
})
}
}

fn shmem_from_id_and_size(id: ShMemId, map_size: usize) -> Result<Self, Error> {
let fd = i32::from(id);
unsafe {
let mut stat = std::mem::zeroed();
if fstat(fd, &mut stat) == -1 {
return Err(Error::unknown(
"Failed to map the memfd mapping".to_string(),
));
}
#[allow(clippy::cast_sign_loss)]
if stat.st_size as usize != map_size {
return Err(Error::unknown(
"The mapping's size differs from the requested size".to_string(),
));
}
let map = mmap(
ptr::null_mut(),
map_size,
PROT_READ | PROT_WRITE,
MAP_SHARED,
fd,
0,
);
if map == usize::MAX as *mut c_void {
return Err(Error::last_os_error(format!(
"mmap() failed for map with fd {fd:?}"
)));
}
Ok(Self {
id: ShMemId::from_int(fd),
map: map as *mut u8,
map_size,
})
}
}
}

#[cfg(unix)]
impl ShMem for MemfdShMem {
fn id(&self) -> ShMemId {
self.id
}
}

impl Deref for MemfdShMem {
type Target = [u8];

fn deref(&self) -> &[u8] {
unsafe { slice::from_raw_parts(self.map, self.map_size) }
}
}

impl DerefMut for MemfdShMem {
fn deref_mut(&mut self) -> &mut [u8] {
unsafe { slice::from_raw_parts_mut(self.map, self.map_size) }
}
}

/// [`Drop`] implementation for [`MemfdShMem`], which cleans up the mapping.
#[cfg(unix)]
impl Drop for MemfdShMem {
#[allow(trivial_numeric_casts)]
fn drop(&mut self) {
let fd = i32::from(self.id);

unsafe {
munmap(self.map as *mut _, self.map_size);
close(fd);
}
}
}

/// A [`ShMemProvider`] which uses memfd to provide shared memory mappings.
#[cfg(unix)]
#[derive(Clone, Debug)]
pub struct MemfdShMemProvider {}

unsafe impl Send for MemfdShMemProvider {}

#[cfg(unix)]
impl Default for MemfdShMemProvider {
fn default() -> Self {
Self::new().unwrap()
}
}

/// Implement [`ShMemProvider`] for [`MemfdShMemProvider`]
#[cfg(unix)]
impl ShMemProvider for MemfdShMemProvider {
type ShMem = MemfdShMem;

fn new() -> Result<Self, Error> {
Ok(Self {})
}

fn new_shmem(&mut self, map_size: usize) -> Result<Self::ShMem, Error> {
let mapping = MemfdShMem::new(map_size)?;
Ok(mapping)
}

fn shmem_from_id_and_size(
&mut self,
id: ShMemId,
size: usize,
) -> Result<Self::ShMem, Error> {
MemfdShMem::shmem_from_id_and_size(id, size)
}
}
}
}

/// Then `win32` implementation for shared memory.
Expand Down

0 comments on commit cfe1240

Please sign in to comment.