tiny refactorings

This commit is contained in:
Dominik Maier 2020-12-19 15:42:15 +01:00
parent ee6d4ab3d2
commit 5c7810763d
2 changed files with 48 additions and 34 deletions

View File

@ -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)
}
}

View File

@ -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<Self, AflError>;
fn existing_from_shm_slice(map_str_bytes: &[u8; 20], map_size: usize)
-> Result<Self, AflError>;
/// Initialize from a shm_str with fixed len of 20
fn existing_from_name(shm_str: &str, map_size: usize) -> Result<Self, AflError> {
fn existing_from_shm_str(shm_str: &str, map_size: usize) -> Result<Self, AflError> {
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<Self, AflError>;
/// 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<Self, AflError> {
let map_shm_str = env::var(env_name)?;
let map_size = str::parse::<usize>(&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, AflError> {
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, AflError> {
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<Self, AflError> {
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<Self, AflError> {
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");
}
}