endpoint_validator/ws/
client.rs

1use anyhow::{Context, Result, anyhow};
2use futures::{SinkExt, StreamExt};
3use serde::Serialize;
4use tokio::net::TcpStream;
5use tokio_tungstenite::tungstenite::http::HeaderValue;
6use tokio_tungstenite::tungstenite::{Message, client::IntoClientRequest};
7use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async};
8
9pub struct WsClient {
10    stream: WebSocketStream<MaybeTlsStream<TcpStream>>,
11    seq: u32,
12}
13
14#[derive(Serialize)]
15struct WsRequest<T: Serialize> {
16    method: u32,
17    seq: u32,
18    params: T,
19}
20
21impl WsClient {
22    pub async fn new(connect_addr: &str, header: &str) -> Result<Self> {
23        let mut req = <&str as IntoClientRequest>::into_client_request(connect_addr)
24            .context("Failed to create client request")?;
25
26        req.headers_mut().insert(
27            "Sec-WebSocket-Protocol",
28            HeaderValue::from_str(header).context("Invalid header value")?,
29        );
30
31        let (ws_stream, _) = connect_async(req)
32            .await
33            .context("Failed to connect to endpoint")?;
34        Ok(Self {
35            stream: ws_stream,
36            seq: 0,
37        })
38    }
39
40    pub async fn send_req(&mut self, method: u32, params: impl Serialize) -> Result<()> {
41        self.seq += 1;
42        let req = serde_json::to_string(&WsRequest {
43            method: method,
44            seq: self.seq,
45            params: params,
46        })
47        .context("Failed to serialize request")?;
48        self.stream
49            .send(Message::text(req))
50            .await
51            .context("Failed to send request")?;
52        Ok(())
53    }
54
55    pub async fn recv_raw(&mut self) -> Result<serde_json::Value> {
56        // Get the next message from the stream, or return an error if the connection is closed
57        let msg = self
58            .stream
59            .next()
60            .await
61            .ok_or_else(|| anyhow!("Connection closed"))?
62            .context("Failed to receive message")?;
63
64        let json_value = match msg {
65            Message::Text(text) => serde_json::from_str(text.as_str())
66                .context("Failed to parse received message as JSON")?,
67            _ => return Err(anyhow!("Received unexpected non-text message")),
68        };
69
70        Ok(json_value)
71    }
72
73    pub async fn close(mut self) -> Result<()> {
74        self.stream
75            .close(None)
76            .await
77            .context("Failed to close connection")?;
78        Ok(())
79    }
80}
endpoint_validator/ws/client.rs

endpoint_validator/ws/
client.rs
