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