Skip to main content

acta_maker_sdk/ws/
client.rs

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