wrpc_websockets/
lib.rs

1//! wRPC WebSocket transport using [`tokio_websockets`]
2
3use core::pin::Pin;
4use core::task::{Context, Poll, ready};
5
6use anyhow::Context as _;
7use bytes::Bytes;
8use futures::stream::{SplitSink, SplitStream};
9use futures::{Sink, Stream, StreamExt as _};
10use tokio::io::{AsyncRead, AsyncWrite};
11use tokio_util::io::{SinkWriter, StreamReader};
12use tokio_websockets::resolver::{Gai, Resolver};
13use tokio_websockets::{Message, WebSocketStream};
14use tracing::instrument;
15
16use wrpc_transport::Invoke;
17use wrpc_transport::frame::invoke;
18
19pub use tokio_websockets;
20#[doc(no_inline)]
21pub use tokio_websockets::{ClientBuilder, ServerBuilder};
22
23/// Outgoing half of a split [`WebSocketStream`].
24pub struct Outgoing<T> {
25    sink: SplitSink<WebSocketStream<T>, Message>,
26    eof_sent: bool,
27}
28
29impl<T: AsyncRead + AsyncWrite + Unpin> Sink<&[u8]> for Outgoing<T> {
30    type Error = std::io::Error;
31
32    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
33        Pin::new(&mut self.sink)
34            .poll_ready(cx)
35            .map_err(std::io::Error::other)
36    }
37
38    fn start_send(mut self: Pin<&mut Self>, buf: &[u8]) -> Result<(), Self::Error> {
39        let buf = Message::binary(Bytes::copy_from_slice(buf));
40        Pin::new(&mut self.sink)
41            .start_send(buf)
42            .map_err(std::io::Error::other)
43    }
44
45    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
46        Pin::new(&mut self.sink)
47            .poll_flush(cx)
48            .map_err(std::io::Error::other)
49    }
50
51    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
52        let this = self.as_mut().get_mut();
53        if !this.eof_sent {
54            ready!(Pin::new(&mut this.sink).poll_ready(cx)).map_err(std::io::Error::other)?;
55            Pin::new(&mut this.sink)
56                .start_send(Message::text(Bytes::new()))
57                .map_err(std::io::Error::other)?;
58            this.eof_sent = true;
59        }
60        Pin::new(&mut this.sink)
61            .poll_flush(cx)
62            .map_err(std::io::Error::other)
63    }
64}
65
66/// Incoming half of a split [`WebSocketStream`].
67#[repr(transparent)]
68pub struct Incoming<T>(SplitStream<WebSocketStream<T>>);
69
70impl<T: AsyncRead + AsyncWrite + Unpin> Stream for Incoming<T> {
71    type Item = Result<Bytes, std::io::Error>;
72
73    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
74        loop {
75            match ready!(Pin::new(&mut self.0).poll_next(cx)) {
76                Some(Ok(msg)) if msg.is_text() => {
77                    if msg.as_payload().is_empty() {
78                        return Poll::Ready(None);
79                    }
80                    return Poll::Ready(Some(Err(std::io::Error::new(
81                        std::io::ErrorKind::InvalidData,
82                        "unexpected non-empty WebSocket text frame",
83                    ))));
84                }
85                Some(Ok(msg)) if msg.is_close() => {
86                    return Poll::Ready(Some(Err(std::io::Error::new(
87                        std::io::ErrorKind::UnexpectedEof,
88                        "WebSocket connection closed before EOF sentinel",
89                    ))));
90                }
91                Some(Ok(msg)) if msg.is_ping() || msg.is_pong() => continue,
92                Some(Ok(msg)) => return Poll::Ready(Some(Ok(Bytes::from(msg.into_payload())))),
93                Some(Err(err)) => return Poll::Ready(Some(Err(std::io::Error::other(err)))),
94                None => return Poll::Ready(None),
95            }
96        }
97    }
98}
99
100/// Splits a [`WebSocketStream`] into wRPC byte-stream halves: an [`AsyncWrite`]
101/// outgoing half and an [`AsyncRead`] incoming half, suitable for use with
102/// [`wrpc_transport::frame`].
103pub fn split<T>(
104    ws: WebSocketStream<T>,
105) -> (SinkWriter<Outgoing<T>>, StreamReader<Incoming<T>, Bytes>)
106where
107    T: AsyncRead + AsyncWrite + Unpin,
108{
109    let (tx, rx) = ws.split();
110    (
111        SinkWriter::new(Outgoing {
112            sink: tx,
113            eof_sent: false,
114        }),
115        StreamReader::new(Incoming(rx)),
116    )
117}
118
119/// [Invoke] implementation in terms of a [`ClientBuilder`].
120#[repr(transparent)]
121pub struct Client<'a, R: Resolver = Gai>(ClientBuilder<'a, R>);
122
123impl<'a, R: Resolver> Client<'a, R> {
124    /// Constructs a new [Client] from a preconfigured [`ClientBuilder`].
125    #[must_use]
126    pub fn from_builder(builder: ClientBuilder<'a, R>) -> Self {
127        Self(builder)
128    }
129}
130
131impl<'a, R: Resolver> From<ClientBuilder<'a, R>> for Client<'a, R> {
132    fn from(builder: ClientBuilder<'a, R>) -> Self {
133        Self::from_builder(builder)
134    }
135}
136
137impl<R: Resolver + Send + Sync> Invoke for Client<'_, R> {
138    type Context = ();
139
140    #[instrument(level = "trace", skip(self, paths, params), fields(params = format!("{params:02x?}")))]
141    async fn invoke<P>(
142        &self,
143        (): Self::Context,
144        instance: &str,
145        func: &str,
146        params: Bytes,
147        paths: impl AsRef<[P]> + Send,
148    ) -> anyhow::Result<(
149        wrpc_transport::frame::Outgoing,
150        wrpc_transport::frame::Incoming,
151    )>
152    where
153        P: AsRef<[Option<usize>]> + Send + Sync,
154    {
155        let (ws, _res) = self
156            .0
157            .connect()
158            .await
159            .context("failed to establish WebSocket connection")?;
160        let (tx, rx) = split(ws);
161        invoke(tx, rx, instance, func, params, paths).await
162    }
163}
wrpc_websockets/lib.rs

wrpc_websockets/
lib.rs
