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
52pub(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 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 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}