jetstream_websocket/
lib.rs

1#![doc(
2    html_logo_url = "https://raw.githubusercontent.com/sevki/jetstream/main/logo/JetStream.png"
3)]
4#![doc(
5    html_favicon_url = "https://raw.githubusercontent.com/sevki/jetstream/main/logo/JetStream.png"
6)]
7#![cfg_attr(docsrs, feature(doc_cfg))]
8
9pub extern crate tokio_tungstenite;
10use futures::{
11    stream::{SplitSink, SplitStream},
12    Sink, SinkExt, Stream, StreamExt,
13};
14use jetstream_rpc::{Error, Frame, Protocol};
15use jetstream_wireformat::wire_format_extensions::ConvertWireFormat;
16use std::{
17    io,
18    marker::PhantomData,
19    pin::Pin,
20    task::{Context, Poll},
21};
22use tokio::net::TcpStream;
23use tokio_tungstenite::WebSocketStream;
24use tokio_tungstenite::{tungstenite::Message, MaybeTlsStream};
25
26pub struct WebSocketTransport<P: Protocol>(
27    SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message>,
28    SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>,
29    PhantomData<P>,
30);
31
32impl<P: Protocol + Unpin> From<WebSocketStream<MaybeTlsStream<TcpStream>>>
33    for WebSocketTransport<P>
34{
35    fn from(ws: WebSocketStream<MaybeTlsStream<TcpStream>>) -> Self {
36        let (read, write) = ws.split();
37        WebSocketTransport(read, write, PhantomData)
38    }
39}
40impl<P: Protocol + Unpin> Sink<jetstream_rpc::Frame<P::Request>>
41    for WebSocketTransport<P>
42{
43    type Error = Error;
44
45    fn poll_ready(
46        self: Pin<&mut Self>,
47        cx: &mut Context<'_>,
48    ) -> Poll<Result<(), Self::Error>> {
49        self.get_mut().0.poll_ready_unpin(cx).map_err(|e| {
50            let err = io::Error::other(e);
51            Error::from(err)
52        })
53    }
54
55    fn start_send(
56        self: Pin<&mut Self>,
57        item: jetstream_rpc::Frame<P::Request>,
58    ) -> Result<(), Self::Error> {
59        self.get_mut()
60            .0
61            .start_send_unpin(Message::Binary(item.to_bytes()))
62            .map_err(|e| Error::from(io::Error::other(e)))
63    }
64
65    fn poll_flush(
66        self: Pin<&mut Self>,
67        cx: &mut Context<'_>,
68    ) -> Poll<Result<(), Self::Error>> {
69        self.get_mut()
70            .0
71            .poll_flush_unpin(cx)
72            .map_err(|e| Error::from(io::Error::other(e)))
73    }
74
75    fn poll_close(
76        self: Pin<&mut Self>,
77        cx: &mut Context<'_>,
78    ) -> Poll<Result<(), Self::Error>> {
79        self.get_mut()
80            .0
81            .poll_close_unpin(cx)
82            .map_err(|e| Error::from(io::Error::other(e)))
83    }
84}
85
86impl<P: Protocol> Stream for WebSocketTransport<P>
87where
88    Self: Unpin,
89{
90    type Item = Result<jetstream_rpc::Frame<P::Response>, Error>;
91
92    fn poll_next(
93        self: Pin<&mut Self>,
94        cx: &mut Context<'_>,
95    ) -> Poll<Option<Self::Item>> {
96        self.get_mut().1.poll_next_unpin(cx).map(|opt| {
97            opt.map(|res| match res {
98                Ok(msg) => {
99                    let data = msg.into_data();
100                    Frame::<P::Response>::from_bytes(&data).map_err(Error::from)
101                }
102                Err(e) => Err(Error::from(io::Error::other(e))),
103            })
104        })
105    }
106}
jetstream_websocket/lib.rs

jetstream_websocket/
lib.rs
