From 9bc22d033111f747d3fae165e37f820b631c29e6 Mon Sep 17 00:00:00 2001 From: Dominik Maier Date: Sun, 11 Apr 2021 02:21:09 +0200 Subject: [PATCH] Llmp Fixes (#51) * fixed llmp --- libafl/src/bolts/llmp.rs | 60 ++++++++++++++++++++++----------------- libafl/src/bolts/shmem.rs | 8 +++--- 2 files changed, 38 insertions(+), 30 deletions(-) diff --git a/libafl/src/bolts/llmp.rs b/libafl/src/bolts/llmp.rs index 57d7573c0e..af745e6d2c 100644 --- a/libafl/src/bolts/llmp.rs +++ b/libafl/src/bolts/llmp.rs @@ -237,7 +237,7 @@ fn msg_offset_from_env(env_name: &str) -> Result, Error> { fn new_map_size(max_alloc: usize) -> usize { max( max_alloc * 2 + EOP_MSG_SIZE + LLMP_PAGE_HEADER_LEN, - LLMP_CFG_INITIAL_MAP_SIZE, + LLMP_CFG_INITIAL_MAP_SIZE - 1, ) .next_power_of_two() } @@ -546,7 +546,7 @@ where last_msg_sent: ptr::null_mut(), out_maps: vec![LlmpSharedMap::new( 0, - SH::new_map(new_map_size(LLMP_CFG_INITIAL_MAP_SIZE))?, + SH::new_map(LLMP_CFG_INITIAL_MAP_SIZE)?, )], // drop pages to the broker if it already read them keep_pages_forever, @@ -659,7 +659,10 @@ where if (*ret).tag == LLMP_TAG_UNINITIALIZED { panic!("Did not call send() on last message!"); } - (*ret).buf_len_padded = size_of::() as u64; + (*ret).buf_len = size_of::() as u64; + + // We don't need to pad the EOP message: it'll always be the last in this page. + (*ret).buf_len_padded = (*ret).buf_len; (*ret).message_id = if !last_msg.is_null() { (*last_msg).message_id + 1 } else { @@ -677,11 +680,13 @@ where let buf_len_padded; let mut complete_msg_size = llmp_align(size_of::() + buf_len); let map = self.out_maps.last_mut().unwrap(); - println!("Allocating on map {:?} (size: {})", map, complete_msg_size); let page = map.page_mut(); - println!("Allocating {} (>={}) bytes on page {:?}", complete_msg_size, buf_len, page); let last_msg = self.last_msg_sent; - println!("last msg: {:?}", last_msg); + #[cfg(all(feature = "llmp_debug", feature = "std"))] + println!( + "Allocating {} (>={}) bytes on page {:?} / map {:?} (last msg: {:?})", + complete_msg_size, buf_len, page, map, last_msg + ); /* DBG("XXX complete_msg_size %lu (h: %lu)\n", complete_msg_size, sizeof(llmp_message)); */ /* In case we don't have enough space, make sure the next page will be large * enough */ @@ -700,11 +705,16 @@ where buf_len_padded = llmp_align(base_addr + complete_msg_size) - base_addr - size_of::(); complete_msg_size = buf_len_padded + size_of::(); - dbg!("LLMP_DEBUG: complete_msg_size NEW {}", complete_msg_size); /* Still space for the new message plus the additional "we're full" message? */ - dbg!((*page).size_used, complete_msg_size, EOP_MSG_SIZE, (*page).size_total); + #[cfg(all(feature = "llmp_debug", feature = "std"))] + dbg!( + (*page).size_used, + complete_msg_size, + EOP_MSG_SIZE, + (*page).size_total + ); if (*page).size_used + complete_msg_size + EOP_MSG_SIZE > (*page).size_total { /* We're full. */ return None; @@ -760,9 +770,7 @@ where (*ret).buf_len = buf_len as u64; /* DBG("Returning new message at %p with len %ld, TAG was %x", ret, ret->buf_len_padded, ret->tag); */ /* Maybe catch some bugs... */ - (dbg!(*_llmp_next_msg_ptr(ret))).tag = LLMP_TAG_UNSET; - dbg!(ret); - dbg!(*ret); + (*_llmp_next_msg_ptr(ret)).tag = LLMP_TAG_UNSET; (*ret).tag = LLMP_TAG_UNINITIALIZED; Some(ret) } @@ -772,7 +780,7 @@ where /// It will be read by the consuming threads (broker->clients or client->broker) #[inline(never)] // Not inlined to make cpu-level reodering (hopefully?) improbable unsafe fn send(&mut self, msg: *mut LlmpMsg) -> Result<(), Error> { - dbg!("Sending msg {:?}", msg); + // dbg!("Sending msg {:?}", msg); if self.last_msg_sent == msg { panic!("Message sent twice!"); @@ -814,6 +822,7 @@ where let old_map = self.out_maps.last_mut().unwrap().page_mut(); + #[cfg(all(feature = "llmp_debug", feature = "std"))] println!("New Map Size {}", new_map_size((*old_map).max_alloc_size)); // Create a new shard page. @@ -823,30 +832,26 @@ where ); let mut new_map = new_map_shmem.page_mut(); + #[cfg(all(feature = "llmp_debug", feature = "std"))] println!("got new map at: {:?}", new_map); ptr::write_volatile(&mut (*new_map).current_msg_id, (*old_map).current_msg_id); + #[cfg(all(feature = "llmp_debug", feature = "std"))] println!("Setting max alloc size: {:?}", (*old_map).max_alloc_size); (*new_map).max_alloc_size = (*old_map).max_alloc_size; /* On the old map, place a last message linking to the new map for the clients * to consume */ let mut out: *mut LlmpMsg = self.alloc_eop()?; - println!("out"); (*out).sender = (*old_map).sender; - println!("sender"); let mut end_of_page_msg = (*out).buf.as_mut_ptr() as *mut LlmpPayloadSharedMapInfo; - println!("size"); (*end_of_page_msg).map_size = new_map_shmem.shmem.map().len(); - println!("str"); (*end_of_page_msg).shm_str = *new_map_shmem.shmem.shm_slice(); - println!("send"); /* Send the last msg on the old buf */ self.send(out)?; - println!("sent"); // Set the new page as current page. self.out_maps.push(new_map_shmem); @@ -855,10 +860,10 @@ where // If we want to get red if old pages, (client to broker), do that now if !self.keep_pages_forever { + #[cfg(all(feature = "llmp_debug", feature = "std"))] println!("pruning"); self.prune_old_pages(); } - println!("Done"); Ok(()) } @@ -874,6 +879,7 @@ where self.handle_out_eop()?; } + #[cfg(all(feature = "llmp_debug", feature = "std"))] println!("Handled out eop"); match unsafe { self.alloc_next_if_space(buf_len) } { @@ -1046,7 +1052,8 @@ where // Handle end of page if (*msg).buf_len < size_of::() as u64 { panic!( - "Illegal message length for EOP (is {}, expected {})", + "Illegal message length for EOP (is {}/{}, expected {})", + (*msg).buf_len, (*msg).buf_len_padded, size_of::() ); @@ -1057,6 +1064,9 @@ where Copy the contents first to be safe (probably fine in rust either way). */ let pageinfo_cpy = *pageinfo; + // Set last msg we received to null (as the map may no longer exist) + self.last_msg_recvd = ptr::null(); + // Mark the old page save to unmap, in case we didn't so earlier. ptr::write_volatile(&mut (*page).save_to_unmap, 1); @@ -1397,7 +1407,8 @@ where If we should need zero copy, we could instead post a link to the original msg with the map_id and offset. */ let actual_size = (*out).buf_len_padded; - msg.copy_to_nonoverlapping(out, size_of::() + (*msg).buf_len_padded as usize); + let complete_size = actual_size as usize + size_of::(); + (msg as *const u8).copy_to_nonoverlapping(out as *mut u8, complete_size); (*out).buf_len_padded = actual_size; /* We need to replace the message ID with our own */ if let Err(e) = self.llmp_out.send(out) { @@ -1507,10 +1518,7 @@ where let llmp_tcp_id = self.llmp_clients.len() as u32; // Tcp out map sends messages from background thread tcp server to foreground client - let tcp_out_map = LlmpSharedMap::new( - llmp_tcp_id, - SH::new_map(new_map_size(LLMP_CFG_INITIAL_MAP_SIZE))?, - ); + let tcp_out_map = LlmpSharedMap::new(llmp_tcp_id, SH::new_map(LLMP_CFG_INITIAL_MAP_SIZE)?); let tcp_out_map_str = tcp_out_map.shmem.shm_str(); let tcp_out_map_size = tcp_out_map.shmem.map().len(); self.register_client(tcp_out_map); @@ -1822,7 +1830,7 @@ where last_msg_sent: ptr::null_mut(), out_maps: vec![LlmpSharedMap::new( 0, - SH::new_map(new_map_size(LLMP_CFG_INITIAL_MAP_SIZE))?, + SH::new_map(LLMP_CFG_INITIAL_MAP_SIZE)?, )], // drop pages to the broker if it already read them keep_pages_forever: false, diff --git a/libafl/src/bolts/shmem.rs b/libafl/src/bolts/shmem.rs index f87b7afbfc..b2f43c4344 100644 --- a/libafl/src/bolts/shmem.rs +++ b/libafl/src/bolts/shmem.rs @@ -344,13 +344,13 @@ pub mod unix_shmem { pub fn new(map_size: usize) -> Result { let mut ret = unix_shmem_unitialized(); let map = unsafe { unix_shmem_init(&mut ret, map_size) }; - if !map.is_null() { - Ok(ret) - } else { + if map.is_null() { Err(Error::Unknown(format!( - "Could not allocate map of size {}", + "Could not allocate map of size {} - check OS limits, (i.e. shmall, shmmax)!", map_size ))) + } else { + Ok(ret) } } }