MGen/src/bin/mgen-client.rs

// Code specific to the client in the client-server mode.

use mgen::updater::Updater;
use mgen::{HandshakeRef, MessageHeader, SerializedMessage};
use rand_xoshiro::{rand_core::SeedableRng, Xoshiro256PlusPlus};
use serde::Deserialize;
use std::result::Result;
use std::sync::Arc;
use tokio::io::{split, AsyncWriteExt, ReadHalf, WriteHalf};
use tokio::net::TcpStream;
use tokio::sync::mpsc;
use tokio::task;
use tokio::time::Duration;
use tokio_rustls::{client::TlsStream, TlsConnector};

mod messenger;

use crate::messenger::dists::{ConfigDistributions, Distributions};
use crate::messenger::error::{FatalError, MessengerError};
use crate::messenger::state::{
    manage_active_conversation, manage_idle_conversation, StateMachine, StateToWriter,
};
use crate::messenger::tcp::{connect, SocksParams};

/// Type for sending messages from the reader thread to the state thread.
type ReaderToState = mpsc::UnboundedSender<MessageHeader>;
/// Type of messages sent to the writer thread.
type MessageHolder = Box<SerializedMessage>;
/// Type for getting messages from the state thread in the writer thread.
type WriterFromState = mpsc::UnboundedReceiver<MessageHolder>;
/// Type for sending the updated read half of the socket.
type ReadSocketUpdaterIn = Updater<ReadHalf<TlsStream<TcpStream>>>;
/// Type for getting the updated read half of the socket.
type ReadSocketUpdaterOut = Updater<ReadHalf<TlsStream<TcpStream>>>;
/// Type for sending the updated write half of the socket.
type WriteSocketUpdaterIn = Updater<WriteHalf<TlsStream<TcpStream>>>;
/// Type for getting the updated write half of the socket.
type WriteSocketUpdaterOut = Updater<WriteHalf<TlsStream<TcpStream>>>;
/// Type for sending errors to other threads.
type ErrorChannelIn = mpsc::UnboundedSender<MessengerError>;
/// Type for getting errors from other threads.
type ErrorChannelOut = mpsc::UnboundedReceiver<MessengerError>;

// we gain a (very) tiny performance win by not bothering to validate the cert
pub struct NoCertificateVerification {}

impl tokio_rustls::rustls::client::ServerCertVerifier for NoCertificateVerification {
    fn verify_server_cert(
        &self,
        _end_entity: &tokio_rustls::rustls::Certificate,
        _intermediates: &[tokio_rustls::rustls::Certificate],
        _server_name: &tokio_rustls::rustls::ServerName,
        _scts: &mut dyn Iterator<Item = &[u8]>,
        _ocsp: &[u8],
        _now: std::time::SystemTime,
    ) -> Result<tokio_rustls::rustls::client::ServerCertVerified, tokio_rustls::rustls::Error> {
        Ok(tokio_rustls::rustls::client::ServerCertVerified::assertion())
    }
}

/// The thread responsible for getting incoming messages,
/// checking for any network errors while doing so,
/// and giving messages to the state thread.
async fn reader(
    message_channel: ReaderToState,
    mut socket_updater: ReadSocketUpdaterOut,
    error_channel: ErrorChannelIn,
) {
    loop {
        let mut stream = socket_updater.recv().await;
        loop {
            let msg = match mgen::get_message(&mut stream).await {
                Ok(msg) => msg,
                Err(e) => {
                    error_channel.send(e.into()).expect("Error channel closed");
                    break;
                }
            };

            message_channel
                .send(msg)
                .expect("Reader message channel closed");
        }
    }
}

/// The thread responsible for sending messages from the state thread,
/// and checking for any network errors while doing so.
async fn writer(
    mut message_channel: WriterFromState,
    mut socket_updater: WriteSocketUpdaterOut,
    error_channel: ErrorChannelIn,
) {
    loop {
        let mut stream = socket_updater.recv().await;
        loop {
            let msg = message_channel
                .recv()
                .await
                .expect("Writer message channel closed");
            if let Err(e) = msg.write_all_to(&mut stream).await {
                error_channel.send(e.into()).expect("Error channel closed");
                break;
            }
        }
    }
}

