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

acta_maker_sdk/ws/
client.rs
