tcp_message_io 1.0.4

#![doc = include_str!("../docs/raw.md")]

use std::future::Future;
#[cfg(feature = "zstd")]
use std::io::Cursor;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;

use anyhow::Result;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::broadcast;
use tokio::sync::broadcast::error::RecvError;
use tokio::sync::broadcast::Sender;
use tokio::time::sleep;
use tokio::{select, spawn};
use tracing::{debug, error, info};
#[cfg(feature = "zstd")]
use zstd::{decode_all, encode_all};

/// Client to make TCP requests in the form of messages.
///
/// Messages can be any sequence of bytes, such as a serialized
/// value. This library uses a 8-byte header to encode the size
/// of the data sent.
pub struct RawTCPClient {
    tcp_stream: TcpStream,
}

impl RawTCPClient {
    /// 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 {
            tcp_stream: TcpStream::connect((host, port)).await?,
        })
    }

    /// Send a message and wait for a response, returning it verbatim.
    ///
    /// Returns an error if there were errors while sending
    /// or receiving data, if the connection was closed by the server
    /// or if there were problems compressing/decompressing the data.
    ///
    /// Server and client must both use compression or not use it.
    pub async fn send(&mut self, message: &[u8]) -> Result<Vec<u8>> {
        write(&mut self.tcp_stream, message).await?;
        read(&mut self.tcp_stream).await
    }
}

/// Response returned by message handlers to indicate which
/// action to take.
#[derive(Debug, Clone)]
pub enum RawTCPResponse {
    /// Causes the server to send the given response message back to the client.
    Message(Vec<u8>),
    /// 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,
}

// This is kept separate from the above, as it's used internally
enum RequestAction {
    CloseConnection,
    StopServer,
    None,
}

/// 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 [`RawTCPResponse`] 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 RawTCPServer<H, F>
where
    H: Fn(Vec<u8>) -> F + Send + Sync + 'static,
    F: Future<Output = Result<RawTCPResponse>> + Send + 'static,
{
    host: String,
    port: u16,
    handler: H,
    inactivity_timeout_ms: AtomicU64,
}

