From ff8a89f0c190a5052af5d1b062e847bca318f808 Mon Sep 17 00:00:00 2001 From: Dominik Maier Date: Tue, 8 Dec 2020 21:12:39 +0100 Subject: [PATCH] added tcp server --- afl/src/events/llmp.rs | 121 ++++++++++++++++++++++++----- afl/src/events/llmp_manager.rs | 49 ------------ afl/src/events/mod.rs | 6 +- afl/src/events/shmem_translated.rs | 1 - 4 files changed, 103 insertions(+), 74 deletions(-) diff --git a/afl/src/events/llmp.rs b/afl/src/events/llmp.rs index 91269d6f6e..18cfcfc1f1 100644 --- a/afl/src/events/llmp.rs +++ b/afl/src/events/llmp.rs @@ -52,7 +52,14 @@ use core::ptr; use core::sync::atomic::{compiler_fence, Ordering}; use core::time::Duration; use libc::{c_uint, c_ulong, c_ushort}; -use std::{cmp::max, ffi::CStr, mem::size_of, thread}; +use std::{ + cmp::max, + ffi::CStr, + io::{Read, Write}, + mem::size_of, + net::TcpListener, + thread, +}; use crate::utils::next_pow2; use crate::AflError; @@ -471,6 +478,28 @@ impl LlmpSender { (*msg).tag = LLMP_TAG_UNSET; (*page).size_used -= (*msg).buf_len_padded as usize + size_of::(); } + + /// Allocates a message of the given size, tags it, and sends it off. + pub fn send_buf(&mut self, tag: u32, buf: &[u8]) -> Result<(), AflError> { + // Make sure we don't reuse already allocated tags + if tag == LLMP_TAG_NEW_SHM_CLIENT + || tag == LLMP_TAG_END_OF_PAGE + || tag == LLMP_TAG_UNINITIALIZED + || tag == LLMP_TAG_UNSET + { + return Err(AflError::Unknown(format!( + "Reserved tag supplied to send_buf ({:#X})", + tag + ))); + } + unsafe { + let msg = self.alloc_next(buf.len())?; + (*msg).tag = tag; + buf.as_ptr() + .copy_to_nonoverlapping((*msg).buf.as_mut_ptr(), buf.len()); + self.send(msg) + } + } } /// Receiving end of an llmp channel @@ -625,7 +654,7 @@ impl LlmpBroker { /// Registers a new client for the given sharedmap str and size. /// Returns the id of the new client in broker.client_map - pub unsafe fn register_client(&mut self, client_page: LlmpSharedMap) { + pub fn register_client(&mut self, client_page: LlmpSharedMap) { let id = self.llmp_clients.len() as u32; self.llmp_clients.push(LlmpReceiver { id, @@ -740,6 +769,75 @@ impl LlmpBroker { } } } + + pub fn launch_tcp_listener(&mut self, port: u16) -> 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(); + + let llmp_tcp_id = self.llmp_clients.len() as u32; + + // Tcp out map sends messages from background thread tcp server to foreground client + let tcp_out_map = LlmpSharedMap::new(llmp_tcp_id, LLMP_PREF_INITIAL_MAP_SIZE)?; + let tcp_out_map_str = tcp_out_map.shmem.shm_str; + let tcp_out_map_size = tcp_out_map.shmem.map_size; + self.register_client(tcp_out_map); + + Ok(thread::spawn(move || { + let mut new_client_sender = LlmpSender { + id: 0, + last_msg_sent: 0 as *mut LlmpMsg, + out_maps: vec![ + LlmpSharedMap::from_name_slice(&tcp_out_map_str, tcp_out_map_size).unwrap(), + ], + // drop pages to the broker if it already read them + keep_pages_forever: false, + }; + + loop { + let (mut stream, addr) = match listener.accept() { + Ok(res) => res, + Err(e) => { + dbg!("Ignoring failed accept", e); + continue; + } + }; + dbg!("New connection", addr, stream.peer_addr().unwrap()); + match stream.write(&broadcast_str_initial) { + Ok(_) => {} // fire & forget + Err(e) => { + dbg!("Could not send to shmap to client", e); + continue; + } + }; + let mut new_client_map_str: [u8; 20] = Default::default(); + let map_str_len = match stream.read(&mut new_client_map_str) { + Ok(res) => res, + Err(e) => { + dbg!("Ignoring failed read from client", e); + continue; + } + }; + if map_str_len < 20 { + dbg!("Didn't receive a complete shmap id str from client. Ignoring."); + continue; + } + + match new_client_sender.send_buf(LLMP_TAG_NEW_SHM_CLIENT, &new_client_map_str) { + Ok(()) => (), + Err(e) => println!("Error forwarding client on map: {:?}", e), + }; + } + })) + } } /// `n` clients connect to a broker. They share an outgoing map with the broker, @@ -770,24 +868,7 @@ impl LlmpClient { /// Allocates a message of the given size, tags it, and sends it off. pub fn send_buf(&mut self, tag: u32, buf: &[u8]) -> Result<(), AflError> { - // Make sure we don't reuse already allocated tags - if tag == LLMP_TAG_NEW_SHM_CLIENT - || tag == LLMP_TAG_END_OF_PAGE - || tag == LLMP_TAG_UNINITIALIZED - || tag == LLMP_TAG_UNSET - { - return Err(AflError::Unknown(format!( - "Reserved tag supplied to send_buf ({:#X})", - tag - ))); - } - unsafe { - let msg = self.alloc_next(buf.len())?; - (*msg).tag = tag; - buf.as_ptr() - .copy_to_nonoverlapping((*msg).buf.as_mut_ptr(), buf.len()); - self.send(msg) - } + self.llmp_out.send_buf(tag, buf) } /// A client receives a broadcast message. diff --git a/afl/src/events/llmp_manager.rs b/afl/src/events/llmp_manager.rs index 4a191c1fe7..78542e509b 100644 --- a/afl/src/events/llmp_manager.rs +++ b/afl/src/events/llmp_manager.rs @@ -9,55 +9,6 @@ use super::{ Event, EventManager, }; -/* -pub unsafe fn llmp_tcp_server_clientloop(client: &mut LlmpClient, _data: *mut c_void) -> ! { - // 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 initial_broadcasts_map_str = client - .as_ref() - .unwrap() - .current_broadcast_map - .as_ref() - .unwrap() - .shm_str; - - let listener = TcpListener::bind("0.0.0.0:3333").unwrap(); - // accept connections and process them, spawning a new thread for each one - println!("Server listening on port 3333"); - loop { - let (mut stream, addr) = match listener.accept() { - Ok(res) => res, - Err(e) => { - dbg!("Ignoring failed accept", e); - continue; - } - }; - dbg!("New connection", addr, stream.peer_addr().unwrap()); - match stream.write(&initial_broadcasts_map_str as &[u8]) { - Ok(_) => {} // fire & forget - Err(e) => { - dbg!("Could not send to shmap to client", e); - continue; - } - }; - let mut new_client_map_str: [u8; 20] = Default::default(); - let map_str_len = match stream.read(&mut new_client_map_str) { - Ok(res) => res, - Err(e) => { - dbg!("Ignoring failed read from client", e); - continue; - } - }; - if map_str_len < 20 { - dbg!("Didn't receive a complete shmap id str from client. Ignoring."); - continue; - } - } -} -*/ - /// Eventmanager for multi-processed application #[cfg(feature = "std")] pub struct LLMPEventManager diff --git a/afl/src/events/mod.rs b/afl/src/events/mod.rs index 42c104a693..a301f2f48d 100644 --- a/afl/src/events/mod.rs +++ b/afl/src/events/mod.rs @@ -6,13 +6,11 @@ use core::marker::PhantomData; use serde::{Deserialize, Serialize}; -#[cfg(feature = "std")] -pub mod llmp_translated; // TODO: Abstract away. #[cfg(feature = "std")] pub mod shmem_translated; -#[cfg(feature = "std")] -pub use crate::events::llmp::LLMPEventManager; +/*#[cfg(feature = "std")] +pub use crate::events::llmp::LLMPEventManager;*/ #[cfg(feature = "std")] use std::io::Write; diff --git a/afl/src/events/shmem_translated.rs b/afl/src/events/shmem_translated.rs index e11c000cd3..dbe195dcdd 100644 --- a/afl/src/events/shmem_translated.rs +++ b/afl/src/events/shmem_translated.rs @@ -59,7 +59,6 @@ const AFL_RET_SUCCESS: c_uint = 0; // too.) #[derive(Clone)] -#[repr(C)] pub struct AflShmem { pub shm_str: [u8; 20], pub shm_id: c_int,