tide-disco 0.9.7

// Copyright (c) 2022 Espresso Systems (espressosys.com)
// This file is part of the tide-disco library.

// You should have received a copy of the MIT License
// along with the tide-disco library. If not, see <https://mit-license.org/>.

//! An interface for asynchronous communication with clients, using WebSockets.

use crate::{
    http::{content::Accept, mime},
    request::{RequestParams, best_response_type},
};
use async_std::sync::Arc;
use futures::{
    FutureExt, Sink, SinkExt, Stream, StreamExt, TryFutureExt,
    future::BoxFuture,
    select, sink,
    stream::BoxStream,
    task::{Context, Poll},
};
use pin_project::pin_project;
use serde::{Serialize, de::DeserializeOwned};
use std::borrow::Cow;
use std::fmt::Display;
use std::marker::PhantomData;
use std::pin::Pin;
use tide_websockets::{
    Message, WebSocketConnection,
    tungstenite::protocol::frame::{CloseFrame, coding::CloseCode},
};
use vbs::{BinarySerializer, Serializer, version::StaticVersionType};

pub use disco_types::error::SocketError;

#[derive(Clone, Copy, Debug)]
enum MessageType {
    Binary,
    Json,
}

/// A connection facilitating bi-directional, asynchronous communication with a client.
///
/// [Connection] implements [Stream], which can be used to receive `FromClient` messages from the
/// client, and [Sink] which can be used to send `ToClient` messages to the client.
#[pin_project]
pub struct Connection<ToClient: ?Sized, FromClient, Error, VER: StaticVersionType> {
    #[pin]
    conn: WebSocketConnection,
    // [Sink] wrapper around `conn`
    sink: Pin<Box<dyn Send + Sink<Message, Error = SocketError<Error>>>>,
    accept: MessageType,
    #[allow(clippy::type_complexity)]
    _phantom: PhantomData<fn(&ToClient, &FromClient, &Error, &VER) -> ()>,
}

impl<ToClient: ?Sized, FromClient: DeserializeOwned, E, VER: StaticVersionType> Stream
    for Connection<ToClient, FromClient, E, VER>
{
    type Item = Result<FromClient, SocketError<E>>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        // Get a `Pin<&mut WebSocketConnection>` for the underlying connection, so we can use the
        // `Stream` implementation of that field.
        match self.project().conn.poll_next(cx) {
            Poll::Ready(None) => Poll::Ready(None),
            Poll::Ready(Some(Err(err))) => {
                Poll::Ready(Some(Err(SocketError::WebSockets(err.to_string()))))
            }
            Poll::Ready(Some(Ok(msg))) => Poll::Ready(Some(match msg {
                Message::Binary(bytes) => {
                    Serializer::<VER>::deserialize(&bytes).map_err(SocketError::from)
                }
                Message::Text(s) => serde_json::from_str(&s).map_err(SocketError::from),
                _ => Err(SocketError::UnsupportedMessageType),
            })),
            Poll::Pending => Poll::Pending,
        }
    }
}

impl<ToClient: Serialize + ?Sized, FromClient, E, VER: StaticVersionType> Sink<&ToClient>
    for Connection<ToClient, FromClient, E, VER>
{
    type Error = SocketError<E>;

    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.sink.as_mut().poll_ready(cx).map_err(SocketError::from)
    }

    fn start_send(mut self: Pin<&mut Self>, item: &ToClient) -> Result<(), Self::Error> {
        let msg = match self.accept {
            MessageType::Binary => Message::Binary(Serializer::<VER>::serialize(item)?),
            MessageType::Json => Message::Text(serde_json::to_string(item)?),
        };
        self.sink.as_mut().start_send(msg)
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.sink.as_mut().poll_flush(cx).map_err(SocketError::from)
    }

    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.sink.as_mut().poll_close(cx).map_err(SocketError::from)
    }
}

impl<ToClient: Serialize, FromClient, E, VER: StaticVersionType> Sink<ToClient>
    for Connection<ToClient, FromClient, E, VER>
{
    type Error = SocketError<E>;

    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        Sink::<&ToClient>::poll_ready(self, cx)
    }

    fn start_send(self: Pin<&mut Self>, item: ToClient) -> Result<(), Self::Error> {
        self.start_send(&item)
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        Sink::<&ToClient>::poll_flush(self, cx)
    }

    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        Sink::<&ToClient>::poll_close(self, cx)
    }
}

impl<ToClient: ?Sized, FromClient, E, VER: StaticVersionType>
    Connection<ToClient, FromClient, E, VER>
{
    #[allow(clippy::result_large_err)]
    fn new(accept: &Accept, conn: WebSocketConnection) -> Result<Self, SocketError<E>> {
        let ty = best_response_type(accept, &[mime::JSON, mime::BYTE_STREAM])?;
        let ty = if ty == mime::JSON {
            MessageType::Json
        } else if ty == mime::BYTE_STREAM {
            MessageType::Binary
        } else {
            unreachable!()
        };
        Ok(Self {
            sink: Self::sink(conn.clone()),
            conn,
            accept: ty,
            _phantom: Default::default(),
        })
    }

    /// Wrap a `WebSocketConnection` in a type that implements `Sink<Message>`.
    fn sink(
        conn: WebSocketConnection,
    ) -> Pin<Box<dyn Send + Sink<Message, Error = SocketError<E>>>> {
        Box::pin(sink::unfold(conn, |conn, msg| async move {
            conn.send(msg)
                .await
                .map_err(|err| SocketError::WebSockets(err.to_string()))?;
            Ok(conn)
        }))
    }
}

