busylib 0.0.10

BUSY Bar Rust HTTP client
Documentation
use futures_util::{SinkExt as _, StreamExt as _};
use http::Request;
use tokio::io::{AsyncRead, AsyncWrite};
use tokio_tungstenite::WebSocketStream;
use tokio_tungstenite::tungstenite::client::IntoClientRequest as _;
use tokio_tungstenite::tungstenite::{Error as TungsteniteError, Message};

use super::{BoxFuture, WsConnection, WsMessage, WsTransport, WsTransportError, WsTransportResult};

#[derive(Debug, Clone, Default)]
pub struct TungsteniteWsTransport;

impl TungsteniteWsTransport {
    pub fn new() -> Self {
        Self
    }
}

impl WsTransport for TungsteniteWsTransport {
    fn connect(&self, request: Request<()>) -> BoxFuture<'_, WsTransportResult> {
        Box::pin(async move {
            let handshake = handshake_request(request)?;

            let (stream, _response) = tokio_tungstenite::connect_async(handshake)
                .await
                .map_err(connect_error)?;

            Ok(Box::new(TungsteniteConnection { stream }) as Box<dyn WsConnection>)
        })
    }
}

fn handshake_request(request: Request<()>) -> Result<Request<()>, WsTransportError> {
    let (parts, ()) = request.into_parts();

    let mut handshake = parts
        .uri
        .into_client_request()
        .map_err(WsTransportError::new)?;

    handshake.headers_mut().extend(parts.headers);

    Ok(handshake)
}

fn connect_error(error: TungsteniteError) -> WsTransportError {
    match error {
        TungsteniteError::Http(response) => {
            let status = response.status();
            WsTransportError::rejected(
                status,
                format!("the device answered the upgrade with {status}"),
            )
        }
        other => WsTransportError::new(other),
    }
}

struct TungsteniteConnection<S> {
    stream: WebSocketStream<S>,
}

impl<S> WsConnection for TungsteniteConnection<S>
where
    S: AsyncRead + AsyncWrite + Unpin + Send,
{
    fn send(&mut self, message: WsMessage) -> BoxFuture<'_, Result<(), WsTransportError>> {
        Box::pin(async move {
            let message = match message {
                WsMessage::Text(text) => Message::text(text),
                WsMessage::Binary(data) => Message::binary(data),
            };

            self.stream
                .send(message)
                .await
                .map_err(WsTransportError::new)
        })
    }

    fn recv(&mut self) -> BoxFuture<'_, Option<Result<WsMessage, WsTransportError>>> {
        Box::pin(async move {
            loop {
                return match self.stream.next().await? {
                    Ok(Message::Text(text)) => Some(Ok(WsMessage::Text(text.as_str().to_owned()))),
                    Ok(Message::Binary(data)) => Some(Ok(WsMessage::Binary(data))),
                    Ok(Message::Ping(_) | Message::Pong(_) | Message::Frame(_)) => continue,
                    Ok(Message::Close(_)) => None,
                    Err(TungsteniteError::ConnectionClosed) => None,
                    Err(error) => Some(Err(WsTransportError::new(error))),
                };
            }
        })
    }

    fn close(&mut self) -> BoxFuture<'_, Result<(), WsTransportError>> {
        Box::pin(async move {
            match self.stream.close(None).await {
                Ok(()) | Err(TungsteniteError::ConnectionClosed) => Ok(()),
                Err(error) => Err(WsTransportError::new(error)),
            }
        })
    }
}