FRET-LibAFL/libafl/examples/llmp_test/main.rs

/*!
This shows how llmp can be used directly, without libafl abstractions
*/
extern crate alloc;

use alloc::rc::Rc;
#[cfg(all(unix, feature = "std"))]
use core::{cell::RefCell, convert::TryInto, time::Duration};
#[cfg(all(unix, feature = "std"))]
use std::{thread, time};

use libafl::bolts::llmp::Tag;
#[cfg(all(unix, feature = "std"))]
use libafl::{
    bolts::{
        llmp,
        shmem::{ShMemProvider, StdShMemProvider},
    },
    Error,
};

const _TAG_SIMPLE_U32_V1: Tag = 0x51300321;
const _TAG_MATH_RESULT_V1: Tag = 0x77474331;
const _TAG_1MEG_V1: Tag = 0xB1111161;

#[cfg(all(unix, feature = "std"))]
fn adder_loop(port: u16) -> ! {
    let shmem_provider = StdShMemProvider::new().unwrap();
    let mut client = llmp::LlmpClient::create_attach_to_tcp(shmem_provider, port).unwrap();
    let mut last_result: u32 = 0;
    let mut current_result: u32 = 0;
    loop {
        let mut msg_counter = 0;
        loop {
            let (_sender, tag, buf) = match client.recv_buf().unwrap() {
                None => break,
                Some(msg) => msg,
            };
            msg_counter += 1;
            match tag {
                _TAG_SIMPLE_U32_V1 => {
                    current_result =
                        current_result.wrapping_add(u32::from_le_bytes(buf.try_into().unwrap()));
                }
                _ => println!(
                    "Adder Client ignored unknown message {} with {} bytes",
                    tag,
                    buf.len()
                ),
            };
        }

        if current_result != last_result {
            println!(
                "Adder handled {} messages, reporting {} to broker",
                msg_counter, current_result
            );

            client
                .send_buf(_TAG_MATH_RESULT_V1, &current_result.to_le_bytes())
                .unwrap();
            last_result = current_result;
        }

        thread::sleep(time::Duration::from_millis(100));
    }
}

#[cfg(all(unix, feature = "std"))]
fn large_msg_loop(port: u16) -> ! {
    let mut client =
        llmp::LlmpClient::create_attach_to_tcp(StdShMemProvider::new().unwrap(), port).unwrap();

    let meg_buf = [1u8; 1 << 20];

    loop {
        client.send_buf(_TAG_1MEG_V1, &meg_buf).unwrap();
        println!("Sending the next megabyte");
        thread::sleep(time::Duration::from_millis(100))
    }
}

#[cfg(all(unix, feature = "std"))]
fn broker_message_hook(
    client_id: u32,
    tag: llmp::Tag,
    message: &[u8],
) -> Result<llmp::LlmpMsgHookResult, Error> {
    match tag {
        _TAG_SIMPLE_U32_V1 => {
            println!(
                "Client {:?} sent message: {:?}",
                client_id,
                u32::from_le_bytes(message.try_into().unwrap())
            );
            Ok(llmp::LlmpMsgHookResult::ForwardToClients)
        }
        _TAG_MATH_RESULT_V1 => {
            println!(
                "Adder Client has this current result: {:?}",
                u32::from_le_bytes(message.try_into().unwrap())
            );
            Ok(llmp::LlmpMsgHookResult::Handled)
        }
        _ => {
            println!("Unknwon message id received!");
            Ok(llmp::LlmpMsgHookResult::ForwardToClients)
        }
    }
}

#[cfg(not(unix))]
fn main() {
    todo!("LLMP is not yet supported on this platform.");
}

#[cfg(unix)]
fn main() {
    /* The main node has a broker, and a few worker threads */

    let mode = std::env::args()
        .nth(1)
        .expect("no mode specified, chose 'broker', 'ctr', 'adder', or 'large'");
    let port: u16 = std::env::args()
        .nth(2)
        .unwrap_or("1337".into())
        .parse::<u16>()
        .unwrap();
    println!("Launching in mode {} on port {}", mode, port);

    match mode.as_str() {
        "broker" => {
            let mut broker = llmp::LlmpBroker::new(StdShMemProvider::new().unwrap()).unwrap();
            broker
                .launch_listener(llmp::Listener::Tcp(
                    std::net::TcpListener::bind(format!("127.0.0.1:{}", port)).unwrap(),
                ))
                .unwrap();
            broker.loop_forever(&mut broker_message_hook, Some(Duration::from_millis(5)))
        }
        "ctr" => {
            let mut client =
                llmp::LlmpClient::create_attach_to_tcp(StdShMemProvider::new().unwrap(), port)
                    .unwrap();
            let mut counter: u32 = 0;
            loop {
                counter = counter.wrapping_add(1);
                client
                    .send_buf(_TAG_SIMPLE_U32_V1, &counter.to_le_bytes())
                    .unwrap();
                println!("CTR Client writing {}", counter);
                thread::sleep(Duration::from_secs(1))
            }
        }
        "adder" => {
            adder_loop(port);
        }
        "large" => {
            large_msg_loop(port);
        }
        _ => {
            println!("No valid mode supplied");
        }
    }
}