1use std::{
22 num::NonZeroU32,
23 str::FromStr,
24 sync::{
25 Arc, LazyLock,
26 atomic::{AtomicBool, AtomicU8, Ordering},
27 },
28 time::Duration,
29};
30
31use arc_swap::ArcSwap;
32#[cfg(test)]
33use nautilus_core::string::secret::REDACTED;
34use nautilus_core::{AtomicMap, string::secret::SecretString};
35use nautilus_live::{
36 SocketControl,
37 task::{TaskJoinOutcome, TaskSlot, finish_task},
38};
39use nautilus_model::{
40 data::BarType,
41 identifiers::{AccountId, InstrumentId},
42 instruments::{Instrument, InstrumentAny},
43};
44use nautilus_network::{
45 http::create_standard_nautilus_headers,
46 mode::ConnectionMode,
47 ratelimiter::quota::Quota,
48 websocket::{
49 SubscriptionState, TransportBackend, WebSocketClient, WebSocketConfig,
50 channel_message_handler,
51 },
52};
53use ustr::Ustr;
54
55use crate::{
56 common::{
57 consts::{
58 RECONNECT_BACKOFF_FACTOR, RECONNECT_BASE_BACKOFF, RECONNECT_JITTER_MS,
59 RECONNECT_MAX_BACKOFF, RECONNECT_TIMEOUT, WS_DISCONNECT_TIMEOUT, WS_HEARTBEAT_SECS,
60 },
61 credential::CoinbaseCredential,
62 enums::CoinbaseWsChannel,
63 },
64 websocket::{
65 handler::{FeedHandler, HandlerCommand, NautilusWsMessage},
66 messages::{CoinbaseWsAction, CoinbaseWsSubscription},
67 },
68};
69
70pub static COINBASE_WS_CONNECTION_QUOTA: LazyLock<Quota> = LazyLock::new(|| {
72 Quota::per_second(NonZeroU32::new(8).expect("non-zero")).expect("valid constant")
73});
74
75pub static COINBASE_WS_SUBSCRIPTION_QUOTA: LazyLock<Quota> = LazyLock::new(|| {
77 Quota::per_second(NonZeroU32::new(8).expect("non-zero")).expect("valid constant")
78});
79
80pub const COINBASE_RATE_LIMIT_KEY_SUBSCRIPTION: &str = "subscription";
82
83pub static COINBASE_WS_SUBSCRIPTION_KEYS: LazyLock<[Ustr; 1]> =
85 LazyLock::new(|| [Ustr::from(COINBASE_RATE_LIMIT_KEY_SUBSCRIPTION)]);
86
87#[derive(Debug)]
92pub struct CoinbaseWebSocketClient {
93 url: String,
94 connection_mode: Arc<ArcSwap<AtomicU8>>,
95 signal: Arc<AtomicBool>,
96 cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<HandlerCommand>>>,
97 out_rx: Option<tokio::sync::mpsc::UnboundedReceiver<NautilusWsMessage>>,
98 instruments: Arc<AtomicMap<InstrumentId, InstrumentAny>>,
99 subscription_aliases: Arc<AtomicMap<Ustr, Ustr>>,
104 bar_types: ahash::AHashMap<String, BarType>,
105 subscriptions: SubscriptionState,
106 credential: Option<CoinbaseCredential>,
107 account_id: Option<AccountId>,
108 task_handle: TaskSlot<()>,
109 shutdown_errors: Vec<String>,
110 transport_backend: TransportBackend,
111 proxy_url: Option<SecretString>,
112 socket_control: Option<SocketControl>,
113}
114
115impl Clone for CoinbaseWebSocketClient {
116 fn clone(&self) -> Self {
117 Self {
118 url: self.url.clone(),
119 connection_mode: Arc::clone(&self.connection_mode),
120 signal: Arc::clone(&self.signal),
121 cmd_tx: Arc::clone(&self.cmd_tx),
122 out_rx: None,
123 instruments: Arc::clone(&self.instruments),
124 subscription_aliases: Arc::clone(&self.subscription_aliases),
125 bar_types: self.bar_types.clone(),
126 subscriptions: self.subscriptions.clone(),
127 credential: self.credential.clone(),
128 account_id: self.account_id,
129 task_handle: TaskSlot::new(),
130 shutdown_errors: Vec::new(),
131 transport_backend: self.transport_backend,
132 proxy_url: self.proxy_url.clone(),
133 socket_control: self.socket_control.clone(),
134 }
135 }
136}
137
138impl CoinbaseWebSocketClient {
139 pub fn new(url: &str, transport_backend: TransportBackend, proxy_url: Option<String>) -> Self {
141 let (placeholder_tx, _) = tokio::sync::mpsc::unbounded_channel();
142
143 Self {
144 url: url.to_string(),
145 connection_mode: Arc::new(ArcSwap::from_pointee(AtomicU8::new(
146 ConnectionMode::Closed.as_u8(),
147 ))),
148 signal: Arc::new(AtomicBool::new(false)),
149 cmd_tx: Arc::new(tokio::sync::RwLock::new(placeholder_tx)),
150 out_rx: None,
151 instruments: Arc::new(AtomicMap::new()),
152 subscription_aliases: Arc::new(AtomicMap::new()),
153 bar_types: ahash::AHashMap::new(),
154 subscriptions: SubscriptionState::new('|'),
155 credential: None,
156 account_id: None,
157 task_handle: TaskSlot::new(),
158 shutdown_errors: Vec::new(),
159 transport_backend,
160 proxy_url: proxy_url.map(SecretString::from),
161 socket_control: None,
162 }
163 }
164
165 #[must_use]
167 pub fn with_socket_control(mut self, control: SocketControl) -> Self {
168 self.socket_control = Some(control);
169 self
170 }
171
172 pub fn with_credential(
174 url: &str,
175 credential: CoinbaseCredential,
176 transport_backend: TransportBackend,
177 proxy_url: Option<String>,
178 ) -> Self {
179 let mut client = Self::new(url, transport_backend, proxy_url);
180 client.credential = Some(credential);
181 client
182 }
183
184 pub async fn set_account_id(&mut self, account_id: AccountId) {
189 self.account_id = Some(account_id);
190
191 let cmd_tx = self.cmd_tx.read().await;
192 if let Err(e) = cmd_tx.send(HandlerCommand::SetAccountId(account_id)) {
193 log::debug!("Failed to send SetAccountId: {e}");
194 }
195 }
196
197 pub async fn initialize_instruments(&self, instruments: Vec<InstrumentAny>) {
204 for instrument in &instruments {
205 self.instruments.insert(instrument.id(), instrument.clone());
206 }
207
208 let cmd_tx = self.cmd_tx.read().await;
209 if let Err(e) = cmd_tx.send(HandlerCommand::InitializeInstruments(instruments)) {
210 log::debug!("Failed to send InitializeInstruments: {e}");
211 }
212 }
213
214 fn prime_default_subscriptions(&self) {
218 self.subscriptions
219 .mark_subscribe(CoinbaseWsChannel::Heartbeats.as_ref());
220 }
221
222 pub async fn connect(&mut self) -> anyhow::Result<()> {
224 if self.is_active() || self.is_reconnecting() {
225 log::warn!("WebSocket already connected or reconnecting");
226 return Ok(());
227 }
228
229 if let Some(outcome) = finish_task(
230 &mut self.task_handle,
231 WS_DISCONNECT_TIMEOUT,
232 WS_DISCONNECT_TIMEOUT,
233 )
234 .await
235 {
236 match outcome {
237 TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted => {}
238 TaskJoinOutcome::Failed(error) => {
239 anyhow::bail!("Coinbase WebSocket handler failed: {error}");
240 }
241 TaskJoinOutcome::Incomplete => {
242 anyhow::bail!("Coinbase WebSocket handler did not stop after abort");
243 }
244 }
245 }
246
247 self.signal.store(false, Ordering::Relaxed);
249
250 let (message_handler, raw_rx) = channel_message_handler();
251 let headers = create_standard_nautilus_headers();
252
253 let cfg = WebSocketConfig {
254 url: self.url.clone(),
255 headers,
256 heartbeat_interval_secs: Some(WS_HEARTBEAT_SECS),
259 heartbeat_payload: None,
260 connect_timeout_ms: Some(RECONNECT_TIMEOUT.as_millis() as u64),
261 reconnect_delay_initial_ms: Some(RECONNECT_BASE_BACKOFF.as_millis() as u64),
262 reconnect_delay_max_ms: Some(RECONNECT_MAX_BACKOFF.as_millis() as u64),
263 reconnect_backoff_factor: Some(RECONNECT_BACKOFF_FACTOR),
264 reconnect_jitter_ms: Some(RECONNECT_JITTER_MS),
265 reconnect_max_attempts: None,
266 heartbeat_timeout_secs: None,
267 idle_timeout_ms: None,
268 backend: self.transport_backend,
269 proxy_url: self
270 .proxy_url
271 .as_ref()
272 .map(|value| value.expose_secret().to_owned()),
273 };
274
275 let keyed_quotas = vec![(
276 COINBASE_RATE_LIMIT_KEY_SUBSCRIPTION.to_string(),
277 *COINBASE_WS_SUBSCRIPTION_QUOTA,
278 )];
279
280 let client = WebSocketClient::builder()
281 .config(cfg)
282 .message_handler(message_handler)
283 .keyed_quotas(keyed_quotas)
284 .default_quota(*COINBASE_WS_CONNECTION_QUOTA)
285 .maybe_state_sink(self.socket_control.as_ref().map(SocketControl::sink))
286 .connect()
287 .await?;
288
289 let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
290 let (out_tx, out_rx) = tokio::sync::mpsc::unbounded_channel::<NautilusWsMessage>();
291
292 *self.cmd_tx.write().await = cmd_tx.clone();
293 self.out_rx = Some(out_rx);
294 self.connection_mode.store(client.connection_mode_atomic());
295 let reconnect_handle = client.reconnect_handle();
296 log::debug!("Coinbase WebSocket connected: {}", self.url);
297
298 if let Err(e) = cmd_tx.send(HandlerCommand::SetClient(client)) {
299 anyhow::bail!("Failed to send SetClient command: {e}");
300 }
301
302 if let Some(control) = &self.socket_control {
303 control.register(move || reconnect_handle.request_reconnect());
304 }
305
306 let instruments_vec: Vec<InstrumentAny> =
307 self.instruments.load().values().cloned().collect();
308
309 if !instruments_vec.is_empty()
310 && let Err(e) = cmd_tx.send(HandlerCommand::InitializeInstruments(instruments_vec))
311 {
312 log::error!("Failed to send InitializeInstruments: {e}");
313 }
314
315 for (key, bar_type) in &self.bar_types {
317 if let Err(e) = cmd_tx.send(HandlerCommand::AddBarType {
318 key: key.clone(),
319 bar_type: *bar_type,
320 }) {
321 log::error!("Failed to restore bar type {key}: {e}");
322 }
323 }
324
325 if let Some(account_id) = self.account_id
326 && let Err(e) = cmd_tx.send(HandlerCommand::SetAccountId(account_id))
327 {
328 log::error!("Failed to restore account_id: {e}");
329 }
330
331 self.prime_default_subscriptions();
332
333 resubscribe_all(
335 &self.subscriptions,
336 &self.credential,
337 &cmd_tx,
338 Some(&out_tx),
339 );
340
341 let signal = Arc::clone(&self.signal);
342 let subscriptions = self.subscriptions.clone();
343 let credential = self.credential.clone();
344 let cmd_tx_reconnect = cmd_tx.clone();
345 let aliases_for_handler = Arc::clone(&self.subscription_aliases);
346
347 if let Err(e) = self.task_handle.spawn(async move {
348 let mut handler = FeedHandler::new(signal, cmd_rx, raw_rx, aliases_for_handler);
349
350 loop {
351 match handler.next().await {
352 Some(NautilusWsMessage::Reconnected) => {
353 subscriptions.reset_after_reconnect();
354 resubscribe_all(
355 &subscriptions,
356 &credential,
357 &cmd_tx_reconnect,
358 Some(&out_tx),
359 );
360
361 if let Err(e) = out_tx.send(NautilusWsMessage::Reconnected) {
362 log::debug!("Output channel closed: {e}");
363 break;
364 }
365 }
366 Some(msg) => {
367 if let Err(e) = out_tx.send(msg) {
368 log::debug!("Output channel closed: {e}");
369 break;
370 }
371 }
372 None => {
373 log::debug!("Feed handler stopped");
374 break;
375 }
376 }
377 }
378 }) {
379 self.out_rx = None;
380 anyhow::bail!("Failed to start Coinbase WebSocket handler task: {e}");
381 }
382
383 Ok(())
384 }
385
386 pub async fn subscribe(
388 &self,
389 channel: CoinbaseWsChannel,
390 product_ids: &[Ustr],
391 ) -> anyhow::Result<()> {
392 let jwt = if channel.requires_auth() {
393 let credential = self
394 .credential
395 .as_ref()
396 .ok_or_else(|| anyhow::anyhow!("Credentials required for {channel}"))?;
397 Some(credential.build_ws_jwt()?)
398 } else {
399 self.credential.as_ref().and_then(|c| c.build_ws_jwt().ok())
400 };
401
402 let sub = protect_subscription(CoinbaseWsSubscription {
403 msg_type: CoinbaseWsAction::Subscribe,
404 product_ids: product_ids.to_vec(),
405 channel,
406 jwt,
407 })?;
408
409 let channel_str = channel.as_ref();
410
411 if product_ids.is_empty() {
412 self.subscriptions.mark_subscribe(channel_str);
413 } else {
414 for product_id in product_ids {
415 let topic = format!("{channel_str}|{product_id}");
416 self.subscriptions.mark_subscribe(&topic);
417 }
418 }
419
420 let cmd_tx = self.cmd_tx.read().await;
421 cmd_tx
422 .send(HandlerCommand::Subscribe {
423 channel: sub.channel,
424 product_ids: sub.product_ids,
425 payload: sub.payload,
426 })
427 .map_err(|e| anyhow::anyhow!("Failed to send Subscribe command: {e}"))
428 }
429
430 pub async fn unsubscribe(
432 &self,
433 channel: CoinbaseWsChannel,
434 product_ids: &[Ustr],
435 ) -> anyhow::Result<()> {
436 let jwt = self.credential.as_ref().and_then(|c| c.build_ws_jwt().ok());
437
438 let unsub = protect_subscription(CoinbaseWsSubscription {
439 msg_type: CoinbaseWsAction::Unsubscribe,
440 product_ids: product_ids.to_vec(),
441 channel,
442 jwt,
443 })?;
444
445 let channel_str = channel.as_ref();
446
447 if product_ids.is_empty() {
448 self.subscriptions.mark_unsubscribe(channel_str);
449 } else {
450 for product_id in product_ids {
451 let topic = format!("{channel_str}|{product_id}");
452 self.subscriptions.mark_unsubscribe(&topic);
453 }
454 }
455
456 let cmd_tx = self.cmd_tx.read().await;
457 cmd_tx
458 .send(HandlerCommand::Unsubscribe {
459 channel: unsub.channel,
460 product_ids: unsub.product_ids,
461 payload: unsub.payload,
462 })
463 .map_err(|e| anyhow::anyhow!("Failed to send Unsubscribe command: {e}"))
464 }
465
466 pub async fn next_message(&mut self) -> Option<NautilusWsMessage> {
468 self.out_rx.as_mut()?.recv().await
469 }
470
471 pub(crate) fn begin_shutdown(&self) {
473 self.signal.store(true, Ordering::Release);
474 }
475
476 pub async fn disconnect(&mut self) -> anyhow::Result<()> {
478 let cmd_tx = self.cmd_tx.read().await;
481
482 if let Err(e) = cmd_tx.send(HandlerCommand::Disconnect) {
483 log::debug!("Failed to send Disconnect command: {e}");
484 }
485 drop(cmd_tx);
486
487 self.begin_shutdown();
490
491 if let Some(outcome) = finish_task(
492 &mut self.task_handle,
493 WS_DISCONNECT_TIMEOUT,
494 WS_DISCONNECT_TIMEOUT,
495 )
496 .await
497 {
498 match outcome {
499 TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted => {}
500 TaskJoinOutcome::Failed(error) => {
501 self.shutdown_errors
502 .push(format!("Coinbase WebSocket handler failed: {error}"));
503 }
504 TaskJoinOutcome::Incomplete => {
505 self.shutdown_errors
506 .push("Coinbase WebSocket handler did not stop after abort".to_string());
507 }
508 }
509 }
510
511 let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
516
517 loop {
518 let mode_ptr = self.connection_mode.load();
519
520 if ConnectionMode::from_u8(mode_ptr.load(Ordering::Relaxed)).is_closed() {
521 break;
522 }
523
524 if tokio::time::Instant::now() >= deadline {
525 self.shutdown_errors
526 .push("Timed out waiting for WebSocket to reach Closed state".to_string());
527 break;
528 }
529
530 tokio::time::sleep(Duration::from_millis(20)).await;
531 }
532
533 if let Some(control) = &self.socket_control {
534 control.deregister();
535 }
536
537 if self.shutdown_errors.is_empty() {
538 Ok(())
539 } else {
540 let errors = std::mem::take(&mut self.shutdown_errors);
541 anyhow::bail!(errors.join("; "))
542 }
543 }
544
545 #[must_use]
547 pub fn is_active(&self) -> bool {
548 let mode_ptr = self.connection_mode.load();
549 let mode_val = mode_ptr.load(Ordering::Relaxed);
550 ConnectionMode::from_u8(mode_val).is_active()
551 }
552
553 #[must_use]
555 pub fn is_reconnecting(&self) -> bool {
556 let mode_ptr = self.connection_mode.load();
557 let mode_val = mode_ptr.load(Ordering::Relaxed);
558 ConnectionMode::from_u8(mode_val).is_reconnect()
559 }
560
561 #[must_use]
563 pub fn instruments(&self) -> &Arc<AtomicMap<InstrumentId, InstrumentAny>> {
564 &self.instruments
565 }
566
567 #[must_use]
569 pub fn subscription_aliases(&self) -> &Arc<AtomicMap<Ustr, Ustr>> {
570 &self.subscription_aliases
571 }
572
573 pub fn register_subscription_alias(&self, canonical: Ustr, subscribed: Ustr) {
577 self.subscription_aliases.insert(canonical, subscribed);
578 }
579
580 pub fn unregister_subscription_alias(&self, canonical: &Ustr) {
582 self.subscription_aliases.remove(canonical);
583 }
584
585 #[must_use]
587 pub fn subscriptions(&self) -> &SubscriptionState {
588 &self.subscriptions
589 }
590
591 pub async fn update_instrument(&self, instrument: InstrumentAny) {
593 let id = instrument.id();
594 self.instruments.insert(id, instrument.clone());
595
596 let cmd_tx = self.cmd_tx.read().await;
597
598 if let Err(e) = cmd_tx.send(HandlerCommand::UpdateInstrument(Box::new(instrument))) {
599 log::debug!("Failed to send UpdateInstrument: {e}");
600 }
601 }
602
603 pub fn take_out_rx(
607 &mut self,
608 ) -> Option<tokio::sync::mpsc::UnboundedReceiver<NautilusWsMessage>> {
609 self.out_rx.take()
610 }
611
612 pub fn register_bar_type(&mut self, key: String, bar_type: BarType) {
617 self.bar_types.insert(key, bar_type);
618 }
619
620 pub async fn add_bar_type(&mut self, key: String, bar_type: BarType) {
622 self.bar_types.insert(key.clone(), bar_type);
623
624 let cmd_tx = self.cmd_tx.read().await;
625
626 if let Err(e) = cmd_tx.send(HandlerCommand::AddBarType { key, bar_type }) {
627 log::debug!("Failed to send AddBarType: {e}");
628 }
629 }
630}
631
632impl Drop for CoinbaseWebSocketClient {
633 fn drop(&mut self) {
634 if let Some(handle) = self.task_handle.as_ref() {
635 self.signal.store(true, Ordering::Release);
636 handle.abort();
637 }
638 }
639}
640
641struct ProtectedSubscription {
642 channel: CoinbaseWsChannel,
643 product_ids: Vec<Ustr>,
644 payload: SecretString,
645}
646
647fn protect_subscription(
648 mut subscription: CoinbaseWsSubscription,
649) -> Result<ProtectedSubscription, serde_json::Error> {
650 let payload = SecretString::from(serde_json::to_string(&subscription)?);
651
652 Ok(ProtectedSubscription {
653 channel: subscription.channel,
654 product_ids: std::mem::take(&mut subscription.product_ids),
655 payload,
656 })
657}
658
659fn resubscribe_all(
660 subscriptions: &SubscriptionState,
661 credential: &Option<CoinbaseCredential>,
662 cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
663 out_tx: Option<&tokio::sync::mpsc::UnboundedSender<NautilusWsMessage>>,
664) {
665 let topics = subscriptions.all_topics();
666
667 if topics.is_empty() {
668 log::debug!("No active subscriptions to restore");
669 return;
670 }
671
672 log::info!(
673 "Resubscribing to {} topics after reconnection",
674 topics.len()
675 );
676
677 for topic in topics {
678 let (channel, product_id) = match topic.split_once('|') {
679 Some((ch, pid)) => (ch, Some(pid)),
680 None => (topic.as_str(), None),
681 };
682
683 let channel_enum = match CoinbaseWsChannel::from_str(channel) {
684 Ok(ch) => ch,
685 Err(_) => {
686 log::warn!("Unknown channel in topic: {topic}");
687 continue;
688 }
689 };
690
691 let jwt = match credential.as_ref() {
692 Some(c) => match c.build_ws_jwt() {
693 Ok(token) => Some(token),
694 Err(e) => {
695 if channel_enum.requires_auth() {
696 let msg = format!(
697 "JWT required for {channel} but build failed: {e}; topic {topic} not restored"
698 );
699 log::error!("{msg}");
700 if let Some(tx) = out_tx {
701 let _ = tx.send(NautilusWsMessage::Error(msg));
702 }
703 continue;
704 }
705 None
706 }
707 },
708 None => {
709 if channel_enum.requires_auth() {
710 let msg = format!(
711 "JWT required for {channel} but no credentials configured; topic {topic} not restored"
712 );
713 log::error!("{msg}");
714 if let Some(tx) = out_tx {
715 let _ = tx.send(NautilusWsMessage::Error(msg));
716 }
717 continue;
718 }
719 None
720 }
721 };
722
723 let product_ids = match product_id {
724 Some(pid) => vec![Ustr::from(pid)],
725 None => vec![],
726 };
727
728 let sub = CoinbaseWsSubscription {
729 msg_type: CoinbaseWsAction::Subscribe,
730 product_ids,
731 channel: channel_enum,
732 jwt,
733 };
734
735 let sub = match protect_subscription(sub) {
736 Ok(sub) => sub,
737 Err(e) => {
738 log::error!("Failed to serialize subscription for {topic}: {e}");
739 continue;
740 }
741 };
742
743 if let Err(e) = cmd_tx.send(HandlerCommand::Subscribe {
744 channel: sub.channel,
745 product_ids: sub.product_ids,
746 payload: sub.payload,
747 }) {
748 log::error!("Failed to resubscribe {topic}: {e}");
749 }
750 }
751}
752
753#[cfg(test)]
754mod tests {
755 use nautilus_network::websocket::SubscriptionState;
756 use rstest::rstest;
757
758 use super::*;
759
760 #[rstest]
761 fn test_debug_redacts_proxy_url() {
762 let proxy_url = "http://user:password@proxy.example:8080";
763 let client = CoinbaseWebSocketClient::new(
764 "wss://test",
765 TransportBackend::default(),
766 Some(proxy_url.to_string()),
767 );
768
769 let debug = format!("{client:?}");
770
771 assert!(debug.contains(REDACTED));
772 assert!(!debug.contains(proxy_url));
773 }
774
775 #[rstest]
776 fn test_drop_clone_does_not_signal_handler() {
777 let client = CoinbaseWebSocketClient::new("wss://test", TransportBackend::default(), None);
778 let clone = client.clone();
779
780 drop(clone);
781
782 assert!(!client.signal.load(Ordering::Acquire));
783 }
784
785 #[rstest]
786 fn test_resubscribe_all_product_level_topic() {
787 let subs = SubscriptionState::new('|');
788 subs.mark_subscribe("level2|BTC-USD");
789
790 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
791 resubscribe_all(&subs, &None, &tx, None);
792
793 let cmd = rx.try_recv().unwrap();
794
795 match cmd {
796 HandlerCommand::Subscribe {
797 channel,
798 product_ids,
799 payload,
800 } => {
801 let payload: serde_json::Value =
802 serde_json::from_str(payload.expose_secret()).unwrap();
803 assert_eq!(channel, CoinbaseWsChannel::Level2);
804 assert_eq!(product_ids, vec![Ustr::from("BTC-USD")]);
805 assert!(payload.get("jwt").is_none());
806 }
807 other => panic!("Expected Subscribe, was {other:?}"),
808 }
809 }
810
811 #[rstest]
812 fn test_resubscribe_all_channel_level_topic() {
813 let subs = SubscriptionState::new('|');
814 subs.mark_subscribe("heartbeats");
815
816 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
817 resubscribe_all(&subs, &None, &tx, None);
818
819 let cmd = rx.try_recv().unwrap();
820
821 match cmd {
822 HandlerCommand::Subscribe {
823 channel,
824 product_ids,
825 ..
826 } => {
827 assert_eq!(channel, CoinbaseWsChannel::Heartbeats);
828 assert!(product_ids.is_empty());
829 }
830 other => panic!("Expected Subscribe, was {other:?}"),
831 }
832 }
833
834 #[rstest]
835 fn test_resubscribe_all_multiple_topics() {
836 let subs = SubscriptionState::new('|');
837 subs.mark_subscribe("market_trades|BTC-USD");
838 subs.mark_subscribe("ticker|ETH-USD");
839
840 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
841 resubscribe_all(&subs, &None, &tx, None);
842
843 let cmd1 = rx.try_recv().unwrap();
844 let cmd2 = rx.try_recv().unwrap();
845
846 assert!(matches!(cmd1, HandlerCommand::Subscribe { .. }));
847 assert!(matches!(cmd2, HandlerCommand::Subscribe { .. }));
848 assert!(rx.try_recv().is_err());
849 }
850
851 #[rstest]
852 fn test_resubscribe_all_empty_subscriptions() {
853 let subs = SubscriptionState::new('|');
854
855 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
856 resubscribe_all(&subs, &None, &tx, None);
857
858 assert!(rx.try_recv().is_err());
859 }
860
861 #[rstest]
862 fn test_resubscribe_all_unknown_channel_skipped() {
863 let subs = SubscriptionState::new('|');
864 subs.mark_subscribe("nonexistent_channel|BTC-USD");
865
866 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
867 resubscribe_all(&subs, &None, &tx, None);
868
869 assert!(rx.try_recv().is_err());
870 }
871
872 #[rstest]
873 #[case("level2|BTC-USD", CoinbaseWsChannel::Level2)]
874 #[case("market_trades|ETH-USD", CoinbaseWsChannel::MarketTrades)]
875 #[case("ticker|BTC-USD", CoinbaseWsChannel::Ticker)]
876 #[case("ticker_batch|BTC-USD", CoinbaseWsChannel::TickerBatch)]
877 #[case("candles|BTC-USD", CoinbaseWsChannel::Candles)]
878 #[case("heartbeats", CoinbaseWsChannel::Heartbeats)]
879 #[case("status", CoinbaseWsChannel::Status)]
880 fn test_resubscribe_all_channel_mapping(
881 #[case] topic: &str,
882 #[case] expected_channel: CoinbaseWsChannel,
883 ) {
884 let subs = SubscriptionState::new('|');
885 subs.mark_subscribe(topic);
886
887 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
888 resubscribe_all(&subs, &None, &tx, None);
889
890 let cmd = rx.try_recv().unwrap();
891
892 match cmd {
893 HandlerCommand::Subscribe { channel, .. } => {
894 assert_eq!(channel, expected_channel);
895 }
896 other => panic!("Expected Subscribe, was {other:?}"),
897 }
898 }
899
900 #[rstest]
901 fn test_protected_subscription_preserves_wire_jwt_and_redacts_command_debug() {
902 let jwt = "jwt-secret-value";
903 let subscription = CoinbaseWsSubscription {
904 msg_type: CoinbaseWsAction::Subscribe,
905 product_ids: vec![Ustr::from("BTC-USD")],
906 channel: CoinbaseWsChannel::User,
907 jwt: Some(SecretString::from(jwt)),
908 };
909 let subscription_debug = format!("{subscription:?}");
910 let protected = protect_subscription(subscription).unwrap();
911 let payload: serde_json::Value =
912 serde_json::from_str(protected.payload.expose_secret()).unwrap();
913 let command = HandlerCommand::Subscribe {
914 channel: protected.channel,
915 product_ids: protected.product_ids,
916 payload: protected.payload,
917 };
918 let debug = format!("{command:?}");
919
920 assert_eq!(payload["jwt"], jwt);
921 assert!(subscription_debug.contains(REDACTED));
922 assert!(!subscription_debug.contains(jwt));
923 assert!(!debug.contains(jwt));
924 }
925
926 #[rstest]
927 #[case("user|BTC-USD")]
928 #[case("futures_balance_summary")]
929 fn test_resubscribe_all_auth_channel_skipped_without_credentials(#[case] topic: &str) {
930 let subs = SubscriptionState::new('|');
931 subs.mark_subscribe(topic);
932
933 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
934 resubscribe_all(&subs, &None, &tx, None);
935
936 assert!(rx.try_recv().is_err());
938 }
939
940 #[rstest]
941 #[case("user|BTC-USD", "user")]
942 #[case("futures_balance_summary", "futures_balance_summary")]
943 fn test_resubscribe_all_emits_error_for_auth_channel_without_credentials(
944 #[case] topic: &str,
945 #[case] channel: &str,
946 ) {
947 let subs = SubscriptionState::new('|');
948 subs.mark_subscribe(topic);
949
950 let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
951 let (out_tx, mut out_rx) = tokio::sync::mpsc::unbounded_channel();
952 resubscribe_all(&subs, &None, &cmd_tx, Some(&out_tx));
953
954 assert!(cmd_rx.try_recv().is_err());
956
957 let msg = out_rx
958 .try_recv()
959 .expect("Error event must be emitted when auth channel cannot resubscribe");
960 match msg {
961 NautilusWsMessage::Error(text) => {
962 assert!(
963 text.contains(channel),
964 "error must mention the channel, was: {text}"
965 );
966 assert!(
967 text.contains(topic),
968 "error must mention the topic, was: {text}"
969 );
970 }
971 other => panic!("expected Error variant, was {other:?}"),
972 }
973 }
974
975 #[rstest]
976 fn test_resubscribe_all_emits_error_when_jwt_build_fails() {
977 let subs = SubscriptionState::new('|');
978 let topic = "user|BTC-USD";
979 subs.mark_subscribe(topic);
980
981 let bad_credential = Some(CoinbaseCredential::new(
984 "organizations/test/apiKeys/test".to_string(),
985 "not-a-pem-key".to_string(),
986 ));
987
988 let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
989 let (out_tx, mut out_rx) = tokio::sync::mpsc::unbounded_channel();
990 resubscribe_all(&subs, &bad_credential, &cmd_tx, Some(&out_tx));
991
992 assert!(cmd_rx.try_recv().is_err(), "no subscribe should be sent");
993 let msg = out_rx
994 .try_recv()
995 .expect("Error event must be emitted when JWT build fails for an auth channel");
996 match msg {
997 NautilusWsMessage::Error(text) => {
998 assert!(text.contains("user"), "error must mention channel: {text}");
999 assert!(text.contains(topic), "error must mention topic: {text}");
1000 }
1001 other => panic!("expected Error variant, was {other:?}"),
1002 }
1003 }
1004
1005 #[rstest]
1006 fn test_prime_default_subscriptions_marks_heartbeats() {
1007 let client = CoinbaseWebSocketClient::new("wss://test", TransportBackend::default(), None);
1008 assert!(client.subscriptions.all_topics().is_empty());
1009
1010 client.prime_default_subscriptions();
1011
1012 let topics = client.subscriptions.all_topics();
1013 assert!(topics.iter().any(|t| t == "heartbeats"), "{topics:?}");
1014 }
1015
1016 #[rstest]
1017 fn test_ws_quotas_match_documented_limits() {
1018 assert_eq!(COINBASE_WS_CONNECTION_QUOTA.burst_size().get(), 8);
1019 assert_eq!(COINBASE_WS_SUBSCRIPTION_QUOTA.burst_size().get(), 8);
1020 }
1021
1022 #[rstest]
1023 fn test_ws_subscription_rate_limit_key_is_stable() {
1024 assert_eq!(COINBASE_RATE_LIMIT_KEY_SUBSCRIPTION, "subscription");
1025 assert_eq!(
1026 COINBASE_WS_SUBSCRIPTION_KEYS[0],
1027 COINBASE_RATE_LIMIT_KEY_SUBSCRIPTION,
1028 );
1029 }
1030}