diff --git a/afl/src/events/llmp.rs b/afl/src/events/llmp.rs index 9a2616e4e8..e27001c4d2 100644 --- a/afl/src/events/llmp.rs +++ b/afl/src/events/llmp.rs @@ -594,7 +594,7 @@ where let mut end_of_page_msg = (*out).buf.as_mut_ptr() as *mut LlmpPayloadSharedMapInfo; (*end_of_page_msg).map_size = new_map_shmem.shmem.map().len(); - (*end_of_page_msg).shm_str = *new_map_shmem.shmem.shm_str_buf(); + (*end_of_page_msg).shm_str = *new_map_shmem.shmem.shm_slice(); // We never sent a msg on the new buf */ self.last_msg_sent = 0 as *mut LlmpMsg; @@ -727,7 +727,7 @@ where ptr::write_volatile(&mut (*page).save_to_unmap, 1); // Map the new page. The old one should be unmapped by Drop self.current_recv_map = - LlmpSharedMap::existing(SH::existing_map_by_shm_bytes( + LlmpSharedMap::existing(SH::existing_from_shm_slice( &pageinfo_cpy.shm_str, pageinfo_cpy.map_size, )?); @@ -962,7 +962,7 @@ where // to read from the initial map id. let client_out_map_mem = &self.llmp_out.out_maps.first().unwrap().shmem; - let broadcast_str_initial = client_out_map_mem.shm_str_buf().clone(); + let broadcast_str_initial = client_out_map_mem.shm_slice().clone(); let llmp_tcp_id = self.llmp_clients.len() as u32; @@ -980,7 +980,7 @@ where id: 0, last_msg_sent: 0 as *mut LlmpMsg, out_maps: vec![LlmpSharedMap::existing( - SH::existing_from_name(&tcp_out_map_str, tcp_out_map_size).unwrap(), + SH::existing_from_shm_str(&tcp_out_map_str, tcp_out_map_size).unwrap(), )], // drop pages to the broker if it already read them keep_pages_forever: false, @@ -1070,8 +1070,7 @@ where } else { let pageinfo = (*msg).buf.as_mut_ptr() as *mut LlmpPayloadSharedMapInfo; - match SH::existing_map_by_shm_bytes(&(*pageinfo).shm_str, (*pageinfo).map_size) - { + match SH::existing_from_shm_slice(&(*pageinfo).shm_str, (*pageinfo).map_size) { Ok(new_map) => { let new_page = LlmpSharedMap::existing(new_map); let id = next_id; @@ -1215,12 +1214,12 @@ where let mut new_broker_map_str: [u8; 20] = Default::default(); stream.read_exact(&mut new_broker_map_str)?; - let ret = Self::new(LlmpSharedMap::existing(SH::existing_map_by_shm_bytes( + let ret = Self::new(LlmpSharedMap::existing(SH::existing_from_shm_slice( &new_broker_map_str, LLMP_PREF_INITIAL_MAP_SIZE, )?))?; - stream.write(ret.llmp_out.out_maps.first().unwrap().shmem.shm_str_buf())?; + stream.write(ret.llmp_out.out_maps.first().unwrap().shmem.shm_slice())?; Ok(ret) } } diff --git a/afl/src/events/shmem.rs b/afl/src/events/shmem.rs index 17f6bda842..bcd92b269c 100644 --- a/afl/src/events/shmem.rs +++ b/afl/src/events/shmem.rs @@ -60,28 +60,32 @@ struct shmid_ds { /// A Shared map pub trait ShMem: Sized { /// Creates a nes variable with the given name, strigified to 20 bytes. - fn existing_map_by_shm_bytes( - map_str_bytes: &[u8; 20], - map_size: usize, - ) -> Result; + fn existing_from_shm_slice(map_str_bytes: &[u8; 20], map_size: usize) + -> Result; /// Initialize from a shm_str with fixed len of 20 - fn existing_from_name(shm_str: &str, map_size: usize) -> Result { + fn existing_from_shm_str(shm_str: &str, map_size: usize) -> Result { let mut slice: [u8; 20] = [0; 20]; for (i, val) in shm_str.as_bytes().iter().enumerate() { slice[i] = *val; } - Self::existing_map_by_shm_bytes(&slice, map_size) + Self::existing_from_shm_slice(&slice, map_size) } /// Creates a new map with the given size fn new_map(map_size: usize) -> Result; /// The string to identify this shm - fn shm_str(&self) -> String; + fn shm_str(&self) -> String { + let bytes = self.shm_slice(); + let eof_pos = bytes.iter().position(|&c| c == 0).unwrap(); + alloc::str::from_utf8(&bytes[..eof_pos]) + .unwrap() + .to_string() + } /// Let's just fix this to a large enough buf - fn shm_str_buf(&self) -> &[u8; 20]; + fn shm_slice(&self) -> &[u8; 20]; /// The actual shared map, in memory fn map(&self) -> &[u8]; @@ -104,7 +108,7 @@ pub trait ShMem: Sized { fn existing_from_env(env_name: &str) -> Result { let map_shm_str = env::var(env_name)?; let map_size = str::parse::(&env::var(format!("{}_SIZE", env_name))?)?; - Self::existing_from_name(&map_shm_str, map_size) + Self::existing_from_shm_str(&map_shm_str, map_size) } } @@ -119,24 +123,21 @@ pub struct AflShmem { #[cfg(feature = "std")] impl ShMem for AflShmem { - fn existing_map_by_shm_bytes( + fn existing_from_shm_slice( map_str_bytes: &[u8; 20], map_size: usize, ) -> Result { - Self::from_name_slice(map_str_bytes, map_size) + unsafe { + let str_bytes = map_str_bytes as *const [u8; 20] as *const libc::c_char; + Self::from_str(CStr::from_ptr(str_bytes), map_size) + } } fn new_map(map_size: usize) -> Result { Self::new(map_size) } - fn shm_str(&self) -> String { - unsafe { CStr::from_ptr(self.shm_str.as_ptr() as *const i8) } - .to_string_lossy() - .into() - } - - fn shm_str_buf(&self) -> &[u8; 20] { + fn shm_slice(&self) -> &[u8; 20] { &self.shm_str } @@ -185,14 +186,6 @@ impl AflShmem { } } - /// Generate a shared map with a fixed byte array of 20 - pub fn from_name_slice(shm_str: &[u8; 20], map_size: usize) -> Result { - unsafe { - let str_bytes = shm_str as *const [u8; 20] as *const libc::c_char; - Self::from_str(CStr::from_ptr(str_bytes), map_size) - } - } - pub fn new(map_size: usize) -> Result { let mut ret = afl_shmem_unitialized(); let map = unsafe { afl_shmem_init(&mut ret, map_size) }; @@ -281,3 +274,25 @@ unsafe fn afl_shmem_by_str(shm: *mut AflShmem, shm_str: &CStr, map_size: usize) } return (*shm).map; } + +#[cfg(test)] +mod tests { + use super::{AflShmem, ShMem}; + + #[cfg(feature = "std")] + #[test] + fn test_str_conversions() { + let mut shm_str: [u8; 20] = [0; 20]; + shm_str[0] = 'A' as u8; + shm_str[1] = 'B' as u8; + shm_str[2] = 'C' as u8; + let faux_shmem = AflShmem { + shm_id: 0, + shm_str, + map: 0 as *mut u8, + map_size: 20, + }; + let str = faux_shmem.shm_str(); + assert_eq!(str, "ABC"); + } +}