impl<ToClient: ?Sized, FromClient, E, VER: StaticVersionType> Clone
    for Connection<ToClient, FromClient, E, VER>
{
    fn clone(&self) -> Self {
        Self {
            sink: Self::sink(self.conn.clone()),
            conn: self.conn.clone(),
            accept: self.accept,
            _phantom: Default::default(),
        }
    }
}

pub(crate) type Handler<State, Error> = Box<
    dyn 'static
        + Send
        + Sync
        + Fn(RequestParams, WebSocketConnection, &State) -> BoxFuture<Result<(), SocketError<Error>>>,
>;

pub(crate) fn handler<State, Error, ToClient, FromClient, F, VER: StaticVersionType>(
    f: F,
) -> Handler<State, Error>
where
    F: 'static
        + Send
        + Sync
        + Fn(
            RequestParams,
            Connection<ToClient, FromClient, Error, VER>,
            &State,
        ) -> BoxFuture<Result<(), Error>>,
    State: 'static + Send + Sync,
    ToClient: 'static + Serialize + ?Sized,
    FromClient: 'static + DeserializeOwned,
    Error: 'static + Send + Display,
{
    raw_handler(move |req, conn, state| {
        f(req, conn, state)
            .map_err(SocketError::AppSpecific)
            .boxed()
    })
}

struct StreamHandler<F, VER: StaticVersionType>(F, PhantomData<VER>);

impl<F, VER: StaticVersionType> StreamHandler<F, VER> {
    fn handle<'a, State, Error, Msg>(
        &self,
        req: RequestParams,
        conn: Connection<Msg, (), Error, VER>,
        state: &'a State,
    ) -> BoxFuture<'a, Result<(), SocketError<Error>>>
    where
        F: 'static + Send + Sync + Fn(RequestParams, &State) -> BoxStream<Result<Msg, Error>>,
        State: 'static + Send + Sync,
        Msg: 'static + Serialize + Send + Sync,
        Error: 'static + Send,
        VER: 'static + Send + Sync,
    {
        let mut stream = (self.0)(req, state).fuse();
        async move {
            // Appease the borrow checker, this is a cheap clone
            let (mut send, mut recv) = (conn.clone(), conn);

            // Neither stream is documented to be cancel-safe, so we store the futures outside select
            let mut item_fut = stream.next();
            let mut client_fut = recv.next().fuse();

            loop {
                select! {
                    item = item_fut => {
                        match item {
                            Some(msg) => {
                                send.send(&msg.map_err(SocketError::AppSpecific)?).await?;
                                item_fut = stream.next();
                            }
                            None => {
                                break;
                            }
                        }
                    }
                    // We don't actually expect to receive anything from the client,
                    // it is being polled only to handle connection closure by the client
                    client_msg = client_fut => {
                        client_fut = recv.next().fuse();
                        match client_msg {
                            None => return Ok(()),
                            Some(Err(e)) => return Err(e),
                            _ => {}
                        }
                    }
                };
            }
            Ok(())
        }
        .boxed()
    }
}

pub(crate) fn stream_handler<State, Error, Msg, F, VER>(f: F) -> Handler<State, Error>
where
    F: 'static + Send + Sync + Fn(RequestParams, &State) -> BoxStream<Result<Msg, Error>>,
    State: 'static + Send + Sync,
    Msg: 'static + Serialize + Send + Sync,
    Error: 'static + Send + Display,
    VER: 'static + Send + Sync + StaticVersionType,
{
    let handler: StreamHandler<F, VER> = StreamHandler(f, Default::default());
    raw_handler(move |req, conn, state| handler.handle(req, conn, state))
}

fn raw_handler<State, Error, ToClient, FromClient, F, VER>(f: F) -> Handler<State, Error>
where
    F: 'static
        + Send
        + Sync
        + Fn(
            RequestParams,
            Connection<ToClient, FromClient, Error, VER>,
            &State,
        ) -> BoxFuture<Result<(), SocketError<Error>>>,
    State: 'static + Send + Sync,
    ToClient: 'static + Serialize + ?Sized,
    FromClient: 'static + DeserializeOwned,
    Error: 'static + Send + Display,
    VER: StaticVersionType,
{
    let close = |conn: WebSocketConnection, res: Result<(), SocketError<Error>>| async move {
        // When the handler finishes, send a close message. If there was an error, include the error
        // message.
        let msg = res.as_ref().err().map(|err| CloseFrame {
            code: CloseCode::Error,
            reason: Cow::Owned(err.to_string()),
        });
        conn.send(Message::Close(msg))
            .await
            .map_err(|err| SocketError::WebSockets(err.to_string()))?;
        res
    };
    Box::new(move |req, raw_conn, state| {
        let accept = match req.accept() {
            Ok(accept) => accept,
            Err(err) => return close(raw_conn, Err(err.into())).boxed(),
        };
        let conn = match Connection::new(&accept, raw_conn.clone()) {
            Ok(conn) => conn,
            Err(err) => return close(raw_conn, Err(err)).boxed(),
        };
        f(req, conn, state)
            .then(move |res| close(raw_conn, res))
            .boxed()
    })
}

