diff --git a/afl/src/events/llmp.rs b/afl/src/events/llmp.rs index 44a5bd06f1..25530467d0 100644 --- a/afl/src/events/llmp.rs +++ b/afl/src/events/llmp.rs @@ -193,6 +193,47 @@ impl LlmpMsg { } } +/// An Llmp instance +pub enum LlmpConnection { + /// A broker and a thread using this tcp background thread + IsBroker { + broker: LlmpBroker, + listener_thread: thread::JoinHandle<()>, + }, + /// A client, connected to the port + IsClient { client: LlmpClient }, +} + +impl LlmpConnection { + /// Creates either a broker, if the tcp port is not bound, or a client, connected to this port. + pub fn on_port(port: u16) -> Result { + match TcpListener::bind(format!("127.0.0.1:{}", port)) { + Ok(listener) => { + // We got the port. We are the broker! :) + dbg!("We're the broker"); + let mut broker = LlmpBroker::new()?; + let listener_thread = broker.launch_tcp_listener(listener)?; + Ok(LlmpConnection::IsBroker { + broker, + listener_thread, + }) + } + Err(e) => { + match e.kind() { + std::io::ErrorKind::AddrInUse => { + // We are the client :) + dbg!("We're the client", e); + Ok(LlmpConnection::IsClient { + client: LlmpClient::create_attach_to_tcp(port)?, + }) + } + _ => Err(AflError::File(e)), + } + } + } + } +} + /// Contents of the share mem pages, used by llmp internally #[derive(Copy, Clone)] #[repr(C, packed)] @@ -869,16 +910,28 @@ impl LlmpBroker { } } - pub fn launch_tcp_listener(&mut self, port: u16) -> Result, AflError> { + /// Launches a thread using a tcp listener socket, on which new clients may connect to this broker + /// Does so on the given port. + pub fn launch_tcp_listener_on( + &mut self, + port: u16, + ) -> Result, AflError> { + let listener = TcpListener::bind(format!("127.0.0.1:{}", port))?; + // accept connections and process them, spawning a new thread for each one + println!("Server listening on port {}", port); + return self.launch_tcp_listener(listener); + } + + /// Launches a thread using a tcp listener socket, on which new clients may connect to this broker + pub fn launch_tcp_listener( + &mut self, + listener: TcpListener, + ) -> Result, AflError> { // Later in the execution, after the initial map filled up, // the current broacast map will will point to a different map. // However, the original map is (as of now) never freed, new clients will start // to read from the initial map id. - let listener = TcpListener::bind(format!("127.0.0.1:{}", port))?; - // accept connections and process them, spawning a new thread for each one - println!("Server listening on port {}", port); - let client_out_map_mem = &self.llmp_out.out_maps.first().unwrap().shmem; let broadcast_str_initial = client_out_map_mem.shm_str.clone(); diff --git a/afl/src/events/mod.rs b/afl/src/events/mod.rs index ba677d2269..c981a06b08 100644 --- a/afl/src/events/mod.rs +++ b/afl/src/events/mod.rs @@ -533,12 +533,16 @@ where #[cfg(test)] mod tests { + use std::{thread, time::Duration}; + use crate::events::Event; use crate::inputs::bytes::BytesInput; use crate::observers::observer_serde::NamedSerdeAnyMap; use crate::observers::{Observer, StdMapObserver}; use crate::serde_anymap::{Ptr, PtrMut}; + use super::llmp::{LlmpConnection, Tag}; + static mut MAP: [u32; 4] = [0; 4]; #[test] @@ -572,4 +576,38 @@ mod tests { _ => panic!("mistmatch".to_string()), }; } + + use crate::events::tests::LlmpConnection::{IsBroker, IsClient}; + + #[test] + pub fn llmp_connection() { + let mut broker = match LlmpConnection::on_port(1337).unwrap() { + IsClient { client } => panic!("Could not bind to port as broker"), + IsBroker { + broker, + listener_thread, + } => broker, + }; + let mut client = match LlmpConnection::on_port(1337).unwrap() { + IsBroker { + broker, + listener_thread, + } => panic!("Second connect should be a client!"), + IsClient { client } => client, + }; + // Add the first client (2nd, actually, because of the tcp listener client) + broker.once().unwrap(); + assert_eq!(broker.llmp_clients.len(), 2); + + let tag: Tag = 0x1337; + let arr: [u8; 1] = [1u8]; + // Send stuff + client.send_buf(tag, &arr).unwrap(); + // Forward stuff to clients + broker.once().unwrap(); + broker.once().unwrap(); + let (tag2, arr2) = client.recv_buf_blocking().unwrap(); + assert_eq!(tag, tag2); + assert_eq!(arr[0], arr2[0]); + } } diff --git a/fuzzers/libfuzzer/src/lib.rs b/fuzzers/libfuzzer/src/lib.rs index 0b373e4ca9..662311784b 100644 --- a/fuzzers/libfuzzer/src/lib.rs +++ b/fuzzers/libfuzzer/src/lib.rs @@ -3,8 +3,6 @@ extern crate alloc; use alloc::boxed::Box; -use alloc::rc::Rc; -use core::cell::RefCell; #[cfg(feature = "std")] use std::io::stderr;