cards/server/src/websocket.rs

use crate::user_handler::*;
use crate::AppState;
use crate::Game;
use crate::NewGameManifest;
use crate::User;
use anyhow::Result;
use axum::extract::ws::CloseFrame;
use axum::{
    extract::{
        ws::{Message, WebSocket},
        ConnectInfo, State, WebSocketUpgrade,
    },
    response::IntoResponse,
};
use futures::{SinkExt, StreamExt};
use lib::*;
use rand::prelude::SliceRandom;
use serde_json::{from_str, to_string};
use std::collections::HashMap;
use std::{
    net::SocketAddr,
    sync::{Arc, RwLock},
};
use tokio::sync::{broadcast::Sender, mpsc};

/// Establish the WebSocket connection
pub async fn websocket_connection_handler(
    ws: WebSocketUpgrade,
    // user_agent: Option<TypedHeader<headers::UserAgent>>,
    ConnectInfo(addr): ConnectInfo<SocketAddr>,
    State(state): State<Arc<AppState>>,
) -> impl IntoResponse {
    tracing::debug!("New connection from {}", &addr);
    ws.on_upgrade(move |socket| websocket_on_connection(socket, state, addr))
}

/// This runs right after a WebSocket connection is established
pub async fn websocket_on_connection(stream: WebSocket, state: Arc<AppState>, addr: SocketAddr) {
    // Split channels to send and receive asynchronously.
    let (mut sender, mut receiver) = stream.split();

    // Create channel for direct messages
    let (dm_tx, mut dm_rx) = mpsc::channel(30);

    let mut map = HashMap::new();
    map.insert(addr, dm_tx.clone());

    let _ = state
        .users_tx
        // add tx
        .send(UserHandlerMessage::NewUser {
            user: User::new(
                format!(
                    "{} {}",
                    state.first_names.choose(&mut rand::thread_rng()).unwrap(),
                    state.last_names.choose(&mut rand::thread_rng()).unwrap(),
                ),
                dm_tx.clone(),
            ),
            addr,
        })
        .await;

    // Subscribe to receive from global broadcast channel
    let mut rx = state.broadcast_tx.subscribe();

    // Send messages to this client
    let mut send_task = tokio::spawn(async move {
        let mut broadcast = None;
        let mut dm = None;
        loop {
            tokio::select! {
                b = rx.recv() => broadcast = Some(b.unwrap()),
                d = dm_rx.recv() => dm = d,
            };

            if let Some(msg) = &dm {
                if sender.send(Message::Text(msg.to_string())).await.is_err() {
                    break;
                } else {
                    dm = Option::None;
                }
            } else if let Some(msg) = &broadcast {
                if sender.send(Message::Text(msg.to_string())).await.is_err() {
                } else {
                    broadcast = Option::None;
                }
            }
        }
    });

    // Receive messages from this client
    let mut recv_task = tokio::spawn(async move {
        while let Some(Ok(message)) = receiver.next().await {
            state
                .messages_tx
                .send((addr.clone(), message.clone()))
                .await
                .unwrap();
            websocket_message_handler(state.clone(), addr, message)
                .await
                .expect("Message Handler exploded!")
        }
    });

    // If either task completes then abort the other
    tokio::select! {
        _ = (&mut send_task) => recv_task.abort(),
        _ = (&mut recv_task) => send_task.abort(),
    };
}

/// Handle incoming messages over the WebSocket
pub async fn websocket_message_handler(
    state: Arc<AppState>,
    addr: SocketAddr,
    message: Message,
) -> Result<()> {
    let tx = &state.broadcast_tx;

    match message {
        Message::Text(text) => match text {
            _new_game if let Ok(_new_game) = from_str::<NewGameRequest>(&text) => {
                tracing::debug!("New game request received.");
                game_handle_new_game(_new_game, &state, tx, addr)?;
            }
            _ => {
                tracing::debug!("Unhandled text message: {}", &text);
            }
        },
        Message::Binary(data) => {
            tracing::debug!("Binary: {:?}", data)
        }
        Message::Close(_close_frame) => {
            // websocket_handle_close(close_frame, &state, tx, addr)?;
        }
        Message::Pong(ping) => {
            tracing::debug!("Pong received with: {:?}", ping);
        }
        Message::Ping(pong) => {
            tracing::debug!("Pong received with: {:?}", pong);
        }
    }

    Ok(())
}

/// This runs when a NewGameRequest is received
fn game_handle_new_game(
    new_game: NewGameRequest,
    state: &Arc<AppState>,
    tx: &Sender<String>,
    addr: SocketAddr,
) -> Result<()> {
    let manifest = NewGameManifest {
        name: new_game.name,
        host: state
            .online_users
            .read()
            .unwrap()
            .get(&addr)
            .unwrap()
            .clone(),
    };
    tracing::debug!("Game Packs {:?}", new_game.packs);

    // create game
    if let Ok(new_game_object) = Game::new(manifest) {
        state
            .games
            .write()
            .unwrap()
            .insert(new_game_object.name.clone(), RwLock::new(new_game_object));
        tx.send(meta_games_browser_update(state))?;
        tx.send(meta_server_summary_update(state))?;
    }

    Ok(())
}