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:
parent
ac43997950
commit
4f6f76e857
@ -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.
|
||||
|
@ -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.
|
||||
|
@ -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(
|
||||
|
@ -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();
|
||||
|
||||
|
@ -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)
|
||||
|
@ -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),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
@ -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);
|
||||
}
|
||||
}
|
||||
|
@ -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();
|
||||
|
@ -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();
|
||||
|
@ -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)?;
|
||||
|
||||
|
@ -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
|
||||
|
||||
|
@ -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>>;
|
||||
|
@ -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);
|
||||
}
|
||||
|
@ -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() {
|
||||
|
@ -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,
|
||||
)
|
||||
|
Loading…
x
Reference in New Issue
Block a user