rust_webrtc/server/src/signal.rs

use std::{
    collections::VecDeque,
    sync::{Arc, LazyLock},
    time::Duration,
};

use axum::{Router, http::StatusCode, response::IntoResponse, routing::get};
use fastwebsockets::{
    Frame, OpCode, WebSocketError,
    upgrade::{IncomingUpgrade, UpgradeFut},
};
use protocol::{Error, Signal, SignalType};
use tokio::{
    net::TcpListener,
    sync::{RwLock, broadcast},
    time::sleep,
};
use tower_http::cors::CorsLayer;

const SERVER_ADDRESS: &str = "192.168.1.3:4546";

static ONLINE_USERS: LazyLock<RwLock<VecDeque<UserMessages>>> =
    LazyLock::new(|| VecDeque::new().into());

#[derive(Debug)]
struct UserMessages {
    user: String,
    message_receiver: broadcast::Receiver<Signal>,
}

pub async fn start_signalling() {
    let router = Router::new()
        .route("/", get(alive))
        .route("/signal", get(signal))
        .layer(CorsLayer::permissive());

    let listener = TcpListener::bind(SERVER_ADDRESS).await.unwrap();

    println!("{}", SERVER_ADDRESS);
    axum::serve(listener, router).await.unwrap();
}

async fn alive() -> impl IntoResponse {
    StatusCode::OK
}

async fn signal(websocket: IncomingUpgrade) -> impl IntoResponse {
    let (response, websocket) = websocket.upgrade().unwrap();
    tokio::spawn(websocket_handler(websocket));
    response
}

async fn remove_user_from_online_users(user: &String) -> Result<(), Error> {
    let mut target_index = None;
    let mut online_users = ONLINE_USERS.write().await;
    for (index, online_user) in online_users.iter().enumerate() {
        if online_user.user == *user {
            target_index = Some(index);
        }
    }
    if let Some(target_index) = target_index {
        online_users.remove(target_index).expect("Should Not");
        Ok(())
    } else {
        Err(Error::UnregisteredUser)
    }
}

async fn websocket_handler(websocket: UpgradeFut) {
    let mut websocket = websocket.await.unwrap();
    websocket.set_auto_pong(true);
    websocket.set_writev(false);
    websocket.set_auto_close(true);

    let (mut websocket_receiver, mut websocker_sender) = websocket.split(tokio::io::split);
    let user = Arc::new(RwLock::new(String::default()));
    let (message_sender, message_receiver) = broadcast::channel(100);

    let user_for_receiver_first_connection_disconnect_check = user.clone();
    if let Ok(received_frame) = websocket_receiver
        .read_frame::<_, WebSocketError>(&mut move |_| {
            let user_for_receiver_disconnect_check =
                user_for_receiver_first_connection_disconnect_check.clone();
            async move {
                let _ = remove_user_from_online_users(
                    &*user_for_receiver_disconnect_check.read().await,
                )
                .await;
                Ok(())
            }
        })
        .await
    {
        if let OpCode::Text = received_frame.opcode {
            let signal =
                serde_json::from_slice::<Signal>(&received_frame.payload.to_vec()).unwrap();

            if signal.get_signal_type() == SignalType::Auth
                && *user.read().await == String::default()
            {
                let new_user = UserMessages {
                    user: signal.get_data(),
                    message_receiver,
                };
                *user.write().await = new_user.user.to_owned();
                ONLINE_USERS.write().await.push_back(new_user);
            } else {
                return;
            }
        }
    }

    let user_for_sender = user.clone();
    tokio::spawn(async move {
        while ONLINE_USERS.read().await.len() < 2 {
            sleep(Duration::from_secs(1)).await;
        }
        loop {
            let mut user_messages = ONLINE_USERS.write().await;
            for user_message in user_messages.iter_mut() {
                if user_message.user != *user_for_sender.read().await
                    && user_message.message_receiver.len() > 0
                {
                    while let Ok(message) = user_message.message_receiver.recv().await {
                        if let Err(err_val) = websocker_sender
                            .write_frame(Frame::text(fastwebsockets::Payload::Owned(
                                serde_json::to_vec(&message).unwrap(),
                            )))
                            .await
                        {
                            eprintln!("Error: WebSocket Send | {}", err_val);
                            let _ = remove_user_from_online_users(
                                &user_for_sender.read().await.to_owned(),
                            )
                            .await;
                            break;
                        }
                        if user_message.message_receiver.len() < 1 {
                            break;
                        }
                    }
                }
            }
            sleep(Duration::from_secs(1)).await;
        }
    });

    while let Ok(received_frame) = websocket_receiver
        .read_frame::<_, WebSocketError>(&mut |_| async {
            let _ = remove_user_from_online_users(&user.read().await.to_owned()).await;
            Ok(())
        })
        .await
    {
        if let OpCode::Text = received_frame.opcode {
            let signal =
                serde_json::from_slice::<Signal>(&received_frame.payload.to_vec()).unwrap();

            if signal.get_signal_type() != SignalType::Auth {
                if let Err(err_val) = message_sender.send(signal) {
                    eprintln!("Error: WebSocket Channel Send | {}", err_val);
                    let _ = remove_user_from_online_users(&user.read().await.to_owned()).await;
                }
            } else {
                let _ = remove_user_from_online_users(&user.read().await.to_owned()).await;
                return;
            }
        }
    }
}