tcp_message_io 1.0.4

use std::cell::RefCell;
use std::fmt::Debug;
use std::future::Future;
use std::marker::PhantomData;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};

use anyhow::{bail, Result};
#[cfg(feature = "postcard")]
use postcard::{from_bytes, to_allocvec};
#[cfg(feature = "postcard")]
use serde::{Deserialize, Serialize};

use crate::raw::{RawTCPClient, RawTCPResponse, RawTCPServer};

/// Trait a message type needs to implement to be usable with the
/// TCPClient/TCPServer implementation. This allows to customize the
/// serialization format.
pub trait SerializeMessage: Sized + Send + Sync + 'static {
    fn serialize(&self) -> Result<Vec<u8>>;
    fn deserialize(message: &[u8]) -> Result<Self>;
}

#[cfg(feature = "postcard")]
impl<T> SerializeMessage for T
where
    T: Serialize + for<'a> Deserialize<'a> + Send + Sync + 'static,
{
    fn serialize(&self) -> Result<Vec<u8>> {
        Ok(to_allocvec(self)?)
    }

    fn deserialize(message: &[u8]) -> Result<Self> {
        Ok(from_bytes(message)?)
    }
}

/// Client to make TCP requests in the form of messages.
///
/// Messages can be any serializable object. This library uses
/// a 9-byte header to encode the size of the serialized data sent
/// and the result.
///
/// By default, postcard is used as wire serialization format
pub struct TCPClient<Q, A>
where
    // These might be 2 different types
    Q: SerializeMessage,
    A: SerializeMessage,
{
    raw_client: RawTCPClient,
    phantom_q: PhantomData<Q>,
    phantom_a: PhantomData<A>,
}

impl<Q, A> TCPClient<Q, A>
where
    Q: SerializeMessage,
    A: SerializeMessage,
{
    /// Connect to a TCP server at the given `host` and `port`.
    ///
    /// Returns an error if the connection cannot be established.
    pub async fn connect(host: &str, port: u16) -> Result<Self> {
        Ok(Self {
            raw_client: RawTCPClient::connect(host, port).await?,
            phantom_q: PhantomData,
            phantom_a: PhantomData,
        })
    }

    /// Serialized and sends a message, waits for the response,
    /// returning the deserialized content.
    ///
    /// Returns an error if there were errors while sending
    /// or receiving data, or if the connection was closed by the server.
    pub async fn send(&mut self, message: Q) -> Result<Option<A>> {
        let raw_message = message.serialize()?;
        let raw_response = self.raw_client.send(&raw_message).await?;
        if raw_response.is_empty() {
            return Ok(None);
        }
        Ok(Some(A::deserialize(&raw_response)?))
    }
}

/// Response returned by message handlers to indicate which
/// action to take.
#[derive(Debug)]
pub enum TCPResponse<A>
where
    A: SerializeMessage,
{
    /// Causes the server to send the given response message back to the client.
    Message(A),
    /// Instructs the server to close the connection. To avoid that the client
    /// receives an error because the connection was closed, this will also
    /// send an empty response back to the client.
    CloseConnection,
    /// Causes the server to stop. To avoid that the client
    /// receives an error because the connection was closed, this will also
    /// send an empty response back to the client.
    StopServer,
}

/// Server to handle TCP requests.
///
/// When the `.listen()` method is called the server starts accepting
/// (possibly simultaneous) connections from any number of clients.
///
/// When a client sends a message, the server will invoke the given request
/// handle to process the request, and take an action based on the handler
/// return type (see [`TCPResponse`] for a list of possible actions).
///
/// If the client closes the connection without warning, the server will
/// simply drop that connection.
///
/// If the client sends a malformed message (without or with a wrong length
/// header), the server can potentially read too little data, or hang waiting
/// for data to read.
pub struct TCPServer<Q, A, H, F>
where
    Q: SerializeMessage,
    A: SerializeMessage,
    H: Fn(Q) -> F + Send + Sync + 'static,
    F: Future<Output = Result<TCPResponse<A>>> + Send + 'static,
{
    host: String,
    port: u16,
    handler: H,
    bad_request_response: Mutex<RefCell<Option<fn() -> TCPResponse<A>>>>,
    inactivity_timeout_ms: AtomicU64,
    phantom_q: PhantomData<Q>,
}

