radioxide/back/src/streaming.rs

use std::{
    fs::File,
    io::{self, BufReader},
    sync::Arc,
    time::Duration,
};

use futures_util::{SinkExt, StreamExt};
use rustls_pemfile::{certs, pkcs8_private_keys};

use tokio::{
    net::TcpListener,
    sync::broadcast::{channel, Receiver, Sender},
    time::Instant,
};
use tokio_rustls::{
    rustls::pki_types::{CertificateDer, PrivateKeyDer},
    TlsAcceptor,
};
use tokio_tungstenite::tungstenite::{Error, Message};

use crate::{Config, Listener, Streamer};

const BUFFER_LENGTH: usize = 1000000;
const MAX_TOLERATED_MESSAGE_COUNT: usize = 10;
pub async fn start(relay_configs: Config) {
    //need to move them for multi streamer
    let listener_socket = TcpListener::bind(relay_configs.listener_address)
        .await
        .unwrap();
    let (record_producer, record_consumer) = channel(BUFFER_LENGTH);
    let streamer_socket = TcpListener::bind(relay_configs.streamer_address)
        .await
        .unwrap();
    let timer = Instant::now();

    let fullchain: io::Result<Vec<CertificateDer<'static>>> = certs(&mut BufReader::new(
        File::open("certificates/fullchain.pem").unwrap(),
    ))
    .collect();
    let fullchain = fullchain.unwrap();
    let privkey: io::Result<PrivateKeyDer<'static>> = pkcs8_private_keys(&mut BufReader::new(
        File::open("certificates/privkey.pem").unwrap(),
    ))
    .next()
    .unwrap()
    .map(Into::into);
    let privkey = privkey.unwrap();

    let server_tls_config = tokio_rustls::rustls::ServerConfig::builder()
        .with_no_client_auth()
        .with_single_cert(fullchain, privkey)
        .unwrap();
    let acceptor = TlsAcceptor::from(Arc::new(server_tls_config));
    loop {
        match streamer_socket.accept().await {
            Ok((streamer_tcp, streamer_info)) => {
                let new_streamer = Streamer {
                    ip: streamer_info.ip(),
                    port: streamer_info.port(),
                };
                println!(
                    "New Streamer: {:#?} | {:#?}",
                    streamer_info,
                    timer.elapsed()
                );
                if relay_configs.tls {
                    let streamer_tcp_tls = acceptor.accept(streamer_tcp).await.unwrap();
                    match tokio_tungstenite::accept_async(streamer_tcp_tls).await {
                        Ok(ws_stream) => {
                            tokio::spawn(streamer_stream(
                                new_streamer,
                                record_producer,
                                ws_stream,
                                timer,
                            ));
                            break;
                        }
                        Err(err_val) => eprintln!("Error: TCP to WS Transform | {}", err_val),
                    }
                } else {
                    match tokio_tungstenite::accept_async(streamer_tcp).await {
                        Ok(ws_stream) => {
                            tokio::spawn(streamer_stream(
                                new_streamer,
                                record_producer,
                                ws_stream,
                                timer,
                            ));
                            break;
                        }
                        Err(err_val) => eprintln!("Error: TCP to WS Transform | {}", err_val),
                    }
                }
            }
            Err(err_val) => eprintln!("Error: TCP Accept Connection | {}", err_val),
        }
    }
    let (message_producer, message_consumer) = channel(BUFFER_LENGTH);
    let (buffered_producer, _) = channel(BUFFER_LENGTH);
    tokio::spawn(message_organizer(
        message_producer.clone(),
        record_consumer,
        relay_configs.latency,
    ));
    tokio::spawn(buffer_layer(
        message_consumer,
        buffered_producer.clone(),
        relay_configs.latency,
    ));
    tokio::spawn(status_checker(buffered_producer.clone(), timer));
    while let Ok((tcp_stream, listener_info)) = listener_socket.accept().await {
        let new_listener = Listener {
            ip: listener_info.ip(),
            port: listener_info.port(),
        };
        if relay_configs.tls {
            let streamer_tcp_tls = acceptor.accept(tcp_stream).await.unwrap();
            let wss_stream = tokio_tungstenite::accept_async(streamer_tcp_tls)
                .await
                .unwrap();
            tokio::spawn(stream(
                new_listener,
                wss_stream,
                buffered_producer.subscribe(),
            ));
        } else {
            let ws_stream = tokio_tungstenite::accept_async(tcp_stream).await.unwrap();
            tokio::spawn(stream(
                new_listener,
                ws_stream,
                buffered_producer.subscribe(),
            ));
        }
        println!("New Listener: {} | {:#?}", listener_info, timer.elapsed());
    }
}
async fn status_checker(buffered_producer: Sender<Message>, timer: Instant) {
    let mut listener_counter = buffered_producer.receiver_count();
    let mut bottleneck_flag = false;
    //let mut buffer_len = buffered_producer.len();
    loop {
        tokio::time::sleep(Duration::from_secs(3)).await;
        if buffered_producer.receiver_count() != 0 {
            if buffered_producer.len() > 2 {
                bottleneck_flag = true;
                println!(
                    "Bottleneck: {} | {:#?}",
                    buffered_producer.len(),
                    timer.elapsed()
                );
            }
            if bottleneck_flag && buffered_producer.len() < 2 {
                bottleneck_flag = false;
                println!("Flawless Again");
            }
            if listener_counter != buffered_producer.receiver_count() {
                listener_counter = buffered_producer.receiver_count();
                println!("Listener(s): {}", listener_counter);
            }
        }
    }
}
async fn buffer_layer(
    mut message_consumer: Receiver<Message>,
    buffered_producer: Sender<Message>,
    delay: u16,
) {
    loop {
        tokio::time::sleep(Duration::from_millis(delay.into())).await;
        while message_consumer.len() > 0 {
            match message_consumer.recv().await {
                Ok(message) => match buffered_producer.send(message) {
                    Ok(_) => {}
                    Err(_) => {}
                },
                Err(_) => {}
            }
        }
    }
}
async fn streamer_stream<
    T: futures_util::Stream<Item = Result<Message, Error>> + std::marker::Unpin,
