reach_websocket/client/

mod.rs

// SPDX-FileCopyrightText: 2024-2025 Michael Goldenberg <m@mgoldenberg.net>
// SPDX-FileCopyrightText: 2024-2025 eaon <eaon@posteo.net>
// SPDX-License-Identifier: EUPL-1.2

use futures::channel::{mpsc, oneshot};
use futures::stream::{Fuse, FusedStream};
use futures::{Sink, SinkExt, StreamExt};
use pin_project::pin_project;
use rand_core::CryptoRngCore;
use thiserror::Error;
use tokio::io::{AsyncRead, AsyncWrite};
use tokio::task::JoinHandle;
use tokio_tungstenite::{
    WebSocketStream,
    tungstenite::{Error as TungsteniteError, protocol::Message},
};

use reach_core::communication::{
    Communicable, CommunicableOwned, CommunicableTypes, Communication, MatchResponse,
};

#[cfg(feature = "local-client")]
mod local;
#[cfg(feature = "local-client")]
pub use local::*;

#[cfg(feature = "remote-client")]
mod remote;
#[cfg(feature = "remote-client")]
pub use remote::*;

mod sealed {
    use super::*;

    pub struct WebSocketClientContext<C, S, R>
    where
        C: CommunicableTypes,
        C::Variant: Responders<C> + WebSocketClientExtension,
    {
        pub request_stream: Fuse<mpsc::Receiver<WebSocketSinkItem<C>>>,
        pub web_socket: Fuse<WebSocketStream<S>>,
        pub extension: <C::Variant as WebSocketClientExtension>::Extension,
        pub responders: <C::Variant as Responders<C>>::Responders,
        pub csprng: R,
    }

    pub trait WebSocketClientHandlers<C>
    where
        C: CommunicableTypes + Send + 'static,
        C::Variant: Responders<C> + WebSocketClientExtension,
    {
        fn handle_request<S, R>(
            item: WebSocketSinkItem<C>,
            context: &mut WebSocketClientContext<C, S, R>,
        ) -> impl Future<Output = Result<(), HandleRequestError<C::Req>>> + Send
        where
            S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
            R: CryptoRngCore + Send + 'static;

        fn handle_response<S, R>(
            item: Result<Message, TungsteniteError>,
            context: &mut WebSocketClientContext<C, S, R>,
        ) -> impl Future<Output = Result<(), HandleResponseError<C::Resp>>> + Send
        where
            S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
            R: CryptoRngCore + Send + 'static;
    }
}

use sealed::*;