/// The thread responsible for (re-)establishing connections to the server,
/// and determining how to handle errors this or other threads receive.
async fn socket_updater(
    str_params: SocksParams,
    retry: f64,
    mut error_channel: ErrorChannelOut,
    reader_channel: ReadSocketUpdaterIn,
    writer_channel: WriteSocketUpdaterIn,
) -> FatalError {
    let retry = Duration::from_secs_f64(retry);

    let tls_config = tokio_rustls::rustls::ClientConfig::builder()
        .with_safe_defaults()
        .with_custom_certificate_verifier(Arc::new(NoCertificateVerification {}))
        .with_no_client_auth();
    let connector = TlsConnector::from(Arc::new(tls_config));

    // unwrap is safe, split always returns at least one element
    let tls_server_str = str_params.target.split(':').next().unwrap();
    let tls_server_name =
        tokio_rustls::rustls::ServerName::try_from(tls_server_str).expect("invalid server name");
    loop {
        let stream: TcpStream = match connect(&str_params).await {
            Ok(stream) => stream,
            Err(MessengerError::Recoverable(_)) => {
                tokio::time::sleep(retry).await;
                continue;
            }
            Err(MessengerError::Fatal(e)) => return e,
        };

        let mut stream = match connector.connect(tls_server_name.clone(), stream).await {
            Ok(stream) => stream,
            Err(_) => {
                tokio::time::sleep(retry).await;
                continue;
            }
        };

        let handshake = HandshakeRef {
            sender: &str_params.user,
            group: &str_params.recipient,
        };

        if stream.write_all(&handshake.serialize()).await.is_err() {
            continue;
        }

        let (rd, wr) = split(stream);
        reader_channel.send(rd);
        writer_channel.send(wr);

        let res = error_channel.recv().await.expect("Error channel closed");
        if let MessengerError::Fatal(e) = res {
            return e;
        }
    }
}

/// The thread responsible for handling the conversation state
/// (i.e., whether the user is active or idle, and when to send messages).
/// Spawns all other threads for this conversation.
async fn manage_conversation(
    user: String,
    socks: Option<String>,
    config: ConversationConfig,
) -> Result<(), MessengerError> {
    let mut rng = Xoshiro256PlusPlus::from_entropy();
    let distributions: Distributions = config.distributions.try_into()?;

    let str_params = SocksParams {
        socks,
        target: config.server,
        user: user.clone(),
        recipient: config.group.clone(),
    };

    let mut state_machine = StateMachine::start(distributions, &mut rng);

    let (reader_to_state, mut state_from_reader) = mpsc::unbounded_channel();
    let (state_to_writer, writer_from_state) = mpsc::unbounded_channel();
    let read_socket_updater_in = Updater::new();
    let read_socket_updater_out = read_socket_updater_in.clone();
    let write_socket_updater_in = Updater::new();
    let write_socket_updater_out = write_socket_updater_in.clone();
    let (errs_in, errs_out) = mpsc::unbounded_channel();

    tokio::spawn(reader(
        reader_to_state,
        read_socket_updater_out,
        errs_in.clone(),
    ));
    tokio::spawn(writer(writer_from_state, write_socket_updater_out, errs_in));
    tokio::spawn(socket_updater(
        str_params,
        config.retry,
        errs_out,
        read_socket_updater_in,
        write_socket_updater_in,
    ));

    tokio::time::sleep(Duration::from_secs_f64(config.bootstrap)).await;

    let mut state_to_writer = StateToWriter {
        channel: state_to_writer,
    };
    loop {
        state_machine = match state_machine {
            StateMachine::Idle(conversation) => {
                manage_idle_conversation::<false, _, _, _>(
                    conversation,
                    &mut state_from_reader,
                    &mut state_to_writer,
                    &user,
                    &config.group,
                    &mut rng,
                )
                .await
            }
            StateMachine::Active(conversation) => {
                manage_active_conversation(
                    conversation,
                    &mut state_from_reader,
                    &mut state_to_writer,
                    &user,
                    &config.group,
                    false,
                    &mut rng,
                )
                .await
            }
        };
    }
}

#[derive(Debug, Deserialize)]
struct ConversationConfig {
    group: String,
    server: String,
    bootstrap: f64,
    retry: f64,
    distributions: ConfigDistributions,
}

#[derive(Debug, Deserialize)]
struct Config {
    user: String,
    socks: Option<String>,
    conversations: Vec<ConversationConfig>,
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let mut args = std::env::args();
    let _ = args.next();
    let mut handles = vec![];
    for config_file in args.flat_map(|a| glob::glob(a.as_str()).unwrap()) {
        let yaml_s = std::fs::read_to_string(config_file?)?;
        let config: Config = serde_yaml::from_str(&yaml_s)?;
        for conversation in config.conversations.into_iter() {
            let handle: task::JoinHandle<Result<(), MessengerError>> = tokio::spawn(
                manage_conversation(config.user.clone(), config.socks.clone(), conversation),
            );
            handles.push(handle);
        }
    }

    let handles: futures::stream::FuturesUnordered<_> = handles.into_iter().collect();
    for handle in handles {
        handle.await??;
    }
    Ok(())
}