impl<Q, A, H, F> TCPServer<Q, A, H, F>
where
    A: SerializeMessage,
    Q: SerializeMessage,
    H: Fn(Q) -> F + Send + Sync + 'static,
    F: Future<Output = Result<TCPResponse<A>>> + Send + 'static,
{
    /// Create a new server listening to the given `host` and `port`,
    /// processing requests with `handler`.
    ///
    /// Does not actually start listening, for that you need to call [`listen`][Self::listen].
    ///
    /// The socket will be freed when the struct is dropped.
    pub fn new(host: impl Into<String>, port: u16, handler: H) -> Arc<Self> {
        Arc::new(Self {
            host: host.into(),
            port,
            handler,
            bad_request_response: Mutex::new(RefCell::new(None)),
            inactivity_timeout_ms: AtomicU64::new(0),
            phantom_q: PhantomData,
        })
    }

    pub fn with_bad_request_handler(
        self: Arc<Self>,
        bad_request_response: fn() -> TCPResponse<A>,
    ) -> Arc<Self> {
        *self.bad_request_response.lock().unwrap().borrow_mut() = Some(bad_request_response);
        self
    }

    /// Instructs the server to quit after the given amount of time without any request.
    ///
    /// Note that requests that take longer than the timeout amount, will be dropped
    /// in the middle of processing
    pub fn with_inactivity_timeout(self: Arc<Self>, timeout_ms: u64) -> Arc<Self> {
        self.inactivity_timeout_ms
            .store(timeout_ms, Ordering::Relaxed);
        self
    }

    /// Start accepting connections from clients, processing and answering messages.
    ///
    /// If multiple servers are started on the same port, this function will panic.
    pub async fn listen(self: Arc<Self>) {
        let cloned_self = self.clone();
        RawTCPServer::new(self.host.clone(), self.port, move |req| {
            cloned_self.clone().handle_raw_message(req)
        })
        .with_inactivity_timeout(self.inactivity_timeout_ms.load(Ordering::Relaxed))
        .listen()
        .await;
    }

    async fn handle_raw_message(self: Arc<Self>, raw_request: Vec<u8>) -> Result<RawTCPResponse> {
        let action = match Q::deserialize(&raw_request) {
            Ok(request) => (self.handler)(request).await?,
            Err(err) => {
                if let Some(bad_request_handler) =
                    self.bad_request_response.lock().unwrap().borrow().clone()
                {
                    bad_request_handler()
                } else {
                    bail!(
                        "Bad request, unable to deserialize (this might be caused \
                        by incompatible client version or mismatch in compression configuration): {}",
                        err
                    );
                }
            }
        };
        Ok(match action {
            TCPResponse::Message(response) => RawTCPResponse::Message(response.serialize()?),
            TCPResponse::CloseConnection => RawTCPResponse::CloseConnection,
            TCPResponse::StopServer => RawTCPResponse::StopServer,
        })
    }
}

#[cfg(test)]
mod tests {
    use std::time::{Duration, Instant};

    use serde::{Deserialize, Serialize};
    use serial_test::serial;
    use tokio::spawn;
    use tokio::time::sleep;

    use super::*;

    const HOST: &str = "127.0.0.1";
    const PORT: u16 = 12345;

    #[derive(Debug, Serialize, Deserialize)]
    enum Request {
        Hello,
        Double { num: u64 },
        Sum { a: u64, b: u64 },
        Close,
        CauseError,
        Stop,
    }

    #[derive(Debug, Serialize, Deserialize)]
    enum OtherRequest {
        One,
        Two,
        Three,
        Four,
        Five,
        Six,
        // Other enum above has only 6 elements,
        // so 7th will not deserialize. I tried with
        // tagged unions and I could not get it to work.
        ImDifferent,
    }

    #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
    enum Response {
        World,
        Result(u64),
    }

    async fn while_server_running(fut: impl Future) {
        let handle = spawn(
            TCPServer::new(HOST, PORT, handle_requests)
                .with_inactivity_timeout(15)
                .listen(),
        );
        // Give some time to server to bind socket
        sleep(Duration::from_millis(10)).await;
        fut.await;
        handle.await.unwrap();
    }

    async fn while_server_running_custom_bad_request(fut: impl Future) {
        let handle = spawn(
            TCPServer::new(HOST, PORT, handle_requests)
                .with_inactivity_timeout(15)
                .with_bad_request_handler(|| TCPResponse::Message(Response::World))
                .listen(),
        );
        // Give some time to server to bind socket
        sleep(Duration::from_millis(10)).await;
        fut.await;
        handle.await.unwrap();
    }

