acta_maker_sdk/ws/
client.rs

1use futures_util::stream::{SplitSink, SplitStream};
2use futures_util::{SinkExt, StreamExt};
3use std::sync::{Arc, OnceLock};
4use std::time::Duration;
5use tokio::net::TcpStream;
6
7use tokio_tungstenite::{
8    MaybeTlsStream, WebSocketStream, connect_async_with_config,
9    tungstenite::{
10        Bytes, Message, Utf8Bytes,
11        protocol::{CloseFrame, WebSocketConfig},
12    },
13};
14
15use crate::ws::error::{WsClientError, WsResult, WsTransportConfigError};
16use crate::ws::types::{
17    AcceptQuoteMessage, AddChannelsData, AddMintsData, AuthChallengeData, BatchQuotesMessage,
18    CancelAllQuotesMessage, CancelQuoteData, CancelRfqData, ClaimReferralCodeData, ClientMessage,
19    CloseCode, GetActiveRfqsMessage, GetExpiriesMessage, GetIndicativePricesMessage,
20    GetMakerPositionsMessage, GetMarketDescriptorsMessage, GetMarketsForMakerMessage,
21    GetMarketsMessage, GetMmSummaryMessage, GetMyActiveRfqsMessage, GetMyQuotesMessage,
22    GetMyReferralInfoData, GetOrderStatusMessage, GetPositionsMessage, GetSubscriptionsMessage,
23    GetTokensMessage, HelloData, IndicativePricesResponseMessage, Lane, QuoteMessage,
24    RedeemInviteData, RemoveChannelsData, RemoveMintsData, ReplaceQuoteMessage, ResumeAuthData,
25    RfqRequestMessage, ServerMessage, StartAuthData, SubmitSignedSponsoredTxData, SubscribeData,
26    UnsubscribeData, WsChannel, parse_server_message,
27};
28use uuid::Uuid;
29
30type Socket = WebSocketStream<MaybeTlsStream<TcpStream>>;
31pub(crate) type WsSink = SplitSink<Socket, Message>;
32
33macro_rules! ws_method {
34    ($name:ident, $variant:ident, $data:ty) => {
35        pub async fn $name(&mut self, data: $data) -> WsResult<()> {
36            self.send(&ClientMessage::$variant(data)).await
37        }
38    };
39}
40
41macro_rules! ws_method_request_id {
42    ($name:ident, $variant:ident, $msg:ident) => {
43        pub async fn $name(&mut self) -> WsResult<()> {
44            self.send(&ClientMessage::$variant($msg {
45                request_id: Uuid::new_v4(),
46            }))
47            .await
48        }
49    };
50}
51
52pub struct WsClient {
53    stream: Socket,
54    server_close: Arc<OnceLock<ServerClose>>,
55}
56
57/// Read half of a split [`WsClient`]. Liveness frames are returned, not
58/// answered, so the owner decides which lane the reply takes.
59pub(crate) struct WsReader {
60    stream: SplitStream<Socket>,
61    server_close: Arc<OnceLock<ServerClose>>,
62}
63
64/// The close frame the server sent.
65#[derive(Debug, Clone, PartialEq, Eq)]
66pub struct ServerClose {
67    pub code: u16,
68    pub reason: String,
69}
70
71impl ServerClose {
72    /// `None` for codes outside [`CloseCode`].
73    #[must_use]
74    pub const fn close_code(&self) -> Option<CloseCode> {
75        CloseCode::from_code(self.code)
76    }
77}
78
79fn record_close(slot: &OnceLock<ServerClose>, frame: Option<CloseFrame>) {
80    if let Some(frame) = frame {
81        let _ = slot.set(ServerClose {
82            code: u16::from(frame.code),
83            reason: frame.reason.to_string(),
84        });
85    }
86}
87
88#[allow(clippy::large_enum_variant)]
89pub(crate) enum WsFrame {
90    Message(ServerMessage),
91    Ping(Bytes),
92    Pong,
93}
94
95#[derive(Debug, Clone)]
96pub struct PreparedClientMessage {
97    text: Utf8Bytes,
98    lane: Lane,
99    subscription: Option<Box<crate::ws::types::SubscriptionChange>>,
100}
101
102impl PreparedClientMessage {
103    pub fn new(message: &ClientMessage) -> Result<Self, serde_json::Error> {
104        Ok(Self {
105            text: serde_json::to_string(message)?.into(),
106            lane: message.lane(),
107            subscription: crate::ws::types::SubscriptionChange::from_message(message).map(Box::new),
108        })
109    }
110
111    pub(crate) fn take_subscription(
112        &mut self,
113    ) -> Option<Box<crate::ws::types::SubscriptionChange>> {
114        self.subscription.take()
115    }
116
117    pub(crate) const fn lane(&self) -> Lane {
118        self.lane
119    }
120
121    pub(crate) fn into_text(self) -> Utf8Bytes {
122        self.text
123    }
124
125    pub fn as_str(&self) -> &str {
126        &self.text
127    }
128
129    #[must_use]
130    pub fn len(&self) -> usize {
131        self.text.len()
132    }
133
134    #[must_use]
135    pub fn is_empty(&self) -> bool {
136        self.text.is_empty()
137    }
138}
139
140#[derive(Debug, Clone, Copy)]
141pub struct WsTransportConfig {
142    pub connect_timeout: Duration,
143    pub read_buffer_size: usize,
144    pub write_buffer_size: usize,
145    pub max_write_buffer_size: usize,
146    pub max_message_size: usize,
147    pub max_frame_size: usize,
148    pub tcp_nodelay: bool,
149}
150
151impl Default for WsTransportConfig {
152    fn default() -> Self {
153        Self {
154            connect_timeout: Duration::from_secs(10),
155            read_buffer_size: 64 * 1024,
156            write_buffer_size: 0,
157            max_write_buffer_size: 1024 * 1024,
158            max_message_size: 2 * 1024 * 1024,
159            max_frame_size: 2 * 1024 * 1024,
160            tcp_nodelay: true,
161        }
162    }
163}
164
165impl WsTransportConfig {
166    /// Validate limits before a connection task is spawned.
167    ///
168    /// # Errors
169    /// Returns a typed error for zero deadlines/limits or an invalid write-buffer range.
170    pub fn validate(self) -> Result<(), WsTransportConfigError> {
171        if self.connect_timeout.is_zero() {
172            return Err(WsTransportConfigError::ZeroConnectTimeout);
173        }
174        if self.max_write_buffer_size <= self.write_buffer_size {
175            return Err(WsTransportConfigError::InvalidWriteBufferRange);
176        }
177        if self.max_message_size == 0 || self.max_frame_size == 0 {
178            return Err(WsTransportConfigError::ZeroMessageOrFrameLimit);
179        }
180        Ok(())
181    }
182}
183
184impl WsClient {
185    pub async fn connect(url: &str) -> WsResult<Self> {
186        Self::connect_with_config(url, WsTransportConfig::default()).await
187    }
188
189    pub async fn connect_with_config(url: &str, config: WsTransportConfig) -> WsResult<Self> {
190        config.validate()?;
191        let websocket_config = WebSocketConfig::default()
192            .read_buffer_size(config.read_buffer_size)
193            .write_buffer_size(config.write_buffer_size)
194            .max_write_buffer_size(config.max_write_buffer_size)
195            .max_message_size(Some(config.max_message_size))
196            .max_frame_size(Some(config.max_frame_size));
197        let connect = connect_async_with_config(url, Some(websocket_config), config.tcp_nodelay);
198        let (stream, _response) = tokio::time::timeout(config.connect_timeout, connect)
199            .await
200            .map_err(|_| WsClientError::ConnectTimeout)??;
201        Ok(Self {
202            stream,
203            server_close: Arc::default(),
204        })
205    }
206
207    pub async fn send(&mut self, msg: &ClientMessage) -> WsResult<()> {
208        let prepared = PreparedClientMessage::new(msg)?;
209        self.send_prepared(&prepared).await?;
210        Ok(())
211    }
212
213    pub async fn send_prepared(&mut self, message: &PreparedClientMessage) -> WsResult<()> {
214        self.stream
215            .send(Message::Text(message.text.clone()))
216            .await?;
217        Ok(())
218    }
219
220    /// Send a prepared message by transferring its buffer into the socket sink.
221    pub async fn send_prepared_owned(&mut self, message: PreparedClientMessage) -> WsResult<()> {
222        self.stream.send(Message::Text(message.text)).await?;
223        Ok(())
224    }
225
226    ws_method!(send_hello, Hello, HelloData);
227    ws_method!(resume_auth, ResumeAuth, ResumeAuthData);
228    ws_method!(auth_challenge, AuthChallenge, AuthChallengeData);
229    ws_method!(quote, Quote, QuoteMessage);
230    ws_method!(replace_quote, ReplaceQuote, ReplaceQuoteMessage);
231    ws_method!(batch_quotes, BatchQuotes, BatchQuotesMessage);
232    ws_method!(subscribe, Subscribe, SubscribeData);
233    ws_method!(unsubscribe, Unsubscribe, UnsubscribeData);
234    ws_method!(get_my_quotes, GetMyQuotes, GetMyQuotesMessage);
235    ws_method!(
236        get_maker_positions,
237        GetMakerPositions,
238        GetMakerPositionsMessage
239    );
240    ws_method!(
241        get_markets_for_maker,
242        GetMarketsForMaker,
243        GetMarketsForMakerMessage
244    );
245    ws_method!(rfq_request, RfqRequest, RfqRequestMessage);
246    ws_method!(accept_quote, AcceptQuote, AcceptQuoteMessage);
247    ws_method!(
248        submit_signed_sponsored_tx,
249        SubmitSignedSponsoredTx,
250        SubmitSignedSponsoredTxData
251    );
252    ws_method!(
253        indicative_prices_response,
254        IndicativePricesResponse,
255        IndicativePricesResponseMessage
256    );
257    ws_method!(
258        get_indicative_prices,
259        GetIndicativePrices,
260        GetIndicativePricesMessage
261    );
262    ws_method!(
263        get_market_descriptors,
264        GetMarketDescriptors,
265        GetMarketDescriptorsMessage
266    );
267    ws_method!(get_tokens, GetTokens, GetTokensMessage);
268    ws_method!(get_expiries, GetExpiries, GetExpiriesMessage);
269    ws_method!(get_order_status, GetOrderStatus, GetOrderStatusMessage);
270    ws_method!(redeem_invite, RedeemInvite, RedeemInviteData);
271    ws_method!(
272        claim_referral_code,
273        ClaimReferralCode,
274        ClaimReferralCodeData
275    );
276    ws_method!(
277        get_my_referral_info,
278        GetMyReferralInfo,
279        GetMyReferralInfoData
280    );
281
282    ws_method_request_id!(get_mm_summary, GetMmSummary, GetMmSummaryMessage);
283    ws_method_request_id!(get_subscriptions, GetSubscriptions, GetSubscriptionsMessage);
284    ws_method_request_id!(get_markets, GetMarkets, GetMarketsMessage);
285
286    pub async fn get_positions(&mut self) -> WsResult<()> {
287        self.send(&ClientMessage::GetPositions(GetPositionsMessage {
288            request_id: Uuid::new_v4(),
289            ..Default::default()
290        }))
291        .await
292    }
293    ws_method_request_id!(get_my_active_rfqs, GetMyActiveRfqs, GetMyActiveRfqsMessage);
294    ws_method_request_id!(get_active_rfqs, GetActiveRfqs, GetActiveRfqsMessage);
295
296    pub async fn ping(&mut self) -> WsResult<()> {
297        self.send(&ClientMessage::Ping).await
298    }
299
300    pub async fn logout(&mut self) -> WsResult<()> {
301        self.send(&ClientMessage::Logout).await
302    }
303
304    pub async fn start_auth(&mut self, pubkey: String) -> WsResult<()> {
305        self.send(&ClientMessage::StartAuth(StartAuthData { pubkey }))
306            .await
307    }
308
309    pub async fn cancel_quote(&mut self, rfq_id: Uuid) -> WsResult<()> {
310        self.send(&ClientMessage::CancelQuote(CancelQuoteData {
311            rfq_id,
312            request_id: Uuid::new_v4(),
313        }))
314        .await
315    }
316
317    pub async fn cancel_all_quotes(&mut self, market: Option<String>) -> WsResult<()> {
318        self.send(&ClientMessage::CancelAllQuotes(CancelAllQuotesMessage {
319            request_id: Uuid::new_v4(),
320            market,
321        }))
322        .await
323    }
324
325    pub async fn cancel_rfq(&mut self, rfq_id: Uuid) -> WsResult<()> {
326        self.send(&ClientMessage::CancelRfq(CancelRfqData {
327            rfq_id,
328            request_id: Uuid::new_v4(),
329        }))
330        .await
331    }
332
333    pub async fn add_mints(
334        &mut self,
335        underlying_mints: Option<Vec<String>>,
336        quote_mints: Option<Vec<String>>,
337    ) -> WsResult<()> {
338        self.send(&ClientMessage::AddMints(AddMintsData {
339            request_id: Uuid::new_v4(),
340            underlying_mints,
341            quote_mints,
342        }))
343        .await
344    }
345
346    pub async fn remove_mints(
347        &mut self,
348        underlying_mints: Option<Vec<String>>,
349        quote_mints: Option<Vec<String>>,
350    ) -> WsResult<()> {
351        self.send(&ClientMessage::RemoveMints(RemoveMintsData {
352            request_id: Uuid::new_v4(),
353            underlying_mints,
354            quote_mints,
355        }))
356        .await
357    }
358
359    pub async fn add_channels(&mut self, channels: Vec<WsChannel>) -> WsResult<()> {
360        self.send(&ClientMessage::AddChannels(AddChannelsData {
361            request_id: Uuid::new_v4(),
362            channels,
363        }))
364        .await
365    }
366
367    pub async fn remove_channels(&mut self, channels: Vec<WsChannel>) -> WsResult<()> {
368        self.send(&ClientMessage::RemoveChannels(RemoveChannelsData {
369            request_id: Uuid::new_v4(),
370            channels,
371        }))
372        .await
373    }
374
375    pub async fn send_text(&mut self, text: impl Into<String>) -> WsResult<()> {
376        self.stream.send(Message::Text(text.into().into())).await?;
377        Ok(())
378    }
379
380    /// The server's close frame once [`Self::next`] has returned `None` after one.
381    #[must_use]
382    pub fn server_close(&self) -> Option<&ServerClose> {
383        self.server_close.get()
384    }
385
386    pub(crate) fn server_close_slot(&self) -> Arc<OnceLock<ServerClose>> {
387        Arc::clone(&self.server_close)
388    }
389
390    /// Read until the server's close frame or end of stream, at most `grace`.
391    pub(crate) async fn await_server_close(&mut self, grace: Duration) {
392        let _ = tokio::time::timeout(grace, async {
393            while let Some(Ok(_)) = self.next().await {}
394        })
395        .await;
396    }
397
398    pub async fn next(&mut self) -> Option<Result<ServerMessage, WsClientError>> {
399        loop {
400            let msg = match self.stream.next().await {
401                Some(Ok(msg)) => msg,
402                Some(Err(err)) => return Some(Err(err.into())),
403                None => return None,
404            };
405
406            match msg {
407                Message::Text(text) => {
408                    let parsed = parse_server_message(&text);
409                    return Some(parsed.map_err(WsClientError::from));
410                }
411                Message::Binary(bin) => {
412                    return Some(parse_binary(&bin).map_err(WsClientError::from));
413                }
414                Message::Ping(payload) => {
415                    if let Err(err) = self.stream.send(Message::Pong(payload)).await {
416                        return Some(Err(err.into()));
417                    }
418                }
419                Message::Pong(_) => {}
420                Message::Close(frame) => {
421                    record_close(&self.server_close, frame);
422                    return None;
423                }
424                Message::Frame(_) => {}
425            }
426        }
427    }
428
429    pub async fn close(mut self) -> WsResult<()> {
430        self.stream.close(None).await?;
431        Ok(())
432    }
433
434    pub(crate) fn split(self) -> (WsSink, WsReader) {
435        let (sink, stream) = self.stream.split();
436        (
437            sink,
438            WsReader {
439                stream,
440                server_close: self.server_close,
441            },
442        )
443    }
444}
445
446impl WsReader {
447    /// Read until the server's close frame or end of stream, at most `grace`.
448    pub(crate) async fn await_server_close(&mut self, grace: Duration) {
449        let _ = tokio::time::timeout(grace, async {
450            while let Some(Ok(_)) = self.next().await {}
451        })
452        .await;
453    }
454
455    pub(crate) async fn next(&mut self) -> Option<WsResult<WsFrame>> {
456        loop {
457            let msg = match self.stream.next().await {
458                Some(Ok(msg)) => msg,
459                Some(Err(err)) => return Some(Err(err.into())),
460                None => return None,
461            };
462            let frame = match msg {
463                Message::Text(text) => parse_server_message(&text).map(WsFrame::Message),
464                Message::Binary(bin) => parse_binary(&bin).map(WsFrame::Message),
465                Message::Ping(payload) => Ok(WsFrame::Ping(payload)),
466                Message::Pong(_) => Ok(WsFrame::Pong),
467                Message::Close(frame) => {
468                    record_close(&self.server_close, frame);
469                    return None;
470                }
471                Message::Frame(_) => continue,
472            };
473            return Some(frame.map_err(WsClientError::from));
474        }
475    }
476}
477
478fn parse_binary(bin: &[u8]) -> Result<ServerMessage, serde_json::Error> {
479    match std::str::from_utf8(bin) {
480        Ok(text) => parse_server_message(text),
481        Err(_) => serde_json::from_slice::<ServerMessage>(bin),
482    }
483}
acta_maker_sdk/ws/client.rs

acta_maker_sdk/ws/
client.rs
