diff --git a/llmp_test/Cargo.toml b/llmp_test/Cargo.toml index c7035f5fba..5ce6bfad10 100644 --- a/llmp_test/Cargo.toml +++ b/llmp_test/Cargo.toml @@ -7,4 +7,5 @@ edition = "2018" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [dependencies] +num_cpus = "1.0" afl = { path = "../" } \ No newline at end of file diff --git a/llmp_test/src/main.rs b/llmp_test/src/main.rs index 964960db01..988fa56555 100644 --- a/llmp_test/src/main.rs +++ b/llmp_test/src/main.rs @@ -1,12 +1,13 @@ +use core::convert::TryInto; use core::ffi::c_void; - -use std::env::args; -use std::ptr; +use core::mem::size_of; +use core::ptr; +use std::thread; +use std::time; use afl::events::llmp_translated::*; -use std::convert::TryInto; -use std::{thread, time}; +const TAG_SIMPLE_U32_V1: u32 = 0x51300321; fn llmp_test_clientloop(client: *mut llmp_client, _data: *mut c_void) -> ! { let mut counter: u32 = 0; @@ -14,13 +15,13 @@ fn llmp_test_clientloop(client: *mut llmp_client, _data: *mut c_void) -> ! { counter += 1; unsafe { - let llmp_message = llmp_client_alloc_next(client, 10); + 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(), - 4, + size_of::(), ); - (*llmp_message).tag = 1; + (*llmp_message).tag = TAG_SIMPLE_U32_V1; llmp_client_send(client, llmp_message); } @@ -36,12 +37,12 @@ fn broker_message_hook( ) -> LlmpMessageHookResult { unsafe { match (*message).tag { - 1 => { + TAG_SIMPLE_U32_V1 => { println!( "Client {:?} sent message: {:?}", (*client_metadata).pid, u32::from_be_bytes( - std::slice::from_raw_parts((*message).buf.as_ptr(), 4) + std::slice::from_raw_parts((*message).buf.as_ptr(), size_of::()) .try_into() .unwrap() ), @@ -57,7 +58,10 @@ fn broker_message_hook( } fn main() { - /* The main node has a broker, a tcp server, and a few worker threads */ + /* The main node has a broker, and a few worker threads */ + + let thread_count = num_cpus::get() - 1; + println!("Running with 1 broker and {} clients", thread_count); let mut broker = llmp_broker_state { last_msg_sent: ptr::null_mut(), @@ -68,7 +72,6 @@ fn main() { llmp_client_count: 0, llmp_clients: ptr::null_mut(), }; - let thread_count = 3; unsafe { llmp_broker_init(&mut broker).expect("Could not init"); for i in 0..thread_count { diff --git a/src/events/llmp_translated.rs b/src/events/llmp_translated.rs index dc0ee56e97..a85ec8230d 100644 --- a/src/events/llmp_translated.rs +++ b/src/events/llmp_translated.rs @@ -1151,7 +1151,7 @@ pub unsafe fn llmp_client_recv_blocking(client: *mut llmp_client) -> *mut llmp_m } /* The current page could have changed in recv (EOP) */ /* Alloc the next message, internally handling end of page by allocating a new one. */ -pub unsafe fn llmp_client_alloc_next(client: *mut llmp_client, size: c_ulong) -> *mut llmp_message { +pub unsafe fn llmp_client_alloc_next(client: *mut llmp_client, size: usize) -> *mut llmp_message { if client.is_null() { panic!("Client is NULL"); } @@ -1159,10 +1159,10 @@ pub unsafe fn llmp_client_alloc_next(client: *mut llmp_client, size: c_ulong) -> shmem2page( &mut *(*client) .out_maps - .offset((*client).out_map_count.wrapping_sub(1 as c_int as c_ulong) as isize), + .offset((*client).out_map_count.wrapping_sub(1) as isize), ), (*client).last_msg_sent, - size, + size as c_ulong, ); if msg.is_null() { let last_map_count: c_ulong = (*client).out_map_count; @@ -1175,7 +1175,7 @@ pub unsafe fn llmp_client_alloc_next(client: *mut llmp_client, size: c_ulong) -> || (*(*shmem2page( &mut *(*client) .out_maps - .offset((*client).out_map_count.wrapping_sub(1 as c_int as c_ulong) as isize), + .offset((*client).out_map_count.wrapping_sub(1) as isize), )) .messages .as_mut_ptr()) @@ -1190,10 +1190,10 @@ pub unsafe fn llmp_client_alloc_next(client: *mut llmp_client, size: c_ulong) -> shmem2page( &mut *(*client) .out_maps - .offset((*client).out_map_count.wrapping_sub(1 as c_int as c_ulong) as isize), + .offset((*client).out_map_count.wrapping_sub(1) as isize), ), 0 as *mut llmp_message, - size, + size as c_ulong, ); if msg.is_null() { return 0 as *mut llmp_message; @@ -1203,7 +1203,7 @@ pub unsafe fn llmp_client_alloc_next(client: *mut llmp_client, size: c_ulong) -> (*msg).message_id = if !(*client).last_msg_sent.is_null() { (*(*client).last_msg_sent).message_id.wrapping_add(1) } else { - 1 as c_int as c_uint + 1 as c_uint }; /* DBG("Allocated message at loc %p with buflen %ld", msg, msg->buf_len_padded); */ return msg;