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)]
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 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 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}