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))]
8use futures::{Sink, Stream};
9use jetstream_rpc::{Frame, Framer, Protocol};
10use jetstream_wireformat::{
11    wire_format_extensions::ConvertWireFormat, WireFormat,
12};
13use std::{
14    io,
15    marker::PhantomData,
16    pin::Pin,
17    task::{Context, Poll},
18};
19use tungstenite::{Message, WebSocket};
20
21pub struct WebSocketTransport<P: Protocol>(
22    WebSocket<tungstenite::stream::MaybeTlsStream<std::net::TcpStream>>,
23    PhantomData<P>,
24);
25
26impl<P: Protocol>
27    From<WebSocket<tungstenite::stream::MaybeTlsStream<std::net::TcpStream>>>
28    for WebSocketTransport<P>
29{
30    fn from(
31        value: WebSocket<
32            tungstenite::stream::MaybeTlsStream<std::net::TcpStream>,
33        >,
34    ) -> Self {
35        Self(value, PhantomData)
36    }
37}
38
39impl<P: Protocol> Sink<jetstream_rpc::Frame<P::Request>>
40    for WebSocketTransport<P>
41where
42    Self: Unpin,
43{
44    type Error = io::Error;
45
46    fn poll_ready(
47        self: Pin<&mut Self>,
48        _cx: &mut Context<'_>,
49    ) -> Poll<Result<(), Self::Error>> {
50        match self.0.can_write() {
51            true => Poll::Ready(Ok(())),
52            false => Poll::Pending,
53        }
54    }
55
56    fn start_send(
57        self: Pin<&mut Self>,
58        item: jetstream_rpc::Frame<P::Request>,
59    ) -> Result<(), Self::Error> {
60        self.get_mut()
61            .0
62            .send(WebsocketFrame(item).into())
63            .map_err(io::Error::other)?;
64        Ok(())
65    }
66
67    fn poll_flush(
68        self: Pin<&mut Self>,
69        _cx: &mut Context<'_>,
70    ) -> Poll<Result<(), Self::Error>> {
71        Poll::Ready(self.get_mut().0.flush().map_err(io::Error::other))
72    }
73
74    fn poll_close(
75        self: Pin<&mut Self>,
76        _cx: &mut Context<'_>,
77    ) -> Poll<Result<(), Self::Error>> {
78        Poll::Ready(self.get_mut().0.close(None).map_err(io::Error::other))
79    }
80}
81
82impl<P: Protocol> Stream for WebSocketTransport<P>
83where
84    Self: Unpin,
85{
86    type Item = Result<jetstream_rpc::Frame<P::Response>, io::Error>;
87
88    fn poll_next(
89        self: Pin<&mut Self>,
90        _cx: &mut Context<'_>,
91    ) -> Poll<Option<Self::Item>> {
92        match self.get_mut().0.read() {
93            Ok(Message::Binary(bytes)) => {
94                let mut reader = io::Cursor::new(bytes);
95                let frame = Frame::<P::Response>::decode(&mut reader).unwrap();
96
97                Poll::Ready(Some(Ok(frame)))
98            }
99            Err(e) => {
100                eprintln!("Error reading from websocket: {:?}", e);
101                Poll::Ready(None)
102            }
103            _ => {
104                eprintln!("Unexpected message type from websocket");
105                Poll::Ready(None)
106            }
107        }
108    }
109}
110
111pub struct WebsocketFrame<F: Framer>(Frame<F>);
112
113impl<F: Framer> From<WebsocketFrame<F>> for tungstenite::protocol::Message {
114    fn from(value: WebsocketFrame<F>) -> Self {
115        Message::Binary(value.0.to_bytes())
116    }
117}
jetstream_websocket/lib.rs

jetstream_websocket/
lib.rs
