From b7c49db3b1f87f50319bfe5c7298acb97ce0fe52 Mon Sep 17 00:00:00 2001 From: Dominik Maier Date: Fri, 27 Nov 2020 12:49:41 +0100 Subject: [PATCH] rewrote llmp broker to be more rusty --- afl/llmp_test/src/main.rs | 10 +- afl/src/events/llmp_translated.rs | 818 +++++++++++++++--------------- afl/src/events/mod.rs | 29 +- 3 files changed, 443 insertions(+), 414 deletions(-) diff --git a/afl/llmp_test/src/main.rs b/afl/llmp_test/src/main.rs index 9a72ccd955..425a69c429 100644 --- a/afl/llmp_test/src/main.rs +++ b/afl/llmp_test/src/main.rs @@ -131,16 +131,14 @@ fn main() { llmp_broker_init(&mut broker).expect("Could not init"); for i in 0..counter_thread_count { println!("Adding client {}", i); - llmp_broker_register_childprocess_clientloop( - &mut broker, + broker.register_childprocess_clientloop( llmp_test_clientloop, ptr::null_mut(), ) .expect("could not add child clientloop"); } - llmp_broker_register_childprocess_clientloop( - &mut broker, + broker.register_childprocess_clientloop( test_adder_clientloop, ptr::null_mut(), ) @@ -148,8 +146,8 @@ fn main() { println!("Spawning broker"); - llmp_broker_add_message_hook(&mut broker, broker_message_hook, ptr::null_mut()); + broker.add_message_hook(broker_message_hook, ptr::null_mut()); - llmp_broker_run(&mut broker); + broker.run(); } } diff --git a/afl/src/events/llmp_translated.rs b/afl/src/events/llmp_translated.rs index 694c3c0f07..55190c6290 100644 --- a/afl/src/events/llmp_translated.rs +++ b/afl/src/events/llmp_translated.rs @@ -125,7 +125,7 @@ pub struct llmp_message { #[derive(Clone)] #[repr(C)] -pub struct llmp_broker_state { +pub struct llmp_broker { pub last_msg_sent: *mut llmp_message, pub broadcast_map_count: c_ulong, pub broadcast_maps: *mut afl_shmem, @@ -171,7 +171,7 @@ pub enum LlmpMessageHookResult { } pub type LlmpMessageHookFn = unsafe fn( - _: *mut llmp_broker_state, + _: *mut llmp_broker, _: *mut llmp_broker_client_metadata, _: *mut llmp_message, _: *mut c_void, @@ -529,7 +529,7 @@ pub unsafe fn llmp_send(page: *mut llmp_page, msg: *mut llmp_message) -> Result< #[inline] unsafe fn _llmp_broker_current_broadcast_map( - broker_state: *mut llmp_broker_state, + broker_state: *mut llmp_broker, ) -> *mut afl_shmem { return &mut *(*broker_state).broadcast_maps.offset( (*broker_state) @@ -626,7 +626,7 @@ unsafe fn llmp_handle_out_eop( } } /* no more space left! We'll have to start a new page */ -pub unsafe fn llmp_broker_handle_out_eop(broker: *mut llmp_broker_state) -> AflRet { +pub unsafe fn llmp_broker_handle_out_eop(broker: *mut llmp_broker) -> AflRet { (*broker).broadcast_maps = llmp_handle_out_eop( (*broker).broadcast_maps, &mut (*broker).broadcast_map_count, @@ -639,7 +639,7 @@ pub unsafe fn llmp_broker_handle_out_eop(broker: *mut llmp_broker_state) -> AflR } as AflRet; } pub unsafe fn llmp_broker_alloc_next( - broker: *mut llmp_broker_state, + broker: *mut llmp_broker, len: c_ulong, ) -> *mut llmp_message { let mut broadcast_page: *mut llmp_page = shmem2page(_llmp_broker_current_broadcast_map(broker)); @@ -666,235 +666,398 @@ pub unsafe fn llmp_broker_alloc_next( } return out; } -/* Registers a new client for the given sharedmap str and size. -Be careful: Intenral realloc may change the location of the client map */ -unsafe fn llmp_broker_register_client( - broker: *mut llmp_broker_state, - shm_str: &CStr, - map_size: c_ulong, -) -> *mut llmp_broker_client_metadata { - /* make space for a new client and calculate its id */ - (*broker).llmp_clients = afl_realloc( - (*broker).llmp_clients as *mut c_void, - (*broker) - .llmp_client_count - .wrapping_add(1 as c_ulong) - .wrapping_mul(::std::mem::size_of::() as c_ulong), - ) as *mut llmp_broker_client_metadata; - if (*broker).llmp_clients.is_null() { - return 0 as *mut llmp_broker_client_metadata; - } - let mut client: *mut llmp_broker_client_metadata = &mut *(*broker) - .llmp_clients - .offset((*broker).llmp_client_count as isize) - as *mut llmp_broker_client_metadata; - memset( - client as *mut c_void, - 0 as c_int, - ::std::mem::size_of::() as c_ulong, - ); - (*client).client_state = calloc( - 1 as c_ulong, - ::std::mem::size_of::() as c_ulong, - ) as *mut llmp_client; - if (*client).client_state.is_null() { - return 0 as *mut llmp_broker_client_metadata; - } - (*(*client).client_state).id = (*broker).llmp_client_count as u32; - (*client).cur_client_map = - calloc(1 as c_ulong, ::std::mem::size_of::() as c_ulong) as *mut afl_shmem; - if (*client).cur_client_map.is_null() { - return 0 as *mut llmp_broker_client_metadata; - } - if afl_shmem_by_str((*client).cur_client_map, shm_str, map_size).is_null() { - return 0 as *mut llmp_broker_client_metadata; - } - (*broker).llmp_client_count = (*broker).llmp_client_count.wrapping_add(1); - // tODO: Add client map - return client; -} -/* broker broadcast to its own page for all others to read */ -#[inline] -unsafe fn llmp_broker_handle_new_msgs( - broker: *mut llmp_broker_state, - mut client: *mut llmp_broker_client_metadata, -) { - // TODO: We could memcpy a range of pending messages, instead of one by one. - /* DBG("llmp_broker_handle_new_msgs %p %p->%u\n", broker, client, client->client_state->id); */ - let incoming: *mut llmp_page = shmem2page((*client).cur_client_map); - let mut current_message_id: u32 = if !(*client).last_msg_broker_read.is_null() { - (*(*client).last_msg_broker_read).message_id - } else { - 0 as c_uint - }; - while current_message_id as c_ulong != (*incoming).current_msg_id { - let msg: *mut llmp_message = llmp_recv(incoming, (*client).last_msg_broker_read); - if msg.is_null() { - panic!("No message received but not all message ids receved! Data out of sync?"); + +impl llmp_broker { + /* Registers a new client for the given sharedmap str and size. + Be careful: Intenral realloc may change the location of the client map */ + unsafe fn register_client( + &mut self, + shm_str: &CStr, + map_size: c_ulong, + ) -> *mut llmp_broker_client_metadata { + /* make space for a new client and calculate its id */ + self.llmp_clients = afl_realloc( + self.llmp_clients as *mut c_void, + self.llmp_client_count + .wrapping_add(1 as c_ulong) + .wrapping_mul(::std::mem::size_of::() as c_ulong), + ) as *mut llmp_broker_client_metadata; + if self.llmp_clients.is_null() { + return 0 as *mut llmp_broker_client_metadata; } - if (*msg).tag == 0xaf1e0f1 as c_uint { - let pageinfo: *mut llmp_payload_new_page = { - let mut _msg: *mut llmp_message = msg; - (if (*_msg).buf_len >= ::std::mem::size_of::() as c_ulong { - (*_msg).buf.as_mut_ptr() - } else { - 0 as *mut u8 - }) as *mut llmp_payload_new_page - }; - if pageinfo.is_null() { - panic!(format!( - "Illegal message length for EOP (is {}, expected {})", - (*msg).buf_len_padded, - ::std::mem::size_of::() as c_ulong - )); - } - /* We can reuse the map mem space, no need to free and calloc. - However, the pageinfo points to the map we're about to unmap. - Copy the contents first. */ - let mut pageinfo_cpy: llmp_payload_new_page = llmp_payload_new_page { - map_size: 0, - shm_str: [0; 20], - }; - memcpy( - &mut pageinfo_cpy as *mut llmp_payload_new_page as *mut c_void, - pageinfo as *const c_void, - ::std::mem::size_of::() as c_ulong, - ); - let client_map: *mut afl_shmem = (*client).cur_client_map; - ::std::ptr::write_volatile( - &mut (*shmem2page(client_map)).save_to_unmap as *mut u16, - 1 as u16, - ); - afl_shmem_deinit(client_map); - if afl_shmem_by_str( - client_map, - CStr::from_bytes_with_nul(&(*pageinfo).shm_str).expect("Illegal shm_str"), - (*pageinfo).map_size, - ) - .is_null() - { - panic!(format!( - "Could not get shmem by str for map {:?} of size {:?}", - (*pageinfo).shm_str.as_mut_ptr(), - (*pageinfo).map_size - )); - } - } else if (*msg).tag == 0xc11e471 as c_uint { - /* This client informs us about yet another new client - add it to the list! Also, no need to forward this msg. */ - let pageinfo: *mut llmp_payload_new_page = { - let mut _msg: *mut llmp_message = msg; - (if (*_msg).buf_len >= ::std::mem::size_of::() as c_ulong { - (*_msg).buf.as_mut_ptr() - } else { - 0 as *mut u8 - }) as *mut llmp_payload_new_page - }; - if pageinfo.is_null() { - println!("Ignoring broken CLIENT_ADDED msg due to incorrect size. Expected {:?} but got {:?}", - ::std::mem::size_of::() as - c_ulong, (*msg).buf_len_padded); - } - /* register_client may realloc the clients, we need to find ours again */ - let client_id: u32 = (*(*client).client_state).id; - if llmp_broker_register_client( - broker, - CStr::from_bytes_with_nul(&(*pageinfo).shm_str).expect("Illegal shm_str"), - (*pageinfo).map_size, - ) - .is_null() - { - panic!(format!( - "Could not register clientprocess with shm_str {:?}", - (*pageinfo).shm_str.as_mut_ptr() - )); - } - (*client).client_type = LLMP_CLIENT_TYPE_FOREIGN_PROCESS; - /* find client again */ - client = &mut *(*broker).llmp_clients.offset(client_id as isize) - as *mut llmp_broker_client_metadata + let mut client: *mut llmp_broker_client_metadata = + self.llmp_clients.offset(self.llmp_client_count as isize) + as *mut llmp_broker_client_metadata; + memset( + client as *mut c_void, + 0 as c_int, + ::std::mem::size_of::() as c_ulong, + ); + (*client).client_state = calloc( + 1 as c_ulong, + ::std::mem::size_of::() as c_ulong, + ) as *mut llmp_client; + if (*client).client_state.is_null() { + return 0 as *mut llmp_broker_client_metadata; + } + (*(*client).client_state).id = (*self).llmp_client_count as u32; + (*client).cur_client_map = + calloc(1 as c_ulong, ::std::mem::size_of::() as c_ulong) as *mut afl_shmem; + if (*client).cur_client_map.is_null() { + return 0 as *mut llmp_broker_client_metadata; + } + if afl_shmem_by_str((*client).cur_client_map, shm_str, map_size).is_null() { + return 0 as *mut llmp_broker_client_metadata; + } + self.llmp_client_count = self.llmp_client_count.wrapping_add(1); + // TODO: Add client map + return client; + } + + /* Adds a hook that gets called in the broker for each new message the broker touches. + if the callback returns false, the message is not forwarded to the clients. */ + pub unsafe fn add_message_hook( + &mut self, + hook: LlmpMessageHookFn, + data: *mut c_void, + ) -> AflRet { + return llmp_add_hook_generic( + &mut self.msg_hooks, + &mut self.msg_hook_count, + ::std::mem::transmute::, *mut c_void>(Some(hook)), + data, + ); + } + + /* broker broadcast to its own page for all others to read */ + #[inline] + unsafe fn handle_new_msgs(&mut self, mut client: *mut llmp_broker_client_metadata) { + // TODO: We could memcpy a range of pending messages, instead of one by one. + /* DBG("llmp_broker_handle_new_msgs %p %p->%u\n", broker, client, client->client_state->id); */ + let incoming: *mut llmp_page = shmem2page((*client).cur_client_map); + let mut current_message_id: u32 = if !(*client).last_msg_broker_read.is_null() { + (*(*client).last_msg_broker_read).message_id } else { - let mut forward_msg: bool = 1 as c_int != 0; - let mut i: c_ulong = 0; - while i < (*broker).msg_hook_count { - let msg_hook: *mut llmp_hookdata_generic = - &mut *(*broker).msg_hooks.offset(i as isize) as *mut llmp_hookdata_generic; - forward_msg = forward_msg as c_int != 0 - && ::std::mem::transmute::<*mut c_void, Option>( - (*msg_hook).func, - ) - .expect("non-null function pointer")( - broker, client, msg, (*msg_hook).data - ) as c_int - != 0; - if !llmp_msg_in_page(shmem2page((*client).cur_client_map), msg) { - /* Special handling in case the client got exchanged inside the message_hook, for example after a crash. */ - return; - } - i = i.wrapping_add(1) + 0 as c_uint + }; + while current_message_id as c_ulong != (*incoming).current_msg_id { + let msg: *mut llmp_message = llmp_recv(incoming, (*client).last_msg_broker_read); + if msg.is_null() { + panic!("No message received but not all message ids receved! Data out of sync?"); } - if forward_msg { - let mut out: *mut llmp_message = - llmp_broker_alloc_next(broker, (*msg).buf_len_padded); - if out.is_null() { + if (*msg).tag == 0xaf1e0f1 as c_uint { + let pageinfo: *mut llmp_payload_new_page = { + let mut _msg: *mut llmp_message = msg; + (if (*_msg).buf_len >= ::std::mem::size_of::() as c_ulong + { + (*_msg).buf.as_mut_ptr() + } else { + 0 as *mut u8 + }) as *mut llmp_payload_new_page + }; + if pageinfo.is_null() { panic!(format!( - "Error allocating {} bytes in shmap {:?}", + "Illegal message length for EOP (is {}, expected {})", (*msg).buf_len_padded, - (*_llmp_broker_current_broadcast_map(broker)) - .shm_str - .as_mut_ptr(), + ::std::mem::size_of::() as c_ulong )); } - /* Copy over the whole message. - 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: c_ulong = (*out).buf_len_padded; - memcpy( - out as *mut c_void, - msg as *const c_void, - (::std::mem::size_of::() as c_ulong) - .wrapping_add((*msg).buf_len_padded), - ); - (*out).buf_len_padded = actual_size; - /* We need to replace the message ID with our own */ - let out_page: *mut llmp_page = - shmem2page(_llmp_broker_current_broadcast_map(broker)); - (*out).message_id = (*out_page).current_msg_id.wrapping_add(1 as c_ulong) as u32; - match llmp_send(out_page, out) { - Err(e) => panic!(format!("Error sending msg: {:?}", e)), - _ => (), + /* We can reuse the map mem space, no need to free and calloc. + However, the pageinfo points to the map we're about to unmap. + Copy the contents first. */ + let mut pageinfo_cpy: llmp_payload_new_page = llmp_payload_new_page { + map_size: 0, + shm_str: [0; 20], }; - (*broker).last_msg_sent = out + memcpy( + &mut pageinfo_cpy as *mut llmp_payload_new_page as *mut c_void, + pageinfo as *const c_void, + ::std::mem::size_of::() as c_ulong, + ); + let client_map: *mut afl_shmem = (*client).cur_client_map; + ::std::ptr::write_volatile( + &mut (*shmem2page(client_map)).save_to_unmap as *mut u16, + 1 as u16, + ); + afl_shmem_deinit(client_map); + if afl_shmem_by_str( + client_map, + CStr::from_bytes_with_nul(&(*pageinfo).shm_str).expect("Illegal shm_str"), + (*pageinfo).map_size, + ) + .is_null() + { + panic!(format!( + "Could not get shmem by str for map {:?} of size {:?}", + (*pageinfo).shm_str.as_mut_ptr(), + (*pageinfo).map_size + )); + } + } else if (*msg).tag == 0xc11e471 as c_uint { + /* This client informs us about yet another new client + add it to the list! Also, no need to forward this msg. */ + let pageinfo: *mut llmp_payload_new_page = { + let mut _msg: *mut llmp_message = msg; + (if (*_msg).buf_len >= ::std::mem::size_of::() as c_ulong + { + (*_msg).buf.as_mut_ptr() + } else { + 0 as *mut u8 + }) as *mut llmp_payload_new_page + }; + if pageinfo.is_null() { + println!("Ignoring broken CLIENT_ADDED msg due to incorrect size. Expected {:?} but got {:?}", + ::std::mem::size_of::() as + c_ulong, (*msg).buf_len_padded); + } + /* register_client may realloc the clients, we need to find ours again */ + let client_id: u32 = (*(*client).client_state).id; + if self + .register_client( + CStr::from_bytes_with_nul(&(*pageinfo).shm_str).expect("Illegal shm_str"), + (*pageinfo).map_size, + ) + .is_null() + { + panic!(format!( + "Could not register clientprocess with shm_str {:?}", + (*pageinfo).shm_str.as_mut_ptr() + )); + } + (*client).client_type = LLMP_CLIENT_TYPE_FOREIGN_PROCESS; + /* find client again */ + client = + self.llmp_clients.offset(client_id as isize) as *mut llmp_broker_client_metadata + } else { + let mut forward_msg: bool = 1 as c_int != 0; + let mut i: c_ulong = 0; + while i < self.msg_hook_count { + let msg_hook: *mut llmp_hookdata_generic = + self.msg_hooks.offset(i as isize) as *mut llmp_hookdata_generic; + forward_msg = forward_msg as c_int != 0 + && ::std::mem::transmute::<*mut c_void, Option>( + (*msg_hook).func, + ) + .expect("non-null function pointer")( + self, client, msg, (*msg_hook).data + ) as c_int + != 0; + if !llmp_msg_in_page(shmem2page((*client).cur_client_map), msg) { + /* Special handling in case the client got exchanged inside the message_hook, for example after a crash. */ + return; + } + i = i.wrapping_add(1) + } + if forward_msg { + let mut out: *mut llmp_message = + llmp_broker_alloc_next(self, (*msg).buf_len_padded); + if out.is_null() { + panic!(format!( + "Error allocating {} bytes in shmap {:?}", + (*msg).buf_len_padded, + (*_llmp_broker_current_broadcast_map(self)) + .shm_str + .as_mut_ptr(), + )); + } + /* Copy over the whole message. + 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: c_ulong = (*out).buf_len_padded; + memcpy( + out as *mut c_void, + msg as *const c_void, + (::std::mem::size_of::() as c_ulong) + .wrapping_add((*msg).buf_len_padded), + ); + (*out).buf_len_padded = actual_size; + /* We need to replace the message ID with our own */ + let out_page: *mut llmp_page = + shmem2page(_llmp_broker_current_broadcast_map(self)); + (*out).message_id = + (*out_page).current_msg_id.wrapping_add(1 as c_ulong) as u32; + match llmp_send(out_page, out) { + Err(e) => panic!(format!("Error sending msg: {:?}", e)), + _ => (), + }; + self.last_msg_sent = out + } } + (*client).last_msg_broker_read = msg; + current_message_id = (*msg).message_id } - (*client).last_msg_broker_read = msg; - current_message_id = (*msg).message_id } -} -/* The broker walks all pages and looks for changes, then broadcasts them on - * its own shared page, once. */ -/* The broker walks all pages and looks for changes, then broadcasts them on - * its own shared page, once. */ -pub unsafe fn llmp_broker_once(broker: *mut llmp_broker_state) { - compiler_fence(Ordering::SeqCst); - let mut i: u32 = 0; - while (i as c_ulong) < (*broker).llmp_client_count { - let client: *mut llmp_broker_client_metadata = - &mut *(*broker).llmp_clients.offset(i as isize) as *mut llmp_broker_client_metadata; - llmp_broker_handle_new_msgs(broker, client); - i = i.wrapping_add(1) - } -} -/* The broker walks all pages and looks for changes, then broadcasts them on - * its own shared page */ -pub unsafe fn llmp_broker_loop(broker: *mut llmp_broker_state) -> ! { - loop { + + /* The broker walks all pages and looks for changes, then broadcasts them on + * its own shared page, once. */ + /* The broker walks all pages and looks for changes, then broadcasts them on + * its own shared page, once. */ + pub unsafe fn once(&mut self) { compiler_fence(Ordering::SeqCst); - llmp_broker_once(broker); - /* 5 milis of sleep for now to not busywait at 100% */ - usleep((5 as c_int * 1000 as c_int) as c_uint); + let mut i: u32 = 0; + while (i as c_ulong) < self.llmp_client_count { + let client: *mut llmp_broker_client_metadata = + self.llmp_clients.offset(i as isize) as *mut llmp_broker_client_metadata; + self.handle_new_msgs(client); + i = i.wrapping_add(1) + } + } + + /* The broker walks all pages and looks for changes, then broadcasts them on + * its own shared page */ + pub unsafe fn broker_loop(&mut self) -> ! { + loop { + compiler_fence(Ordering::SeqCst); + self.once(); + /* 5 milis of sleep for now to not busywait at 100% */ + usleep((5 as c_int * 1000 as c_int) as c_uint); + } + } + + /* launch a specific client. This function is rarely needed - all registered clients will get launched at broker_run */ + pub unsafe fn launch_client( + &mut self, + mut clientdata: *mut llmp_broker_client_metadata, + ) -> bool { + if clientdata < self.llmp_clients + || clientdata + > self + .llmp_clients + .offset(self.llmp_client_count.wrapping_sub(1 as c_ulong) as isize) + as *mut llmp_broker_client_metadata + { + println!( + "[!] WARNING: Illegal client specified at ptr {:?} (instead of {:?} to {:?})", + clientdata, + self.llmp_clients, + self.llmp_clients + .offset(self.llmp_client_count.wrapping_sub(1 as c_ulong) as isize,) + as *mut llmp_broker_client_metadata, + ); + return 0 as c_int != 0; + } + if (*clientdata).client_type as c_uint == LLMP_CLIENT_TYPE_CHILD_PROCESS as c_int as c_uint + { + if (*clientdata).pid != 0 { + println!("[!] WARNING: Tried to relaunch already running client. Set ->pid to 0 if this is what you want."); + return 0 as c_int != 0; + } + let child_id: c_int = fork(); + if child_id < 0 as c_int { + println!("[!] WARNING: Could not fork"); + return 0 as c_int != 0; + } else { + if child_id == 0 as c_int { + /* child */ + _llmp_client_wrapped_loop(clientdata as *mut c_void); + } + } + /* parent */ + (*clientdata).pid = child_id; + return 1 as c_int != 0; + } else { + println!("[!] WARNING: Tried to spawn llmp child with unknown thread type."); + return 0 as c_int != 0; + } + //return 1 as c_int != 0; + } + + pub unsafe fn launch_clientloops(&mut self) -> Result<(), AflError> { + let mut i: c_ulong = 0; + while i < self.llmp_client_count { + if (*self.llmp_clients.offset(i as isize)).client_type as c_uint + == LLMP_CLIENT_TYPE_CHILD_PROCESS as c_uint + { + if !self.launch_client(self.llmp_clients.offset(i as isize)) { + println!("[!] WARNING: Could not launch all clients"); + return Err(AflError::Unknown("Failed to launch clients".into())); + } + } + i = i.wrapping_add(1) + } + Ok(()) + } + + /* The broker walks all pages and looks for changes, then broadcasts them on + its own shared page. + Never returns. */ + /* Start all threads and the main broker. + Same as llmp_broker_launch_threaded clients(); + Never returns. */ + /* Start all threads and the main broker. Never returns. */ + pub unsafe fn run(&mut self) -> ! { + self.launch_clientloops().expect("Failed to launch clients"); + self.broker_loop(); + } + + /* Register a new forked/child client. + Client thread will be called with llmp_client client, containing + the data in ->data. This will register a client to be spawned up as soon as + broker_loop() starts. Clients can also be added later via + llmp_broker_register_remote(..) or the local_tcp_client + */ + pub unsafe fn register_childprocess_clientloop( + &mut self, + clientloop: LlmpClientloopFn, + data: *mut c_void, + ) -> Result<(), AflError> { + let mut client_map: afl_shmem = { + let init = afl_shmem { + shm_str: [0; 20], + shm_id: 0, + map: 0 as *mut u8, + map_size: 0, + }; + init + }; + if llmp_new_page_shmem( + &mut client_map, + self.llmp_client_count, + ((1 as c_int) << 28 as c_int) as c_ulong, + ) + .is_null() + { + return Err(AflError::Unknown("Alloc".into())); + } + let mut client: *mut llmp_broker_client_metadata = self.register_client( + CStr::from_ptr(&client_map.shm_str as *const i8), + client_map.map_size, + ); + if client.is_null() { + afl_shmem_deinit(&mut client_map); + return Err(AflError::Unknown("Something in clients failed".into())); + } + (*client).clientloop = Some(clientloop); + (*client).data = data; + (*client).client_type = LLMP_CLIENT_TYPE_CHILD_PROCESS; + /* Copy the already allocated shmem to the client state */ + (*(*client).client_state).out_maps = afl_realloc( + (*(*client).client_state).out_maps as *mut c_void, + ::std::mem::size_of::() as c_ulong, + ) as *mut afl_shmem; + if (*(*client).client_state).out_maps.is_null() { + afl_shmem_deinit(&mut client_map); + afl_shmem_deinit((*client).cur_client_map); + /* "Unregister" by subtracting the client from count */ + self.llmp_client_count = self.llmp_client_count.wrapping_sub(1); + return Err(AflError::Unknown("Something in clients failed".into())); + } + memcpy( + (*(*client).client_state).out_maps as *mut c_void, + &mut client_map as *mut afl_shmem as *const c_void, + ::std::mem::size_of::() as c_ulong, + ); + (*(*client).client_state).out_map_count = 1 as c_ulong; + /* Each client starts with the very first map. + They should then iterate through all maps once and work on all old messages. + */ + (*(*client).client_state).current_broadcast_map = + self.broadcast_maps.offset(0 as isize) as *mut afl_shmem; + (*(*client).client_state).out_map_count = 1 as c_ulong; + return Ok(()); } } + /* A new page will be used. Notify each registered hook in the client about this fact. */ unsafe fn llmp_clientrigger_new_out_page_hooks(client: *mut llmp_client) { let mut i: c_ulong = 0; @@ -927,84 +1090,6 @@ unsafe fn _llmp_client_wrapped_loop(llmp_client_broker_metadata_ptr: *mut c_void ); } -/* launch a specific client. This function is rarely needed - all registered clients will get launched at broker_run */ -pub unsafe fn llmp_broker_launch_client( - broker: *mut llmp_broker_state, - mut clientdata: *mut llmp_broker_client_metadata, -) -> bool { - if clientdata < (*broker).llmp_clients - || clientdata - > &mut *(*broker) - .llmp_clients - .offset((*broker).llmp_client_count.wrapping_sub(1 as c_ulong) as isize) - as *mut llmp_broker_client_metadata - { - println!( - "[!] WARNING: Illegal client specified at ptr {:?} (instead of {:?} to {:?})", - clientdata, - (*broker).llmp_clients, - &mut *(*broker) - .llmp_clients - .offset((*broker).llmp_client_count.wrapping_sub(1 as c_ulong) as isize,) - as *mut llmp_broker_client_metadata, - ); - return 0 as c_int != 0; - } - if (*clientdata).client_type as c_uint == LLMP_CLIENT_TYPE_CHILD_PROCESS as c_int as c_uint { - if (*clientdata).pid != 0 { - println!("[!] WARNING: Tried to relaunch already running client. Set ->pid to 0 if this is what you want."); - return 0 as c_int != 0; - } - let child_id: c_int = fork(); - if child_id < 0 as c_int { - println!("[!] WARNING: Could not fork"); - return 0 as c_int != 0; - } else { - if child_id == 0 as c_int { - /* child */ - _llmp_client_wrapped_loop(clientdata as *mut c_void); - } - } - /* parent */ - (*clientdata).pid = child_id; - return 1 as c_int != 0; - } else { - println!("[!] WARNING: Tried to spawn llmp child with unknown thread type."); - return 0 as c_int != 0; - } - //return 1 as c_int != 0; -} - -pub unsafe fn llmp_broker_launch_clientloops( - broker: *mut llmp_broker_state, -) -> Result<(), AflError> { - let mut i: c_ulong = 0; - while i < (*broker).llmp_client_count { - if (*(*broker).llmp_clients.offset(i as isize)).client_type as c_uint - == LLMP_CLIENT_TYPE_CHILD_PROCESS as c_uint - { - if !llmp_broker_launch_client(broker, &mut *(*broker).llmp_clients.offset(i as isize)) { - println!("[!] WARNING: Could not launch all clients"); - return Err(AflError::Unknown("Failed to launch clients".into())); - } - } - i = i.wrapping_add(1) - } - Ok(()) -} - -/* The broker walks all pages and looks for changes, then broadcasts them on -its own shared page. -Never returns. */ -/* Start all threads and the main broker. -Same as llmp_broker_launch_threaded clients(); -Never returns. */ -/* Start all threads and the main broker. Never returns. */ -pub unsafe fn llmp_broker_run(broker: *mut llmp_broker_state) -> ! { - llmp_broker_launch_clientloops(broker).expect("Failed to launch clients"); - llmp_broker_loop(broker); -} - /* For non zero-copy, we want to get rid of old pages with duplicate messages eventually. This function This funtion sees if we can unallocate older pages. @@ -1293,74 +1378,6 @@ impl Drop for llmp_client { } } -/* Register a new forked/child client. -Client thread will be called with llmp_client client, containing -the data in ->data. This will register a client to be spawned up as soon as -broker_loop() starts. Clients can also be added later via -llmp_broker_register_remote(..) or the local_tcp_client -*/ -pub unsafe fn llmp_broker_register_childprocess_clientloop( - broker: *mut llmp_broker_state, - clientloop: LlmpClientloopFn, - data: *mut c_void, -) -> Result<(), AflError> { - let mut client_map: afl_shmem = { - let init = afl_shmem { - shm_str: [0; 20], - shm_id: 0, - map: 0 as *mut u8, - map_size: 0, - }; - init - }; - if llmp_new_page_shmem( - &mut client_map, - (*broker).llmp_client_count, - ((1 as c_int) << 28 as c_int) as c_ulong, - ) - .is_null() - { - return Err(AflError::Unknown("Alloc".into())); - } - let mut client: *mut llmp_broker_client_metadata = llmp_broker_register_client( - broker, - CStr::from_ptr(&client_map.shm_str as *const i8), - client_map.map_size, - ); - if client.is_null() { - afl_shmem_deinit(&mut client_map); - return Err(AflError::Unknown("Something in clients failed".into())); - } - (*client).clientloop = Some(clientloop); - (*client).data = data; - (*client).client_type = LLMP_CLIENT_TYPE_CHILD_PROCESS; - /* Copy the already allocated shmem to the client state */ - (*(*client).client_state).out_maps = afl_realloc( - (*(*client).client_state).out_maps as *mut c_void, - ::std::mem::size_of::() as c_ulong, - ) as *mut afl_shmem; - if (*(*client).client_state).out_maps.is_null() { - afl_shmem_deinit(&mut client_map); - afl_shmem_deinit((*client).cur_client_map); - /* "Unregister" by subtracting the client from count */ - (*broker).llmp_client_count = (*broker).llmp_client_count.wrapping_sub(1); - return Err(AflError::Unknown("Something in clients failed".into())); - } - memcpy( - (*(*client).client_state).out_maps as *mut c_void, - &mut client_map as *mut afl_shmem as *const c_void, - ::std::mem::size_of::() as c_ulong, - ); - (*(*client).client_state).out_map_count = 1 as c_ulong; - /* Each client starts with the very first map. - They should then iterate through all maps once and work on all old messages. - */ - (*(*client).client_state).current_broadcast_map = - &mut *(*broker).broadcast_maps.offset(0 as isize) as *mut afl_shmem; - (*(*client).client_state).out_map_count = 1 as c_ulong; - return Ok(()); -} - /* Generic function to add a hook to the mem pointed to by hooks_p, using afl_realloc on the mem area, and increasing * hooks_count_p */ pub unsafe fn llmp_add_hook_generic( @@ -1403,28 +1420,41 @@ pub unsafe fn llmp_client_add_new_out_page_hook( ); } -/* Adds a hook that gets called in the broker for each new message the broker touches. -if the callback returns false, the message is not forwarded to the clients. */ -pub unsafe fn llmp_broker_add_message_hook( - broker: *mut llmp_broker_state, - hook: LlmpMessageHookFn, - data: *mut c_void, -) -> AflRet { - return llmp_add_hook_generic( - &mut (*broker).msg_hooks, - &mut (*broker).msg_hook_count, - ::std::mem::transmute::, *mut c_void>(Some(hook)), - data, - ); +/* Clean up the broker instance */ +pub unsafe fn llmp_broker_deinit(broker: *mut llmp_broker) { + let mut i: c_ulong; + i = 0 as c_ulong; + while i < (*broker).broadcast_map_count { + afl_shmem_deinit(&mut *(*broker).broadcast_maps.offset(i as isize)); + i = i.wrapping_add(1) + } + i = 0 as c_ulong; + while i < (*broker).llmp_client_count { + afl_shmem_deinit((*(*broker).llmp_clients.offset(i as isize)).cur_client_map); + free((*(*broker).llmp_clients.offset(i as isize)).cur_client_map as *mut c_void); + i = i.wrapping_add(1) + // TODO: Properly clean up the client + } + afl_free((*broker).broadcast_maps as *mut c_void); + (*broker).broadcast_map_count = 0 as c_ulong; + afl_free((*broker).llmp_clients as *mut c_void); + (*broker).llmp_client_count = 0 as c_ulong; } + +impl Drop for llmp_broker { + fn drop(&mut self) { + unsafe { llmp_broker_deinit(self) }; + } +} + /* Allocate and set up the new broker instance. Afterwards, run with * broker_run. */ -pub unsafe fn llmp_broker_init(broker: *mut llmp_broker_state) -> Result<(), AflError> { +pub unsafe fn llmp_broker_init(broker: *mut llmp_broker) -> Result<(), AflError> { memset( broker as *mut c_void, 0 as c_int, - ::std::mem::size_of::() as c_ulong, + ::std::mem::size_of::() as c_ulong, ); /* let's create some space for outgoing maps */ (*broker).broadcast_maps = afl_realloc( @@ -1449,29 +1479,3 @@ pub unsafe fn llmp_broker_init(broker: *mut llmp_broker_state) -> Result<(), Afl } return Ok(()); } -/* Clean up the broker instance */ -pub unsafe fn llmp_broker_deinit(broker: *mut llmp_broker_state) { - let mut i: c_ulong; - i = 0 as c_ulong; - while i < (*broker).broadcast_map_count { - afl_shmem_deinit(&mut *(*broker).broadcast_maps.offset(i as isize)); - i = i.wrapping_add(1) - } - i = 0 as c_ulong; - while i < (*broker).llmp_client_count { - afl_shmem_deinit((*(*broker).llmp_clients.offset(i as isize)).cur_client_map); - free((*(*broker).llmp_clients.offset(i as isize)).cur_client_map as *mut c_void); - i = i.wrapping_add(1) - // TODO: Properly clean up the client - } - afl_free((*broker).broadcast_maps as *mut c_void); - (*broker).broadcast_map_count = 0 as c_ulong; - afl_free((*broker).llmp_clients as *mut c_void); - (*broker).llmp_client_count = 0 as c_ulong; -} - -impl Drop for llmp_broker_state { - fn drop(&mut self) { - unsafe { llmp_broker_deinit(self) }; - } -} diff --git a/afl/src/events/mod.rs b/afl/src/events/mod.rs index 634555e720..9ea756d6eb 100644 --- a/afl/src/events/mod.rs +++ b/afl/src/events/mod.rs @@ -429,7 +429,34 @@ where R: Rand, //TODO CE: CustomEvent, { - pub fn new() -> Self { + /// Forks n processes, never returns in broker. + pub fn spawn(process_count: usize) -> Self { + /* + llmp_broker_init(&mut broker).expect("Could not init"); + for i in 0..counter_thread_count { + println!("Adding client {}", i); + llmp_broker_register_childprocess_clientloop( + &mut broker, + llmp_test_clientloop, + ptr::null_mut(), + ) + .expect("could not add child clientloop"); + } + + llmp_broker_register_childprocess_clientloop( + &mut broker, + test_adder_clientloop, + ptr::null_mut(), + ) + .expect("Error registering childprocess"); + + println!("Spawning broker"); + + llmp_broker_add_message_hook(&mut broker, broker_message_hook, ptr::null_mut()); + + llmp_broker_run(&mut broker); + */ + Self { _marker: PhantomData, }