>(
    streamer: Streamer,
    record_producer: Sender<Message>,
    mut ws_stream: T,
    timer: Instant,
) {
    loop {
        match ws_stream.next().await {
            Some(message_with_question) => {
                match message_with_question {
                    Ok(message) => {
                        //println!("{}", message.len());
                        match record_producer.send(message) {
                            Ok(_) => {}
                            Err(_) => {}
                        }
                    }
                    Err(_) => {}
                }
            }
            None => {
                println!(
                    "Streamer Disconnected: {} | {:#?}",
                    format!("{}:{}", streamer.ip, streamer.port),
                    timer.elapsed()
                );
                return;
            }
        }
    }
}

async fn message_organizer(
    message_producer: Sender<Message>,
    mut record_consumer: Receiver<Message>,
    delay: u16,
) {
    loop {
        match record_consumer.recv().await {
            Ok(single_message) => match message_producer.send(single_message) {
                Ok(_) => {}
                Err(_) => {}
            },
            Err(_) => {}
        }
        tokio::time::sleep(Duration::from_millis(delay.into())).await;
    }
}
async fn stream<T: futures_util::Sink<Message> + std::marker::Unpin>(
    listener: Listener,
    mut ws_stream: T,
    mut buffered_consumer: Receiver<Message>,
) {
    while let Ok(message) = buffered_consumer.recv().await {
        if buffered_consumer.len() > MAX_TOLERATED_MESSAGE_COUNT {
            println!(
                "{} Forced to Disconnect | Reason -> Slow Consumer",
                format!("{}:{}", listener.ip, listener.port)
            );
            break;
        }

        match ws_stream.send(message).await {
            Ok(_) => {
                if let Err(_) = ws_stream.flush().await {
                    println!(
                        "{} is Disconnected",
                        format!("{}:{}", listener.ip, listener.port)
                    );
                    break;
                }
            }
            Err(_) => {
                println!(
                    "{} is Disconnected",
                    format!("{}:{}", listener.ip, listener.port)
                );
                break;
            }
        }
    }
}