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

acta_maker_sdk/ws/
client.rs
