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
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155

// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    stream::streaming_client::StreamingClientReceiver, StreamError, StreamResult, USER_AGENT,
};
use futures::{
    stream::{SplitSink, SplitStream},
    SinkExt, StreamExt,
};
use std::{
    str::FromStr,
    sync::atomic::{AtomicU64, Ordering},
};

use diem_json_rpc_types::{
    stream::{
        request::{StreamJsonRpcRequest, StreamMethodRequest},
        response::StreamJsonRpcResponse,
    },
    Id,
};
use reqwest::Method;
use tokio::{sync::mpsc, task::JoinHandle};
use tokio_tungstenite::{
    connect_async_with_config,
    tungstenite::{handshake::client::Request, protocol::WebSocketConfig, Message},
    MaybeTlsStream, WebSocketStream,
};

pub struct WebsocketTransport {
    stream: Option<SplitStream<WebSocketStream<MaybeTlsStream<tokio::net::TcpStream>>>>,
    sink: SplitSink<WebSocketStream<MaybeTlsStream<tokio::net::TcpStream>>, Message>,
    next_id: AtomicU64,
    channel_task: Option<JoinHandle<()>>,
}

impl WebsocketTransport {
    pub async fn new<T: Into<String>>(
        url: T,
        websocket_config: Option<WebSocketConfig>,
    ) -> StreamResult<Self> {
        let request = Request::builder()
            .header(reqwest::header::USER_AGENT, USER_AGENT)
            .header(reqwest::header::CONTENT_LENGTH, 1_000)
            .uri(url.into())
            .method(Method::GET)
            .body(())
            .map_err(StreamError::from_http_error)?;

        let (stream, _) = connect_async_with_config(request, websocket_config)
            .await
            .map_err(StreamError::from_tungstenite_error)?;

        let (sink, stream) = stream.split();

        Ok(Self {
            stream: Some(stream),
            sink,
            next_id: AtomicU64::new(0),
            channel_task: None,
        })
    }

    pub async fn send(&mut self, request_json: String) -> StreamResult<()> {
        self.sink
            .send(Message::text(request_json))
            .await
            .map_err(StreamError::encode)?;
        Ok(())
    }

    pub async fn send_method_request(
        &mut self,
        request: StreamMethodRequest,
        id: Option<Id>,
    ) -> StreamResult<Id> {
        let id = id.unwrap_or_else(|| self.get_next_id());
        let request = StreamJsonRpcRequest::new(request, id.clone());
        self.send_request(&request).await
    }

    pub async fn send_request(&mut self, request: &StreamJsonRpcRequest) -> StreamResult<Id> {
        let json = serde_json::to_string(&request)?;
        self.send(json).await?;
        Ok(request.id.clone())
    }

    pub fn get_next_id(&self) -> Id {
        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
        Id::Number(id)
    }

    /// This consumes the websocket stream (but not the sink)
    /// Possibilities:
    /// 1. `stream.next().await` returns `None`: this happens if a stream is closed
    /// 2. `stream.next().await` returns `Some(Err)`: this is generally a transport error, and the
    ///     stream can be considered closed
    /// 3. `stream.next().await` returns `Some(Ok(Message))`: this may or may not be valid JSON from
    ///     the server, so we need to do a bit of validation:
    ///     1. `msg.is_text()` must be `true`: we ignore `Ping`/`Pong`, `Binary`, etc message types
    ///     2. `msg.to_text()` must be `Ok(String)`: Ensures the message is valid UTF8
    ///     3. `StreamJsonRpcResponse::from_str(msg)` must be `Ok(StreamJsonRpcResponse)`: this is
    ///         serde deserialization, which validates that the JSON is well formed, and matches the
    ///         `StreamJsonRpcResponse` struct
    ///
    pub fn get_stream(mut self) -> (StreamingClientReceiver, Self) {
        let (sender, receiver) = mpsc::channel(100);

        let mut stream = self
            .stream
            .expect("Stream is `None`: it has already been consumed");
        self.stream = None;

        self.channel_task = Some(tokio::task::spawn(async move {
            loop {
                match stream.next().await {
                    None => {
                        sender
                            .send(Err(StreamError::connection_closed(None::<StreamError>)))
                            .await
                            .ok();
                    }
                    Some(msg) => match msg {
                        Ok(msg) => {
                            if msg.is_text() {
                                let msg = match msg
                                    .to_text()
                                    .map_err(StreamError::from_tungstenite_error)
                                {
                                    Ok(msg) => msg,
                                    Err(e) => {
                                        let _ = sender.send(Err(e)).await;
                                        continue;
                                    }
                                };
                                match StreamJsonRpcResponse::from_str(msg) {
                                    Ok(msg) => sender.send(Ok(msg)).await.ok(),
                                    Err(e) => sender.send(Err(StreamError::from(e))).await.ok(),
                                };
                            }
                        }
                        Err(e) => {
                            let _ = sender
                                .send(Err(StreamError::from_tungstenite_error(e)))
                                .await;
                        }
                    },
                };
            }
        }));

        (receiver, self)
    }
}