SeaBattle/rust/cli_player/src/client.rs

use crate::cli_args::cli_args;
use crate::server;
use futures::stream::{SplitSink, SplitStream};
use futures::{SinkExt, StreamExt};
use sea_battle_backend::data::GameRules;
use sea_battle_backend::human_player_ws::{ClientMessage, ServerMessage};
use sea_battle_backend::server::BotPlayQuery;
use sea_battle_backend::utils::{boxed_error, Res};
use std::sync::mpsc;
use std::sync::mpsc::TryRecvError;
use tokio::net::TcpStream;
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};

type WsStream = WebSocketStream<MaybeTlsStream<TcpStream>>;

/// Connection client
///
/// This structure acts as a wrapper around websocket connection that handles automatically parsing
/// of incoming messages and encoding of outgoing messages
pub struct Client {
    sink: SplitSink<WsStream, Message>,
    receiver: mpsc::Receiver<ServerMessage>,
}

impl Client {
    /// Start to play against a bot
    ///
    /// When playing against a bot, local server is always used
    pub async fn start_bot_play(rules: &GameRules) -> Res<Self> {
        server::start_server_if_missing().await;

        Self::connect_url(
            &cli_args().local_server_address(),
            &format!(
                "/play/bot?{}",
                serde_urlencoded::to_string(&BotPlayQuery {
                    rules: rules.clone(),
                    player_name: "Human".to_string()
                })
                .unwrap()
            ),
        )
        .await
    }

    /// Do connect to a server, returning
    async fn connect_url(server: &str, uri: &str) -> Res<Self> {
        let mut url = server.replace("http", "ws");
        url.push_str(uri);
        log::debug!("Connecting to {}", url);

        // Connect to websocket & split streams
        let (socket, _) = tokio_tungstenite::connect_async(url).await?;
        let (sink, mut stream) = socket.split();

        // Receive server message on a separate task
        let (sender, receiver) = mpsc::channel();
        tokio::task::spawn(async move {
            loop {
                match Self::recv_next_msg(&mut stream).await {
                    Ok(msg) => {
                        if let Err(e) = sender.send(msg.clone()) {
                            log::error!("Failed to forward ws message! {} (msg={:?})", e, msg);
                            break;
                        }
                    }
                    Err(e) => {
                        log::error!("Failed receive next message from websocket! {}", e);
                        break;
                    }
                }
            }
        });

        Ok(Self { sink, receiver })
    }

    /// Receive next message from stream
    async fn recv_next_msg(stream: &mut SplitStream<WsStream>) -> Res<ServerMessage> {
        loop {
            let chunk = match stream.next().await {
                None => return Err(boxed_error("No more message in queue!")),
                Some(d) => d,
            };

            match chunk? {
                Message::Text(t) => {
                    log::debug!("TEXT Got a text message from server!");
                    let msg: ServerMessage = serde_json::from_str(&t)?;
                    return Ok(msg);
                }
                Message::Binary(_) => {
                    log::debug!("BINARY Got an unexpected binary message");
                    return Err(boxed_error("Received an unexpected binary message!"));
                }
                Message::Ping(_) => {
                    log::debug!("PING Got a ping message from server");
                }
                Message::Pong(_) => {
                    log::debug!("PONG Got a pong message");
                }
                Message::Close(_) => {
                    log::debug!("CLOSE Got a close websocket message");
                    return Err(boxed_error("Server requested to close connection!"));
                }
                Message::Frame(_) => {
                    log::debug!("FRAME Got an unexpected frame from server!");
                    return Err(boxed_error("Got an unexpected frame!"));
                }
            }
        }
    }

    /// Send a message through the stream
    pub async fn send_message(&mut self, msg: &ClientMessage) -> Res {
        self.sink
            .send(Message::Text(serde_json::to_string(&msg)?))
            .await?;
        Ok(())
    }

    /// Try to receive next message from websocket, in a non-blocking way
    pub async fn try_recv_next_message(&self) -> Res<Option<ServerMessage>> {
        match self.receiver.try_recv() {
            Ok(msg) => Ok(Some(msg)),
            Err(TryRecvError::Empty) => Ok(None),
            Err(TryRecvError::Disconnected) => Err(boxed_error("Receiver channel disconnected!")),
        }
    }

    /// Block until the next message from websocket is availabl
    pub async fn recv_next_message(&self) -> Res<ServerMessage> {
        Ok(self.receiver.recv()?)
    }
}