impl<H, F> RawTCPServer<H, F>
where
    H: Fn(Vec<u8>) -> F + Send + Sync + 'static,
    F: Future<Output = Result<RawTCPResponse>> + 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,
            inactivity_timeout_ms: AtomicU64::new(0),
        })
    }

    /// 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>) {
        // This channel is to communicate the reception of a message,
        // so that we know how long it has been since the last message (activity).
        // It also allows the handler to signal that it wants to shut the server down,
        // likely because it it was requested by the client.
        let (signal_sender, mut signal_receiver) = broadcast::channel(1);
        spawn(self.clone().accept_connections(signal_sender));

        let inactivity_timeout_ms = self.inactivity_timeout_ms.load(Ordering::Relaxed);
        loop {
            let must_exit = if inactivity_timeout_ms == 0 {
                signal_receiver
                    .recv()
                    .await
                    .expect("Unable to read message from channel")
            } else {
                select! {
                    must_exit = signal_receiver.recv() => match must_exit {
                        Ok(must_exit) => must_exit,
                        // This happens when the receiver points to an old value,
                        // and it must be ignored as the next iteration of loop
                        // will fetch the correct value.
                        Err(RecvError::Lagged(_)) => false,
                        // Other errors is mostly channel close
                        Err(_) => true,
                    },
                    _ = sleep(Duration::from_millis(inactivity_timeout_ms)) => {
                        info!(
                            port = self.port,
                            "Not receiving requests since more that {} seconds, stopping request handling...",
                            inactivity_timeout_ms / 1_000
                        );
                        true
                    },
                }
            };

            if must_exit {
                // Stop accepting connections and handling requests
                break;
            }
        }
    }

    async fn accept_connections(self: Arc<Self>, signal_sender: Sender<bool>) {
        let listener = match TcpListener::bind((self.host.as_str(), self.port)).await {
            Ok(listener) => listener,
            Err(err) => {
                error!(
                    port = self.port,
                    error = ?err,
                    "Unable to bind socket listener: {:?}",
                    err,
                );
                panic!("Unable to bind socket listener: {:?}", err);
            }
        };

        let mut signal_receiver = signal_sender.subscribe();

        // Loop to keep accepting new connections
        loop {
            select! {
                result = listener.accept() => match result {
                    Ok((tcp_stream, _)) => {
                        debug!(port = self.port, "Client successfully connected");
                        spawn(
                            self.clone()
                                .handle_connection(tcp_stream, signal_sender.clone())
                        );
                    }
                    // We are resilient to any error that happens at connection time
                    Err(err) => error!(
                        port = self.port,
                        error = ?err,
                        "Error while accepting connection: {:?}",
                        err,
                    ),
                },
                must_exit = signal_receiver.recv() => match must_exit {
                    Ok(must_exit) => if must_exit { break },
                    // This happens when the receiver points to an old value,
                    // and it must be ignored as the next iteration of loop
                    // will fetch the correct value.
                    Err(RecvError::Lagged(_)) => {},
                    // Other errors is mostly channel close
                    Err(_) => break,
                },
            }
        }
    }

    async fn handle_connection(
        self: Arc<Self>,
        mut tcp_stream: TcpStream,
        signal_sender: Sender<bool>,
    ) {
        // Keep accepting new messages from the same client
        loop {
            let action = self.clone().handle_request(&mut tcp_stream).await;
            signal_sender
                .send(matches!(action, RequestAction::StopServer))
                .expect("Unable to send message to channel");
            match action {
                RequestAction::CloseConnection | RequestAction::StopServer => {
                    break;
                }
                RequestAction::None => {}
            }
        }
    }

    async fn handle_request(self: Arc<Self>, tcp_stream: &mut TcpStream) -> RequestAction {
        let request = match read(tcp_stream).await {
            Ok(msg) => msg,
            Err(err) => {
                // This will happen when the client closes the connection
                // or the connection breaks.
                debug!(
                    port = self.port,
                    error = ?err,
                    "Connection closed by client",
                );
                return RequestAction::CloseConnection;
            }
        };

        match (self.handler)(request).await {
            Ok(action) => match action {
                RawTCPResponse::Message(message) => {
                    if !message.is_empty() {
                        if let Err(error) = write(tcp_stream, &message).await {
                            error!(
                                port = self.port,
                                ?error,
                                ?message,
                                "Connection was unexpectedly closed by the client while handling \
                                the request. The handler is unable to return the response \
                                and will be forced to discard it: {:?}",
                                error,
                            );
                            return RequestAction::CloseConnection;
                        }
                    }
                    RequestAction::None
                }
                // In both of the following cases, we answer with empty message,
                // so that we do not simply suddenly close the connection,
                // causing an EOF error on the client.
                //
                // Any error at this point, will be ignored
                // and we simply close the connection.
                RawTCPResponse::CloseConnection => {
                    write(tcp_stream, &[]).await.expect(
                        "Unable to write empty message to socket before closing connection.",
                    );
                    RequestAction::CloseConnection
                }
                RawTCPResponse::StopServer => {
                    write(tcp_stream, &[]).await.expect(
                        "Unable to write empty message to socket before stopping TCP server.",
                    );
                    RequestAction::StopServer
                }
            },
            Err(error) => {
                // We log the error, then keep serving more requests.
                //
                // The handler can decide to send errors to the clients
                // by adding a layer on top of this implementation,
                // to shut the server down using TCPResponse::StopServer
                // or to close the connection using TCPResponse::CloseConnection.
                //
                // That is left specifically to the users of this library
                // as we do not want to impose any behavior they might now want.
                error!(
                    port = self.port,
                    ?error,
                    "Error while handling request: {:?}",
                    error,
                );
                // We need to write something to not leave the client hanging
                write(tcp_stream, &[])
                    .await
                    .expect("Unable to write empty message to socket after handler error.");
                RequestAction::None
            }
        }
    }
}