struct MapErr<State, Error, F> {
    handler: Handler<State, Error>,
    map: Arc<F>,
}

impl<State, Error, F> MapErr<State, Error, F> {
    fn handle<'a, Error2>(
        &self,
        req: RequestParams,
        conn: WebSocketConnection,
        state: &'a State,
    ) -> BoxFuture<'a, Result<(), SocketError<Error2>>>
    where
        F: 'static + Send + Sync + Fn(Error) -> Error2,
        State: 'static + Send + Sync,
        Error: 'static,
    {
        let map = self.map.clone();
        let fut = (self.handler)(req, conn, state);
        async move { fut.await.map_err(|err| err.map_app_specific(&*map)) }.boxed()
    }
}

pub(crate) fn map_err<State, Error, Error2>(
    h: Handler<State, Error>,
    f: impl 'static + Send + Sync + Fn(Error) -> Error2,
) -> Handler<State, Error2>
where
    State: 'static + Send + Sync,
    Error: 'static,
{
    let handler = MapErr {
        handler: h,
        map: Arc::new(f),
    };
    Box::new(move |req, conn, state| handler.handle(req, conn, state))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{Api, App, Url, error::ServerError, testing::test_ws_client};
    use async_std::task::{sleep, spawn};
    use async_tungstenite::tungstenite::Message as TungsteniteMessage;
    use futures::{StreamExt, stream};
    use pin_project::pinned_drop;
    use portpicker::pick_unused_port;
    use std::{
        sync::{
            Arc,
            atomic::{AtomicBool, Ordering},
        },
        time::Duration,
    };
    use vbs::version::StaticVersion;

    type StaticVer01 = StaticVersion<0, 1>;

    #[pin_project(PinnedDrop)]
    struct DropStream<S: Stream> {
        #[pin]
        stream: S,
        dropped: Arc<AtomicBool>,
    }

    impl<S: Stream> Stream for DropStream<S> {
        type Item = S::Item;

        fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
            let stream = self.project().stream;
            stream.poll_next(cx)
        }
    }

    #[pinned_drop]
    impl<S: Stream> PinnedDrop for DropStream<S> {
        fn drop(self: Pin<&mut Self>) {
            self.dropped.store(true, Ordering::SeqCst);
        }
    }

    #[async_std::test]
    async fn test_stream_handler_client_closure() {
        // Setup: Create a simple API with a stream endpoint
        let port = pick_unused_port().expect("No ports available");

        let mut app = App::<(), ServerError>::with_state(());
        let toml_content = r#"
            [meta]
            FORMAT_VERSION = "0.1.0"

            [route.stream_test]
            PATH = ["/stream"]
            METHOD = "SOCKET"
            "#;

        let mut api =
            Api::<(), ServerError, StaticVer01>::new(toml_content.parse::<toml::Value>().unwrap())
                .unwrap();

        // Register a stream handler that sends multiple messages and indicates
        // whether it was dropped
        let dropped = Arc::new(AtomicBool::new(false));
        let _dropped = dropped.clone();
        api.stream("stream_test", move |_req, _state| {
            Box::pin(DropStream {
                stream: stream::iter(0..).map(Result::Ok),
                dropped: _dropped.clone(),
            })
        })
        .unwrap();

        app.register_module("test", api).unwrap();

        // Start the server
        spawn(async move {
            app.serve(format!("127.0.0.1:{}", port), StaticVer01::instance())
                .await
                .unwrap();
        });

        // Give the server time to start
        sleep(Duration::from_millis(500)).await;

        // Connect as a client
        let url = Url::parse(&format!("http://127.0.0.1:{}/test/stream", port)).unwrap();
        let mut ws_stream = test_ws_client(url).await;

        // Receive a few messages
        let mut received_count = 0;
        for _ in 0..5 {
            if let Some(Ok(TungsteniteMessage::Text(msg))) = ws_stream.next().await {
                let parsed: usize = serde_json::from_str(&msg).unwrap();
                assert_eq!(parsed, received_count);
                received_count += 1;
            }
        }

        // Close the client connection
        ws_stream
            .close(None)
            .await
            .expect("Failed to close connection");

        // Wait a bit to ensure the server processes the closure
        sleep(Duration::from_millis(300)).await;

        // The underlying stream should've been dropped
        assert!(dropped.load(Ordering::SeqCst));
    }
}