#[derive(Debug, Error)]
pub enum WebSocketClientError<Req, Resp> {
    #[error("request channel: {0}")]
    RequestChannel(#[from] mpsc::SendError),
    #[error("request handler: {0}")]
    RequestHandler(#[from] HandleRequestError<Req>),
    #[error("response handler: {0}")]
    ResponseHandler(#[from] HandleResponseError<Resp>),
    #[error("response channel disconnected")]
    ResponseChannelDisconnected(#[from] oneshot::Canceled),
    #[error("response decoder: {0}")]
    ResponseDecoder(#[from] reach_core::error::DecodeError),
    #[error("response error: {0}")]
    ResponseError(Resp),
}

type WebSocketClientResult<T, Req, Resp> = Result<T, WebSocketClientError<Req, Resp>>;

#[derive(Debug, Error)]
pub enum HandleRequestError<Req> {
    #[error("permissible communication lengths not set")]
    MissingPermissibleCommunicationLenghts,
    #[error("responder already exists")]
    ResponderAlreadyExists,
    #[error("tag already in use: {0:?}")]
    TagAlreadyInUse(Communication<Req>),
    #[error("response channel disconnected")]
    ResponseChannelDisconnected,
    #[error("web socket: {0}")]
    WebSocket(#[from] TungsteniteError),
}

impl<Req> HandleRequestError<Req> {
    pub fn is_fatal(&self) -> bool {
        matches!(self, Self::WebSocket(_))
    }
}

#[derive(Debug, Error)]
pub enum HandleResponseError<Resp> {
    #[error("response decoder: {0}")]
    ResponseDecoder(#[from] reach_core::error::DecodeError),
    #[error("web socket: {0}")]
    WebSocket(#[from] TungsteniteError),
    #[error("message not supported: {0:?}")]
    MessageNotSupported(Message),
    #[error("response not tagged: {0:?}")]
    ResponseNotTagged(Communication<Resp>),
    #[error("responder not found: {0:?}")]
    ResponderNotFound(Communication<Resp>),
    #[error("response channel disconnected")]
    ResponseChannelDisconnected,
}

impl<Resp> HandleResponseError<Resp> {
    pub fn is_fatal(&self) -> bool {
        matches!(self, Self::WebSocket(_))
    }
}

type WebSocketClientTask<Req, Resp> = JoinHandle<WebSocketClientResult<(), Req, Resp>>;

#[pin_project]
pub struct WebSocketClient<C>
where
    C: CommunicableTypes,
{
    #[pin]
    inner: mpsc::Sender<WebSocketSinkItem<C>>,
    task: WebSocketClientTask<C::Req, C::Resp>,
}

pub struct WebSocketClientOptions<E, S, R, const B: usize> {
    pub web_socket: WebSocketStream<S>,
    pub csprng: R,
    pub extension: E,
}

impl<C> WebSocketClient<C>
where
    C: CommunicableTypes<Req: Send + 'static, Resp: Send + 'static> + Send + 'static,
    C::Variant: Responders<C>,
    C::Variant: WebSocketClientHandlers<C> + WebSocketClientExtension,
{
    async fn send_communication_and_wait<Req>(
        &mut self,
        communication: Communication<C::Req>,
    ) -> WebSocketClientResult<Req::Resp, C::Req, C::Resp>
    where
        Req: MatchResponse<C::Req, Resp: Communicable<C::Resp>>,
    {
        let (responder, receiver) = oneshot::channel();
        let item = WebSocketSinkItem {
            request: communication,
            responder: Responder { inner: responder },
        };
        self.inner.send(item).await?;

        let response = receiver.await??;

        Req::Resp::try_from_communication(&response)
            .map_err(|_| WebSocketClientError::ResponseError(response.r#type))
    }

    pub async fn send_and_wait<Req>(
        &mut self,
        request: &Req,
    ) -> WebSocketClientResult<Req::Resp, C::Req, C::Resp>
    where
        Req: MatchResponse<C::Req, Resp: Communicable<C::Resp>> + Communicable<C::Req>,
    {
        self.send_communication_and_wait::<Req>(request.to_communication())
            .await
    }

    pub async fn send_owned_and_wait<Req>(
        &mut self,
        request: Req,
    ) -> WebSocketClientResult<Req::Resp, C::Req, C::Resp>
    where
        Req: MatchResponse<C::Req, Resp: Communicable<C::Resp>> + CommunicableOwned<C::Req>,
    {
        self.send_communication_and_wait::<Req>(request.to_communication())
            .await
    }

    async fn start<S, R>(
        mut context: WebSocketClientContext<C, S, R>,
    ) -> WebSocketClientResult<(), C::Req, C::Resp>
    where
        S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
        R: CryptoRngCore + Send + 'static,
    {
        loop {
            tokio::select! {
                item = &mut context.request_stream.next() => match item {
                    None => {
                        if context.web_socket.is_terminated() || !context.responders.has_responder()
                        {
                            break;
                        } else {
                            continue;
                        }
                    },
                    Some(request) => {
                        if let Err(e) = C::Variant::handle_request(request, &mut context).await {
                            if e.is_fatal() {
                                return Err(e.into());
                            } else {
                                // TODO: log non-fatal errors here
                                // log::error!("{e}");
                                continue;
                            }
                        }
                    }
                },
                item = &mut context.web_socket.next() => match item {
                    None => break,
                    Some(response) => {
                        if let Err(e) = C::Variant::handle_response(response, &mut context).await {
                            if e.is_fatal() {
                                return Err(e.into());
                            } else {
                                // TODO: log non-fatal errors here
                                // log::error!("{e}");
                                continue;
                            }
                        }
                    }
                }
            }
        }

        Ok(())
    }
}

type ResponseSender<Req, Resp> =
    oneshot::Sender<WebSocketClientResult<Communication<Resp>, Req, Resp>>;

pub struct Responder<C>
where
    C: CommunicableTypes,
{
    inner: ResponseSender<C::Req, C::Resp>,
}

pub struct WebSocketSinkItem<C: CommunicableTypes> {
    request: Communication<C::Req>,
    responder: Responder<C>,
}

impl<C> Sink<WebSocketSinkItem<C>> for WebSocketClient<C>
where
    C: CommunicableTypes,
{
    type Error = WebSocketClientError<C::Req, C::Resp>;

    fn poll_ready(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Result<(), Self::Error>> {
        self.project().inner.poll_ready(cx).map_err(Into::into)
    }

    fn start_send(
        self: std::pin::Pin<&mut Self>,
        item: WebSocketSinkItem<C>,
    ) -> WebSocketClientResult<(), C::Req, C::Resp> {
        self.project().inner.start_send(item).map_err(Into::into)
    }

    fn poll_flush(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Result<(), Self::Error>> {
        self.project().inner.poll_flush(cx).map_err(Into::into)
    }

    fn poll_close(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Result<(), Self::Error>> {
        self.project().inner.poll_close(cx).map_err(Into::into)
    }
}

pub trait Responders<C> {
    type Responders: ContainsResponders;
}

pub trait ContainsResponders {
    fn has_responder(&self) -> bool;
}

pub trait WebSocketClientExtension {
    type Extension;
}