rust_okx/ws/
conn.rs

1//! WebSocket connection traits and the default tungstenite connector.
2
3use std::future::Future;
4
5use bytes::Bytes;
6
7use super::error::WsError;
8
9/// A WebSocket frame handled by the OKX WebSocket client.
10#[derive(Debug, Clone, PartialEq, Eq)]
11#[non_exhaustive]
12pub enum WsFrame {
13    /// A text frame.
14    Text(String),
15    /// A ping frame.
16    Ping(Bytes),
17    /// A pong frame.
18    Pong(Bytes),
19    /// A close frame.
20    Close,
21}
22
23/// A full-duplex WebSocket connection.
24///
25/// This trait is intentionally minimal so tests can use a fake connection
26/// without depending on a concrete WebSocket implementation.
27pub trait WsConn: Send {
28    /// Send a text frame.
29    fn send_text(&mut self, text: String) -> impl Future<Output = Result<(), WsError>> + Send;
30
31    /// Send a pong control frame with the supplied ping payload.
32    fn send_pong(&mut self, payload: Bytes) -> impl Future<Output = Result<(), WsError>> + Send;
33
34    /// Receive the next frame.
35    fn recv(&mut self) -> impl Future<Output = Result<Option<WsFrame>, WsError>> + Send;
36
37    /// Close the connection.
38    fn close(&mut self) -> impl Future<Output = Result<(), WsError>> + Send;
39}
40
41/// Creates WebSocket connections for an OKX WebSocket endpoint URL.
42pub trait WsConnector: Send + Sync {
43    /// The connection type returned by this connector.
44    type Conn: WsConn;
45
46    /// Connect to a WebSocket URL.
47    fn connect(&self, url: &str) -> impl Future<Output = Result<Self::Conn, WsError>> + Send;
48}
49
50#[cfg(feature = "websocket")]
51mod tungstenite_impl {
52    use futures_util::{SinkExt, StreamExt};
53    use tokio::net::TcpStream;
54    use tokio_tungstenite::tungstenite::Message;
55    use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async};
56
57    use super::*;
58    use crate::TransportError;
59
60    /// Default WebSocket connector backed by `tokio-tungstenite`.
61    #[derive(Debug, Clone, Copy, Default)]
62    pub struct TungsteniteConnector;
63
64    /// WebSocket connection backed by `tokio-tungstenite`.
65    pub struct TungsteniteConn {
66        inner: WebSocketStream<MaybeTlsStream<TcpStream>>,
67    }
68
69    impl WsConnector for TungsteniteConnector {
70        type Conn = TungsteniteConn;
71
72        fn connect(&self, url: &str) -> impl Future<Output = Result<Self::Conn, WsError>> + Send {
73            let url = url.to_owned();
74            async move {
75                let (inner, _) = connect_async(url)
76                    .await
77                    .map_err(|e| WsError::from(TransportError::new(e)))?;
78                Ok(TungsteniteConn { inner })
79            }
80        }
81    }
82
83    #[allow(clippy::manual_async_fn)]
84    impl WsConn for TungsteniteConn {
85        fn send_text(&mut self, text: String) -> impl Future<Output = Result<(), WsError>> + Send {
86            async move {
87                self.inner
88                    .send(Message::Text(text))
89                    .await
90                    .map_err(|e| WsError::from(TransportError::new(e)))
91            }
92        }
93
94        fn send_pong(
95            &mut self,
96            payload: Bytes,
97        ) -> impl Future<Output = Result<(), WsError>> + Send {
98            async move {
99                self.inner
100                    .send(Message::Pong(payload.to_vec()))
101                    .await
102                    .map_err(|e| WsError::from(TransportError::new(e)))
103            }
104        }
105
106        fn recv(&mut self) -> impl Future<Output = Result<Option<WsFrame>, WsError>> + Send {
107            async move {
108                loop {
109                    let Some(message) = self.inner.next().await else {
110                        return Ok(None);
111                    };
112                    let message = message.map_err(|e| WsError::from(TransportError::new(e)))?;
113                    match message {
114                        Message::Text(text) => return Ok(Some(WsFrame::Text(text))),
115                        Message::Ping(bytes) => return Ok(Some(WsFrame::Ping(Bytes::from(bytes)))),
116                        Message::Pong(bytes) => return Ok(Some(WsFrame::Pong(Bytes::from(bytes)))),
117                        Message::Close(_) => return Ok(Some(WsFrame::Close)),
118                        Message::Binary(_) | Message::Frame(_) => continue,
119                    }
120                }
121            }
122        }
123
124        fn close(&mut self) -> impl Future<Output = Result<(), WsError>> + Send {
125            async move {
126                self.inner
127                    .close(None)
128                    .await
129                    .map_err(|e| WsError::from(TransportError::new(e)))
130            }
131        }
132    }
133}
134
135#[cfg(feature = "websocket")]
136pub use tungstenite_impl::{TungsteniteConn, TungsteniteConnector};
rust_okx/ws/conn.rs

rust_okx/ws/
conn.rs
