rust_okx/ws/
client.rs

1//! WebSocket client and connection lifecycle.
2
3use std::collections::VecDeque;
4
5use crate::credentials::Credentials;
6use crate::ws::WsError;
7use crate::{Error, signing};
8
9use super::arg::Arg;
10use super::conn::{TungsteniteConnector, WsConn, WsConnector, WsFrame};
11use super::event::{WsEvent, parse_text_event};
12use super::request::{ChannelRequest, LoginArg, LoginRequest};
13
14/// OKX WebSocket channel group.
15#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
16#[non_exhaustive]
17pub enum WsChannelGroup {
18    /// Public market-data channels.
19    Public,
20    /// Private account and order channels.
21    Private,
22    /// Business channels such as candle variants and some advanced feeds.
23    Business,
24}
25
26impl WsChannelGroup {
27    fn path(self) -> &'static str {
28        match self {
29            Self::Public => "public",
30            Self::Private => "private",
31            Self::Business => "business",
32        }
33    }
34}
35
36/// Builder for an [`OkxWs`] client.
37pub struct OkxWsBuilder<C: WsConnector = TungsteniteConnector> {
38    connector: C,
39    group: WsChannelGroup,
40    credentials: Option<Credentials>,
41    demo: bool,
42}
43
44impl<C: WsConnector> OkxWsBuilder<C> {
45    /// Create a builder from a connector, channel group, and region.
46    pub fn new(connector: C, group: WsChannelGroup) -> Self {
47        Self {
48            connector,
49            group,
50            credentials: None,
51            demo: false,
52        }
53    }
54
55    /// Set API credentials for private WebSocket login.
56    pub fn credentials(mut self, credentials: Credentials) -> Self {
57        self.credentials = Some(credentials);
58        self
59    }
60
61    /// Toggle OKX demo trading WebSocket endpoints.
62    pub fn demo_trading(mut self, demo: bool) -> Self {
63        self.demo = demo;
64        self
65    }
66
67    /// Connect and build the WebSocket client.
68    pub async fn connect(self) -> Result<OkxWs<C>, Error> {
69        let url = ws_url(self.group, self.demo);
70        let conn = self.connector.connect(&url).await?;
71        Ok(OkxWs {
72            connector: self.connector,
73            conn,
74            url,
75            group: self.group,
76            credentials: self.credentials,
77            logged_in: false,
78            login_sent: false,
79            subscriptions: Vec::new(),
80            pending_after_login: Vec::new(),
81            pending_payloads_after_login: Vec::new(),
82            queued: VecDeque::new(),
83        })
84    }
85}
86
87impl OkxWsBuilder<TungsteniteConnector> {
88    /// Create a public WebSocket builder using the default connector.
89    pub fn public() -> Self {
90        Self::new(TungsteniteConnector, WsChannelGroup::Public)
91    }
92
93    /// Create a business WebSocket builder using the default connector.
94    pub fn business() -> Self {
95        Self::new(TungsteniteConnector, WsChannelGroup::Business)
96    }
97
98    /// Create a private WebSocket builder using the default connector.
99    pub fn private(credentials: Credentials) -> Self {
100        Self::new(TungsteniteConnector, WsChannelGroup::Private).credentials(credentials)
101    }
102}
103
104/// OKX WebSocket client.
105pub struct OkxWs<C: WsConnector = TungsteniteConnector> {
106    connector: C,
107    conn: C::Conn,
108    url: String,
109    group: WsChannelGroup,
110    credentials: Option<Credentials>,
111    logged_in: bool,
112    login_sent: bool,
113    subscriptions: Vec<Arg>,
114    pending_after_login: Vec<Arg>,
115    pending_payloads_after_login: Vec<String>,
116    queued: VecDeque<WsEvent>,
117}
118
119impl OkxWs<TungsteniteConnector> {
120    /// Start building a public WebSocket client.
121    pub fn public() -> OkxWsBuilder<TungsteniteConnector> {
122        OkxWsBuilder::public()
123    }
124
125    /// Start building a business WebSocket client.
126    pub fn business() -> OkxWsBuilder<TungsteniteConnector> {
127        OkxWsBuilder::business()
128    }
129
130    /// Start building a private WebSocket client.
131    pub fn private(credentials: Credentials) -> OkxWsBuilder<TungsteniteConnector> {
132        OkxWsBuilder::private(credentials)
133    }
134}
135
136impl<C: WsConnector> OkxWs<C> {
137    pub(crate) fn require_channel_group(
138        &self,
139        required: WsChannelGroup,
140        operation: &str,
141    ) -> Result<(), Error> {
142        if self.group != required {
143            return Err(WsError::Configuration(format!(
144                "WebSocket operation `{operation}` requires the {required:?} channel group"
145            ))
146            .into());
147        }
148        Ok(())
149    }
150
151    pub(crate) fn require_credentials(&self, operation: &str) -> Result<(), Error> {
152        if self.credentials.is_none() {
153            return Err(WsError::Configuration(format!(
154                "WebSocket operation `{operation}` requires API credentials"
155            ))
156            .into());
157        }
158        Ok(())
159    }
160
161    /// Subscribe to one or more channel arguments.
162    ///
163    /// OKX docs: <https://www.okx.com/docs-v5/en/#overview-websocket-subscribe>
164    pub async fn subscribe(&mut self, args: &[Arg]) -> Result<(), Error> {
165        self.track_subscriptions(args);
166        if self.needs_login()? && !self.logged_in {
167            self.pending_after_login.extend_from_slice(args);
168            if !self.login_sent {
169                self.send_login().await?;
170            }
171            return Ok(());
172        }
173        self.send_op("subscribe", args).await
174    }
175
176    /// Unsubscribe from one or more channel arguments.
177    ///
178    /// OKX docs: <https://www.okx.com/docs-v5/en/#overview-websocket-unsubscribe>
179    pub async fn unsubscribe(&mut self, args: &[Arg]) -> Result<(), Error> {
180        self.subscriptions.retain(|arg| !args.contains(arg));
181        self.send_op("unsubscribe", args).await
182    }
183
184    /// Receive and decode the next WebSocket event.
185    pub async fn next_event(&mut self) -> Result<Option<WsEvent>, Error> {
186        if let Some(event) = self.queued.pop_front() {
187            return Ok(Some(event));
188        }
189
190        loop {
191            let frame = match self.conn.recv().await {
192                Ok(Some(frame)) => frame,
193                Ok(None) | Err(_) => {
194                    self.reconnect().await?;
195                    return Ok(Some(WsEvent::Reconnected));
196                }
197            };
198
199            match frame {
200                WsFrame::Text(text) if text == "pong" => continue,
201                WsFrame::Text(text) if text == "ping" => {
202                    self.conn.send_text("pong".to_owned()).await?;
203                    continue;
204                }
205                WsFrame::Text(text) => {
206                    if let Some(event) = self.parse_text(&text).await? {
207                        return Ok(Some(event));
208                    }
209                }
210                WsFrame::Ping(payload) => {
211                    self.conn.send_pong(payload).await?;
212                }
213                WsFrame::Pong(_) => continue,
214                WsFrame::Close => {
215                    self.reconnect().await?;
216                    return Ok(Some(WsEvent::Reconnected));
217                }
218            }
219        }
220    }
221
222    /// Close the WebSocket connection.
223    pub async fn close(&mut self) -> Result<(), Error> {
224        self.conn.close().await.map_err(Error::from)
225    }
226
227    pub(crate) async fn send_operation_payload(&mut self, payload: String) -> Result<(), Error> {
228        if self.needs_login()? && !self.logged_in {
229            self.pending_payloads_after_login.push(payload);
230            if !self.login_sent {
231                self.send_login().await?;
232            }
233            return Ok(());
234        }
235        self.conn.send_text(payload).await.map_err(Error::from)
236    }
237
238    fn track_subscriptions(&mut self, args: &[Arg]) {
239        for arg in args {
240            if !self.subscriptions.contains(arg) {
241                self.subscriptions.push(arg.clone());
242            }
243        }
244    }
245
246    fn needs_login(&self) -> Result<bool, Error> {
247        if self.group == WsChannelGroup::Private && self.credentials.is_none() {
248            return Err(WsError::Configuration(
249                "private WebSocket requires credentials".to_owned(),
250            )
251            .into());
252        }
253        Ok(self.group == WsChannelGroup::Private
254            || (self.group == WsChannelGroup::Business && self.credentials.is_some()))
255    }
256
257    async fn reconnect(&mut self) -> Result<(), Error> {
258        self.conn = self.connector.connect(&self.url).await?;
259        self.logged_in = false;
260        self.login_sent = false;
261        self.pending_payloads_after_login.clear();
262
263        if self.needs_login()? {
264            self.pending_after_login = self.subscriptions.clone();
265            self.send_login().await?;
266        } else if !self.subscriptions.is_empty() {
267            let args = self.subscriptions.clone();
268            self.send_op("subscribe", &args).await?;
269        }
270        Ok(())
271    }
272
273    async fn send_login(&mut self) -> Result<(), Error> {
274        let credentials = self.credentials.as_ref().ok_or_else(|| {
275            WsError::Configuration("WebSocket login requires credentials".to_owned())
276        })?;
277        let payload = login_payload(credentials, &signing::ws_timestamp())?;
278        self.conn.send_text(payload).await?;
279        self.login_sent = true;
280        Ok(())
281    }
282
283    async fn send_op(&mut self, op: &str, args: &[Arg]) -> Result<(), Error> {
284        let request = match op {
285            "subscribe" => ChannelRequest::subscribe(args),
286            "unsubscribe" => ChannelRequest::unsubscribe(args),
287            _ => ChannelRequest { id: None, op, args },
288        };
289        let payload = serde_json::to_string(&request).map_err(|e| WsError::Encode { source: e })?;
290        self.conn.send_text(payload).await.map_err(Error::from)
291    }
292
293    async fn parse_text(&mut self, text: &str) -> Result<Option<WsEvent>, Error> {
294        let Some(event) = parse_text_event(text)? else {
295            return Ok(None);
296        };
297
298        if matches!(event, WsEvent::Login) {
299            self.logged_in = true;
300            self.login_sent = false;
301            if !self.pending_after_login.is_empty() {
302                let args = std::mem::take(&mut self.pending_after_login);
303                self.send_op("subscribe", &args).await?;
304            }
305            if !self.pending_payloads_after_login.is_empty() {
306                let payloads = std::mem::take(&mut self.pending_payloads_after_login);
307                for payload in payloads {
308                    self.conn.send_text(payload).await?;
309                }
310            }
311        }
312
313        Ok(Some(event))
314    }
315}
316
317/// Build the authenticated WebSocket login body.
318///
319/// OKX docs: <https://www.okx.com/docs-v5/en/#overview-websocket-login>
320pub(crate) fn login_payload(credentials: &Credentials, timestamp: &str) -> Result<String, Error> {
321    let payload = LoginRequest::new(LoginArg::new(
322        credentials.api_key(),
323        credentials.passphrase(),
324        timestamp,
325        signing::ws_login_sign(timestamp, credentials.secret_key()),
326    ));
327    serde_json::to_string(&payload)
328        .map_err(|e| WsError::Encode { source: e })
329        .map_err(Error::from)
330}
331
332pub(crate) fn ws_url(group: WsChannelGroup, demo: bool) -> String {
333    let host = if demo { "wspap.okx.com" } else { "ws.okx.com" };
334    format!("wss://{host}:8443/ws/v5/{}", group.path())
335}
rust_okx/ws/client.rs

rust_okx/ws/
client.rs
