jetstream_websocket 9.5.0

Jetstream is a RPC framework for Rust, based on the 9P protocol and QUIC.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106

#![doc(
    html_logo_url = "https://raw.githubusercontent.com/sevki/jetstream/main/logo/JetStream.png"
)]
#![doc(
    html_favicon_url = "https://raw.githubusercontent.com/sevki/jetstream/main/logo/JetStream.png"
)]
#![cfg_attr(docsrs, feature(doc_cfg))]

pub extern crate tokio_tungstenite;
use futures::{
    stream::{SplitSink, SplitStream},
    Sink, SinkExt, Stream, StreamExt,
};
use jetstream_rpc::{Frame, Protocol};
use jetstream_wireformat::wire_format_extensions::ConvertWireFormat;
use std::{
    io,
    marker::PhantomData,
    pin::Pin,
    task::{Context, Poll},
};
use tokio::net::TcpStream;
use tokio_tungstenite::WebSocketStream;
use tokio_tungstenite::{tungstenite::Message, MaybeTlsStream};

pub struct WebSocketTransport<P: Protocol>(
    SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message>,
    SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>,
    PhantomData<P>,
);

impl<P: Protocol + Unpin> From<WebSocketStream<MaybeTlsStream<TcpStream>>>
    for WebSocketTransport<P>
{
    fn from(ws: WebSocketStream<MaybeTlsStream<TcpStream>>) -> Self {
        let (read, write) = ws.split();
        WebSocketTransport(read, write, PhantomData)
    }
}
impl<P: Protocol + Unpin> Sink<jetstream_rpc::Frame<P::Request>>
    for WebSocketTransport<P>
{
    type Error = std::io::Error;

    fn poll_ready(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<(), Self::Error>> {
        self.get_mut()
            .0
            .poll_ready_unpin(cx)
            .map_err(io::Error::other)
    }

    fn start_send(
        self: Pin<&mut Self>,
        item: jetstream_rpc::Frame<P::Request>,
    ) -> Result<(), Self::Error> {
        self.get_mut()
            .0
            .start_send_unpin(Message::Binary(item.to_bytes()))
            .map_err(io::Error::other)
    }

    fn poll_flush(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<(), Self::Error>> {
        self.get_mut()
            .0
            .poll_flush_unpin(cx)
            .map_err(io::Error::other)
    }

    fn poll_close(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<(), Self::Error>> {
        self.get_mut()
            .0
            .poll_close_unpin(cx)
            .map_err(io::Error::other)
    }
}

impl<P: Protocol> Stream for WebSocketTransport<P>
where
    Self: Unpin,
{
    type Item = Result<jetstream_rpc::Frame<P::Response>, io::Error>;

    fn poll_next(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Option<Self::Item>> {
        self.get_mut().1.poll_next_unpin(cx).map(|opt| {
            opt.map(|res| match res {
                Ok(msg) => {
                    let data = msg.into_data();
                    Frame::<P::Response>::from_bytes(&data)
                }
                Err(e) => Err(io::Error::other(e)),
            })
        })
    }
}