    async fn handle_requests(req: Request) -> Result<TCPResponse<Response>> {
        Ok(match req {
            Request::Hello => TCPResponse::Message(Response::World),
            Request::Double { num } => TCPResponse::Message(Response::Result(2 * num)),
            Request::Sum { a, b } => TCPResponse::Message(Response::Result(a + b)),
            Request::Close => TCPResponse::CloseConnection,
            Request::CauseError => bail!("An error occurred".to_string()),
            Request::Stop => TCPResponse::StopServer,
        })
    }

    // We send a couple of messages, to the server,
    // check that responses are correct, then stop it.
    #[tokio::test]
    #[serial]
    async fn test_tcp_server() {
        while_server_running(async {
            let mut client = TCPClient::<_, Response>::connect(HOST, PORT).await.unwrap();
            assert_eq!(
                client.send(Request::Hello).await.unwrap().unwrap(),
                Response::World
            );
            assert_eq!(
                client
                    .send(Request::Double { num: 3 })
                    .await
                    .unwrap()
                    .unwrap(),
                Response::Result(6)
            );
            assert_eq!(
                client
                    .send(Request::Sum { a: 3, b: 5 })
                    .await
                    .unwrap()
                    .unwrap(),
                Response::Result(8)
            );
            // Stop server
            assert_eq!(client.send(Request::Stop).await.unwrap(), None);
        })
        .await;
    }

    #[tokio::test]
    #[serial]
    async fn test_tcp_server_handler_error() {
        while_server_running(async {
            let mut client = TCPClient::<_, Response>::connect(HOST, PORT).await.unwrap();
            // Invoke handler route that causes error
            assert_eq!(client.send(Request::CauseError).await.unwrap(), None);
            // Stop server
            assert_eq!(client.send(Request::Stop).await.unwrap(), None);
        })
        .await;
    }

    #[tokio::test]
    #[serial]
    async fn test_tcp_server_bad_request_error() {
        let start = Instant::now();

        while_server_running(async {
            let mut client = TCPClient::<_, Response>::connect(HOST, PORT).await.unwrap();
            // Send wrong message type
            assert_eq!(client.send(OtherRequest::ImDifferent).await.unwrap(), None);
        })
        .await;

        let elapsed_ms = start.elapsed().as_millis();
        // let is be off by 1ms to make the test robust
        assert!(elapsed_ms >= 15, "Elapsed time: {} ms", elapsed_ms);
    }

    #[tokio::test]
    #[serial]
    async fn test_tcp_server_bad_request_error_custom() {
        let start = Instant::now();

        while_server_running_custom_bad_request(async {
            let mut client = TCPClient::<_, Response>::connect(HOST, PORT).await.unwrap();
            // Send wrong message type
            assert_eq!(
                client.send(OtherRequest::ImDifferent).await.unwrap(),
                Some(Response::World)
            );
        })
        .await;

        let elapsed_ms = start.elapsed().as_millis();
        // let is be off by 1ms to make the test robust
        assert!(elapsed_ms >= 15, "Elapsed time: {} ms", elapsed_ms);
    }

    #[tokio::test]
    #[serial]
    #[should_panic]
    async fn test_tcp_server_close_connection() {
        while_server_running(async {
            let mut client1 = TCPClient::<_, Response>::connect(HOST, PORT).await.unwrap();
            let mut client2 = TCPClient::<_, Response>::connect(HOST, PORT).await.unwrap();

            // Invoke handler route that closes connection on client1
            assert_eq!(client1.send(Request::Close).await.unwrap(), None);

            // client2 should still work as expected
            assert_eq!(
                client2.send(Request::Hello).await.unwrap().unwrap(),
                Response::World
            );

            // Invoke any other route, this will panic as the connection
            // has been closed for client1
            client1.send(Request::Hello).await.unwrap();
        })
        .await;
    }

    // Test that the inactivity timeout works as expected
    #[tokio::test]
    #[serial]
    async fn test_handle_tcp_requests_timeout() {
        let start = Instant::now();

        let handle = spawn(
            TCPServer::new(HOST, PORT, handle_requests)
                .with_inactivity_timeout(15)
                .listen(),
        );
        handle.await.unwrap();

        let elapsed_ms = start.elapsed().as_millis();
        // let is be off by 1ms to make the test robust
        assert!(
            14 <= elapsed_ms && elapsed_ms <= 16,
            "Elapsed time: {} ms",
            elapsed_ms
        );
    }
}