Streamline ShMem API (#472)

* from warning

* fix latest clippy

* clippy fixes++

* renamed shmem parameters

* renamed map to shmem

* make forkserver executor work for any (non-system) shmem

* Mem -> ShMem

* rework windows

* fix nit

* fix symbolic
This commit is contained in:
Dominik Maier 2022-01-17 18:28:26 +01:00 committed by GitHub
parent ac43997950
commit 4f6f76e857
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
15 changed files with 422 additions and 351 deletions

View File

@ -26,10 +26,10 @@ Next, we'll take a look at the `ForkserverExecutor`. In this case, it is `afl-cc
As you can see from the forkserver example,
```rust,ignore
//Coverage map shared between observer and executor
let mut shmem = StdShMemProvider::new().unwrap().new_map(MAP_SIZE).unwrap();
let mut shmem = StdShMemProvider::new().unwrap().new_shmem(MAP_SIZE).unwrap();
//let the forkserver know the shmid
shmem.write_to_env("__AFL_SHM_ID").unwrap();
let mut shmem_map = shmem.map_mut();
let mut shmem_buf = shmem.as_mut_slice();
```
Here we make a shared memory region; `shmem`, and write this to environmental variable `__AFL_SHM_ID`. Then the instrumented binary, or the forkserver, finds this shared memory region (from the aforementioned env var) to record its coverage. On your fuzzer side, you can pass this shmem map to your `Observer` to obtain coverage feedbacks combined with any `Feedback`.
@ -47,11 +47,11 @@ On your fuzzer side, you can allocate a shared memory region and make the `EDGES
```rust,ignore
let mut shmem;
unsafe{
shmem = StdShMemProvider::new().unwrap().new_map(MAX_EDGES_NUM).unwrap();
shmem = StdShMemProvider::new().unwrap().new_shmem(MAX_EDGES_NUM).unwrap();
}
let shmem_map = shmem.map_mut();
let shmem_buf = shmem.as_mut_slice();
unsafe{
EDGES_PTR = shmem_map.as_ptr();
EDGES_PTR = shmem_buf.as_ptr();
}
```
Again, you can pass this shmem map to your `Observer` and `Feedback` to obtain coverage feedbacks.

View File

@ -21,16 +21,16 @@ The broker can also intercept and filter the messages it receives instead of for
A common use-case for messages filtered by the broker are the status messages sent from each client to the broker directly.
The broker used this information to paint a simple UI, with up-to-date information about all clients, however the other clients don't need to receive this information.
### Speedy Local Messages via Shared Maps
### Speedy Local Messages via Shared Memory
Throughout LibAFL, we use a wrapper around different operating system's shared maps, called `ShMem`.
Shared maps are the backbone of `LLMP`.
Shared maps, called shared memory for the sake of not colliding with Rust's `map()` functions, are the backbone of `LLMP`.
Each client, usually a fuzzer trying to share stats and new testcases, maps an outgoing `ShMem` map.
With very few exceptions, only this client writes to this map, therefore, we do not run in race conditions and can live without locks.
The broker reads from all client's `ShMem` maps.
It checks all incoming client maps periodically, and then forwards new messages to its outgoing broadcast-`ShMem`, mapped by all connected clients.
It checks all incoming client maps periodically and then forwards new messages to its outgoing broadcast-`ShMem`, mapped by all connected clients.
To send new messages, a client places a new message at the end of their map, and then updates a static field to notify the broker.
To send new messages, a client places a new message at the end of their shared memory and then updates a static field to notify the broker.
Once the outgoing map is full, the sender allocates a new `ShMem` using the respective `ShMemProvider`.
It then sends the information needed to map the newly-allocated page in connected processes to the old page, using an end of page (`EOP`) message.
Once the receiver maps the new page, flags it as safe for unmapping from the sending process (to avoid race conditions if we have more than a single EOP in a short time), and then continues to read from the new `ShMem`.
@ -49,7 +49,7 @@ The schema for client's maps to the broker is as follows:
The broker loops over all incoming maps, and checks for new messages.
On `std` builds, the broker will sleep a few milliseconds after a loop, since we do not need the messages to arrive instantly.
After the broker received a new message from clientN, (`clientN_out->current_id != last_message->message_id`) the broker copies the message content to its own broadcast map.
After the broker received a new message from clientN, (`clientN_out->current_id != last_message->message_id`) the broker copies the message content to its own broadcast shared memory.
The clients periodically, for example after finishing `n` mutations, check for new incoming messages by checking if (`current_broadcast_map->current_id != last_message->message_id`).
While the broker uses the same EOP mechanism to map new `ShMem`s for its outgoing map, it never unmaps old pages.
@ -61,7 +61,7 @@ So the outgoing messages flow like this over the outgoing broadcast `Shmem`:
```text
[broker]
|
[current_broadcast_map]
[current_broadcast_shmem]
|
|___________________________________
|_________________ \
@ -83,10 +83,10 @@ Finally, call `LlmpBroker::loop_forever()`.
### B2B: Connecting Fuzzers via TCP
For `broker2broker` communication, all broadcast messages are additionally forwarded via network sockets.
To facilitate this, we spawn an additional client thread in the broker, that reads the broadcast map, just like any other client would.
To facilitate this, we spawn an additional client thread in the broker, that reads the broadcast shared memory, just like any other client would.
For broker2broker communication, this b2b client listens for TCP connections from other, remote brokers.
It keeps a pool of open sockets to other, remote, b2b brokers around at any time.
When receiving a new message on the local broker map, the b2b client will forward it to all connected remote brokers via TCP.
When receiving a new message on the local broker shared memory, the b2b client will forward it to all connected remote brokers via TCP.
Additionally, the broker can receive messages from all connected (remote) brokers, and forward them to the local broker over a client `ShMem`.
As a sidenote, the tcp listener used for b2b communication is also used for an initial handshake when a new client tries to connect to a broker locally, simply exchanging the initial `ShMem` descriptions.

View File

@ -49,6 +49,12 @@ pub fn main() {
.long("timeout")
.default_value("1200"),
)
.arg(
Arg::new("debug_child")
.help("If not set, the child's stdout and stderror will be redirected to /dev/null")
.short('d')
.long("debug-child"),
)
.arg(
Arg::new("arguments")
.help("Arguments passed to the target")
@ -61,16 +67,18 @@ pub fn main() {
const MAP_SIZE: usize = 65536;
//Coverage map shared between observer and executor
let mut shmem = StdShMemProvider::new().unwrap().new_map(MAP_SIZE).unwrap();
//let the forkserver know the shmid
// The default, OS-specific privider for shared memory
let mut shmem_provider = StdShMemProvider::new().unwrap();
// The coverage map shared between observer and executor
let mut shmem = shmem_provider.new_shmem(MAP_SIZE).unwrap();
// let the forkserver know the shmid
shmem.write_to_env("__AFL_SHM_ID").unwrap();
let shmem_map = shmem.map_mut();
let shmem_buf = shmem.as_mut_slice();
// Create an observation channel using the signals map
let edges_observer = HitcountsMapObserver::new(ConstMapObserver::<_, MAP_SIZE>::new(
"shared_mem",
shmem_map,
shmem_buf,
));
// Create an observation channel to keep track of the execution time
@ -127,6 +135,9 @@ pub fn main() {
// A fuzzer with feedbacks and a corpus scheduler
let mut fuzzer = StdFuzzer::new(scheduler, feedback, objective);
// If we should debug the child
let debug_child = res.value_of("debug_child").is_some();
// Create the executor for the forkserver
let args = match res.values_of("arguments") {
Some(vec) => vec.map(|s| s.to_string()).collect::<Vec<String>>().to_vec(),
@ -134,11 +145,12 @@ pub fn main() {
};
let mut executor = TimeoutForkserverExecutor::new(
ForkserverExecutor::new(
ForkserverExecutor::with_shmem_inputs(
res.value_of("executable").unwrap().to_string(),
&args,
true,
tuple_list!(edges_observer, time_observer),
debug_child,
&mut shmem_provider,
)
.unwrap(),
Duration::from_millis(

View File

@ -206,13 +206,13 @@ fn fuzz(
// The shared memory for the concolic runtime to write its trace to
let mut concolic_shmem = StdShMemProvider::new()
.unwrap()
.new_map(DEFAULT_SIZE)
.new_shmem(DEFAULT_SIZE)
.unwrap();
concolic_shmem.write_to_env(DEFAULT_ENV_NAME).unwrap();
// The concolic observer observers the concolic shared memory map.
let concolic_observer =
ConcolicObserver::new("concolic".to_string(), concolic_shmem.map_mut());
ConcolicObserver::new("concolic".to_string(), concolic_shmem.as_slice_mut());
let concolic_observer_name = concolic_observer.name().to_string();

View File

@ -2,7 +2,7 @@
A library for low level message passing
To send new messages, the clients place a new message at the end of their
`client_out_map`. If the current map is filled up, they place an end of page (`EOP`)
`client_out_mem`. If the current map is filled up, they place an end of page (`EOP`)
msg and alloc a new [`ShMem`].
Once the broker mapped this same page, it flags it as safe for unmapping.
@ -21,7 +21,7 @@ After the broker received a new message for clientN, (`clientN_out->current_id
!= last_message->message_id`) the broker will copy the message content to its
own, centralized page.
The clients periodically check (`current_broadcast_map->current_id !=
The clients periodically check (`current_broadcast_shmem->current_id !=
last_message->message_id`) for new incoming messages. If the page is filled up,
the broker instead creates a new page and places an end of page (`EOP`)
message in its queue. The `EOP` buf contains the new description to
@ -31,7 +31,7 @@ current map.
```text
[broker]
|
[current_broadcast_map]
[current_broadcast_shmem]
|
|___________________________________
|_________________ \
@ -41,8 +41,8 @@ current map.
[client0] [client1] ... [clientN]
```
In the future, if we would need zero copy, the `current_broadcast_map` could instead
list the `client_out_map` ID an offset for each message. In that case, the clients
In the future, if we would need zero copy, the `current_broadcast_shmem` could instead
list the `client_out_shmem` ID an offset for each message. In that case, the clients
also need to create a new [`ShMem`] each time their bufs are filled up.
@ -227,7 +227,7 @@ pub enum TcpResponse {
/// After receiving a new connection, the broker immediately sends a Hello.
BrokerConnectHello {
/// The broker page a new local client can listen on
broker_map_description: ShMemDescription,
broker_shmem_description: ShMemDescription,
/// This broker's hostname
hostname: String,
},
@ -294,14 +294,14 @@ impl Listener {
#[inline]
#[allow(clippy::cast_ptr_alignment)]
unsafe fn shmem2page_mut<SHM: ShMem>(afl_shmem: &mut SHM) -> *mut LlmpPage {
afl_shmem.map_mut().as_mut_ptr() as *mut LlmpPage
afl_shmem.as_mut_slice().as_mut_ptr() as *mut LlmpPage
}
/// Get sharedmem from a page
#[inline]
#[allow(clippy::cast_ptr_alignment)]
unsafe fn shmem2page<SHM: ShMem>(afl_shmem: &SHM) -> *const LlmpPage {
afl_shmem.map().as_ptr() as *const LlmpPage
afl_shmem.as_slice().as_ptr() as *const LlmpPage
}
/// Return, if a msg is contained in the current page
@ -411,7 +411,7 @@ fn recv_tcp_msg(stream: &mut TcpStream) -> Result<Vec<u8>, Error> {
/// enough. For now, we want to have at least enough space to store 2 of the
/// largest messages we encountered (plus message one `new_page` message).
#[inline]
fn new_map_size(max_alloc: usize) -> usize {
fn next_shmem_size(max_alloc: usize) -> usize {
max(
max_alloc * 2 + EOP_MSG_SIZE + LLMP_PAGE_HEADER_LEN,
LLMP_CFG_INITIAL_MAP_SIZE - 1,
@ -460,7 +460,7 @@ unsafe fn llmp_next_msg_ptr_checked<SHM: ShMem>(
alloc_size: usize,
) -> Result<*mut LlmpMsg, Error> {
let page = map.page_mut();
let map_size = map.shmem.map().len();
let map_size = map.shmem.as_slice().len();
let msg_begin_min = (page as *const u8).add(size_of::<LlmpPage>());
// We still need space for this msg (alloc_size).
let msg_begin_max = (page as *const u8).add(map_size - alloc_size);
@ -544,7 +544,7 @@ impl LlmpMsg {
#[inline]
pub fn as_slice<SHM: ShMem>(&self, map: &mut LlmpSharedMap<SHM>) -> Result<&[u8], Error> {
unsafe {
if self.in_map(map) {
if self.in_shmem(map) {
Ok(self.as_slice_unsafe())
} else {
Err(Error::IllegalState("Current message not in page. The sharedmap get tampered with or we have a BUG.".into()))
@ -554,9 +554,9 @@ impl LlmpMsg {
/// Returns true, if the pointer is, indeed, in the page of this shared map.
#[inline]
pub fn in_map<SHM: ShMem>(&self, map: &mut LlmpSharedMap<SHM>) -> bool {
pub fn in_shmem<SHM: ShMem>(&self, map: &mut LlmpSharedMap<SHM>) -> bool {
unsafe {
let map_size = map.shmem.map().len();
let map_size = map.shmem.as_slice().len();
let buf_ptr = self.buf.as_ptr();
if buf_ptr > (map.page_mut() as *const u8).add(size_of::<LlmpPage>())
&& buf_ptr <= (map.page_mut() as *const u8).add(map_size - size_of::<LlmpMsg>())
@ -721,7 +721,7 @@ where
/// If null, a new page (just) started.
pub last_msg_sent: *const LlmpMsg,
/// A vec of page wrappers, each containing an intialized AfShmem
pub out_maps: Vec<LlmpSharedMap<SP::Mem>>,
pub out_shmems: Vec<LlmpSharedMap<SP::ShMem>>,
/// If true, pages will never be pruned.
/// The broker uses this feature.
/// By keeping the message history around,
@ -745,9 +745,9 @@ where
Ok(Self {
id,
last_msg_sent: ptr::null_mut(),
out_maps: vec![LlmpSharedMap::new(
out_shmems: vec![LlmpSharedMap::new(
0,
shmem_provider.new_map(LLMP_CFG_INITIAL_MAP_SIZE)?,
shmem_provider.new_shmem(LLMP_CFG_INITIAL_MAP_SIZE)?,
)],
// drop pages to the broker if it already read them
keep_pages_forever,
@ -762,7 +762,11 @@ where
/// # Safety
/// Only safe if you really really restart the page on everything connected
pub unsafe fn reset(&mut self) {
_llmp_page_init(&mut self.out_maps.last_mut().unwrap().shmem, self.id, true);
_llmp_page_init(
&mut self.out_shmems.last_mut().unwrap().shmem,
self.id,
true,
);
self.last_msg_sent = ptr::null_mut();
}
@ -785,11 +789,11 @@ where
env::set_var(&format!("{}_CLIENT_ID", env_name), &format!("{}", id));
}
/// Reattach to a vacant `out_map`, to with a previous sender stored the information in an env before.
/// Reattach to a vacant `out_shmem`, to with a previous sender stored the information in an env before.
#[cfg(feature = "std")]
pub fn on_existing_from_env(mut shmem_provider: SP, env_name: &str) -> Result<Self, Error> {
let msg_sent_offset = msg_offset_from_env(env_name)?;
let mut ret = Self::on_existing_map(
let mut ret = Self::on_existing_shmem(
shmem_provider.clone(),
shmem_provider.existing_from_env(env_name)?,
msg_sent_offset,
@ -802,10 +806,10 @@ where
/// A new client can reattach to it using [`LlmpSender::on_existing_from_env()`].
#[cfg(feature = "std")]
pub fn to_env(&self, env_name: &str) -> Result<(), Error> {
let current_out_map = self.out_maps.last().unwrap();
current_out_map.shmem.write_to_env(env_name)?;
let current_out_shmem = self.out_shmems.last().unwrap();
current_out_shmem.shmem.write_to_env(env_name)?;
Self::client_id_to_env(env_name, self.id);
unsafe { current_out_map.msg_to_env(self.last_msg_sent, env_name) }
unsafe { current_out_shmem.msg_to_env(self.last_msg_sent, env_name) }
}
/// Waits for this sender to be save to unmap.
@ -831,10 +835,10 @@ where
/// If we are allowed to unmap this client
pub fn safe_to_unmap(&self) -> bool {
let current_out_map = self.out_maps.last().unwrap();
let current_out_shmem = self.out_shmems.last().unwrap();
unsafe {
// println!("Reading safe_to_unmap from {:?}", current_out_map.page() as *const _);
(*current_out_map.page())
// println!("Reading safe_to_unmap from {:?}", current_out_shmem.page() as *const _);
(*current_out_shmem.page())
.safe_to_unmap
.load(Ordering::Relaxed)
!= 0
@ -845,29 +849,29 @@ where
/// # Safety
/// If this method is called, the page may be unmapped before it is read by any receiver.
pub unsafe fn mark_safe_to_unmap(&mut self) {
(*self.out_maps.last_mut().unwrap().page_mut())
(*self.out_shmems.last_mut().unwrap().page_mut())
.safe_to_unmap
.store(1, Ordering::Relaxed);
}
/// Reattach to a vacant `out_map`.
/// Reattach to a vacant `out_shmem`.
/// It is essential, that the receiver (or someone else) keeps a pointer to this map
/// else reattach will get a new, empty page, from the OS, or fail.
pub fn on_existing_map(
pub fn on_existing_shmem(
shmem_provider: SP,
current_out_map: SP::Mem,
current_out_shmem: SP::ShMem,
last_msg_sent_offset: Option<u64>,
) -> Result<Self, Error> {
let mut out_map = LlmpSharedMap::existing(current_out_map);
let mut out_shmem = LlmpSharedMap::existing(current_out_shmem);
let last_msg_sent = match last_msg_sent_offset {
Some(offset) => out_map.msg_from_offset(offset)?,
Some(offset) => out_shmem.msg_from_offset(offset)?,
None => ptr::null_mut(),
};
Ok(Self {
id: 0,
last_msg_sent,
out_maps: vec![out_map],
out_shmems: vec![out_shmem],
// drop pages to the broker if it already read them
keep_pages_forever: false,
has_unsent_message: false,
@ -881,7 +885,7 @@ where
unsafe fn prune_old_pages(&mut self) {
// Exclude the current page by splitting of the last element for this iter
let mut unmap_until_excl = 0;
for map in self.out_maps.split_last_mut().unwrap().1 {
for map in self.out_shmems.split_last_mut().unwrap().1 {
if (*map.page()).safe_to_unmap.load(Ordering::Relaxed) == 0 {
// The broker didn't read this page yet, no more pages to unmap.
break;
@ -889,7 +893,7 @@ where
unmap_until_excl += 1;
}
if unmap_until_excl == 0 && self.out_maps.len() > LLMP_CFG_MAX_PENDING_UNREAD_PAGES {
if unmap_until_excl == 0 && self.out_shmems.len() > LLMP_CFG_MAX_PENDING_UNREAD_PAGES {
// We send one last information to the broker before quitting.
self.send_buf(LLMP_SLOW_RECEIVER_PANIC, &[]).unwrap();
panic!("The receiver/broker could not process our sent llmp messages in time. Either we're sending too many messages too fast, the broker got stuck, or it crashed. Giving up.");
@ -897,7 +901,7 @@ where
// Remove all maps that the broker already mapped
// simply removing them from the vec should then call drop and unmap them.
self.out_maps.drain(0..unmap_until_excl);
self.out_shmems.drain(0..unmap_until_excl);
}
/// Intern: Special allocation function for `EOP` messages (and nothing else!)
@ -905,7 +909,7 @@ where
/// So if [`alloc_next`] fails, create new page if necessary, use this function,
/// place `EOP`, commit `EOP`, reset, alloc again on the new space.
unsafe fn alloc_eop(&mut self) -> Result<*mut LlmpMsg, Error> {
let map = self.out_maps.last_mut().unwrap();
let map = self.out_shmems.last_mut().unwrap();
let page = map.page_mut();
let last_msg = self.last_msg_sent;
assert!((*page).size_used + EOP_MSG_SIZE <= (*page).size_total,
@ -940,7 +944,7 @@ where
/// Never call [`alloc_next`] without either sending or cancelling the last allocated message for this page!
/// There can only ever be up to one message allocated per page at each given time.
unsafe fn alloc_next_if_space(&mut self, buf_len: usize) -> Option<*mut LlmpMsg> {
let map = self.out_maps.last_mut().unwrap();
let map = self.out_shmems.last_mut().unwrap();
let page = map.page_mut();
let last_msg = self.last_msg_sent;
@ -1032,7 +1036,7 @@ where
if overwrite_client_id {
(*msg).sender = self.id;
}
let page = self.out_maps.last_mut().unwrap().page_mut();
let page = self.out_shmems.last_mut().unwrap().page_mut();
if msg.is_null() || !llmp_msg_in_page(page, msg) {
return Err(Error::Unknown(format!(
"Llmp Message {:?} is null or not in current page",
@ -1060,7 +1064,7 @@ where
let bt = Backtrace::new();
#[cfg(not(debug_assertions))]
let bt = "<n/a (release)>";
let shm = self.out_maps.last().unwrap();
let shm = self.out_shmems.last().unwrap();
println!(
"LLMP_DEBUG: End of page reached for map {} with len {}, sending EOP, bt: {:?}",
shm.shmem.id(),
@ -1069,16 +1073,19 @@ where
);
}
let old_map = self.out_maps.last_mut().unwrap().page_mut();
let old_map = self.out_shmems.last_mut().unwrap().page_mut();
#[cfg(all(feature = "llmp_debug", feature = "std"))]
println!("New Map Size {}", new_map_size((*old_map).max_alloc_size));
println!(
"Next ShMem Size {}",
next_shmem_size((*old_map).max_alloc_size)
);
// Create a new shard page.
let mut new_map_shmem = LlmpSharedMap::new(
(*old_map).sender,
self.shmem_provider
.new_map(new_map_size((*old_map).max_alloc_size))?,
.new_shmem(next_shmem_size((*old_map).max_alloc_size))?,
);
let mut new_map = new_map_shmem.page_mut();
@ -1109,7 +1116,7 @@ where
self.send(out, true)?;
// Set the new page as current page.
self.out_maps.push(new_map_shmem);
self.out_shmems.push(new_map_shmem);
// We never sent a msg on the new buf */
self.last_msg_sent = ptr::null_mut();
@ -1152,7 +1159,7 @@ where
pub unsafe fn cancel_send(&mut self, msg: *mut LlmpMsg) {
/* DBG("Client %d cancels send of msg at %p with tag 0x%X and size %ld", client->id, msg, msg->tag,
* msg->buf_len_padded); */
let page = self.out_maps.last_mut().unwrap().page_mut();
let page = self.out_shmems.last_mut().unwrap().page_mut();
(*msg).tag = LLMP_TAG_UNSET;
(*page).size_used -= (*msg).buf_len_padded as usize + size_of::<LlmpMsg>();
}
@ -1191,7 +1198,7 @@ where
(*msg).buf_len = shrinked_len as u64;
(*msg).buf_len_padded = buf_len_padded as u64;
let page = self.out_maps.last_mut().unwrap().page_mut();
let page = self.out_shmems.last_mut().unwrap().page_mut();
// Doing this step by step will catch underflows in debug builds :)
(*page).size_used -= old_len_padded as usize;
@ -1252,7 +1259,7 @@ where
/// Describe this [`LlmpClient`] in a way that it can be restored later, using [`Self::on_existing_from_description`].
pub fn describe(&self) -> Result<LlmpDescription, Error> {
let map = self.out_maps.last().unwrap();
let map = self.out_shmems.last().unwrap();
let last_message_offset = if self.last_msg_sent.is_null() {
None
} else {
@ -1270,9 +1277,9 @@ where
mut shmem_provider: SP,
description: &LlmpDescription,
) -> Result<Self, Error> {
Self::on_existing_map(
Self::on_existing_shmem(
shmem_provider.clone(),
shmem_provider.map_from_decription(description.shmem)?,
shmem_provider.shmem_from_description(description.shmem)?,
description.last_message_offset,
)
}
@ -1291,7 +1298,7 @@ where
/// The shmem provider
pub shmem_provider: SP,
/// current page. After EOP, this gets replaced with the new one
pub current_recv_map: LlmpSharedMap<SP::Mem>,
pub current_recv_shmem: LlmpSharedMap<SP::ShMem>,
/// Caches the highest msg id we've seen so far
highest_msg_id: u64,
}
@ -1301,10 +1308,10 @@ impl<SP> LlmpReceiver<SP>
where
SP: ShMemProvider,
{
/// Reattach to a vacant `recv_map`, to with a previous sender stored the information in an env before.
/// Reattach to a vacant `recv_shmem`, to with a previous sender stored the information in an env before.
#[cfg(feature = "std")]
pub fn on_existing_from_env(mut shmem_provider: SP, env_name: &str) -> Result<Self, Error> {
Self::on_existing_map(
Self::on_existing_shmem(
shmem_provider.clone(),
shmem_provider.existing_from_env(env_name)?,
msg_offset_from_env(env_name)?,
@ -1315,28 +1322,28 @@ where
/// A new client can reattach to it using [`LlmpReceiver::on_existing_from_env()`]
#[cfg(feature = "std")]
pub fn to_env(&self, env_name: &str) -> Result<(), Error> {
let current_out_map = &self.current_recv_map;
current_out_map.shmem.write_to_env(env_name)?;
unsafe { current_out_map.msg_to_env(self.last_msg_recvd, env_name) }
let current_out_shmem = &self.current_recv_shmem;
current_out_shmem.shmem.write_to_env(env_name)?;
unsafe { current_out_shmem.msg_to_env(self.last_msg_recvd, env_name) }
}
/// Create a Receiver, reattaching to an existing sender map.
/// It is essential, that the sender (or someone else) keeps a pointer to the `sender_map`
/// It is essential, that the sender (or someone else) keeps a pointer to the `sender_shmem`
/// else reattach will get a new, empty page, from the OS, or fail.
pub fn on_existing_map(
pub fn on_existing_shmem(
shmem_provider: SP,
current_sender_map: SP::Mem,
current_sender_shmem: SP::ShMem,
last_msg_recvd_offset: Option<u64>,
) -> Result<Self, Error> {
let mut current_recv_map = LlmpSharedMap::existing(current_sender_map);
let mut current_recv_shmem = LlmpSharedMap::existing(current_sender_shmem);
let last_msg_recvd = match last_msg_recvd_offset {
Some(offset) => current_recv_map.msg_from_offset(offset)?,
Some(offset) => current_recv_shmem.msg_from_offset(offset)?,
None => ptr::null_mut(),
};
Ok(Self {
id: 0,
current_recv_map,
current_recv_shmem,
last_msg_recvd,
shmem_provider,
highest_msg_id: 0,
@ -1348,7 +1355,7 @@ where
#[inline(never)]
unsafe fn recv(&mut self) -> Result<Option<*mut LlmpMsg>, Error> {
/* DBG("recv %p %p\n", page, last_msg); */
let mut page = self.current_recv_map.page_mut();
let mut page = self.current_recv_shmem.page_mut();
let last_msg = self.last_msg_recvd;
let (current_msg_id, loaded) =
@ -1380,7 +1387,7 @@ where
}
// We don't know how big the msg wants to be, assert at least the header has space.
Some(llmp_next_msg_ptr_checked(
&mut self.current_recv_map,
&mut self.current_recv_shmem,
last_msg,
size_of::<LlmpMsg>(),
)?)
@ -1388,7 +1395,7 @@ where
// Let's see what we got.
if let Some(msg) = ret {
if !(*msg).in_map(&mut self.current_recv_map) {
if !(*msg).in_shmem(&mut self.current_recv_shmem) {
return Err(Error::IllegalState("Unexpected message in map (out of map bounds) - bugy client or tampered shared map detedted!".into()));
}
// Handle special, LLMP internal, messages.
@ -1430,20 +1437,20 @@ where
(*page).safe_to_unmap.store(1, Ordering::Relaxed);
// Map the new page. The old one should be unmapped by Drop
self.current_recv_map =
LlmpSharedMap::existing(self.shmem_provider.map_from_id_and_size(
self.current_recv_shmem =
LlmpSharedMap::existing(self.shmem_provider.shmem_from_id_and_size(
ShMemId::from_slice(&pageinfo_cpy.shm_str),
pageinfo_cpy.map_size,
)?);
page = self.current_recv_map.page_mut();
page = self.current_recv_shmem.page_mut();
// Mark the new page save to unmap also (it's mapped by us, the broker now)
(*page).safe_to_unmap.store(1, Ordering::Relaxed);
#[cfg(all(feature = "llmp_debug", feature = "std"))]
println!(
"LLMP_DEBUG: Got a new recv map {} with len {:?}",
self.current_recv_map.shmem.id(),
self.current_recv_map.shmem.len()
self.current_recv_shmem.shmem.id(),
self.current_recv_shmem.shmem.len()
);
// After we mapped the new page, return the next message, if available
return self.recv();
@ -1463,7 +1470,7 @@ where
/// Returns a raw ptr, on the recv map. Should be safe in general
pub unsafe fn recv_blocking(&mut self) -> Result<*mut LlmpMsg, Error> {
let mut current_msg_id = 0;
let page = self.current_recv_map.page_mut();
let page = self.current_recv_shmem.page_mut();
let last_msg = self.last_msg_recvd;
if !last_msg.is_null() {
assert!(
@ -1505,7 +1512,7 @@ where
(*msg).sender,
(*msg).tag,
(*msg).flags,
(*msg).as_slice(&mut self.current_recv_map)?,
(*msg).as_slice(&mut self.current_recv_shmem)?,
)),
None => None,
})
@ -1520,14 +1527,14 @@ where
Ok((
(*msg).sender,
(*msg).tag,
(*msg).as_slice(&mut self.current_recv_map)?,
(*msg).as_slice(&mut self.current_recv_shmem)?,
))
}
}
/// Describe this client in a way, that it can be restored later with [`Self::on_existing_from_description`]
pub fn describe(&self) -> Result<LlmpDescription, Error> {
let map = &self.current_recv_map;
let map = &self.current_recv_shmem;
let last_message_offset = if self.last_msg_recvd.is_null() {
None
} else {
@ -1544,9 +1551,9 @@ where
mut shmem_provider: SP,
description: &LlmpDescription,
) -> Result<Self, Error> {
Self::on_existing_map(
Self::on_existing_shmem(
shmem_provider.clone(),
shmem_provider.map_from_decription(description.shmem)?,
shmem_provider.shmem_from_description(description.shmem)?,
description.last_message_offset,
)
}
@ -1571,22 +1578,22 @@ where
SHM: ShMem,
{
/// Creates a new page, initializing the passed shared mem struct
pub fn new(sender: ClientId, mut new_map: SHM) -> Self {
pub fn new(sender: ClientId, mut new_shmem: SHM) -> Self {
#[cfg(all(feature = "llmp_debug", feature = "std"))]
println!(
"LLMP_DEBUG: Initializing map on {} with size {}",
new_map.id(),
new_map.len()
new_shmem.id(),
new_shmem.len()
);
unsafe {
_llmp_page_init(&mut new_map, sender, false);
_llmp_page_init(&mut new_shmem, sender, false);
}
Self { shmem: new_map }
Self { shmem: new_shmem }
}
/// Maps and wraps an existing
pub fn existing(existing_map: SHM) -> Self {
pub fn existing(existing_shmem: SHM) -> Self {
#[cfg(all(feature = "llmp_debug", feature = "std"))]
//{
//#[cfg(debug_assertions)]
@ -1595,14 +1602,14 @@ where
//let bt = "<n/a (release)>";
println!(
"LLMP_DEBUG: Using existing map {} with size {}",
existing_map.id(),
existing_map.len(),
existing_shmem.id(),
existing_shmem.len(),
//bt
);
//}
let ret = Self {
shmem: existing_map,
shmem: existing_shmem,
};
unsafe {
assert!(
@ -1692,7 +1699,7 @@ where
let offset = offset as usize;
unsafe {
let page = self.page_mut();
let page_size = self.shmem.map().len() - size_of::<LlmpPage>();
let page_size = self.shmem.as_slice().len() - size_of::<LlmpPage>();
if offset > page_size {
Err(Error::IllegalArgument(format!(
"Msg offset out of bounds (size: {}, requested offset: {})",
@ -1753,9 +1760,9 @@ where
llmp_out: LlmpSender {
id: 0,
last_msg_sent: ptr::null_mut(),
out_maps: vec![LlmpSharedMap::new(
out_shmems: vec![LlmpSharedMap::new(
0,
shmem_provider.new_map(new_map_size(0))?,
shmem_provider.new_shmem(next_shmem_size(0))?,
)],
// Broker never cleans up the pages so that new
// clients may join at any time
@ -1787,15 +1794,15 @@ where
}
/// Registers a new client for the given sharedmap str and size.
/// Returns the id of the new client in [`broker.client_map`]
pub fn register_client(&mut self, mut client_page: LlmpSharedMap<SP::Mem>) {
/// Returns the id of the new client in [`broker.client_shmem`]
pub fn register_client(&mut self, mut client_page: LlmpSharedMap<SP::ShMem>) {
// Tell the client it may unmap this page now.
client_page.mark_safe_to_unmap();
let id = self.llmp_clients.len() as u32;
self.llmp_clients.push(LlmpReceiver {
id,
current_recv_map: client_page,
current_recv_shmem: client_page,
last_msg_recvd: ptr::null_mut(),
shmem_provider: self.shmem_provider.clone(),
highest_msg_id: 0,
@ -1815,7 +1822,7 @@ where
match (&recv_tcp_msg(&mut stream)?).try_into()? {
TcpResponse::BrokerConnectHello {
broker_map_description: _,
broker_shmem_description: _,
hostname,
} => println!("B2B: Connected to {}", hostname),
_ => {
@ -1851,14 +1858,22 @@ where
let map_description = Self::b2b_thread_on(
stream,
self.llmp_clients.len() as ClientId,
&self.llmp_out.out_maps.first().unwrap().shmem.description(),
&self
.llmp_out
.out_shmems
.first()
.unwrap()
.shmem
.description(),
)?;
let new_map =
LlmpSharedMap::existing(self.shmem_provider.map_from_decription(map_description)?);
let new_shmem = LlmpSharedMap::existing(
self.shmem_provider
.shmem_from_description(map_description)?,
);
{
self.register_client(new_map);
self.register_client(new_shmem);
}
Ok(())
@ -1998,9 +2013,9 @@ where
fn b2b_thread_on(
mut stream: TcpStream,
b2b_client_id: ClientId,
broker_map_description: &ShMemDescription,
broker_shmem_description: &ShMemDescription,
) -> Result<ShMemDescription, Error> {
let broker_map_description = *broker_map_description;
let broker_shmem_description = *broker_shmem_description;
// A channel to get the new "client's" sharedmap id from
let (send, recv) = channel();
@ -2026,7 +2041,7 @@ where
}
};
send.send(new_sender.out_maps.first().unwrap().shmem.description())
send.send(new_sender.out_shmems.first().unwrap().shmem.description())
.expect("B2B: Error sending map description to channel!");
// the receiver receives from the local broker, and forwards it to the tcp stream.
@ -2034,7 +2049,7 @@ where
shmem_provider_bg,
&LlmpDescription {
last_message_offset: None,
shmem: broker_map_description,
shmem: broker_shmem_description,
},
)
.expect("Failed to map local page in broker 2 broker thread!");
@ -2120,7 +2135,7 @@ where
request: &TcpRequest,
current_client_id: &mut u32,
sender: &mut LlmpSender<SP>,
broker_map_description: &ShMemDescription,
broker_shmem_description: &ShMemDescription,
) {
match request {
TcpRequest::LocalClientHello { shmem_description } => {
@ -2156,7 +2171,7 @@ where
}
if let Ok(shmem_description) =
Self::b2b_thread_on(stream, *current_client_id, broker_map_description)
Self::b2b_thread_on(stream, *current_client_id, broker_shmem_description)
{
if Self::announce_new_client(sender, &shmem_description).is_err() {
println!("B2B: Error announcing client {:?}", shmem_description);
@ -2175,26 +2190,26 @@ where
// However, the original map is (as of now) never freed, new clients will start
// to read from the initial map id.
let client_out_map_mem = &self.llmp_out.out_maps.first().unwrap().shmem;
let broker_map_description = client_out_map_mem.description();
let client_out_shmem_mem = &self.llmp_out.out_shmems.first().unwrap().shmem;
let broker_shmem_description = client_out_shmem_mem.description();
let hostname = hostname::get()
.unwrap_or_else(|_| "<unknown>".into())
.to_string_lossy()
.into();
let broker_hello = TcpResponse::BrokerConnectHello {
broker_map_description,
broker_shmem_description: broker_shmem_description,
hostname,
};
let llmp_tcp_id = self.llmp_clients.len() as ClientId;
// Tcp out map sends messages from background thread tcp server to foreground client
let tcp_out_map = LlmpSharedMap::new(
let tcp_out_shmem = LlmpSharedMap::new(
llmp_tcp_id,
self.shmem_provider.new_map(LLMP_CFG_INITIAL_MAP_SIZE)?,
self.shmem_provider.new_shmem(LLMP_CFG_INITIAL_MAP_SIZE)?,
);
let tcp_out_map_description = tcp_out_map.shmem.description();
self.register_client(tcp_out_map);
let tcp_out_shmem_description = tcp_out_shmem.shmem.description();
self.register_client(tcp_out_shmem);
let ret = thread::spawn(move || {
// Create a new ShMemProvider for this background thread.
@ -2205,9 +2220,9 @@ where
let mut tcp_incoming_sender = LlmpSender {
id: llmp_tcp_id,
last_msg_sent: ptr::null_mut(),
out_maps: vec![LlmpSharedMap::existing(
out_shmems: vec![LlmpSharedMap::existing(
shmem_provider_bg
.map_from_decription(tcp_out_map_description)
.shmem_from_description(tcp_out_shmem_description)
.unwrap(),
)],
// drop pages to the broker, if it already read them.
@ -2255,7 +2270,7 @@ where
&req,
&mut current_client_id,
&mut tcp_incoming_sender,
&broker_map_description,
&broker_shmem_description,
);
}
ListenerStream::Empty() => {
@ -2313,18 +2328,18 @@ where
} else {
let pageinfo = (*msg).buf.as_mut_ptr() as *mut LlmpPayloadSharedMapInfo;
match self.shmem_provider.map_from_id_and_size(
match self.shmem_provider.shmem_from_id_and_size(
ShMemId::from_slice(&(*pageinfo).shm_str),
(*pageinfo).map_size,
) {
Ok(new_map) => {
let mut new_page = LlmpSharedMap::existing(new_map);
Ok(new_shmem) => {
let mut new_page = LlmpSharedMap::existing(new_shmem);
let id = next_id;
next_id += 1;
new_page.mark_safe_to_unmap();
self.llmp_clients.push(LlmpReceiver {
id,
current_recv_map: new_page,
current_recv_shmem: new_page,
last_msg_recvd: ptr::null_mut(),
shmem_provider: self.shmem_provider.clone(),
highest_msg_id: 0,
@ -2347,7 +2362,7 @@ where
// The message is not specifically for use. Let the user handle it, then forward it to the clients, if necessary.
let mut should_forward_msg = true;
let map = &mut self.llmp_clients[client_id as usize].current_recv_map;
let map = &mut self.llmp_clients[client_id as usize].current_recv_shmem;
let msg_buf = (*msg).as_slice(map)?;
if let LlmpMsgHookResult::Handled =
(on_new_msg)(client_id, (*msg).tag, (*msg).flags, msg_buf)?
@ -2391,25 +2406,25 @@ where
SP: ShMemProvider,
{
/// Reattach to a vacant client map.
/// It is essential, that the broker (or someone else) kept a pointer to the `out_map`
/// It is essential, that the broker (or someone else) kept a pointer to the `out_shmem`
/// else reattach will get a new, empty page, from the OS, or fail
#[allow(clippy::needless_pass_by_value)]
pub fn on_existing_map(
pub fn on_existing_shmem(
shmem_provider: SP,
_current_out_map: SP::Mem,
_current_out_shmem: SP::ShMem,
_last_msg_sent_offset: Option<u64>,
current_broker_map: SP::Mem,
current_broker_shmem: SP::ShMem,
last_msg_recvd_offset: Option<u64>,
) -> Result<Self, Error> {
Ok(Self {
receiver: LlmpReceiver::on_existing_map(
receiver: LlmpReceiver::on_existing_shmem(
shmem_provider.clone(),
current_broker_map.clone(),
current_broker_shmem.clone(),
last_msg_recvd_offset,
)?,
sender: LlmpSender::on_existing_map(
sender: LlmpSender::on_existing_shmem(
shmem_provider,
current_broker_map,
current_broker_shmem,
last_msg_recvd_offset,
)?,
})
@ -2431,7 +2446,7 @@ where
}
/// Write the current state to env.
/// A new client can attach to exactly the same state by calling [`LlmpClient::on_existing_map()`].
/// A new client can attach to exactly the same state by calling [`LlmpClient::on_existing_shmem()`].
#[cfg(feature = "std")]
pub fn to_env(&self, env_name: &str) -> Result<(), Error> {
self.sender.to_env(&format!("{}_SENDER", env_name))?;
@ -2487,14 +2502,14 @@ where
/// Creates a new [`LlmpClient`]
pub fn new(
mut shmem_provider: SP,
initial_broker_map: LlmpSharedMap<SP::Mem>,
initial_broker_shmem: LlmpSharedMap<SP::ShMem>,
) -> Result<Self, Error> {
Ok(Self {
sender: LlmpSender {
id: 0,
last_msg_sent: ptr::null_mut(),
out_maps: vec![LlmpSharedMap::new(0, {
shmem_provider.new_map(LLMP_CFG_INITIAL_MAP_SIZE)?
out_shmems: vec![LlmpSharedMap::new(0, {
shmem_provider.new_shmem(LLMP_CFG_INITIAL_MAP_SIZE)?
})],
// drop pages to the broker if it already read them
keep_pages_forever: false,
@ -2504,7 +2519,7 @@ where
receiver: LlmpReceiver {
id: 0,
current_recv_map: initial_broker_map,
current_recv_shmem: initial_broker_shmem,
last_msg_recvd: ptr::null_mut(),
shmem_provider,
highest_msg_id: 0,
@ -2626,24 +2641,25 @@ where
};
println!("Connected to port {}", port);
let broker_map_description = if let TcpResponse::BrokerConnectHello {
broker_map_description,
let broker_shmem_description = if let TcpResponse::BrokerConnectHello {
broker_shmem_description,
hostname: _,
} = (&recv_tcp_msg(&mut stream)?).try_into()?
{
broker_map_description
broker_shmem_description
} else {
return Err(Error::IllegalState(
"Received unexpected Broker Hello".to_string(),
));
};
let map =
LlmpSharedMap::existing(shmem_provider.map_from_decription(broker_map_description)?);
let map = LlmpSharedMap::existing(
shmem_provider.shmem_from_description(broker_shmem_description)?,
);
let mut ret = Self::new(shmem_provider, map)?;
let client_hello_req = TcpRequest::LocalClientHello {
shmem_description: ret.sender.out_maps.first().unwrap().shmem.description(),
shmem_description: ret.sender.out_shmems.first().unwrap().shmem.description(),
};
send_tcp_msg(&mut stream, &client_hello_req)?;
@ -2663,7 +2679,7 @@ where
ret.sender.id = client_id;
// Also set the sender on our initial llmp map correctly.
unsafe {
(*ret.sender.out_maps.first_mut().unwrap().page_mut()).sender = client_id;
(*ret.sender.out_shmems.first_mut().unwrap().page_mut()).sender = client_id;
}
Ok(ret)

View File

@ -89,12 +89,12 @@ where
self.inner.len()
}
fn map(&self) -> &[u8] {
self.inner.map()
fn as_slice(&self) -> &[u8] {
self.inner.as_slice()
}
fn map_mut(&mut self) -> &mut [u8] {
self.inner.map_mut()
fn as_mut_slice(&mut self) -> &mut [u8] {
self.inner.as_mut_slice()
}
}
@ -154,7 +154,7 @@ impl<SP> ShMemProvider for ServedShMemProvider<SP>
where
SP: ShMemProvider,
{
type Mem = ServedShMem<SP::Mem>;
type ShMem = ServedShMem<SP::ShMem>;
/// Connect to the server and return a new [`ServedShMemProvider`]
/// Will try to spawn a [`ShMemService`]. This will only work for the first try.
@ -172,21 +172,19 @@ where
res.id = id;
Ok(res)
}
fn new_map(&mut self, map_size: usize) -> Result<Self::Mem, Error> {
fn new_shmem(&mut self, map_size: usize) -> Result<Self::ShMem, Error> {
let (server_fd, client_fd) = self.send_receive(ServedShMemRequest::NewMap(map_size))?;
Ok(ServedShMem {
inner: ManuallyDrop::new(
self.inner.map_from_id_and_size(
ShMemId::from_string(&format!("{}", client_fd)),
map_size,
)?,
),
inner: ManuallyDrop::new(self.inner.shmem_from_id_and_size(
ShMemId::from_string(&format!("{}", client_fd)),
map_size,
)?),
server_fd,
})
}
fn map_from_id_and_size(&mut self, id: ShMemId, size: usize) -> Result<Self::Mem, Error> {
fn shmem_from_id_and_size(&mut self, id: ShMemId, size: usize) -> Result<Self::ShMem, Error> {
let parts = id.as_str().split(':').collect::<Vec<&str>>();
let server_id_str = parts.get(0).unwrap();
let (server_fd, client_fd) = self.send_receive(ServedShMemRequest::ExistingMap(
@ -194,8 +192,10 @@ where
))?;
Ok(ServedShMem {
inner: ManuallyDrop::new(
self.inner
.map_from_id_and_size(ShMemId::from_string(&format!("{}", client_fd)), size)?,
self.inner.shmem_from_id_and_size(
ShMemId::from_string(&format!("{}", client_fd)),
size,
)?,
),
server_fd,
})
@ -224,7 +224,7 @@ where
Ok(())
}
fn release_map(&mut self, map: &mut Self::Mem) {
fn release_shmem(&mut self, map: &mut Self::ShMem) {
let (refcount, _) = self
.send_receive(ServedShMemRequest::Deregister(map.server_fd))
.expect("Could not communicate with ServedShMem server!");
@ -284,7 +284,7 @@ enum ServedShMemResponse<SP>
where
SP: ShMemProvider,
{
Mapping(Rc<RefCell<SP::Mem>>),
Mapping(Rc<RefCell<SP::ShMem>>),
Id(i32),
RefCount(u32),
}
@ -446,10 +446,10 @@ where
SP: ShMemProvider,
{
provider: SP,
clients: HashMap<RawFd, SharedShMemClient<SP::Mem>>,
clients: HashMap<RawFd, SharedShMemClient<SP::ShMem>>,
/// Maps from a pre-fork (parent) client id to its cloned maps.
forking_clients: HashMap<RawFd, HashMap<i32, Vec<Rc<RefCell<SP::Mem>>>>>,
all_maps: HashMap<i32, Weak<RefCell<SP::Mem>>>,
forking_clients: HashMap<RawFd, HashMap<i32, Vec<Rc<RefCell<SP::ShMem>>>>>,
all_shmems: HashMap<i32, Weak<RefCell<SP::ShMem>>>,
}
impl<SP> ServedShMemServiceWorker<SP>
@ -461,13 +461,13 @@ where
Ok(Self {
provider: SP::new()?,
clients: HashMap::new(),
all_maps: HashMap::new(),
all_shmems: HashMap::new(),
forking_clients: HashMap::new(),
})
}
fn upgrade_map_with_id(&mut self, description_id: i32) -> Rc<RefCell<SP::Mem>> {
self.all_maps
fn upgrade_shmem_with_id(&mut self, description_id: i32) -> Rc<RefCell<SP::ShMem>> {
self.all_shmems
.get_mut(&description_id)
.unwrap()
.clone()
@ -496,11 +496,11 @@ where
/*
// remove temporarily
let client = self.clients.remove(&client_id);
let mut forking_maps = HashMap::new();
let mut forking_shmems = HashMap::new();
for (id, map) in client.as_ref().unwrap().maps.iter() {
forking_maps.insert(*id, map.clone());
forking_shmems.insert(*id, map.clone());
}
self.forking_clients.insert(client_id, forking_maps);
self.forking_clients.insert(client_id, forking_shmems);
self.clients.insert(client_id, client.unwrap());
*/
@ -512,10 +512,10 @@ where
Ok(ServedShMemResponse::Id(client_id))
}
ServedShMemRequest::NewMap(map_size) => {
let new_map = self.provider.new_map(map_size)?;
let description = new_map.description();
let new_rc = Rc::new(RefCell::new(new_map));
self.all_maps
let new_shmem = self.provider.new_shmem(map_size)?;
let description = new_shmem.description();
let new_rc = Rc::new(RefCell::new(new_shmem));
self.all_shmems
.insert(description.id.into(), Rc::downgrade(&new_rc));
Ok(ServedShMemResponse::Mapping(new_rc))
}
@ -536,12 +536,12 @@ where
{
map.clone()
} else {
self.upgrade_map_with_id(description_id)
self.upgrade_shmem_with_id(description_id)
},
))
} else {
Ok(ServedShMemResponse::Mapping(
self.upgrade_map_with_id(description_id),
self.upgrade_shmem_with_id(description_id),
))
}
}

View File

@ -24,39 +24,28 @@ pub use unix_shmem::{MmapShMem, MmapShMemProvider};
#[cfg(all(feature = "std", unix))]
pub use unix_shmem::{UnixShMem, UnixShMemProvider};
#[cfg(all(windows, feature = "std"))]
pub use win32_shmem::{Win32ShMem, Win32ShMemProvider};
#[cfg(all(feature = "std", unix))]
pub use crate::bolts::os::unix_shmem_server::{ServedShMemProvider, ShMemService};
#[cfg(all(windows, feature = "std"))]
pub use win32_shmem::{Win32ShMem, Win32ShMemProvider};
/// The standard sharedmem provider
#[cfg(all(windows, feature = "std"))]
pub type StdShMemProvider = Win32ShMemProvider;
/// The standard sharedmem type
#[cfg(all(windows, feature = "std"))]
pub type StdShMem = Win32ShMem;
/// The standard sharedmem provider
#[cfg(all(target_os = "android", feature = "std"))]
pub type StdShMemProvider =
RcShMemProvider<ServedShMemProvider<unix_shmem::ashmem::AshmemShMemProvider>>;
/// The standard sharedmem type
#[cfg(all(target_os = "android", feature = "std"))]
pub type StdShMem = RcShMem<ServedShMemProvider<unix_shmem::ashmem::AshmemShMemProvider>>;
/// The standard sharedmem service
#[cfg(all(target_os = "android", feature = "std"))]
pub type StdShMemService = ShMemService<unix_shmem::ashmem::AshmemShMemProvider>;
/// The standard sharedmem provider
#[cfg(all(feature = "std", target_vendor = "apple"))]
pub type StdShMemProvider = RcShMemProvider<ServedShMemProvider<MmapShMemProvider>>;
/// The standard sharedmem type
#[cfg(all(feature = "std", target_vendor = "apple"))]
pub type StdShMem = RcShMem<ServedShMemProvider<MmapShMemProvider>>;
#[cfg(all(feature = "std", target_vendor = "apple"))]
/// The standard sharedmem service
pub type StdShMemService = ShMemService<MmapShMemProvider>;
/// The default [`ShMemProvider`] for this os.
#[cfg(all(
feature = "std",
@ -64,14 +53,6 @@ pub type StdShMemService = ShMemService<MmapShMemProvider>;
not(any(target_os = "android", target_vendor = "apple"))
))]
pub type StdShMemProvider = UnixShMemProvider;
/// The default [`ShMemProvider`] for this os.
#[cfg(all(
feature = "std",
unix,
not(any(target_os = "android", target_vendor = "apple"))
))]
pub type StdShMem = UnixShMem;
/// The standard sharedmem service
#[cfg(any(
not(any(target_os = "android", target_vendor = "apple")),
@ -185,10 +166,10 @@ pub trait ShMem: Sized + Debug + Clone {
}
/// The actual shared map, in memory
fn map(&self) -> &[u8];
fn as_slice(&self) -> &[u8];
/// The actual shared map, mutable
fn map_mut(&mut self) -> &mut [u8];
fn as_mut_slice(&mut self) -> &mut [u8];
/// Write this map's config to env
#[cfg(feature = "std")]
@ -206,33 +187,36 @@ pub trait ShMem: Sized + Debug + Clone {
/// All you need for scaling on a new target is to implement this interface, as well as the respective [`ShMem`].
pub trait ShMemProvider: Clone + Default + Debug {
/// The actual shared map handed out by this [`ShMemProvider`].
type Mem: ShMem;
type ShMem: ShMem;
/// Create a new instance of the provider
fn new() -> Result<Self, Error>;
/// Create a new shared memory mapping
fn new_map(&mut self, map_size: usize) -> Result<Self::Mem, Error>;
fn new_shmem(&mut self, map_size: usize) -> Result<Self::ShMem, Error>;
/// Get a mapping given its id and size
fn map_from_id_and_size(&mut self, id: ShMemId, size: usize) -> Result<Self::Mem, Error>;
fn shmem_from_id_and_size(&mut self, id: ShMemId, size: usize) -> Result<Self::ShMem, Error>;
/// Get a mapping given a description
fn map_from_decription(&mut self, description: ShMemDescription) -> Result<Self::Mem, Error> {
self.map_from_id_and_size(description.id, description.size)
fn shmem_from_description(
&mut self,
description: ShMemDescription,
) -> Result<Self::ShMem, Error> {
self.shmem_from_id_and_size(description.id, description.size)
}
/// Create a new sharedmap reference from an existing `id` and `len`
fn clone_ref(&mut self, mapping: &Self::Mem) -> Result<Self::Mem, Error> {
self.map_from_id_and_size(mapping.id(), mapping.len())
fn clone_ref(&mut self, mapping: &Self::ShMem) -> Result<Self::ShMem, Error> {
self.shmem_from_id_and_size(mapping.id(), mapping.len())
}
/// Reads an existing map config from env vars, then maps it
#[cfg(feature = "std")]
fn existing_from_env(&mut self, env_name: &str) -> Result<Self::Mem, Error> {
fn existing_from_env(&mut self, env_name: &str) -> Result<Self::ShMem, Error> {
let map_shm_str = env::var(env_name)?;
let map_size = str::parse::<usize>(&env::var(format!("{}_SIZE", env_name))?)?;
self.map_from_decription(ShMemDescription::from_string_and_size(
self.shmem_from_description(ShMemDescription::from_string_and_size(
&map_shm_str,
map_size,
))
@ -255,7 +239,7 @@ pub trait ShMemProvider: Clone + Default + Debug {
}
/// Release the resources associated with the given [`ShMem`]
fn release_map(&mut self, _map: &mut Self::Mem) {
fn release_shmem(&mut self, _shmem: &mut Self::ShMem) {
// do nothing
}
}
@ -265,7 +249,7 @@ pub trait ShMemProvider: Clone + Default + Debug {
/// Useful if the `ShMemProvider` needs to keep local state.
#[derive(Debug, Clone)]
pub struct RcShMem<T: ShMemProvider> {
internal: ManuallyDrop<T::Mem>,
internal: ManuallyDrop<T::ShMem>,
provider: Rc<RefCell<T>>,
}
@ -281,18 +265,18 @@ where
self.internal.len()
}
fn map(&self) -> &[u8] {
self.internal.map()
fn as_slice(&self) -> &[u8] {
self.internal.as_slice()
}
fn map_mut(&mut self) -> &mut [u8] {
self.internal.map_mut()
fn as_mut_slice(&mut self) -> &mut [u8] {
self.internal.as_mut_slice()
}
}
impl<T: ShMemProvider> Drop for RcShMem<T> {
fn drop(&mut self) {
self.provider.borrow_mut().release_map(&mut self.internal);
self.provider.borrow_mut().release_shmem(&mut self.internal);
}
}
@ -325,7 +309,7 @@ impl<SP> ShMemProvider for RcShMemProvider<SP>
where
SP: ShMemProvider + Debug,
{
type Mem = RcShMem<SP>;
type ShMem = RcShMem<SP>;
fn new() -> Result<Self, Error> {
Ok(Self {
@ -335,26 +319,30 @@ where
})
}
fn new_map(&mut self, map_size: usize) -> Result<Self::Mem, Error> {
Ok(Self::Mem {
internal: ManuallyDrop::new(self.internal.borrow_mut().new_map(map_size)?),
fn new_shmem(&mut self, map_size: usize) -> Result<Self::ShMem, Error> {
Ok(Self::ShMem {
internal: ManuallyDrop::new(self.internal.borrow_mut().new_shmem(map_size)?),
provider: self.internal.clone(),
})
}
fn map_from_id_and_size(&mut self, id: ShMemId, size: usize) -> Result<Self::Mem, Error> {
Ok(Self::Mem {
internal: ManuallyDrop::new(self.internal.borrow_mut().map_from_id_and_size(id, size)?),
fn shmem_from_id_and_size(&mut self, id: ShMemId, size: usize) -> Result<Self::ShMem, Error> {
Ok(Self::ShMem {
internal: ManuallyDrop::new(
self.internal
.borrow_mut()
.shmem_from_id_and_size(id, size)?,
),
provider: self.internal.clone(),
})
}
fn release_map(&mut self, map: &mut Self::Mem) {
self.internal.borrow_mut().release_map(&mut map.internal);
fn release_shmem(&mut self, map: &mut Self::ShMem) {
self.internal.borrow_mut().release_shmem(&mut map.internal);
}
fn clone_ref(&mut self, mapping: &Self::Mem) -> Result<Self::Mem, Error> {
Ok(Self::Mem {
fn clone_ref(&mut self, mapping: &Self::ShMem) -> Result<Self::ShMem, Error> {
Ok(Self::ShMem {
internal: ManuallyDrop::new(self.internal.borrow_mut().clone_ref(&mapping.internal)?),
provider: self.internal.clone(),
})
@ -622,7 +610,7 @@ pub mod unix_shmem {
}
}
fn map_from_id_and_size(id: ShMemId, map_size: usize) -> Result<Self, Error> {
fn shmem_from_id_and_size(id: ShMemId, map_size: usize) -> Result<Self, Error> {
unsafe {
let shm_fd: i32 = id.to_string().parse().unwrap();
@ -659,7 +647,7 @@ pub mod unix_shmem {
#[cfg(unix)]
#[derive(Clone, Debug)]
pub struct MmapShMemProvider {
current_map_id: usize,
current_shmem_id: usize,
}
unsafe impl Send for MmapShMemProvider {}
@ -674,22 +662,24 @@ pub mod unix_shmem {
/// Implement [`ShMemProvider`] for [`MmapShMemProvider`].
#[cfg(unix)]
impl ShMemProvider for MmapShMemProvider {
type Mem = MmapShMem;
type ShMem = MmapShMem;
fn new() -> Result<Self, Error> {
Ok(Self { current_map_id: 0 })
Ok(Self {
current_shmem_id: 0,
})
}
fn new_map(&mut self, map_size: usize) -> Result<Self::Mem, Error> {
self.current_map_id += 1;
MmapShMem::new(map_size, self.current_map_id)
fn new_shmem(&mut self, map_size: usize) -> Result<Self::ShMem, Error> {
self.current_shmem_id += 1;
MmapShMem::new(map_size, self.current_shmem_id)
}
fn map_from_id_and_size(
fn shmem_from_id_and_size(
&mut self,
id: ShMemId,
size: usize,
) -> Result<Self::Mem, Error> {
MmapShMem::map_from_id_and_size(id, size)
) -> Result<Self::ShMem, Error> {
MmapShMem::shmem_from_id_and_size(id, size)
}
}
@ -702,11 +692,11 @@ pub mod unix_shmem {
self.map_size
}
fn map(&self) -> &[u8] {
fn as_slice(&self) -> &[u8] {
unsafe { slice::from_raw_parts(self.map, self.map_size) }
}
fn map_mut(&mut self) -> &mut [u8] {
fn as_mut_slice(&mut self) -> &mut [u8] {
unsafe { slice::from_raw_parts_mut(self.map, self.map_size) }
}
}
@ -775,7 +765,7 @@ pub mod unix_shmem {
}
/// Get a [`UnixShMem`] of the existing shared memory mapping identified by id
pub fn map_from_id_and_size(id: ShMemId, map_size: usize) -> Result<Self, Error> {
pub fn shmem_from_id_and_size(id: ShMemId, map_size: usize) -> Result<Self, Error> {
unsafe {
let id_int: i32 = id.into();
let map = shmat(id_int, ptr::null(), 0) as *mut c_uchar;
@ -801,11 +791,11 @@ pub mod unix_shmem {
self.map_size
}
fn map(&self) -> &[u8] {
fn as_slice(&self) -> &[u8] {
unsafe { slice::from_raw_parts(self.map, self.map_size) }
}
fn map_mut(&mut self) -> &mut [u8] {
fn as_mut_slice(&mut self) -> &mut [u8] {
unsafe { slice::from_raw_parts_mut(self.map, self.map_size) }
}
}
@ -838,21 +828,21 @@ pub mod unix_shmem {
/// Implement [`ShMemProvider`] for [`UnixShMemProvider`].
#[cfg(unix)]
impl ShMemProvider for CommonUnixShMemProvider {
type Mem = CommonUnixShMem;
type ShMem = CommonUnixShMem;
fn new() -> Result<Self, Error> {
Ok(Self {})
}
fn new_map(&mut self, map_size: usize) -> Result<Self::Mem, Error> {
fn new_shmem(&mut self, map_size: usize) -> Result<Self::ShMem, Error> {
CommonUnixShMem::new(map_size)
}
fn map_from_id_and_size(
fn shmem_from_id_and_size(
&mut self,
id: ShMemId,
size: usize,
) -> Result<Self::Mem, Error> {
CommonUnixShMem::map_from_id_and_size(id, size)
) -> Result<Self::ShMem, Error> {
CommonUnixShMem::shmem_from_id_and_size(id, size)
}
}
}
@ -959,7 +949,7 @@ pub mod unix_shmem {
}
/// Get a [`crate::bolts::shmem::unix_shmem::UnixShMem`] of the existing [`ShMem`] mapping identified by id.
pub fn map_from_id_and_size(id: ShMemId, map_size: usize) -> Result<Self, Error> {
pub fn shmem_from_id_and_size(id: ShMemId, map_size: usize) -> Result<Self, Error> {
unsafe {
let fd: i32 = id.to_string().parse().unwrap();
#[allow(trivial_numeric_casts, clippy::cast_sign_loss)]
@ -1003,11 +993,11 @@ pub mod unix_shmem {
self.map_size
}
fn map(&self) -> &[u8] {
fn as_slice(&self) -> &[u8] {
unsafe { slice::from_raw_parts(self.map, self.map_size) }
}
fn map_mut(&mut self) -> &mut [u8] {
fn as_mut_slice(&mut self) -> &mut [u8] {
unsafe { slice::from_raw_parts_mut(self.map, self.map_size) }
}
}
@ -1052,23 +1042,23 @@ pub mod unix_shmem {
/// Implement [`ShMemProvider`] for [`AshmemShMemProvider`], for the Android `ShMem`.
#[cfg(unix)]
impl ShMemProvider for AshmemShMemProvider {
type Mem = AshmemShMem;
type ShMem = AshmemShMem;
fn new() -> Result<Self, Error> {
Ok(Self {})
}
fn new_map(&mut self, map_size: usize) -> Result<Self::Mem, Error> {
fn new_shmem(&mut self, map_size: usize) -> Result<Self::ShMem, Error> {
let mapping = AshmemShMem::new(map_size)?;
Ok(mapping)
}
fn map_from_id_and_size(
fn shmem_from_id_and_size(
&mut self,
id: ShMemId,
size: usize,
) -> Result<Self::Mem, Error> {
AshmemShMem::map_from_id_and_size(id, size)
) -> Result<Self::ShMem, Error> {
AshmemShMem::shmem_from_id_and_size(id, size)
}
}
}
@ -1122,7 +1112,7 @@ pub mod win32_shmem {
}
impl Win32ShMem {
fn new_map(map_size: usize) -> Result<Self, Error> {
fn new_shmem(map_size: usize) -> Result<Self, Error> {
unsafe {
let uuid = Uuid::new_v4();
let mut map_str = format!("libafl_{}", uuid.to_simple());
@ -1159,7 +1149,7 @@ pub mod win32_shmem {
}
}
fn map_from_id_and_size(id: ShMemId, map_size: usize) -> Result<Self, Error> {
fn shmem_from_id_and_size(id: ShMemId, map_size: usize) -> Result<Self, Error> {
unsafe {
let map_str_bytes = id.id;
// Unlike MapViewOfFile this one needs u32
@ -1200,11 +1190,11 @@ pub mod win32_shmem {
self.map_size
}
fn map(&self) -> &[u8] {
fn as_slice(&self) -> &[u8] {
unsafe { slice::from_raw_parts(self.map, self.map_size) }
}
fn map_mut(&mut self) -> &mut [u8] {
fn as_mut_slice(&mut self) -> &mut [u8] {
unsafe { slice::from_raw_parts_mut(self.map, self.map_size) }
}
}
@ -1231,17 +1221,21 @@ pub mod win32_shmem {
/// Implement [`ShMemProvider`] for [`Win32ShMemProvider`]
impl ShMemProvider for Win32ShMemProvider {
type Mem = Win32ShMem;
type ShMem = Win32ShMem;
fn new() -> Result<Self, Error> {
Ok(Self {})
}
fn new_map(&mut self, map_size: usize) -> Result<Self::Mem, Error> {
Win32ShMem::new_map(map_size)
fn new_shmem(&mut self, map_size: usize) -> Result<Self::ShMem, Error> {
Win32ShMem::new_shmem(map_size)
}
fn map_from_id_and_size(&mut self, id: ShMemId, size: usize) -> Result<Self::Mem, Error> {
Win32ShMem::map_from_id_and_size(id, size)
fn shmem_from_id_and_size(
&mut self,
id: ShMemId,
size: usize,
) -> Result<Self::ShMem, Error> {
Win32ShMem::shmem_from_id_and_size(id, size)
}
}
}
@ -1280,7 +1274,7 @@ impl<T: ShMem> ShMemCursor<T> {
/// Slice from the current location on this map to the end, mutable
fn empty_slice_mut(&mut self) -> &mut [u8] {
&mut (self.inner.map_mut()[self.pos..])
&mut (self.inner.as_mut_slice()[self.pos..])
}
}
@ -1327,7 +1321,7 @@ impl<T: ShMem> std::io::Seek for ShMemCursor<T> {
let effective_new_pos = match pos {
std::io::SeekFrom::Start(s) => s,
std::io::SeekFrom::End(offset) => {
let map_len = self.inner.map().len();
let map_len = self.inner.as_slice().len();
assert!(i64::try_from(map_len).is_ok());
let signed_pos = map_len as i64;
let effective = signed_pos.checked_add(offset).unwrap();
@ -1360,8 +1354,8 @@ mod tests {
#[serial]
fn test_shmem_service() {
let mut provider = StdShMemProvider::new().unwrap();
let mut map = provider.new_map(1024).unwrap();
map.map_mut()[0] = 1;
assert!(map.map()[0] == 1);
let mut map = provider.new_shmem(1024).unwrap();
map.as_mut_slice()[0] = 1;
assert!(map.as_slice()[0] == 1);
}
}

View File

@ -58,7 +58,7 @@ pub struct StateRestorer<SP>
where
SP: ShMemProvider,
{
shmem: SP::Mem,
shmem: SP::ShMem,
phantom: PhantomData<*const SP>,
}
@ -85,7 +85,7 @@ where
}
/// Create a new [`StateRestorer`].
pub fn new(shmem: SP::Mem) -> Self {
pub fn new(shmem: SP::ShMem) -> Self {
let mut ret = Self {
shmem,
phantom: PhantomData,
@ -178,7 +178,7 @@ where
}
fn content_mut(&mut self) -> &mut StateShMemContent {
let ptr = self.shmem.map().as_ptr();
let ptr = self.shmem.as_slice().as_ptr();
#[allow(clippy::cast_ptr_alignment)] // Beginning of the page will always be aligned
unsafe {
&mut *(ptr as *mut StateShMemContent)
@ -188,7 +188,7 @@ where
/// The content is either the name of the tmpfile, or the serialized bytes directly, if they fit on a single page.
fn content(&self) -> &StateShMemContent {
#[allow(clippy::cast_ptr_alignment)] // Beginning of the page will always be aligned
let ptr = self.shmem.map().as_ptr() as *const StateShMemContent;
let ptr = self.shmem.as_slice().as_ptr() as *const StateShMemContent;
unsafe { &*(ptr) }
}
@ -251,7 +251,7 @@ mod tests {
const TESTMAP_SIZE: usize = 1024;
let mut shmem_provider = StdShMemProvider::new().unwrap();
let shmem = shmem_provider.new_map(TESTMAP_SIZE).unwrap();
let shmem = shmem_provider.new_shmem(TESTMAP_SIZE).unwrap();
let mut state_restorer = StateRestorer::<StdShMemProvider>::new(shmem);
let state = "hello world".to_string();

View File

@ -835,7 +835,7 @@ where
// First, create a channel from the current fuzzer to the next to store state between restarts.
let staterestorer: StateRestorer<SP> =
StateRestorer::new(self.shmem_provider.new_map(256 * 1024 * 1024)?);
StateRestorer::new(self.shmem_provider.new_shmem(256 * 1024 * 1024)?);
// Store the information to a map.
staterestorer.write_to_env(_ENV_FUZZER_SENDER)?;
@ -976,7 +976,7 @@ mod tests {
let mut llmp_client = LlmpClient::new(
shmem_provider.clone(),
LlmpSharedMap::new(0, shmem_provider.new_map(1024).unwrap()),
LlmpSharedMap::new(0, shmem_provider.new_shmem(1024).unwrap()),
)
.unwrap();
@ -1007,7 +1007,7 @@ mod tests {
// First, create a channel from the current fuzzer to the next to store state between restarts.
let mut staterestorer = StateRestorer::<StdShMemProvider>::new(
shmem_provider.new_map(256 * 1024 * 1024).unwrap(),
shmem_provider.new_shmem(256 * 1024 * 1024).unwrap(),
);
staterestorer.reset();

View File

@ -333,7 +333,7 @@ where
let mut staterestorer = if std::env::var(_ENV_FUZZER_SENDER).is_err() {
// First, create a place to store state in, for restarts.
let staterestorer: StateRestorer<SP> =
StateRestorer::new(shmem_provider.new_map(256 * 1024 * 1024)?);
StateRestorer::new(shmem_provider.new_shmem(256 * 1024 * 1024)?);
//let staterestorer = { LlmpSender::new(shmem_provider.clone(), 0, false)? };
staterestorer.write_to_env(_ENV_FUZZER_SENDER)?;

View File

@ -18,7 +18,7 @@ use std::{
use crate::{
bolts::{
os::{dup2, pipes::Pipe},
shmem::{ShMem, ShMemProvider, StdShMem, StdShMemProvider},
shmem::{ShMem, ShMemProvider, StdShMemProvider},
},
executors::{Executor, ExitKind, HasObservers},
inputs::{HasTargetBytes, Input},
@ -354,6 +354,9 @@ impl Forkserver {
/// A struct that has a forkserver
pub trait HasForkserver {
/// The [`ShMemProvider`] used for this forkserver's map
type SP: ShMemProvider;
/// The forkserver
fn forkserver(&self) -> &Forkserver;
@ -367,10 +370,10 @@ pub trait HasForkserver {
fn out_file_mut(&mut self) -> &mut OutFile;
/// The map of the fuzzer
fn map(&self) -> &Option<StdShMem>;
fn shmem(&self) -> &Option<<<Self as HasForkserver>::SP as ShMemProvider>::ShMem>;
/// The map of the fuzzer, mutable
fn map_mut(&mut self) -> &mut Option<StdShMem>;
fn shmem_mut(&mut self) -> &mut Option<<<Self as HasForkserver>::SP as ShMemProvider>::ShMem>;
}
/// The timeout forkserver executor that wraps around the standard forkserver executor and sets a timeout before each run.
@ -417,14 +420,14 @@ where
let last_run_timed_out = self.executor.forkserver().last_run_timed_out();
match &mut self.executor.map_mut() {
Some(map) => {
match &mut self.executor.shmem_mut() {
Some(shmem) => {
let target_bytes = input.target_bytes();
let size = target_bytes.as_slice().len();
let size_in_bytes = size.to_ne_bytes();
// The first four bytes tells the size of the shmem.
map.map_mut()[..4].copy_from_slice(&size_in_bytes[..4]);
map.map_mut()[SHMEM_FUZZ_HDR_SIZE..(SHMEM_FUZZ_HDR_SIZE + size)]
shmem.as_mut_slice()[..4].copy_from_slice(&size_in_bytes[..4]);
shmem.as_mut_slice()[SHMEM_FUZZ_HDR_SIZE..(SHMEM_FUZZ_HDR_SIZE + size)]
.copy_from_slice(target_bytes.as_slice());
}
None => {
@ -498,24 +501,26 @@ where
/// This [`Executor`] can run binaries compiled for AFL/AFL++ that make use of a forkserver.
/// Shared memory feature is also available, but you have to set things up in your code.
/// Please refer to AFL++'s docs. <https://github.com/AFLplusplus/AFLplusplus/blob/stable/instrumentation/README.persistent_mode.md>
pub struct ForkserverExecutor<I, OT, S>
pub struct ForkserverExecutor<I, OT, S, SP>
where
I: Input + HasTargetBytes,
OT: ObserversTuple<I, S>,
SP: ShMemProvider,
{
target: String,
args: Vec<String>,
out_file: OutFile,
forkserver: Forkserver,
observers: OT,
map: Option<StdShMem>,
map: Option<SP::ShMem>,
phantom: PhantomData<(I, S)>,
}
impl<I, OT, S> Debug for ForkserverExecutor<I, OT, S>
impl<I, OT, S, SP> Debug for ForkserverExecutor<I, OT, S, SP>
where
I: Input + HasTargetBytes,
OT: ObserversTuple<I, S>,
SP: ShMemProvider,
{
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
f.debug_struct("ForkserverExecutor")
@ -529,28 +534,54 @@ where
}
}
impl<I, OT, S> ForkserverExecutor<I, OT, S>
impl<I, OT, S> ForkserverExecutor<I, OT, S, StdShMemProvider>
where
I: Input + HasTargetBytes,
OT: ObserversTuple<I, S>,
{
/// Creates a new [`ForkserverExecutor`] with the given target, arguments and observers.
/// Creates a new `AFL`-style [`ForkserverExecutor`] with the given target, arguments and observers.
/// This Forserver won't attempt to provide inputs over shared mem but write them to an iput file
/// If `debug_child` is set, the child will print to `stdout`/`stderr`.
pub fn new(
target: String,
arguments: &[String],
use_shmem_testcase: bool,
observers: OT,
debug_child: bool,
) -> Result<Self, Error> {
Self::with_debug(target, arguments, use_shmem_testcase, observers, false)
Self::new_internal(target, arguments, observers, debug_child, None)
}
}
impl<I, OT, S, SP> ForkserverExecutor<I, OT, S, SP>
where
I: Input + HasTargetBytes,
OT: ObserversTuple<I, S>,
SP: ShMemProvider,
{
/// Creates a new [`ForkserverExecutor`] with the given target, arguments and observers.
pub fn with_shmem_inputs(
target: String,
arguments: &[String],
observers: OT,
debug_child: bool,
shmem_provider: &mut SP,
) -> Result<Self, Error> {
Self::new_internal(
target,
arguments,
observers,
debug_child,
Some(shmem_provider),
)
}
/// Creates a new [`ForkserverExecutor`] with the given target, arguments and observers, with debug mode
pub fn with_debug(
fn new_internal(
target: String,
arguments: &[String],
use_shmem_testcase: bool,
observers: OT,
debug_output: bool,
debug_child: bool,
shmem_provider: Option<&mut SP>,
) -> Result<Self, Error> {
let mut args = Vec::<String>::new();
let mut use_stdin = true;
@ -567,17 +598,18 @@ where
let out_file = OutFile::new(&out_filename)?;
let mut map = None;
if use_shmem_testcase {
// setup shared memory
let mut provider = StdShMemProvider::new()?;
let mut shmem = provider.new_map(MAX_FILE + SHMEM_FUZZ_HDR_SIZE)?;
shmem.write_to_env("__AFL_SHM_FUZZ_ID")?;
let map = match shmem_provider {
None => None,
Some(provider) => {
// setup shared memory
let mut shmem = provider.new_shmem(MAX_FILE + SHMEM_FUZZ_HDR_SIZE)?;
shmem.write_to_env("__AFL_SHM_FUZZ_ID")?;
let size_in_bytes = (MAX_FILE + SHMEM_FUZZ_HDR_SIZE).to_ne_bytes();
shmem.map_mut()[..4].clone_from_slice(&size_in_bytes[..4]);
map = Some(shmem);
}
let size_in_bytes = (MAX_FILE + SHMEM_FUZZ_HDR_SIZE).to_ne_bytes();
shmem.as_mut_slice()[..4].clone_from_slice(&size_in_bytes[..4]);
Some(shmem)
}
};
let mut forkserver = Forkserver::new(
target.clone(),
@ -585,7 +617,7 @@ where
out_file.as_raw_fd(),
use_stdin,
0,
debug_output,
debug_child,
)?;
let (rlen, status) = forkserver.read_st()?; // Initial handshake, read 4-bytes hello message from the forkserver.
@ -598,7 +630,7 @@ where
println!("All right - fork server is up.");
// If forkserver is responding, we then check if there's any option enabled.
if status & FS_OPT_ENABLED == FS_OPT_ENABLED {
if (status & FS_OPT_SHDMEM_FUZZ == FS_OPT_SHDMEM_FUZZ) & use_shmem_testcase {
if (status & FS_OPT_SHDMEM_FUZZ == FS_OPT_SHDMEM_FUZZ) & map.is_some() {
println!("Using SHARED MEMORY FUZZING feature.");
let send_status = FS_OPT_ENABLED | FS_OPT_SHDMEM_FUZZ;
@ -645,10 +677,11 @@ where
}
}
impl<EM, I, OT, S, Z> Executor<EM, I, S, Z> for ForkserverExecutor<I, OT, S>
impl<EM, I, OT, S, SP, Z> Executor<EM, I, S, Z> for ForkserverExecutor<I, OT, S, SP>
where
I: Input + HasTargetBytes,
OT: ObserversTuple<I, S>,
SP: ShMemProvider,
{
#[inline]
fn run_target(
@ -667,8 +700,8 @@ where
let size = target_bytes.as_slice().len();
let size_in_bytes = size.to_ne_bytes();
// The first four bytes tells the size of the shmem.
map.map_mut()[..4].copy_from_slice(&size_in_bytes[..4]);
map.map_mut()[SHMEM_FUZZ_HDR_SIZE..(SHMEM_FUZZ_HDR_SIZE + size)]
map.as_mut_slice()[..4].copy_from_slice(&size_in_bytes[..4]);
map.as_mut_slice()[SHMEM_FUZZ_HDR_SIZE..(SHMEM_FUZZ_HDR_SIZE + size)]
.copy_from_slice(target_bytes.as_slice());
}
None => {
@ -719,10 +752,11 @@ where
}
}
impl<I, OT, S> HasObservers<I, OT, S> for ForkserverExecutor<I, OT, S>
impl<I, OT, S, SP> HasObservers<I, OT, S> for ForkserverExecutor<I, OT, S, SP>
where
I: Input + HasTargetBytes,
OT: ObserversTuple<I, S>,
SP: ShMemProvider,
{
#[inline]
fn observers(&self) -> &OT {
@ -735,11 +769,14 @@ where
}
}
impl<I, OT, S> HasForkserver for ForkserverExecutor<I, OT, S>
impl<I, OT, S, SP> HasForkserver for ForkserverExecutor<I, OT, S, SP>
where
I: Input + HasTargetBytes,
OT: ObserversTuple<I, S>,
SP: ShMemProvider,
{
type SP = SP;
#[inline]
fn forkserver(&self) -> &Forkserver {
&self.forkserver
@ -761,12 +798,12 @@ where
}
#[inline]
fn map(&self) -> &Option<StdShMem> {
fn shmem(&self) -> &Option<SP::ShMem> {
&self.map
}
#[inline]
fn map_mut(&mut self) -> &mut Option<StdShMem> {
fn shmem_mut(&mut self) -> &mut Option<SP::ShMem> {
&mut self.map
}
}
@ -808,20 +845,23 @@ mod tests {
let bin = "echo";
let args = vec![String::from("@@")];
let mut shmem = StdShMemProvider::new().unwrap().new_map(MAP_SIZE).unwrap();
let mut shmem_provider = StdShMemProvider::new().unwrap();
let mut shmem = shmem_provider.new_shmem(MAP_SIZE).unwrap();
shmem.write_to_env("__AFL_SHM_ID").unwrap();
let shmem_map = shmem.map_mut();
let shmem_buf = shmem.as_mut_slice();
let edges_observer = HitcountsMapObserver::new(ConstMapObserver::<_, MAP_SIZE>::new(
"shared_mem",
shmem_map,
shmem_buf,
));
let executor = ForkserverExecutor::<NopInput, _, ()>::new(
let executor = ForkserverExecutor::<NopInput, _, (), _>::with_shmem_inputs(
bin.to_string(),
&args,
false,
tuple_list!(edges_observer),
false,
&mut shmem_provider,
);
// Since /usr/bin/echo is not a instrumented binary file, the test will just check if the forkserver has failed at the initial handshake

View File

@ -488,7 +488,7 @@ impl<T: ShMem> MessageFileWriter<ShMemCursor<T>> {
}
}
impl MessageFileWriter<ShMemCursor<<StdShMemProvider as ShMemProvider>::Mem>> {
impl MessageFileWriter<ShMemCursor<<StdShMemProvider as ShMemProvider>::ShMem>> {
/// Creates a new `MessageFileWriter` by reading a [`ShMem`] from the given environment variable.
pub fn from_stdshmem_env_with_name(env_name: impl AsRef<str>) -> io::Result<Self> {
Self::from_shmem(
@ -507,4 +507,4 @@ impl MessageFileWriter<ShMemCursor<<StdShMemProvider as ShMemProvider>::Mem>> {
/// A writer that will write messages to a shared memory buffer.
pub type StdShMemMessageFileWriter =
MessageFileWriter<ShMemCursor<<StdShMemProvider as ShMemProvider>::Mem>>;
MessageFileWriter<ShMemCursor<<StdShMemProvider as ShMemProvider>::ShMem>>;

View File

@ -194,7 +194,9 @@ where
let val = unsafe {
// SAFETY: the index is modulo by the length, therefore it is always in bounds
let len = self.hitcounts_map.len();
self.hitcounts_map.map_mut().get_unchecked_mut(hash % len)
self.hitcounts_map
.as_mut_slice()
.get_unchecked_mut(hash % len)
};
*val = val.saturating_add(1);
}

View File

@ -62,7 +62,7 @@ fn main() {
let mut shmemprovider = StdShMemProvider::default();
let concolic_shmem = shmemprovider
.new_map(1024 * 1024 * 1024)
.new_shmem(1024 * 1024 * 1024)
.expect("unable to create shared mapping");
concolic_shmem
.write_to_env(DEFAULT_ENV_NAME)
@ -70,7 +70,7 @@ fn main() {
let coverage_map = StdShMemProvider::new()
.unwrap()
.new_map(COVERAGE_MAP_SIZE)
.new_shmem(COVERAGE_MAP_SIZE)
.unwrap();
//let the forkserver know the shmid
coverage_map.write_to_env(HITMAP_ENV_NAME).unwrap();
@ -104,7 +104,7 @@ fn main() {
File::create(coverage_file_path).expect("unable to open coverage file"),
);
for (index, count) in coverage_map
.map()
.as_slice()
.iter()
.enumerate()
.filter(|(_, &v)| v != 0)
@ -117,7 +117,7 @@ fn main() {
let output_file_path = opt.output.unwrap_or_else(|| "trace".into());
let mut output_file =
BufWriter::new(File::create(output_file_path).expect("unable to open output file"));
let mut reader = MessageFileReader::from_length_prefixed_buffer(concolic_shmem.map())
let mut reader = MessageFileReader::from_length_prefixed_buffer(concolic_shmem.as_slice())
.expect("unable to create trace reader");
if opt.plain_text {
while let Some(message) = reader.next_message() {

View File

@ -101,17 +101,15 @@ impl<'a, const MAP_SIZE: usize> ForkserverBytesCoverageSugar<'a, MAP_SIZE> {
out_dir.push("queue");
let shmem_provider = StdShMemProvider::new().expect("Failed to init shared memory");
let mut shmem_provider_client = shmem_provider.clone();
let monitor = MultiMonitor::new(|s| println!("{}", s));
let mut run_client = |state: Option<StdState<_, _, _, _, _>>, mut mgr, _core_id| {
// Coverage map shared between target and fuzzer
let mut shmem = StdShMemProvider::new()
.expect("Failed to init shared memory")
.new_map(MAP_SIZE)
.unwrap();
let mut shmem = shmem_provider_client.new_shmem(MAP_SIZE).unwrap();
shmem.write_to_env("__AFL_SHM_ID").unwrap();
let shmem_map = shmem.map_mut();
let shmem_map = shmem.as_mut_slice();
// Create an observation channel using the coverage map
let edges_observer = unsafe {
@ -171,13 +169,22 @@ impl<'a, const MAP_SIZE: usize> ForkserverBytesCoverageSugar<'a, MAP_SIZE> {
// Create the executor for an in-process function with one observer for edge coverage and one for the execution time
let mut executor = TimeoutForkserverExecutor::new(
ForkserverExecutor::with_debug(
self.program.clone(),
self.arguments,
self.shmem_testcase,
tuple_list!(edges_observer, time_observer),
self.debug_output,
)
if self.shmem_testcase {
ForkserverExecutor::with_shmem_inputs(
self.program.clone(),
self.arguments,
tuple_list!(edges_observer, time_observer),
self.debug_output,
&mut shmem_provider_client,
)
} else {
ForkserverExecutor::new(
self.program.clone(),
self.arguments,
tuple_list!(edges_observer, time_observer),
self.debug_output,
)
}
.expect("Failed to create the executor."),
timeout,
)