diff --git a/llmp_test/src/main.rs b/llmp_test/src/main.rs index 988fa56555..4fffc2bc09 100644 --- a/llmp_test/src/main.rs +++ b/llmp_test/src/main.rs @@ -8,60 +8,105 @@ use std::time; use afl::events::llmp_translated::*; const TAG_SIMPLE_U32_V1: u32 = 0x51300321; +const TAG_MATH_RESULT_V1: u32 = 0x77474331; -fn llmp_test_clientloop(client: *mut llmp_client, _data: *mut c_void) -> ! { +unsafe fn llmp_test_clientloop(client: *mut llmp_client, _data: *mut c_void) -> ! { let mut counter: u32 = 0; loop { counter += 1; - unsafe { + let llmp_message = llmp_client_alloc_next(client, size_of::()); + std::ptr::copy( + counter.to_be_bytes().as_ptr(), + (*llmp_message).buf.as_mut_ptr(), + size_of::(), + ); + (*llmp_message).tag = TAG_SIMPLE_U32_V1; + llmp_client_send(client, llmp_message); + + thread::sleep(time::Duration::from_millis(100)); + } +} + +unsafe fn u32_from_msg(message: *const llmp_message) -> u32 { + u32::from_be_bytes( + std::slice::from_raw_parts((*message).buf.as_ptr(), size_of::()) + .try_into() + .unwrap(), + ) +} + +unsafe fn test_adder_clientloop(client: *mut llmp_client, _data: *mut c_void) -> ! { + let mut last_result: u32 = 0; + let mut current_result: u32 = 0; + loop { + loop { + let last_msg = llmp_client_recv(client); + if last_msg == 0 as *mut llmp_message { + break; + } + match (*last_msg).tag { + TAG_SIMPLE_U32_V1 => { + current_result = last_result.wrapping_add(u32_from_msg(last_msg)); + } + _ => println!("Adder Client ignored unknown message {}", (*last_msg).tag), + }; + } + + if current_result != last_result { let llmp_message = llmp_client_alloc_next(client, size_of::()); std::ptr::copy( - counter.to_be_bytes().as_ptr(), + current_result.to_be_bytes().as_ptr(), (*llmp_message).buf.as_mut_ptr(), size_of::(), ); - (*llmp_message).tag = TAG_SIMPLE_U32_V1; + (*llmp_message).tag = TAG_MATH_RESULT_V1; llmp_client_send(client, llmp_message); + last_result = current_result; } thread::sleep(time::Duration::from_millis(100)); } } -fn broker_message_hook( +unsafe fn broker_message_hook( _broker: *mut llmp_broker_state, client_metadata: *mut llmp_broker_client_metadata, message: *mut llmp_message, _data: *mut c_void, ) -> LlmpMessageHookResult { - unsafe { - match (*message).tag { - TAG_SIMPLE_U32_V1 => { - println!( - "Client {:?} sent message: {:?}", - (*client_metadata).pid, - u32::from_be_bytes( - std::slice::from_raw_parts((*message).buf.as_ptr(), size_of::()) - .try_into() - .unwrap() - ), - ); - LlmpMessageHookResult::Handled - } - _ => { - println!("Unknwon message id received!"); - LlmpMessageHookResult::ForwardToClients - } + match (*message).tag { + TAG_SIMPLE_U32_V1 => { + println!( + "Client {:?} sent message: {:?}", + (*client_metadata).pid, + u32_from_msg(message) + ); + LlmpMessageHookResult::ForwardToClients + } + TAG_MATH_RESULT_V1 => { + println!( + "Adder Client has this current result: {:?}", + u32_from_msg(message) + ); + LlmpMessageHookResult::Handled + } + _ => { + println!("Unknwon message id received!"); + LlmpMessageHookResult::ForwardToClients } } } fn main() { /* The main node has a broker, and a few worker threads */ + let threads_total = num_cpus::get(); - let thread_count = num_cpus::get() - 1; - println!("Running with 1 broker and {} clients", thread_count); + let counter_thread_count = threads_total - 2; + println!( + "Running with 1 broker, 1 adder, and {} counter clients", + counter_thread_count + ); let mut broker = llmp_broker_state { last_msg_sent: ptr::null_mut(), @@ -74,7 +119,7 @@ fn main() { }; unsafe { llmp_broker_init(&mut broker).expect("Could not init"); - for i in 0..thread_count { + for i in 0..counter_thread_count { println!("Adding client {}", i); llmp_broker_register_childprocess_clientloop( &mut broker, @@ -84,6 +129,13 @@ fn main() { .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()); diff --git a/src/events/llmp_translated.rs b/src/events/llmp_translated.rs index a85ec8230d..d33eed5b58 100644 --- a/src/events/llmp_translated.rs +++ b/src/events/llmp_translated.rs @@ -147,7 +147,7 @@ pub struct llmp_broker_client_metadata { pub data: *mut c_void, } -pub type LlmpClientloopFn = fn(_: *mut llmp_client, _: *mut c_void) -> !; +pub type LlmpClientloopFn = unsafe fn(_: *mut llmp_client, _: *mut c_void) -> !; pub type LlmpClientType = c_uint; pub const LLMP_CLIENT_TYPE_FOREIGN_PROCESS: LlmpClientType = 3; pub const LLMP_CLIENT_TYPE_CHILD_PROCESS: LlmpClientType = 2; @@ -170,7 +170,7 @@ pub enum LlmpMessageHookResult { ForwardToClients, } -pub type LlmpMessageHookFn = fn( +pub type LlmpMessageHookFn = unsafe fn( _: *mut llmp_broker_state, _: *mut llmp_broker_client_metadata, _: *mut llmp_message, @@ -1075,7 +1075,7 @@ pub unsafe fn llmp_client_recv(mut client: *mut llmp_client) -> *mut llmp_messag if (*msg).tag == 0xdeadaf as c_uint { panic!("BUG: Read unallocated msg"); } else { - if (*msg).tag == 0xaf1e0f1 as c_int as c_uint { + if (*msg).tag == 0xaf1e0f1 as c_uint { /* we reached the end of the current page. We'll init a new page but can reuse the mem are of the current map. However, we cannot use the message if we deinit its page, so let's copy */ diff --git a/src/events/shmem_translated.rs b/src/events/shmem_translated.rs index a9da60f9a5..cf12b29698 100644 --- a/src/events/shmem_translated.rs +++ b/src/events/shmem_translated.rs @@ -1,4 +1,3 @@ -use ::libc; use libc::{c_char, c_int, c_long, c_uchar, c_uint, c_ulong, c_ushort, c_void}; use std::ffi::CStr;