#[cfg(not(feature = "zstd"))]
async fn read(tcp_stream: &mut TcpStream) -> Result<Vec<u8>> {
    raw_read(tcp_stream).await
}

#[cfg(feature = "zstd")]
async fn read(tcp_stream: &mut TcpStream) -> Result<Vec<u8>> {
    let compressed_response = raw_read(tcp_stream).await?;
    Ok(decode_all(Cursor::new(compressed_response))?)
}

#[cfg(not(feature = "zstd"))]
async fn write(tcp_stream: &mut TcpStream, message: &[u8]) -> Result<()> {
    raw_write(tcp_stream, message).await
}

#[cfg(feature = "zstd")]
async fn write(tcp_stream: &mut TcpStream, message: &[u8]) -> Result<()> {
    let compressed_message = encode_all(message, 0)?;
    raw_write(tcp_stream, &compressed_message).await
}

async fn raw_read(tcp_stream: &mut TcpStream) -> Result<Vec<u8>> {
    tcp_stream.readable().await?;

    let len = tcp_stream.read_u64_le().await?;

    let mut buffer = Vec::with_capacity(len as usize);
    tcp_stream.take(len).read_to_end(&mut buffer).await?;

    Ok(buffer)
}

async fn raw_write(tcp_stream: &mut TcpStream, message: &[u8]) -> Result<()> {
    tcp_stream.writable().await?;

    tcp_stream.write_u64_le(message.len() as u64).await?;
    tcp_stream.write_all(message).await?;
    tcp_stream.flush().await?;

    Ok(())
}

#[cfg(test)]
mod tests {
    use anyhow::bail;
    use serial_test::serial;
    use tokio::time::Instant;

    use super::*;

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

    async fn while_server_running(fut: impl Future) {
        let handle = spawn(RawTCPServer::new(HOST, PORT, handle_requests).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: Vec<u8>) -> Result<RawTCPResponse> {
        Ok(match &req[..] {
            &[1] => RawTCPResponse::Message(vec![10]),
            &[2] => RawTCPResponse::Message(vec![200]),
            &[3] => bail!("Test error"),
            &[4] => RawTCPResponse::CloseConnection,
            _ => RawTCPResponse::StopServer,
        })
    }

    // We send a couple of messages, to the server,
    // check that responses are correct, then stop it.
    #[tokio::test]
    #[serial]
    async fn test_raw_tcp_server() {
        while_server_running(async {
            let mut client = RawTCPClient::connect(HOST, PORT).await.unwrap();
            assert_eq!(client.send(&[1]).await.unwrap(), vec![10]);
            assert_eq!(client.send(&[2]).await.unwrap(), vec![200]);
            // Stop server
            assert_eq!(client.send(&[]).await.unwrap(), vec![]);
        })
        .await;
    }

    #[tokio::test]
    #[serial]
    async fn test_raw_tcp_server_handler_error() {
        while_server_running(async {
            let mut client = RawTCPClient::connect(HOST, PORT).await.unwrap();
            // Invoke handler route that causes error
            assert_eq!(client.send(&[3]).await.unwrap(), vec![]);
            // Stop server
            client.send(&[]).await.unwrap();
        })
        .await;
    }

    #[tokio::test]
    #[serial]
    #[should_panic]
    async fn test_raw_tcp_server_close_connection() {
        while_server_running(async {
            let mut client1 = RawTCPClient::connect(HOST, PORT).await.unwrap();
            let mut client2 = RawTCPClient::connect(HOST, PORT).await.unwrap();

            // Invoke handler route that closes connection on client1
            assert_eq!(client1.send(&[4]).await.unwrap(), vec![]);

            // client2 should still work as expected
            assert_eq!(client2.send(&[1]).await.unwrap(), vec![10]);

            // Invoke any other route, this will panic as the connection
            // has been closed for client1
            client1.send(&[1]).await.unwrap();
        })
        .await;
    }

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

        let handle = spawn(
            RawTCPServer::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);
    }
}