Skip to content

Commit

Permalink
Add memfd shmem backend
Browse files Browse the repository at this point in the history
  • Loading branch information
bernhl committed Nov 1, 2024
1 parent 89cff63 commit 43c3a37
Show file tree
Hide file tree
Showing 2 changed files with 174 additions and 1 deletion.
2 changes: 1 addition & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,7 @@ libc = "0.2.159" # For (*nix) libc
log = "0.4.22"
meminterval = "0.4.1"
mimalloc = { version = "0.1.43", default-features = false }
nix = { version = "0.29.0", default-features = false }
nix = { version = "0.29.0", default-features = false, features = ["fs"] }
num_enum = { version = "0.7.3", default-features = false }
num-traits = { version = "0.2.19", default-features = false }
paste = "1.0.15"
Expand Down
173 changes: 173 additions & 0 deletions libafl_bolts/src/shmem.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1328,6 +1328,179 @@ 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::AsRawFd};

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(format!("Failed to create memfd")));
};
let fd = fd.as_raw_fd();
if ftruncate(fd, map_size as i64) == -1 {
close(fd);
return Err(Error::last_os_error(format!("Failed to ftruncate memfd")));
}
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_string(&format!("{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 = id.to_string().parse().unwrap();
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(),
));
}
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_string(&format!("{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 = self.id.to_string().parse().unwrap();

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 43c3a37

Please sign in to comment.