omikron/src/main.rs

mod auth;
mod calls;
mod data;
mod omega;
mod rho;
mod util;

use async_tungstenite::accept_hdr_async;
use futures::StreamExt;
use std::sync::Arc;
use tokio::net::TcpListener;
use tokio_util::compat::TokioAsyncReadCompatExt;
use tungstenite::handshake::server::{Request, Response};

use crate::{
    calls::call_connection::CallConnection,
    omega::omega_connection::OmegaConnection,
    rho::{client_connection::ClientConnection, iota_connection::IotaConnection},
    util::config_util::CONFIG,
    util::print::{PrintType, line, line_err, print_start_message},
};

#[tokio::main]
async fn main() {
    print_start_message();

    tokio::spawn(async move {
        OmegaConnection::new().connect().await;
    });
    let address = format!("{}:{}", &CONFIG.read().await.ip, &CONFIG.read().await.port);
    let listener = TcpListener::bind(&address).await.unwrap();

    line(
        PrintType::General,
        &format!("WebSocket server listening on {}", &address),
    );

    while let Ok((stream, _)) = listener.accept().await {
        tokio::spawn(async move {
            let mut path: String = "/".to_string();

            let callback = |req: &Request, response: Response| {
                path = req.uri().path().to_string(); // Extract URI path
                Ok(response)
            };
            let ws_stream = match accept_hdr_async(stream.compat(), callback).await {
                Ok(ws) => ws,
                Err(e) => {
                    line_err(
                        PrintType::General,
                        &format!("WebSocket upgrade failed: {}", e),
                    );
                    return;
                }
            };
            let (sender, receiver) = ws_stream.split();
            if path == "/ws/client/" {
                line(PrintType::ClientIn, "New Client connection");
                let client_conn: Arc<ClientConnection> =
                    Arc::from(ClientConnection::new(sender, receiver));
                loop {
                    let msg_result = {
                        let mut session_lock = client_conn.receiver.write().await;
                        session_lock.next().await
                    };

                    match msg_result {
                        Some(Ok(msg)) => {
                            if msg.is_text() {
                                let text = msg.into_text().unwrap();
                                client_conn.clone().handle_message(text).await;
                            } else if msg.is_close() {
                                line(PrintType::ClientIn, "Client disconnected");
                                client_conn.handle_close().await;
                                return;
                            }
                        }
                        Some(Err(e)) => {
                            line_err(PrintType::ClientIn, &format!("WebSocket error: {}", e));
                            client_conn.handle_close().await;
                            return;
                        }
                        None => {
                            line(PrintType::ClientIn, "Client stream ended");
                            client_conn.handle_close().await;
                            return;
                        }
                    }
                }
            } else if path == "/ws/iota/" {
                line(PrintType::IotaIn, "New Iota connection");
                let iota_conn: Arc<IotaConnection> =
                    Arc::from(IotaConnection::new(sender, receiver));
                loop {
                    let msg_result = {
                        let mut session_lock = iota_conn.receiver.write().await;
                        session_lock.next().await
                    };

                    match msg_result {
                        Some(Ok(msg)) => {
                            if msg.is_text() {
                                let text = msg.into_text().unwrap();
                                iota_conn.clone().handle_message(text).await;
                            } else if msg.is_close() {
                                line(PrintType::IotaIn, "Iota disconnected");
                                iota_conn.handle_close().await;
                                return;
                            }
                        }
                        Some(Err(e)) => {
                            line_err(PrintType::IotaIn, &format!("WebSocket error: {}", e));
                            iota_conn.handle_close().await;
                            return;
                        }
                        None => {
                            // Stream ended
                            line(PrintType::IotaIn, "Iota stream ended");
                            iota_conn.handle_close().await;
                            return;
                        }
                    }
                }
            } else if path == "/ws/call/" {
                line(PrintType::CallIn, "New Call connection");
                let call_conn: Arc<CallConnection> =
                    Arc::from(CallConnection::new(sender, receiver).await);
                loop {
                    let msg_result = {
                        let mut session_lock = call_conn.receiver.write().await;
                        session_lock.next().await
                    };

                    match msg_result {
                        Some(Ok(msg)) => {
                            if msg.is_text() {
                                let text = msg.into_text().unwrap();
                                call_conn.clone().handle_message(text).await;
                            } else if msg.is_close() {
                                line(PrintType::CallIn, "Call disconnected");
                                call_conn.handle_close().await;
                                return;
                            }
                        }
                        Some(Err(e)) => {
                            line_err(PrintType::CallIn, &format!("WebSocket error: {}", e));
                            call_conn.handle_close().await;
                            return;
                        }
                        None => {
                            line(PrintType::CallIn, "Call stream ended");
                            call_conn.handle_close().await;
                            return;
                        }
                    }
                }
            }
        });
    }
}