1use std::{
57 collections::VecDeque,
58 fmt::Debug,
59 future::Future,
60 pin::pin,
61 sync::{
62 Arc, OnceLock,
63 atomic::{AtomicBool, AtomicU8, AtomicU32, AtomicU64, Ordering},
64 },
65 time::Duration,
66};
67
68use futures_util::{SinkExt, StreamExt};
69use http::HeaderName;
70use nautilus_core::string::secret::REDACTED;
71use nautilus_cryptography::providers::install_cryptographic_provider;
72use parking_lot::RwLock;
73#[cfg(any(feature = "turmoil", feature = "transport-sockudo"))]
74use rustls::ClientConfig;
75#[cfg(feature = "transport-sockudo")]
76use sockudo_ws::{
77 Config as SockudoConfig, Http1, Role, Stream as SockudoStream,
78 WebSocketStream as SockudoWebSocketStream,
79};
80#[cfg(feature = "transport-sockudo")]
81use tokio::io::{AsyncRead, AsyncWrite};
82#[cfg(any(feature = "turmoil", feature = "transport-sockudo"))]
83use tokio_rustls::TlsConnector;
84#[cfg(feature = "turmoil")]
85use tokio_tungstenite::MaybeTlsStream;
86#[cfg(feature = "turmoil")]
87use tokio_tungstenite::client_async_with_config;
88#[cfg(all(not(feature = "turmoil"), not(all(feature = "simulation", madsim))))]
89use tokio_tungstenite::connect_async_with_config;
90use tokio_tungstenite::tungstenite::{
91 client::IntoClientRequest, handshake::client::Request, http::HeaderValue,
92};
93use tokio_util::sync::CancellationToken;
94use ustr::Ustr;
95
96#[cfg(not(feature = "turmoil"))]
97use super::proxy::{ProxyKind, WsTarget, tunnel_via_proxy};
98use super::{
99 auth::{AuthState, AuthTracker},
100 config::{InitialConnectRetryPolicy, TransportBackend, WebSocketConfig},
101 consts::{
102 CONNECTION_STATE_CHECK_INTERVAL_MS, GRACEFUL_SHUTDOWN_DELAY_MS,
103 GRACEFUL_SHUTDOWN_TIMEOUT_SECS,
104 },
105 types::{
106 CloseHandler, EpochMessageHandler, EpochPingHandler, MessageHandler, MessageReader,
107 MessageWriter, PingHandler, WriterCommand,
108 },
109};
110#[cfg(feature = "turmoil")]
111use crate::net::TcpConnector;
112#[cfg(feature = "transport-sockudo")]
113use crate::net::TcpStream;
114#[cfg(feature = "transport-sockudo")]
115use crate::transport::sockudo::{
116 PrefixedIo, SockudoTransport, client_handshake_with_headers, validate_extra_headers,
117};
118use crate::{
119 RECONNECTED, SocketState, SocketStateSink,
120 backoff::{
121 ExponentialBackoff, RECONNECT_STABILITY_THRESHOLD, ReconnectThrottle, wait_reconnect_delay,
122 },
123 dst,
124 error::{SendError, is_connection_drop_io_error},
125 logging::{log_task_aborted, log_task_started, log_task_stopped},
126 mode::{
127 ConnectionMode, ControllerLifecycle, ReadSessionFence, ReconnectOutcome,
128 ReconnectRequestOutcome,
129 },
130 ratelimiter::{RateLimiter, clock::MonotonicClock, quota::Quota},
131 retry::{RetryConfig, RetryError, RetryManager},
132 transport::{
133 BoxedWsTransport, CloseFrame, Message, TransportError, error::retryable_status,
134 tungstenite::TungsteniteTransport,
135 },
136 writer::{DEFAULT_WRITER_CAPACITY, WriterReceiver, WriterSender},
137};
138
139const WRITE_TIMEOUT_SECS: u64 = 5;
140const CONTROLLER_FALLBACK_INTERVAL_MS: u64 = 100;
141const INITIAL_CONNECT_OPERATION: &str = "WebSocket initial connection";
142
143const MAX_CONTROL_FRAME_PAYLOAD_BYTES: usize = 125;
145
146pub struct WebSocketClientInner {
171 config: WebSocketConfig,
172 reconnect_headers: ReconnectHeaders,
173 handler: Option<IncomingHandler>,
174 ping_handler: Option<IncomingPingHandler>,
175 close_handler: Option<CloseHandler>,
176 read_task: Option<dst::task::JoinHandle<()>>,
177 read_fence: Option<ReadSessionFence>,
178 write_task: dst::task::JoinHandle<()>,
179 writer_tx: WriterSender<WriterCommand>,
180 heartbeat_task: Option<dst::task::JoinHandle<()>>,
181 connection_mode: Arc<AtomicU8>,
182 connection_epoch: Arc<AtomicU64>,
183 state_notify: Arc<tokio::sync::Notify>,
184 controller_notify: Arc<tokio::sync::Notify>,
185 reconnect_published: Arc<AtomicBool>,
186 connect_timeout: Duration,
187 heartbeat_timeout: Option<Duration>,
188 backoff: ExponentialBackoff,
189 reconnect_throttle: ReconnectThrottle,
190 reconnect_max_attempts: Option<u32>,
191 reconnection_attempt_count: u32,
192 auth_tracker: Arc<OnceLock<AuthTracker>>,
193 reconnect_buffer_waits_for_auth: Arc<AtomicBool>,
194 state_sink: Option<SocketStateSink>,
195 connection_rate_limit: Option<ConnectionRateLimit>,
196}
197
198#[derive(Clone, Debug)]
199struct ConnectionRateLimit {
200 limiter: Arc<RateLimiter<Ustr, MonotonicClock>>,
201 keys: Arc<[Ustr]>,
202}
203
204#[derive(Default)]
205struct InitialConnectOptions {
206 retry_policy: Option<InitialConnectRetryPolicy>,
207 cancellation_token: Option<CancellationToken>,
208}
209
210impl WebSocketClientInner {
211 #[allow(unknown_lints, reason = "Clippy lint is unavailable on Rust 1.97")]
219 #[expect(
220 clippy::unused_async,
221 clippy::unused_async_trait_impl,
222 reason = "async signature for consistency with connect-based constructors"
223 )]
224 pub async fn new_with_writer(
225 config: WebSocketConfig,
226 writer: MessageWriter,
227 ) -> Result<Self, TransportError> {
228 Self::new_with_writer_and_state_sink(config, writer, None)
229 }
230
231 fn new_with_writer_and_state_sink(
232 mut config: WebSocketConfig,
233 writer: MessageWriter,
234 state_sink: Option<SocketStateSink>,
235 ) -> Result<Self, TransportError> {
236 install_cryptographic_provider();
237
238 crate::writer::validate_capacity(config.writer_capacity).map_err(|e| {
239 TransportError::Io(std::io::Error::new(std::io::ErrorKind::InvalidInput, e))
240 })?;
241
242 if config.heartbeat_interval_secs == Some(0) {
243 return Err(TransportError::Io(std::io::Error::new(
244 std::io::ErrorKind::InvalidInput,
245 "Heartbeat interval cannot be zero",
246 )));
247 }
248
249 let connection_mode = Arc::new(AtomicU8::new(ConnectionMode::Reconnect.as_u8()));
250 let connection_epoch = Arc::new(AtomicU64::new(0));
251 let state_notify = Arc::new(tokio::sync::Notify::new());
252 let controller_notify = Arc::new(tokio::sync::Notify::new());
253 let reconnect_published = Arc::new(AtomicBool::new(true));
254 let outcome =
255 ConnectionMode::complete_reconnect_with_sink(&connection_mode, state_sink.as_ref());
256 debug_assert_eq!(outcome, ReconnectOutcome::Reconnected);
257
258 let read_task = None;
260 let read_fence = None;
261
262 let backoff = ExponentialBackoff::new(
264 Duration::from_secs(2),
265 Duration::from_secs(30),
266 1.5,
267 100,
268 true,
269 )
270 .map_err(|e| {
271 TransportError::Io(std::io::Error::new(std::io::ErrorKind::InvalidInput, e))
272 })?;
273
274 let auth_tracker = Arc::new(OnceLock::new());
275 let reconnect_buffer_waits_for_auth = Arc::new(AtomicBool::new(false));
276
277 let (writer_tx, writer_rx) = crate::writer::channel::<WriterCommand>(
278 config.writer_capacity.unwrap_or(DEFAULT_WRITER_CAPACITY),
279 );
280 let write_task = Self::spawn_write_task(
281 Arc::clone(&connection_mode),
282 Arc::clone(&controller_notify),
283 Arc::clone(&reconnect_published),
284 writer,
285 writer_rx,
286 Arc::clone(&connection_epoch),
287 Arc::clone(&auth_tracker),
288 Arc::clone(&reconnect_buffer_waits_for_auth),
289 state_sink.clone(),
290 );
291
292 let heartbeat_task = config.heartbeat_interval_secs.map(|heartbeat_interval| {
293 Self::spawn_heartbeat_task(
294 Arc::clone(&connection_mode),
295 heartbeat_interval,
296 config.heartbeat_payload.clone(),
297 writer_tx.clone(),
298 )
299 });
300
301 let reconnect_max_attempts = None; let connect_timeout = Duration::from_secs(10);
303
304 let reconnect_headers = ReconnectHeaders::new(std::mem::take(&mut config.headers));
305
306 Ok(Self {
307 config,
308 reconnect_headers,
309 handler: None, ping_handler: None,
311 close_handler: None,
312 writer_tx,
313 connection_mode,
314 connection_epoch,
315 state_notify,
316 controller_notify,
317 reconnect_published,
318 connect_timeout,
319 heartbeat_timeout: None,
320 heartbeat_task,
321 read_task,
322 read_fence,
323 write_task,
324 backoff,
325 reconnect_throttle: ReconnectThrottle::default(),
326 reconnect_max_attempts,
327 reconnection_attempt_count: 0,
328 auth_tracker,
329 reconnect_buffer_waits_for_auth,
330 state_sink,
331 connection_rate_limit: None,
332 })
333 }
334
335 pub async fn connect_url(
344 config: WebSocketConfig,
345 message_handler: Option<MessageHandler>,
346 ping_handler: Option<PingHandler>,
347 ) -> Result<Self, TransportError> {
348 Self::connect_url_with_handler(
349 config,
350 message_handler.map(IncomingHandler::Message),
351 ping_handler.map(IncomingPingHandler::Ping),
352 None,
353 None,
354 None,
355 InitialConnectOptions::default(),
356 )
357 .await
358 }
359
360 async fn connect_url_with_handler(
361 config: WebSocketConfig,
362 handler: Option<IncomingHandler>,
363 ping_handler: Option<IncomingPingHandler>,
364 close_handler: Option<CloseHandler>,
365 state_sink: Option<SocketStateSink>,
366 connection_rate_limit: Option<ConnectionRateLimit>,
367 initial_connect_options: InitialConnectOptions,
368 ) -> Result<Self, TransportError> {
369 install_cryptographic_provider();
370
371 let is_stream_mode = handler.is_none();
372
373 if is_stream_mode {
377 crate::writer::validate_capacity(config.writer_capacity).map_err(|e| {
378 TransportError::Io(std::io::Error::new(std::io::ErrorKind::InvalidInput, e))
379 })?;
380
381 if config.heartbeat_interval_secs == Some(0) {
382 return Err(TransportError::Io(std::io::Error::new(
383 std::io::ErrorKind::InvalidInput,
384 "Heartbeat interval cannot be zero",
385 )));
386 }
387
388 validate_honored_inbound_limits(&config)?;
389 } else {
390 config.validate().map_err(|e| {
391 TransportError::Io(std::io::Error::new(std::io::ErrorKind::InvalidInput, e))
392 })?;
393 }
394
395 let heartbeat_timeout = config.resolved_heartbeat_timeout().map(Duration::from_secs);
396 let reconnect_max_attempts = config.reconnect_max_attempts;
397
398 let connect_timeout = if is_stream_mode {
400 Duration::from_secs(10)
401 } else {
402 Duration::from_millis(config.connect_timeout_ms.unwrap_or(10_000))
403 };
404
405 let backoff = ExponentialBackoff::new(
406 Duration::from_millis(config.reconnect_delay_initial_ms.unwrap_or(2_000)),
407 Duration::from_millis(config.reconnect_delay_max_ms.unwrap_or(30_000)),
408 config.reconnect_backoff_factor.unwrap_or(1.5),
409 config.reconnect_jitter_ms.unwrap_or(100),
410 true, )
412 .map_err(|e| {
413 TransportError::Io(std::io::Error::new(std::io::ErrorKind::InvalidInput, e))
414 })?;
415
416 let cancellation_token = initial_connect_options
417 .cancellation_token
418 .unwrap_or_default();
419 let retry_policy = initial_connect_options.retry_policy;
420 let max_attempts = retry_policy
421 .as_ref()
422 .map_or(1, |policy| policy.max_attempts.get());
423 let retry_manager = RetryManager::new(initial_connect_retry_config(retry_policy.as_ref())?);
424 let attempt = AtomicU32::new(0);
425 let operation = || {
426 attempt.fetch_add(1, Ordering::Relaxed);
427
428 await_initial_connect_attempt(&cancellation_token, async {
429 if let Some(rate_limit) = &connection_rate_limit {
430 rate_limit
431 .limiter
432 .await_keys_ready(Some(&rate_limit.keys))
433 .await;
434 }
435
436 dst::time::timeout(
438 connect_timeout,
439 Box::pin(Self::connect_with_server(
440 &config.url,
441 config.headers.clone(),
442 config.backend,
443 config.proxy_url.as_deref(),
444 config.max_message_size_bytes,
445 config.max_frame_size_bytes,
446 )),
447 )
448 .await
449 .unwrap_or_else(|_| {
450 Err(TransportError::Io(std::io::Error::new(
451 std::io::ErrorKind::TimedOut,
452 format!(
453 "connection timed out after {}s",
454 connect_timeout.as_secs_f64()
455 ),
456 )))
457 })
458 })
459 };
460
461 let classify = |error: &TransportError| {
462 let retryable = is_retryable_initial_connect_error(error);
463 let attempt = attempt.load(Ordering::Relaxed);
464 if retryable && attempt < max_attempts && !cancellation_token.is_cancelled() {
465 log::warn!(
466 "WebSocket connection attempt {attempt}/{max_attempts} to {REDACTED} failed: {}",
467 error.redacted_message()
468 );
469 }
470
471 retryable
472 };
473
474 let transport = retry_manager
475 .invocation(
476 INITIAL_CONNECT_OPERATION,
477 operation,
478 classify,
479 initial_connect_retry_error,
480 )
481 .cancellation_token(&cancellation_token)
482 .execute()
483 .await?;
484
485 let attempt = attempt.load(Ordering::Relaxed);
486 if attempt > 1 {
487 log::info!("WebSocket connection established after {attempt} attempts");
488 }
489
490 let (writer, reader) = transport;
491 let mut config = config;
492 let reconnect_headers = ReconnectHeaders::new(std::mem::take(&mut config.headers));
493
494 let connection_mode = Arc::new(AtomicU8::new(ConnectionMode::Reconnect.as_u8()));
495 let connection_epoch = Arc::new(AtomicU64::new(0));
496 let state_notify = Arc::new(tokio::sync::Notify::new());
497 let controller_notify = Arc::new(tokio::sync::Notify::new());
498 let reconnect_published = Arc::new(AtomicBool::new(true));
499 let outcome =
500 ConnectionMode::complete_reconnect_with_sink(&connection_mode, state_sink.as_ref());
501 debug_assert_eq!(outcome, ReconnectOutcome::Reconnected);
502
503 let (read_task, read_fence) = if is_stream_mode {
504 (None, None)
505 } else {
506 let read_fence = ReadSessionFence::new();
507 let read_task = Self::spawn_message_handler_task(
508 Arc::clone(&connection_mode),
509 Arc::clone(&state_notify),
510 read_fence.clone(),
511 reader,
512 0,
513 handler.as_ref(),
514 ping_handler.as_ref(),
515 close_handler.as_ref(),
516 state_sink.as_ref(),
517 config.idle_timeout_ms,
518 heartbeat_timeout,
519 );
520 (Some(read_task), Some(read_fence))
521 };
522
523 let auth_tracker = Arc::new(OnceLock::new());
524 let reconnect_buffer_waits_for_auth = Arc::new(AtomicBool::new(false));
525
526 let (writer_tx, writer_rx) = crate::writer::channel::<WriterCommand>(
527 config.writer_capacity.unwrap_or(DEFAULT_WRITER_CAPACITY),
528 );
529 let write_task = Self::spawn_write_task(
530 Arc::clone(&connection_mode),
531 Arc::clone(&controller_notify),
532 Arc::clone(&reconnect_published),
533 writer,
534 writer_rx,
535 Arc::clone(&connection_epoch),
536 Arc::clone(&auth_tracker),
537 Arc::clone(&reconnect_buffer_waits_for_auth),
538 state_sink.clone(),
539 );
540
541 let heartbeat_task = config.heartbeat_interval_secs.map(|heartbeat_secs| {
543 Self::spawn_heartbeat_task(
544 Arc::clone(&connection_mode),
545 heartbeat_secs,
546 config.heartbeat_payload.clone(),
547 writer_tx.clone(),
548 )
549 });
550
551 Ok(Self {
552 config,
553 reconnect_headers,
554 handler,
555 ping_handler,
556 close_handler,
557 read_task,
558 read_fence,
559 write_task,
560 writer_tx,
561 heartbeat_task,
562 connection_mode,
563 connection_epoch,
564 state_notify,
565 controller_notify,
566 reconnect_published,
567 connect_timeout,
568 heartbeat_timeout,
569 backoff,
570 reconnect_throttle: ReconnectThrottle::default(),
571 reconnect_max_attempts,
572 reconnection_attempt_count: 0,
573 auth_tracker,
574 reconnect_buffer_waits_for_auth,
575 state_sink,
576 connection_rate_limit,
577 })
578 }
579
580 #[inline]
602 pub async fn connect_with_server(
603 url: &str,
604 headers: Vec<(String, String)>,
605 backend: TransportBackend,
606 proxy_url: Option<&str>,
607 max_message_size_bytes: Option<usize>,
608 max_frame_size_bytes: Option<usize>,
609 ) -> Result<(MessageWriter, MessageReader), TransportError> {
610 #[cfg(all(feature = "simulation", madsim))]
611 if backend == TransportBackend::Sockudo {
612 return Err(TransportError::Other(
613 "Sockudo timers are unsupported under simulation; select Tungstenite".into(),
614 ));
615 }
616
617 #[cfg(all(feature = "simulation", madsim))]
618 if proxy_url.is_some() || !url.starts_with("ws://") {
619 return Err(TransportError::InvalidUrl(
620 "WebSocket simulation requires plaintext ws:// endpoints without a proxy".into(),
621 ));
622 }
623
624 let limits = InboundSizeLimits {
625 max_message_size_bytes,
626 max_frame_size_bytes,
627 };
628
629 match backend {
630 TransportBackend::Tungstenite => match proxy_url {
631 Some(proxy) => {
632 Box::pin(Self::connect_tungstenite_via_proxy(
633 url, headers, proxy, limits,
634 ))
635 .await
636 }
637 None => Self::connect_tungstenite(url, headers, limits).await,
638 },
639 TransportBackend::Sockudo => {
640 #[cfg(feature = "transport-sockudo")]
641 {
642 match proxy_url {
643 Some(proxy) => {
644 Box::pin(Self::connect_sockudo_via_proxy(url, headers, proxy, limits))
645 .await
646 }
647 None => Self::connect_sockudo(url, headers, limits).await,
648 }
649 }
650
651 #[cfg(not(feature = "transport-sockudo"))]
652 {
653 Err(TransportError::Other(
654 "sockudo backend selected but the transport-sockudo \
655 Cargo feature is not enabled"
656 .to_string(),
657 ))
658 }
659 }
660 }
661 }
662
663 #[inline]
666 #[cfg(not(feature = "turmoil"))]
667 async fn connect_tungstenite(
668 url: &str,
669 headers: Vec<(String, String)>,
670 limits: InboundSizeLimits,
671 ) -> Result<(MessageWriter, MessageReader), TransportError> {
672 let request = tungstenite_request(url, headers)?;
673 let protocol = tungstenite_protocol_config(limits);
674
675 #[cfg(not(all(feature = "simulation", madsim)))]
677 let (stream, _resp) = connect_async_with_config(request, protocol, false)
678 .await
679 .map_err(TransportError::from)?;
680 #[cfg(not(all(feature = "simulation", madsim)))]
681 crate::net::apply_socket_options(stream.get_ref().get_ref());
682
683 #[cfg(all(feature = "simulation", madsim))]
684 let (stream, _resp) = {
685 let target = request.uri();
686 let host = target
687 .host()
688 .ok_or_else(|| TransportError::InvalidUrl("missing hostname".into()))?;
689 let address = format!("{host}:{}", target.port_u16().unwrap_or(80));
690 let socket = dst::net::TcpStream::connect(address.as_str()).await?;
691 tokio_tungstenite::client_async_with_config(request, socket, protocol)
692 .await
693 .map_err(TransportError::from)?
694 };
695
696 let transport: BoxedWsTransport = Box::pin(TungsteniteTransport::new(stream));
697 Ok(transport.split())
698 }
699
700 #[inline]
708 #[cfg(not(feature = "turmoil"))]
709 async fn connect_tungstenite_via_proxy(
710 url: &str,
711 headers: Vec<(String, String)>,
712 proxy_url: &str,
713 limits: InboundSizeLimits,
714 ) -> Result<(MessageWriter, MessageReader), TransportError> {
715 let proxy = match ProxyKind::parse(proxy_url)? {
716 ProxyKind::Http(target) => target,
717 ProxyKind::Unsupported { scheme } => {
718 log::warn!(
719 "WebSocket proxy_url scheme '{scheme}' is not yet supported; \
720 connecting without a WebSocket proxy"
721 );
722 return Self::connect_tungstenite(url, headers, limits).await;
723 }
724 };
725
726 let request = tungstenite_request(url, headers)?;
727
728 let target = WsTarget::parse(url)?;
729 let stream = tunnel_via_proxy(&target, &proxy).await?;
730
731 let transport: BoxedWsTransport =
735 Box::pin(proxied_ws_handshake(request, stream, limits)).await?;
736
737 Ok(transport.split())
738 }
739
740 #[inline]
743 #[cfg(feature = "turmoil")]
744 #[allow(unknown_lints, reason = "Clippy lint is unavailable on Rust 1.97")]
745 #[expect(
746 clippy::unused_async,
747 clippy::unused_async_trait_impl,
748 reason = "signature mirrors the production variant; both are awaited in the dispatcher"
749 )]
750 async fn connect_tungstenite_via_proxy(
751 _url: &str,
752 _headers: Vec<(String, String)>,
753 _proxy_url: &str,
754 _limits: InboundSizeLimits,
755 ) -> Result<(MessageWriter, MessageReader), TransportError> {
756 Err(TransportError::Other(
757 "proxy_url is not supported under the turmoil simulator".to_string(),
758 ))
759 }
760
761 #[inline]
764 #[cfg(feature = "turmoil")]
765 async fn connect_tungstenite(
766 url: &str,
767 headers: Vec<(String, String)>,
768 limits: InboundSizeLimits,
769 ) -> Result<(MessageWriter, MessageReader), TransportError> {
770 let request = tungstenite_request(url, headers)?;
771
772 let uri = request.uri();
773 let scheme = uri.scheme_str().unwrap_or("ws");
774 let host = uri
775 .host()
776 .ok_or_else(|| TransportError::InvalidUrl("missing hostname".to_string()))?;
777
778 let port = uri
780 .port_u16()
781 .unwrap_or_else(|| if scheme == "wss" { 443 } else { 80 });
782
783 let addr = format!("{host}:{port}");
784
785 let connector = crate::net::RealTcpConnector;
787 let tcp_stream = connector.connect(&addr).await?;
788 crate::net::apply_socket_options(&tcp_stream);
789
790 let maybe_tls_stream = if scheme == "wss" {
792 let mut root_store = rustls::RootCertStore::empty();
794 root_store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
795
796 let config = ClientConfig::builder()
797 .with_root_certificates(root_store)
798 .with_no_client_auth();
799
800 let tls_connector = TlsConnector::from(std::sync::Arc::new(config));
801 let domain = rustls::pki_types::ServerName::try_from(host.to_string())
802 .map_err(|e| TransportError::Tls(format!("Invalid DNS name: {e}")))?;
803
804 let tls_stream = tls_connector
805 .connect(domain, tcp_stream)
806 .await
807 .map_err(TransportError::Io)?;
808 MaybeTlsStream::Rustls(tls_stream)
809 } else {
810 MaybeTlsStream::Plain(tcp_stream)
811 };
812
813 let (stream, _resp) = client_async_with_config(
815 request,
816 maybe_tls_stream,
817 tungstenite_protocol_config(limits),
818 )
819 .await
820 .map_err(TransportError::from)?;
821 let transport: BoxedWsTransport = Box::pin(TungsteniteTransport::new(stream));
822 Ok(transport.split())
823 }
824
825 #[inline]
834 #[cfg(feature = "transport-sockudo")]
835 async fn connect_sockudo(
836 url: &str,
837 headers: Vec<(String, String)>,
838 limits: InboundSizeLimits,
839 ) -> Result<(MessageWriter, MessageReader), TransportError> {
840 let target = SockudoTarget::parse(url)?;
841 validate_extra_headers(&headers).map_err(TransportError::from)?;
842
843 #[cfg(feature = "turmoil")]
844 if target.is_tls {
845 return Err(TransportError::Tls(
846 "wss:// is not supported under the turmoil simulator; use ws://".to_string(),
847 ));
848 }
849
850 let tcp_stream = TcpStream::connect((target.host.as_str(), target.port))
851 .await
852 .map_err(TransportError::Io)?;
853
854 crate::net::apply_socket_options(&tcp_stream);
855
856 #[cfg(not(feature = "turmoil"))]
857 if target.is_tls {
858 let mut root_store = rustls::RootCertStore::empty();
859 root_store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
860 let config = ClientConfig::builder()
861 .with_root_certificates(root_store)
862 .with_no_client_auth();
863 let connector = TlsConnector::from(std::sync::Arc::new(config));
864 let domain = rustls::pki_types::ServerName::try_from(target.host.clone())
865 .map_err(|e| TransportError::Tls(format!("Invalid DNS name: {e}")))?;
866 let tls_stream = connector
867 .connect(domain, tcp_stream)
868 .await
869 .map_err(TransportError::Io)?;
870 return Self::finish_sockudo_handshake(tls_stream, &target, &headers, limits).await;
871 }
872
873 Self::finish_sockudo_handshake(tcp_stream, &target, &headers, limits).await
874 }
875
876 #[inline]
882 #[cfg(all(feature = "transport-sockudo", not(feature = "turmoil")))]
883 async fn connect_sockudo_via_proxy(
884 url: &str,
885 headers: Vec<(String, String)>,
886 proxy_url: &str,
887 limits: InboundSizeLimits,
888 ) -> Result<(MessageWriter, MessageReader), TransportError> {
889 let proxy = match ProxyKind::parse(proxy_url)? {
890 ProxyKind::Http(target) => target,
891 ProxyKind::Unsupported { scheme } => {
892 log::warn!(
893 "WebSocket proxy_url scheme '{scheme}' is not yet supported; \
894 connecting without a WebSocket proxy"
895 );
896 return Self::connect_sockudo(url, headers, limits).await;
897 }
898 };
899
900 let target = SockudoTarget::parse(url)?;
901 validate_extra_headers(&headers).map_err(TransportError::from)?;
902
903 let ws_target = WsTarget::parse(url)?;
907 let stream = tunnel_via_proxy(&ws_target, &proxy).await?;
908
909 Self::finish_sockudo_handshake(stream, &target, &headers, limits).await
910 }
911
912 #[inline]
915 #[cfg(all(feature = "transport-sockudo", feature = "turmoil"))]
916 #[allow(unknown_lints, reason = "Clippy lint is unavailable on Rust 1.97")]
917 #[expect(
918 clippy::unused_async,
919 clippy::unused_async_trait_impl,
920 reason = "signature mirrors the production variant; both are awaited in the dispatcher"
921 )]
922 async fn connect_sockudo_via_proxy(
923 _url: &str,
924 _headers: Vec<(String, String)>,
925 _proxy_url: &str,
926 _limits: InboundSizeLimits,
927 ) -> Result<(MessageWriter, MessageReader), TransportError> {
928 Err(TransportError::Other(
929 "proxy_url is not supported under the turmoil simulator".to_string(),
930 ))
931 }
932
933 #[cfg(feature = "transport-sockudo")]
934 async fn finish_sockudo_handshake<S>(
935 mut stream: S,
936 target: &SockudoTarget,
937 headers: &[(String, String)],
938 limits: InboundSizeLimits,
939 ) -> Result<(MessageWriter, MessageReader), TransportError>
940 where
941 S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
942 {
943 let handshake =
947 client_handshake_with_headers(&mut stream, &target.host_header, &target.path, headers)
948 .await?;
949
950 let stream = match handshake.leftover {
953 Some(prefix) => SockudoStream::<Http1>::new(PrefixedIo::new(stream, prefix)),
954 None => SockudoStream::<Http1>::new(stream),
955 };
956
957 let ws =
958 SockudoWebSocketStream::from_raw(stream, Role::Client, sockudo_protocol_config(limits));
959 let transport: BoxedWsTransport = Box::pin(SockudoTransport::new(ws));
960 Ok(transport.split())
961 }
962}
963
964#[derive(Clone, Copy, Debug, Default)]
965struct InboundSizeLimits {
966 max_message_size_bytes: Option<usize>,
967 max_frame_size_bytes: Option<usize>,
968}
969
970fn validate_honored_inbound_limits(config: &WebSocketConfig) -> Result<(), TransportError> {
971 if config.max_message_size_bytes == Some(0) {
972 return Err(TransportError::Io(std::io::Error::new(
973 std::io::ErrorKind::InvalidInput,
974 "max_message_size_bytes must be positive",
975 )));
976 }
977
978 if config.max_frame_size_bytes == Some(0) {
979 return Err(TransportError::Io(std::io::Error::new(
980 std::io::ErrorKind::InvalidInput,
981 "max_frame_size_bytes must be positive",
982 )));
983 }
984
985 Ok(())
986}
987
988fn tungstenite_protocol_config(
989 limits: InboundSizeLimits,
990) -> Option<tokio_tungstenite::tungstenite::protocol::WebSocketConfig> {
991 if limits.max_message_size_bytes.is_none() && limits.max_frame_size_bytes.is_none() {
992 return None;
993 }
994
995 let mut config = tokio_tungstenite::tungstenite::protocol::WebSocketConfig::default();
996 if let Some(size) = limits.max_message_size_bytes {
997 config.max_message_size = Some(size);
998 }
999
1000 if let Some(size) = limits.max_frame_size_bytes {
1001 config.max_frame_size = Some(size);
1002 }
1003
1004 Some(config)
1005}
1006
1007#[cfg(feature = "transport-sockudo")]
1008fn sockudo_protocol_config(limits: InboundSizeLimits) -> SockudoConfig {
1009 let mut config = SockudoConfig::default();
1010 if let Some(size) = limits.max_message_size_bytes {
1011 config.max_message_size = size;
1012 }
1013
1014 if let Some(size) = limits.max_frame_size_bytes {
1015 config.max_frame_size = size;
1016 }
1017
1018 config
1019}
1020
1021fn tungstenite_request(
1022 url: &str,
1023 headers: Vec<(String, String)>,
1024) -> Result<Request, TransportError> {
1025 let mut request = url.into_client_request().map_err(TransportError::from)?;
1026
1027 for (key, value) in headers {
1028 let value = HeaderValue::from_str(&value)
1029 .map_err(|e| TransportError::Handshake(format!("invalid header value: {e}")))?;
1030 let name: HeaderName = key
1031 .parse()
1032 .map_err(|e| TransportError::Handshake(format!("invalid header name: {e}")))?;
1033 request.headers_mut().insert(name, value);
1034 }
1035
1036 Ok(request)
1037}
1038
1039fn is_connection_drop_transport_error(err: &TransportError) -> bool {
1040 err.is_closed() || matches!(err, TransportError::Io(e) if is_connection_drop_io_error(e))
1041}
1042
1043fn is_retryable_initial_connect_error(err: &TransportError) -> bool {
1044 match err {
1045 TransportError::ConnectionClosed
1046 | TransportError::ConnectionReset
1047 | TransportError::ClosedByPeer(_) => true,
1048 TransportError::Io(error) => !matches!(
1049 error.kind(),
1050 std::io::ErrorKind::InvalidInput
1051 | std::io::ErrorKind::InvalidData
1052 | std::io::ErrorKind::Unsupported
1053 | std::io::ErrorKind::PermissionDenied
1054 ),
1055 TransportError::UpgradeRejected(status) | TransportError::ProxyConnectRejected(status) => {
1056 retryable_status(*status)
1057 }
1058 TransportError::InvalidUrl(_)
1059 | TransportError::Handshake(_)
1060 | TransportError::Tls(_)
1061 | TransportError::Protocol(_)
1062 | TransportError::MessageTooLarge
1063 | TransportError::FrameTooLarge
1064 | TransportError::InvalidUtf8
1065 | TransportError::Other(_) => false,
1066 }
1067}
1068
1069fn initial_connect_retry_config(
1070 policy: Option<&InitialConnectRetryPolicy>,
1071) -> Result<RetryConfig, TransportError> {
1072 if let Some(policy) = policy {
1073 ExponentialBackoff::new(
1074 policy.delay_initial,
1075 policy.delay_max,
1076 policy.backoff_factor,
1077 policy.jitter_ms,
1078 false,
1079 )
1080 .map_err(|e| {
1081 TransportError::Io(std::io::Error::new(std::io::ErrorKind::InvalidInput, e))
1082 })?;
1083
1084 Ok(RetryConfig {
1085 max_retries: policy.max_attempts.get() - 1,
1086 initial_delay_ms: duration_to_millis("delay_initial", policy.delay_initial)?,
1087 max_delay_ms: duration_to_millis("delay_max", policy.delay_max)?,
1088 backoff_factor: policy.backoff_factor,
1089 jitter_ms: policy.jitter_ms,
1090 operation_timeout_ms: None,
1091 immediate_first: false,
1092 max_elapsed_ms: None,
1093 })
1094 } else {
1095 Ok(RetryConfig {
1096 max_retries: 0,
1097 initial_delay_ms: 1,
1098 max_delay_ms: 1,
1099 backoff_factor: 1.0,
1100 jitter_ms: 0,
1101 operation_timeout_ms: None,
1102 immediate_first: false,
1103 max_elapsed_ms: None,
1104 })
1105 }
1106}
1107
1108fn duration_to_millis(field: &str, duration: Duration) -> Result<u64, TransportError> {
1109 let nanoseconds = duration.as_nanos();
1110 if nanoseconds > u128::from(u64::MAX) {
1111 return Err(TransportError::Io(std::io::Error::new(
1112 std::io::ErrorKind::InvalidInput,
1113 format!("{field} exceeds the maximum backoff duration"),
1114 )));
1115 }
1116
1117 let milliseconds = nanoseconds
1118 .div_ceil(1_000_000)
1119 .min(u128::from(u64::MAX / 1_000_000));
1120 u64::try_from(milliseconds).map_err(|_| {
1121 TransportError::Io(std::io::Error::new(
1122 std::io::ErrorKind::InvalidInput,
1123 format!("{field} exceeds the maximum backoff duration"),
1124 ))
1125 })
1126}
1127
1128async fn await_initial_connect_attempt<F, T>(
1129 cancellation_token: &CancellationToken,
1130 attempt: F,
1131) -> Result<T, TransportError>
1132where
1133 F: Future<Output = Result<T, TransportError>>,
1134{
1135 tokio::select! {
1136 biased;
1137 () = cancellation_token.cancelled() => Err(initial_connect_cancelled()),
1138 result = attempt => result,
1139 }
1140}
1141
1142fn initial_connect_retry_error(error: RetryError) -> TransportError {
1143 let kind = match error {
1144 RetryError::Canceled => return initial_connect_cancelled(),
1145 RetryError::InvalidConfiguration { .. } => std::io::ErrorKind::InvalidInput,
1146 RetryError::OperationTimeout { .. } | RetryError::ElapsedBudgetExceeded { .. } => {
1147 std::io::ErrorKind::TimedOut
1148 }
1149 };
1150
1151 TransportError::Io(std::io::Error::new(kind, error))
1152}
1153
1154fn initial_connect_cancelled() -> TransportError {
1155 TransportError::Io(std::io::Error::new(
1156 std::io::ErrorKind::Interrupted,
1157 "initial WebSocket connection cancelled",
1158 ))
1159}
1160
1161fn read_termination_log_level(connection_state: &AtomicU8) -> log::Level {
1163 let mode = ConnectionMode::from_atomic(connection_state);
1164 if mode.is_disconnect() || mode.is_closed() {
1165 log::Level::Debug
1166 } else {
1167 log::Level::Warn
1168 }
1169}
1170
1171#[cfg(test)]
1172mod connection_error_tests {
1173 use std::io;
1174
1175 use rstest::rstest;
1176
1177 use super::*;
1178 use crate::transport::CloseFrame;
1179
1180 #[rstest]
1181 #[case(TransportError::ConnectionClosed, true)]
1182 #[case(TransportError::ConnectionReset, true)]
1183 #[case(TransportError::ClosedByPeer(Some(CloseFrame::new(1000, "bye"))), true)]
1184 #[case(TransportError::ClosedByPeer(None), true)]
1185 #[case(TransportError::Io(io::Error::from(io::ErrorKind::BrokenPipe)), true)]
1186 #[case(
1187 TransportError::Io(io::Error::from(io::ErrorKind::ConnectionReset)),
1188 true
1189 )]
1190 #[case(TransportError::Io(io::Error::from(io::ErrorKind::TimedOut)), true)]
1191 #[case(
1192 TransportError::Io(io::Error::from(io::ErrorKind::UnexpectedEof)),
1193 true
1194 )]
1195 #[case(
1196 TransportError::Io(io::Error::from(io::ErrorKind::InvalidInput)),
1197 false
1198 )]
1199 #[case(TransportError::InvalidUrl("http://example.com".into()), false)]
1200 #[case(TransportError::Handshake("bad".into()), false)]
1201 #[case(TransportError::Protocol("bad opcode".into()), false)]
1202 #[case(TransportError::Tls("bad certificate".into()), false)]
1203 #[case(TransportError::MessageTooLarge, false)]
1204 #[case(TransportError::FrameTooLarge, false)]
1205 #[case(TransportError::InvalidUtf8, false)]
1206 #[case(TransportError::Other("backend protocol mismatch".into()), false)]
1207 fn connection_drop_transport_error_classification(
1208 #[case] err: TransportError,
1209 #[case] expected: bool,
1210 ) {
1211 assert_eq!(is_connection_drop_transport_error(&err), expected);
1212 }
1213
1214 #[rstest]
1215 #[case(Duration::ZERO, 0)]
1216 #[case(Duration::from_nanos(1), 1)]
1217 #[case(Duration::from_micros(500), 1)]
1218 #[case(Duration::from_nanos(1_000_001), 2)]
1219 #[case(Duration::from_millis(500), 500)]
1220 #[case(Duration::from_secs(5), 5_000)]
1221 fn duration_to_millis_rounds_up(#[case] duration: Duration, #[case] expected: u64) {
1222 assert_eq!(duration_to_millis("delay", duration).unwrap(), expected);
1223 }
1224
1225 #[rstest]
1226 fn duration_to_millis_caps_at_backoff_maximum() {
1227 assert_eq!(
1228 duration_to_millis("delay", Duration::from_nanos(u64::MAX)).unwrap(),
1229 u64::MAX / 1_000_000
1230 );
1231 }
1232
1233 #[rstest]
1234 fn duration_to_millis_rejects_excessive_duration() {
1235 let error = duration_to_millis("delay", Duration::from_secs(u64::MAX)).unwrap_err();
1236
1237 let TransportError::Io(error) = error else {
1238 panic!("expected an I/O error, was: {error:?}");
1239 };
1240
1241 assert_eq!(error.kind(), io::ErrorKind::InvalidInput);
1242 assert_eq!(
1243 error.to_string(),
1244 "delay exceeds the maximum backoff duration"
1245 );
1246 }
1247
1248 #[rstest]
1249 fn initial_connect_retry_config_preserves_duration_validation() {
1250 let policy = InitialConnectRetryPolicy {
1251 max_attempts: std::num::NonZeroU32::new(2).unwrap(),
1252 delay_initial: Duration::from_micros(900),
1253 delay_max: Duration::from_micros(500),
1254 backoff_factor: 2.0,
1255 jitter_ms: 0,
1256 };
1257
1258 let error = initial_connect_retry_config(Some(&policy)).unwrap_err();
1259
1260 let TransportError::Io(error) = error else {
1261 panic!("expected an I/O error, was: {error:?}");
1262 };
1263
1264 assert_eq!(error.kind(), io::ErrorKind::InvalidInput);
1265 assert!(
1266 error
1267 .to_string()
1268 .contains("delay_max must be >= delay_initial")
1269 );
1270 }
1271
1272 #[rstest]
1273 #[tokio::test]
1274 async fn initial_connect_attempt_prefers_cancellation_over_ready_result() {
1275 let token = CancellationToken::new();
1276 token.cancel();
1277
1278 let error = await_initial_connect_attempt(&token, std::future::ready(Ok(())))
1279 .await
1280 .expect_err("cancellation should take priority");
1281
1282 assert!(
1283 matches!(error, TransportError::Io(ref error) if error.kind() == io::ErrorKind::Interrupted)
1284 );
1285 }
1286
1287 #[rstest]
1288 fn unset_inbound_limits_preserve_backend_defaults() {
1289 let limits = InboundSizeLimits::default();
1290 assert!(tungstenite_protocol_config(limits).is_none());
1291
1292 #[cfg(feature = "transport-sockudo")]
1293 {
1294 let config = sockudo_protocol_config(limits);
1295 let default = SockudoConfig::default();
1296 assert_eq!(config.max_message_size, default.max_message_size);
1297 assert_eq!(config.max_frame_size, default.max_frame_size);
1298 }
1299 }
1300
1301 #[rstest]
1302 fn configured_inbound_limit_overrides_only_that_cap() {
1303 let limits = InboundSizeLimits {
1304 max_message_size_bytes: Some(32),
1305 max_frame_size_bytes: None,
1306 };
1307
1308 let config = tungstenite_protocol_config(limits).expect("override should build a config");
1309 let default = tokio_tungstenite::tungstenite::protocol::WebSocketConfig::default();
1310
1311 assert_eq!(config.max_message_size, Some(32));
1312 assert_eq!(config.max_frame_size, default.max_frame_size);
1313
1314 #[cfg(feature = "transport-sockudo")]
1315 {
1316 let config = sockudo_protocol_config(limits);
1317 let default = SockudoConfig::default();
1318 assert_eq!(config.max_message_size, 32);
1319 assert_eq!(config.max_frame_size, default.max_frame_size);
1320 }
1321 }
1322}
1323
1324#[cfg(not(feature = "turmoil"))]
1329async fn proxied_ws_handshake<S>(
1330 request: tokio_tungstenite::tungstenite::handshake::client::Request,
1331 stream: S,
1332 limits: InboundSizeLimits,
1333) -> Result<BoxedWsTransport, TransportError>
1334where
1335 S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
1336{
1337 let (ws, _resp) = tokio_tungstenite::client_async_with_config(
1338 request,
1339 stream,
1340 tungstenite_protocol_config(limits),
1341 )
1342 .await
1343 .map_err(TransportError::from)?;
1344 Ok(Box::pin(TungsteniteTransport::new(ws)))
1345}
1346
1347#[cfg(feature = "transport-sockudo")]
1354#[derive(Debug, PartialEq, Eq)]
1355struct SockudoTarget {
1356 host: String,
1357 host_header: String,
1358 port: u16,
1359 path: String,
1360 is_tls: bool,
1361}
1362
1363#[cfg(feature = "transport-sockudo")]
1364impl SockudoTarget {
1365 fn parse(url: &str) -> Result<Self, TransportError> {
1366 let parsed = url::Url::parse(url)
1367 .map_err(|e| TransportError::InvalidUrl(format!("invalid WebSocket URL: {e}")))?;
1368
1369 let scheme = parsed.scheme();
1370
1371 let is_tls = match scheme {
1372 "ws" => false,
1373 "wss" => true,
1374 other => {
1375 return Err(TransportError::InvalidUrl(format!(
1376 "expected ws:// or wss:// scheme, was {other}"
1377 )));
1378 }
1379 };
1380
1381 let raw_host = parsed
1382 .host_str()
1383 .ok_or_else(|| TransportError::InvalidUrl("missing hostname".to_string()))?;
1384
1385 let is_bracketed = raw_host.starts_with('[') && raw_host.ends_with(']');
1390 let host = if is_bracketed {
1391 raw_host[1..raw_host.len() - 1].to_string()
1392 } else {
1393 raw_host.to_string()
1394 };
1395
1396 let explicit_port = parsed.port();
1397 let port = explicit_port.unwrap_or(if is_tls { 443 } else { 80 });
1398
1399 let host_header = match explicit_port {
1400 Some(p) => format!("{raw_host}:{p}"),
1401 None => raw_host.to_string(),
1402 };
1403
1404 let path = if parsed.path().is_empty() {
1405 "/".to_string()
1406 } else {
1407 let mut p = parsed.path().to_string();
1408 if let Some(query) = parsed.query() {
1409 p.push('?');
1410 p.push_str(query);
1411 }
1412
1413 p
1414 };
1415
1416 Ok(Self {
1417 host,
1418 host_header,
1419 port,
1420 path,
1421 is_tls,
1422 })
1423 }
1424}
1425
1426impl WebSocketClientInner {
1427 pub async fn reconnect(&mut self) -> Result<(), TransportError> {
1448 Box::pin(self.reconnect_with_outcome()).await.map(|_| ())
1449 }
1450
1451 async fn wait_for_reconnect_publication(&self) -> bool {
1452 let fallback_interval = Duration::from_millis(CONTROLLER_FALLBACK_INTERVAL_MS);
1453
1454 loop {
1455 let mut notified = pin!(self.controller_notify.notified());
1456 notified.as_mut().enable();
1457
1458 if !ConnectionMode::from_atomic(&self.connection_mode).is_reconnect() {
1459 return false;
1460 }
1461
1462 if self.reconnect_published.load(Ordering::SeqCst) {
1463 return true;
1464 }
1465
1466 tokio::select! {
1467 biased;
1468 () = notified => {}
1469 () = dst::time::sleep(fallback_interval) => {}
1470 }
1471 }
1472 }
1473
1474 async fn reconnect_with_outcome(&mut self) -> Result<ReconnectOutcome, TransportError> {
1475 log::info!("Reconnecting");
1476
1477 if self.handler.is_none() {
1478 log::warn!(
1479 "Auto-reconnect disabled for stream-based WebSocket client; \
1480 stream users must manually reconnect by creating a new connection"
1481 );
1482 self.connection_mode
1484 .store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
1485 fail_registered_auth(
1486 self.auth_tracker.as_ref(),
1487 "WebSocket stream mode cannot reconnect",
1488 );
1489 self.state_notify.notify_waiters();
1493 return Ok(ReconnectOutcome::Aborted);
1494 }
1495
1496 if ConnectionMode::from_atomic(&self.connection_mode).is_disconnect() {
1497 log::debug!("Reconnect aborted due to disconnect state");
1498 return Ok(ReconnectOutcome::Aborted);
1499 }
1500
1501 if let Some(rate_limit) = &self.connection_rate_limit {
1502 rate_limit
1503 .limiter
1504 .await_keys_ready(Some(&rate_limit.keys))
1505 .await;
1506 }
1507
1508 if ConnectionMode::from_atomic(&self.connection_mode).is_disconnect() {
1509 log::debug!("Reconnect aborted during connection rate-limit wait");
1510 return Ok(ReconnectOutcome::Aborted);
1511 }
1512
1513 let (new_writer, reader) = dst::time::timeout(
1515 self.connect_timeout,
1516 Box::pin(Self::connect_with_server(
1517 &self.config.url,
1518 self.reconnect_headers.snapshot(),
1519 self.config.backend,
1520 self.config.proxy_url.as_deref(),
1521 self.config.max_message_size_bytes,
1522 self.config.max_frame_size_bytes,
1523 )),
1524 )
1525 .await
1526 .map_err(|_| {
1527 TransportError::Io(std::io::Error::new(
1528 std::io::ErrorKind::TimedOut,
1529 format!(
1530 "reconnection timed out after {}s",
1531 self.connect_timeout.as_secs_f64()
1532 ),
1533 ))
1534 })??;
1535
1536 if ConnectionMode::from_atomic(&self.connection_mode).is_disconnect() {
1537 log::debug!("Reconnect aborted mid-flight (after connect)");
1538 return Ok(ReconnectOutcome::Aborted);
1539 }
1540
1541 let (tx, rx) = tokio::sync::oneshot::channel();
1544
1545 if let Err(e) = self
1546 .writer_tx
1547 .send_update(WriterCommand::Update(new_writer, tx))
1548 {
1549 log::error!("{e}");
1550 return Err(TransportError::Io(std::io::Error::new(
1551 std::io::ErrorKind::BrokenPipe,
1552 format!("Failed to send update command: {e}"),
1553 )));
1554 }
1555
1556 let connection_epoch = match rx.await {
1558 Ok(connection_epoch) => {
1559 log::debug!("Writer confirmed socket update: epoch={connection_epoch}");
1560 connection_epoch
1561 }
1562 Err(e) => {
1563 log::error!("Writer dropped update channel: {e}");
1564 return Err(TransportError::Io(std::io::Error::new(
1565 std::io::ErrorKind::BrokenPipe,
1566 "Writer task dropped response channel",
1567 )));
1568 }
1569 };
1570
1571 dst::time::sleep(Duration::from_millis(GRACEFUL_SHUTDOWN_DELAY_MS)).await;
1573
1574 if ConnectionMode::from_atomic(&self.connection_mode).is_disconnect() {
1575 log::debug!("Reconnect aborted mid-flight (after delay)");
1576 return Ok(ReconnectOutcome::Aborted);
1577 }
1578
1579 if let Some(read_fence) = self.read_fence.take() {
1580 read_fence.invalidate();
1581 }
1582
1583 if let Some(ref read_task) = self.read_task.take()
1584 && !read_task.is_finished()
1585 {
1586 read_task.abort();
1587 log_task_aborted("read");
1588 }
1589
1590 if !self.wait_for_reconnect_publication().await {
1591 log::debug!("Reconnect aborted before state publication completed");
1592 return Ok(ReconnectOutcome::Aborted);
1593 }
1594
1595 if ConnectionMode::complete_reconnect_with_sink(
1598 &self.connection_mode,
1599 self.state_sink.as_ref(),
1600 ) == ReconnectOutcome::Aborted
1601 {
1602 log::debug!("Reconnect aborted (state changed during reconnect)");
1603 return Ok(ReconnectOutcome::Aborted);
1604 }
1605
1606 let read_fence = ReadSessionFence::new();
1607 self.read_task = Some(Self::spawn_message_handler_task(
1608 Arc::clone(&self.connection_mode),
1609 Arc::clone(&self.state_notify),
1610 read_fence.clone(),
1611 reader,
1612 connection_epoch,
1613 self.handler.as_ref(),
1614 self.ping_handler.as_ref(),
1615 self.close_handler.as_ref(),
1616 self.state_sink.as_ref(),
1617 self.config.idle_timeout_ms,
1618 self.heartbeat_timeout,
1619 ));
1620 self.read_fence = Some(read_fence);
1621
1622 log::info!("Reconnect succeeded");
1623 Ok(ReconnectOutcome::Reconnected)
1624 }
1625
1626 #[inline]
1632 #[must_use]
1633 pub fn is_alive(&self) -> bool {
1634 match &self.read_task {
1635 Some(read_task) => !read_task.is_finished() && !self.write_task.is_finished(),
1636 None => !self.write_task.is_finished(),
1637 }
1638 }
1639
1640 #[expect(
1641 clippy::too_many_arguments,
1642 reason = "both handler modes share the same reader lifecycle"
1643 )]
1644 fn spawn_message_handler_task(
1645 connection_state: Arc<AtomicU8>,
1646 state_notify: Arc<tokio::sync::Notify>,
1647 read_fence: ReadSessionFence,
1648 mut reader: MessageReader,
1649 connection_epoch: u64,
1650 handler: Option<&IncomingHandler>,
1651 ping_handler: Option<&IncomingPingHandler>,
1652 close_handler: Option<&CloseHandler>,
1653 state_sink: Option<&SocketStateSink>,
1654 idle_timeout_ms: Option<u64>,
1655 heartbeat_timeout: Option<Duration>,
1656 ) -> dst::task::JoinHandle<()> {
1657 log::debug!("Started message handler task 'read'");
1658
1659 let check_interval = Duration::from_millis(CONNECTION_STATE_CHECK_INTERVAL_MS);
1660 let idle_timeout = idle_timeout_ms.map(Duration::from_millis);
1661
1662 let handler = handler.cloned();
1663 let ping_handler = ping_handler.cloned();
1664 let close_handler = close_handler.cloned();
1665 let state_sink = state_sink.cloned();
1666
1667 dst::task::spawn(async move {
1668 let mut last_data_time = dst::time::Instant::now();
1669 let mut last_frame_time = dst::time::Instant::now();
1670
1671 loop {
1672 if !ConnectionMode::from_atomic(&connection_state).is_active()
1673 || !read_fence.is_valid()
1674 {
1675 break;
1676 }
1677
1678 let read_result = dst::time::timeout(check_interval, reader.next()).await;
1679
1680 if let Ok(Some(Ok(ref message))) = read_result
1681 && (!ConnectionMode::from_atomic(&connection_state).is_active()
1682 || !read_fence.is_valid())
1683 {
1684 log::debug!(
1685 "Dropping WebSocket message with {} bytes after session ended",
1686 message.as_bytes().len()
1687 );
1688 break;
1689 }
1690
1691 if matches!(&read_result, Ok(Some(Ok(_)))) {
1692 last_frame_time = dst::time::Instant::now();
1693 }
1694
1695 match read_result {
1696 Ok(Some(Ok(Message::Binary(data)))) => {
1697 log::trace!("Received message <binary> {} bytes", data.len());
1698 last_data_time = dst::time::Instant::now();
1699
1700 if !ConnectionMode::from_atomic(&connection_state).is_active()
1701 || !read_fence.is_valid()
1702 {
1703 log::debug!(
1704 "Dropping WebSocket message with {} bytes after session ended",
1705 data.len()
1706 );
1707 break;
1708 }
1709
1710 if let Some(ref handler) = handler {
1711 handler.handle(connection_epoch, Message::Binary(data));
1712 }
1713 }
1714 Ok(Some(Ok(Message::Text(data)))) => {
1715 log::trace!("Received text frame ({} bytes)", data.len());
1716 last_data_time = dst::time::Instant::now();
1717
1718 if !ConnectionMode::from_atomic(&connection_state).is_active()
1719 || !read_fence.is_valid()
1720 {
1721 log::debug!(
1722 "Dropping WebSocket message with {} bytes after session ended",
1723 data.len()
1724 );
1725 break;
1726 }
1727
1728 if let Some(ref handler) = handler {
1729 handler.handle(connection_epoch, Message::Text(data));
1730 }
1731 }
1732 Ok(Some(Ok(Message::Ping(ping_data)))) => {
1733 log::trace!("Received ping frame ({} bytes)", ping_data.len());
1734 if let Some(ref handler) = ping_handler {
1739 if !ConnectionMode::from_atomic(&connection_state).is_active()
1740 || !read_fence.is_valid()
1741 {
1742 log::debug!(
1743 "Dropping WebSocket ping with {} bytes after session ended",
1744 ping_data.len()
1745 );
1746 break;
1747 }
1748
1749 handler.handle(connection_epoch, ping_data.to_vec());
1750 }
1751
1752 if idle_timeout_exceeded(last_data_time, idle_timeout) {
1753 break;
1754 }
1755 }
1756 Ok(Some(Ok(Message::Pong(_)))) => {
1757 log::trace!("Received pong");
1758 if idle_timeout_exceeded(last_data_time, idle_timeout) {
1761 break;
1762 }
1763 }
1764 Ok(Some(Ok(Message::Close(frame)))) => {
1765 Self::dispatch_close_frame(
1766 &connection_state,
1767 &read_fence,
1768 close_handler.as_ref(),
1769 state_sink.as_ref(),
1770 frame,
1771 );
1772 break;
1773 }
1774 Ok(Some(Err(e))) => {
1775 if is_connection_drop_transport_error(&e) {
1776 log::warn!("Received connection error, terminating: {e}");
1777 } else {
1778 log::error!("Received transport error, terminating: {e}");
1779 }
1780
1781 break;
1782 }
1783 Ok(None) => {
1784 log::log!(
1785 read_termination_log_level(&connection_state),
1786 "Connection closed by peer (no close frame), terminating"
1787 );
1788 break;
1789 }
1790 Err(_) => {
1791 if heartbeat_timeout_exceeded(last_frame_time, heartbeat_timeout) {
1792 break;
1793 }
1794
1795 if idle_timeout_exceeded(last_data_time, idle_timeout) {
1796 break;
1797 }
1798 }
1799 }
1800 }
1801
1802 state_notify.notify_one();
1804 })
1805 }
1806
1807 fn dispatch_close_frame(
1808 connection_state: &AtomicU8,
1809 read_fence: &ReadSessionFence,
1810 close_handler: Option<&CloseHandler>,
1811 state_sink: Option<&SocketStateSink>,
1812 frame: Option<CloseFrame>,
1813 ) {
1814 match &frame {
1815 Some(frame) => log::log!(
1816 read_termination_log_level(connection_state),
1817 "Received close frame, terminating: code={}, reason={REDACTED}",
1818 frame.code
1819 ),
1820 None => log::log!(
1821 read_termination_log_level(connection_state),
1822 "Received close frame with no code or reason, terminating"
1823 ),
1824 }
1825
1826 if !ConnectionMode::from_atomic(connection_state).is_active() || !read_fence.is_valid() {
1827 return;
1828 }
1829
1830 if let Some(handler) = close_handler
1831 && !handler(frame)
1832 {
1833 ConnectionMode::close_websocket_on_loss(connection_state, state_sink);
1834 }
1835 }
1836
1837 fn buffer_for_replay(
1845 buffer: &mut VecDeque<(Message, tokio::sync::OwnedSemaphorePermit)>,
1846 msg: Message,
1847 permit: tokio::sync::OwnedSemaphorePermit,
1848 ) {
1849 if msg.is_control() {
1850 return;
1851 }
1852
1853 log::debug!(
1854 "Buffering message for replay (buffer size: {})",
1855 buffer.len() + 1
1856 );
1857
1858 buffer.push_back((msg, permit));
1859 }
1860
1861 async fn drain_reconnect_buffer(
1866 buffer: &mut VecDeque<(Message, tokio::sync::OwnedSemaphorePermit)>,
1867 writer: &mut MessageWriter,
1868 connection_state: &AtomicU8,
1869 auth_tracker: &Arc<OnceLock<AuthTracker>>,
1870 reconnect_buffer_waits_for_auth: &AtomicBool,
1871 ) -> bool {
1872 if buffer.is_empty() {
1873 return false;
1874 }
1875
1876 let initial_buffer_len = buffer.len();
1877 log::info!("Sending {initial_buffer_len} buffered messages after reconnection");
1878
1879 while !buffer.is_empty() {
1880 match Self::reconnect_buffer_action(
1881 reconnect_buffer_waits_for_auth,
1882 auth_tracker,
1883 connection_state,
1884 ) {
1885 ReconnectBufferAction::Drain => {}
1886 ReconnectBufferAction::Wait => return false,
1887 ReconnectBufferAction::Discard => {
1888 log::warn!(
1889 "Discarding {} buffered messages after authentication failed",
1890 buffer.len()
1891 );
1892 buffer.clear();
1893 return false;
1894 }
1895 }
1896
1897 let msg_to_send = buffer
1899 .front()
1900 .expect("reconnect buffer should not be empty")
1901 .0
1902 .clone();
1903
1904 if let Err(e) = writer.send(msg_to_send).await {
1905 if is_connection_drop_transport_error(&e) {
1906 log::warn!(
1907 "Failed to send buffered message after reconnection: {e}, {} messages remain in buffer",
1908 buffer.len()
1909 );
1910 } else {
1911 log::error!(
1912 "Failed to send buffered message after reconnection: {e}, {} messages remain in buffer",
1913 buffer.len()
1914 );
1915 }
1916
1917 return true;
1918 }
1919
1920 buffer.pop_front();
1922 }
1923
1924 if buffer.is_empty() {
1925 log::info!("Successfully sent all {initial_buffer_len} buffered messages");
1926 }
1927
1928 false
1929 }
1930
1931 fn can_drain_reconnect_buffer(
1932 reconnect_buffer_waits_for_auth: &AtomicBool,
1933 auth_tracker: &Arc<OnceLock<AuthTracker>>,
1934 ) -> ReconnectBufferAction {
1935 if !reconnect_buffer_waits_for_auth.load(Ordering::Acquire) {
1936 return ReconnectBufferAction::Drain;
1937 }
1938
1939 match auth_tracker.get().map(AuthTracker::auth_state) {
1940 Some(AuthState::Authenticated) => ReconnectBufferAction::Drain,
1941 Some(AuthState::Failed) => ReconnectBufferAction::Discard,
1942 Some(AuthState::Unauthenticated) | None => ReconnectBufferAction::Wait,
1943 }
1944 }
1945
1946 fn reconnect_buffer_action(
1947 reconnect_buffer_waits_for_auth: &AtomicBool,
1948 auth_tracker: &Arc<OnceLock<AuthTracker>>,
1949 connection_state: &AtomicU8,
1950 ) -> ReconnectBufferAction {
1951 let action =
1952 Self::can_drain_reconnect_buffer(reconnect_buffer_waits_for_auth, auth_tracker);
1953
1954 if action == ReconnectBufferAction::Drain
1956 && !ConnectionMode::from_atomic(connection_state).is_active()
1957 {
1958 ReconnectBufferAction::Wait
1959 } else {
1960 action
1961 }
1962 }
1963
1964 #[expect(
1965 clippy::too_many_arguments,
1966 reason = "writer task owns the transport and shared lifecycle coordination state"
1967 )]
1968 fn spawn_write_task(
1969 connection_state: Arc<AtomicU8>,
1970 controller_notify: Arc<tokio::sync::Notify>,
1971 reconnect_published: Arc<AtomicBool>,
1972 writer: MessageWriter,
1973 mut writer_rx: WriterReceiver<WriterCommand>,
1974 connection_epoch: Arc<AtomicU64>,
1975 auth_tracker: Arc<OnceLock<AuthTracker>>,
1976 reconnect_buffer_waits_for_auth: Arc<AtomicBool>,
1977 state_sink: Option<SocketStateSink>,
1978 ) -> dst::task::JoinHandle<()> {
1979 log_task_started("write");
1980
1981 let check_interval = Duration::from_millis(CONNECTION_STATE_CHECK_INTERVAL_MS);
1983
1984 dst::task::spawn(async move {
1985 let mut active_writer = writer;
1986 let mut reconnect_buffer: VecDeque<(Message, tokio::sync::OwnedSemaphorePermit)> =
1989 VecDeque::new();
1990
1991 loop {
1992 let mode = ConnectionMode::from_atomic(&connection_state);
1993
1994 match mode {
1995 ConnectionMode::Disconnect => {
1996 if !reconnect_buffer.is_empty() {
1998 log::warn!(
1999 "Discarding {} buffered messages due to disconnect",
2000 reconnect_buffer.len()
2001 );
2002 reconnect_buffer.clear();
2003 }
2004
2005 _ = dst::time::timeout(
2008 Duration::from_secs(GRACEFUL_SHUTDOWN_TIMEOUT_SECS),
2009 active_writer.close(),
2010 )
2011 .await;
2012 break;
2013 }
2014 ConnectionMode::Closed => {
2015 if !reconnect_buffer.is_empty() {
2017 log::warn!(
2018 "Discarding {} buffered messages due to closed connection",
2019 reconnect_buffer.len()
2020 );
2021 reconnect_buffer.clear();
2022 }
2023
2024 break;
2025 }
2026 _ => {}
2027 }
2028
2029 if mode.is_active() && !reconnect_buffer.is_empty() {
2030 match Self::reconnect_buffer_action(
2031 reconnect_buffer_waits_for_auth.as_ref(),
2032 &auth_tracker,
2033 &connection_state,
2034 ) {
2035 ReconnectBufferAction::Drain => {
2036 let drain_result = dst::time::timeout(
2037 Duration::from_secs(WRITE_TIMEOUT_SECS),
2038 Self::drain_reconnect_buffer(
2039 &mut reconnect_buffer,
2040 &mut active_writer,
2041 &connection_state,
2042 &auth_tracker,
2043 reconnect_buffer_waits_for_auth.as_ref(),
2044 ),
2045 )
2046 .await;
2047
2048 let send_error = drain_result.unwrap_or_else(|_| {
2049 log::warn!(
2050 "Timed out draining reconnect buffer after {WRITE_TIMEOUT_SECS}s, {} messages remain",
2051 reconnect_buffer.len()
2052 );
2053 true
2054 });
2055
2056 if send_error {
2058 _ = request_websocket_reconnect(
2059 &connection_state,
2060 &reconnect_published,
2061 state_sink.as_ref(),
2062 &auth_tracker,
2063 &controller_notify,
2064 || {},
2065 );
2066 }
2067
2068 continue;
2069 }
2070 ReconnectBufferAction::Discard => {
2071 log::warn!(
2072 "Discarding {} buffered messages after authentication failed",
2073 reconnect_buffer.len()
2074 );
2075 reconnect_buffer.clear();
2076 continue;
2077 }
2078 ReconnectBufferAction::Wait => {}
2079 }
2080 }
2081
2082 match dst::time::timeout(check_interval, writer_rx.recv()).await {
2083 Ok(Some((msg, permit))) => {
2084 let mode = ConnectionMode::from_atomic(&connection_state);
2086
2087 if matches!(mode, ConnectionMode::Disconnect | ConnectionMode::Closed) {
2088 break;
2089 }
2090
2091 match msg {
2092 WriterCommand::Update(new_writer, tx) => {
2093 log::debug!("Received new writer");
2094
2095 dst::time::sleep(Duration::from_millis(100)).await;
2097
2098 _ = dst::time::timeout(
2101 Duration::from_secs(GRACEFUL_SHUTDOWN_TIMEOUT_SECS),
2102 active_writer.close(),
2103 )
2104 .await;
2105
2106 active_writer = new_writer;
2107 let epoch = connection_epoch.fetch_add(1, Ordering::AcqRel) + 1;
2108 log::debug!("Updated writer: epoch={epoch}");
2109
2110 if let Err(e) = tx.send(epoch) {
2111 log::error!(
2112 "Failed to report writer update to controller: {e:?}"
2113 );
2114 }
2115 }
2116 WriterCommand::Send(msg)
2117 if mode.is_reconnect()
2118 || (!reconnect_buffer.is_empty() && !msg.is_control()) =>
2119 {
2120 Self::buffer_for_replay(&mut reconnect_buffer, msg, permit);
2121 }
2122 WriterCommand::Heartbeat(_)
2123 | WriterCommand::SendPongOnConnection { .. }
2124 if mode.is_reconnect() => {}
2125 WriterCommand::SendOnConnection { response_tx, .. }
2126 if mode.is_reconnect() =>
2127 {
2128 _ = response_tx.send(Err(SendError::ConnectionChanged));
2129 }
2130 WriterCommand::SendPongOnConnection {
2131 data,
2132 connection_epoch: expected_epoch,
2133 } => {
2134 let epoch = connection_epoch.load(Ordering::Acquire);
2135 if epoch != expected_epoch {
2136 continue;
2137 }
2138
2139 let send_result = dst::time::timeout(
2140 Duration::from_secs(WRITE_TIMEOUT_SECS),
2141 active_writer.send(Message::Pong(data.into())),
2142 )
2143 .await;
2144
2145 let send_failed = match send_result {
2146 Ok(Ok(())) => false,
2147 Ok(Err(e)) => {
2148 if is_connection_drop_transport_error(&e) {
2149 log::warn!("Failed to send pong: {e}");
2150 } else {
2151 log::error!("Failed to send pong: {e}");
2152 }
2153
2154 true
2155 }
2156 Err(_) => {
2157 log::warn!(
2158 "Timed out sending pong after {WRITE_TIMEOUT_SECS}s"
2159 );
2160 true
2161 }
2162 };
2163
2164 if send_failed
2165 && request_websocket_reconnect(
2166 &connection_state,
2167 &reconnect_published,
2168 state_sink.as_ref(),
2169 &auth_tracker,
2170 &controller_notify,
2171 || {},
2172 ) == ReconnectRequestOutcome::Accepted
2173 {
2174 log::warn!("Writer triggering reconnect");
2175 }
2176 }
2177 WriterCommand::SendOnConnection {
2178 message,
2179 connection_epoch: expected_epoch,
2180 response_tx,
2181 } => {
2182 let epoch = connection_epoch.load(Ordering::Acquire);
2183 if epoch != expected_epoch {
2184 _ = response_tx.send(Err(SendError::ConnectionChanged));
2185 continue;
2186 }
2187
2188 let send_result = dst::time::timeout(
2189 Duration::from_secs(WRITE_TIMEOUT_SECS),
2190 active_writer.send(message),
2191 )
2192 .await;
2193
2194 let result = match send_result {
2198 Ok(Ok(())) => Ok(()),
2199 Ok(Err(e)) => {
2200 if is_connection_drop_transport_error(&e) {
2201 log::warn!("Failed to send message: {e}");
2202 } else {
2203 log::error!("Failed to send message: {e}");
2204 }
2205
2206 Err(SendError::BrokenPipe(e.to_string()))
2207 }
2208 Err(_) => {
2209 log::warn!(
2210 "Timed out sending message after {WRITE_TIMEOUT_SECS}s"
2211 );
2212
2213 Err(SendError::WriteTimeout)
2214 }
2215 };
2216
2217 let send_failed = result.is_err();
2218 _ = response_tx.send(result);
2219
2220 if send_failed
2221 && request_websocket_reconnect(
2222 &connection_state,
2223 &reconnect_published,
2224 state_sink.as_ref(),
2225 &auth_tracker,
2226 &controller_notify,
2227 || {},
2228 ) == ReconnectRequestOutcome::Accepted
2229 {
2230 log::warn!("Writer triggering reconnect");
2231 }
2232 }
2233 WriterCommand::Send(msg) => {
2234 let send_failed =
2235 Self::write_outbound(&mut active_writer, msg.clone()).await;
2236 if send_failed {
2237 Self::buffer_for_replay(&mut reconnect_buffer, msg, permit);
2238
2239 if request_websocket_reconnect(
2241 &connection_state,
2242 &reconnect_published,
2243 state_sink.as_ref(),
2244 &auth_tracker,
2245 &controller_notify,
2246 || {},
2247 ) == ReconnectRequestOutcome::Accepted
2248 {
2249 log::warn!("Writer triggering reconnect");
2250 }
2251 }
2252 }
2253 WriterCommand::Heartbeat(msg) => {
2254 let send_failed =
2255 Self::write_outbound(&mut active_writer, msg).await;
2256
2257 if send_failed
2258 && request_websocket_reconnect(
2259 &connection_state,
2260 &reconnect_published,
2261 state_sink.as_ref(),
2262 &auth_tracker,
2263 &controller_notify,
2264 || {},
2265 ) == ReconnectRequestOutcome::Accepted
2266 {
2267 log::warn!("Writer triggering reconnect");
2268 }
2269 }
2270 }
2271 }
2272 Ok(None) => {
2273 log::debug!("Writer channel closed, terminating writer task");
2275 break;
2276 }
2277 Err(_) => {
2278 }
2280 }
2281 }
2282
2283 _ = dst::time::timeout(
2286 Duration::from_secs(GRACEFUL_SHUTDOWN_TIMEOUT_SECS),
2287 active_writer.close(),
2288 )
2289 .await;
2290
2291 log_task_stopped("write");
2292 })
2293 }
2294
2295 async fn write_outbound(writer: &mut MessageWriter, msg: Message) -> bool {
2296 let send_result =
2297 dst::time::timeout(Duration::from_secs(WRITE_TIMEOUT_SECS), writer.send(msg)).await;
2298
2299 match send_result {
2300 Ok(Ok(())) => false,
2301 Ok(Err(e)) => {
2302 if is_connection_drop_transport_error(&e) {
2303 log::warn!("Failed to send message: {e}");
2304 } else {
2305 log::error!("Failed to send message: {e}");
2306 }
2307
2308 true
2309 }
2310 Err(_) => {
2311 log::warn!("Timed out sending message after {WRITE_TIMEOUT_SECS}s");
2312 true
2313 }
2314 }
2315 }
2316
2317 fn spawn_heartbeat_task(
2318 connection_state: Arc<AtomicU8>,
2319 heartbeat_secs: u64,
2320 message: Option<String>,
2321 writer_tx: WriterSender<WriterCommand>,
2322 ) -> dst::task::JoinHandle<()> {
2323 log_task_started("heartbeat");
2324
2325 dst::task::spawn(async move {
2326 let interval = Duration::from_secs(heartbeat_secs);
2327 loop {
2328 dst::time::sleep(interval).await;
2329
2330 match ConnectionMode::from_u8(connection_state.load(Ordering::SeqCst)) {
2331 ConnectionMode::Active => {
2332 let msg = match &message {
2333 Some(text) => {
2334 WriterCommand::Heartbeat(Message::Text(text.clone().into()))
2335 }
2336 None => WriterCommand::Heartbeat(Message::Ping(vec![].into())),
2337 };
2338
2339 match writer_tx.send_control(msg) {
2340 Ok(()) => log::trace!("Sent heartbeat to writer task"),
2341 Err(e) => {
2342 log::warn!("Failed to send heartbeat to writer task: {e}");
2343 }
2344 }
2345 }
2346 ConnectionMode::Reconnect => {}
2347 ConnectionMode::Disconnect | ConnectionMode::Closed => break,
2348 }
2349 }
2350
2351 log_task_stopped("heartbeat");
2352 })
2353 }
2354}
2355
2356fn heartbeat_timeout_exceeded(
2357 last_frame_time: dst::time::Instant,
2358 timeout: Option<Duration>,
2359) -> bool {
2360 if let Some(timeout) = timeout {
2361 let elapsed = last_frame_time.elapsed();
2362 if elapsed >= timeout {
2363 log::warn!(
2364 "Heartbeat timeout: no frame received for {:.1}s",
2365 elapsed.as_secs_f64()
2366 );
2367 return true;
2368 }
2369 }
2370
2371 false
2372}
2373
2374fn idle_timeout_exceeded(
2375 last_data_time: dst::time::Instant,
2376 idle_timeout: Option<Duration>,
2377) -> bool {
2378 if let Some(timeout) = idle_timeout {
2379 let idle_duration = last_data_time.elapsed();
2380 if idle_duration >= timeout {
2381 log::warn!(
2382 "Read idle timeout: no data received for {:.1}s",
2383 idle_duration.as_secs_f64()
2384 );
2385 return true;
2386 }
2387 }
2388
2389 false
2390}
2391
2392impl Drop for WebSocketClientInner {
2393 fn drop(&mut self) {
2394 if let Some(read_fence) = self.read_fence.take() {
2395 read_fence.invalidate();
2396 }
2397
2398 if let Some(ref read_task) = self.read_task.take()
2399 && !read_task.is_finished()
2400 {
2401 read_task.abort();
2402 log_task_aborted("read");
2403 }
2404
2405 if !self.write_task.is_finished() {
2406 self.write_task.abort();
2407 log_task_aborted("write");
2408 }
2409
2410 if let Some(ref handle) = self.heartbeat_task.take()
2411 && !handle.is_finished()
2412 {
2413 handle.abort();
2414 log_task_aborted("heartbeat");
2415 }
2416 }
2417}
2418
2419#[expect(
2420 clippy::missing_fields_in_debug,
2421 reason = "handler closures and internal task handles are intentionally omitted"
2422)]
2423impl Debug for WebSocketClientInner {
2424 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2425 f.debug_struct(stringify!(WebSocketClientInner))
2426 .field("config", &self.config)
2427 .field(
2428 "connection_mode",
2429 &ConnectionMode::from_atomic(&self.connection_mode),
2430 )
2431 .field("connect_timeout", &self.connect_timeout)
2432 .field("is_stream_mode", &self.handler.is_none())
2433 .finish()
2434 }
2435}
2436
2437#[derive(Clone)]
2438enum IncomingHandler {
2439 Message(MessageHandler),
2440 Epoch(EpochMessageHandler),
2441}
2442
2443impl IncomingHandler {
2444 fn handle(&self, connection_epoch: u64, message: Message) {
2445 match self {
2446 Self::Message(handler) => handler(message),
2447 Self::Epoch(handler) => handler(connection_epoch, message),
2448 }
2449 }
2450}
2451
2452#[derive(Clone)]
2453enum IncomingPingHandler {
2454 Ping(PingHandler),
2455 Epoch(EpochPingHandler),
2456}
2457
2458impl IncomingPingHandler {
2459 fn handle(&self, connection_epoch: u64, data: Vec<u8>) {
2460 match self {
2461 Self::Ping(handler) => handler(data),
2462 Self::Epoch(handler) => handler(connection_epoch, data),
2463 }
2464 }
2465}
2466
2467#[derive(Clone, Copy, Debug, Eq, PartialEq)]
2468enum ReconnectBufferAction {
2469 Drain,
2470 Wait,
2471 Discard,
2472}
2473
2474pub struct WebSocketClient {
2481 pub(crate) controller_task: dst::task::JoinHandle<()>,
2482 pub(crate) connection_mode: Arc<AtomicU8>,
2483 pub(crate) connection_epoch: Arc<AtomicU64>,
2484 pub(crate) state_notify: Arc<tokio::sync::Notify>,
2485 pub(crate) connect_timeout: Duration,
2486 pub(crate) rate_limiter: Arc<RateLimiter<Ustr, MonotonicClock>>,
2487 pub(crate) writer_tx: WriterSender<WriterCommand>,
2488 auth_tracker: Arc<OnceLock<AuthTracker>>,
2489 reconnect_buffer_waits_for_auth: Arc<AtomicBool>,
2490 reconnect_headers: ReconnectHeaders,
2491 state_sink: Option<SocketStateSink>,
2492 controller_lifecycle: Arc<ControllerLifecycle>,
2493 controller_notify: Arc<tokio::sync::Notify>,
2494 reconnect_published: Arc<AtomicBool>,
2495 reconnect_supported: bool,
2496}
2497
2498#[derive(Clone)]
2502pub struct ReconnectHeaders {
2503 inner: Arc<RwLock<Vec<(String, String)>>>,
2504}
2505
2506impl ReconnectHeaders {
2507 fn new(headers: Vec<(String, String)>) -> Self {
2508 Self {
2509 inner: Arc::new(RwLock::new(headers)),
2510 }
2511 }
2512
2513 pub fn update(&self, name: &str, value: &str) -> Result<(), TransportError> {
2519 let name = HeaderName::from_bytes(name.as_bytes()).map_err(|e| {
2520 TransportError::Io(std::io::Error::new(
2521 std::io::ErrorKind::InvalidInput,
2522 format!("Invalid WebSocket reconnect header name: {e}"),
2523 ))
2524 })?;
2525
2526 HeaderValue::from_str(value).map_err(|e| {
2527 TransportError::Io(std::io::Error::new(
2528 std::io::ErrorKind::InvalidInput,
2529 format!("Invalid WebSocket reconnect header value: {e}"),
2530 ))
2531 })?;
2532
2533 let name = name.as_str();
2534 let mut headers = self.inner.write();
2535 headers.retain(|(existing, _)| !existing.eq_ignore_ascii_case(name));
2536 headers.push((name.to_string(), value.to_string()));
2537 Ok(())
2538 }
2539
2540 fn snapshot(&self) -> Vec<(String, String)> {
2541 self.inner.read().clone()
2542 }
2543}
2544
2545impl Debug for ReconnectHeaders {
2546 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2547 f.debug_struct(stringify!(ReconnectHeaders))
2548 .finish_non_exhaustive()
2549 }
2550}
2551
2552#[derive(Clone)]
2554pub struct WebSocketReconnectHandle {
2555 connection_mode: Arc<AtomicU8>,
2556 auth_tracker: Arc<OnceLock<AuthTracker>>,
2557 state_sink: Option<SocketStateSink>,
2558 controller_lifecycle: Arc<ControllerLifecycle>,
2559 controller_notify: Arc<tokio::sync::Notify>,
2560 reconnect_published: Arc<AtomicBool>,
2561 supported: bool,
2562}
2563
2564impl Debug for WebSocketReconnectHandle {
2565 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2566 f.debug_struct(stringify!(WebSocketReconnectHandle))
2567 .field(
2568 "connection_mode",
2569 &ConnectionMode::from_atomic(&self.connection_mode),
2570 )
2571 .field("supported", &self.supported)
2572 .finish_non_exhaustive()
2573 }
2574}
2575
2576impl WebSocketReconnectHandle {
2577 #[must_use]
2585 pub fn request_reconnect(&self) -> ReconnectRequestOutcome {
2586 if !self.supported {
2587 return ReconnectRequestOutcome::Unsupported;
2588 }
2589
2590 let Some(request) = self.controller_lifecycle.enter_request() else {
2591 return ReconnectRequestOutcome::Closed;
2592 };
2593
2594 let mut request = Some(request);
2595 request_websocket_reconnect(
2596 &self.connection_mode,
2597 &self.reconnect_published,
2598 self.state_sink.as_ref(),
2599 &self.auth_tracker,
2600 &self.controller_notify,
2601 || drop(request.take()),
2602 )
2603 }
2604}
2605
2606struct ReconnectPublication<'a> {
2607 published: &'a AtomicBool,
2608 controller_notify: &'a tokio::sync::Notify,
2609}
2610
2611impl Drop for ReconnectPublication<'_> {
2612 fn drop(&mut self) {
2613 self.published.store(true, Ordering::SeqCst);
2614 self.controller_notify.notify_one();
2615 }
2616}
2617
2618fn request_websocket_reconnect<F>(
2619 connection_mode: &AtomicU8,
2620 reconnect_published: &AtomicBool,
2621 state_sink: Option<&SocketStateSink>,
2622 auth_tracker: &OnceLock<AuthTracker>,
2623 controller_notify: &tokio::sync::Notify,
2624 on_handoff: F,
2625) -> ReconnectRequestOutcome
2626where
2627 F: FnOnce(),
2628{
2629 if reconnect_published
2630 .compare_exchange(true, false, Ordering::SeqCst, Ordering::SeqCst)
2631 .is_err()
2632 {
2633 return match ConnectionMode::from_atomic(connection_mode) {
2634 ConnectionMode::Active | ConnectionMode::Reconnect => {
2635 ReconnectRequestOutcome::AlreadyReconnecting
2636 }
2637 ConnectionMode::Disconnect => ReconnectRequestOutcome::Disconnected,
2638 ConnectionMode::Closed => ReconnectRequestOutcome::Closed,
2639 };
2640 }
2641
2642 let outcome = ConnectionMode::request_reconnect_outcome(connection_mode);
2643 if outcome != ReconnectRequestOutcome::Accepted {
2644 reconnect_published.store(true, Ordering::SeqCst);
2645 return outcome;
2646 }
2647
2648 let _publication = ReconnectPublication {
2649 published: reconnect_published,
2650 controller_notify,
2651 };
2652
2653 if let Some(tracker) = auth_tracker.get() {
2654 tracker.invalidate();
2655 }
2656
2657 controller_notify.notify_one();
2658 on_handoff();
2659
2660 if let Some(sink) = state_sink {
2661 sink.publish_websocket(SocketState::Disconnected);
2662 }
2663
2664 ReconnectRequestOutcome::Accepted
2665}
2666
2667#[cfg(test)]
2668mod reconnect_request_tests {
2669 use std::sync::{Arc, OnceLock, atomic::AtomicU8};
2670
2671 use rstest::rstest;
2672
2673 use super::*;
2674
2675 fn handle(
2676 mode: ConnectionMode,
2677 supported: bool,
2678 ) -> (
2679 WebSocketReconnectHandle,
2680 AuthTracker,
2681 Arc<tokio::sync::Notify>,
2682 ) {
2683 let tracker = AuthTracker::new();
2684 let _receiver = tracker.begin();
2685 tracker.succeed();
2686 let auth_tracker = Arc::new(OnceLock::new());
2687 auth_tracker
2688 .set(tracker.clone())
2689 .expect("auth tracker should be unset");
2690 let notify = Arc::new(tokio::sync::Notify::new());
2691
2692 let handle = WebSocketReconnectHandle {
2693 connection_mode: Arc::new(AtomicU8::new(mode.as_u8())),
2694 auth_tracker,
2695 state_sink: None,
2696 controller_lifecycle: Arc::new(ControllerLifecycle::new()),
2697 controller_notify: Arc::clone(¬ify),
2698 reconnect_published: Arc::new(AtomicBool::new(true)),
2699 supported,
2700 };
2701
2702 (handle, tracker, notify)
2703 }
2704
2705 #[rstest]
2706 #[tokio::test]
2707 async fn accepted_request_invalidates_auth_and_wakes_controller_once() {
2708 let (handle, tracker, notify) = handle(ConnectionMode::Active, true);
2709
2710 assert_eq!(
2711 handle.request_reconnect(),
2712 ReconnectRequestOutcome::Accepted
2713 );
2714 assert!(!tracker.is_authenticated());
2715 tokio::time::timeout(Duration::from_millis(10), notify.notified())
2716 .await
2717 .expect("accepted request should notify controller");
2718
2719 let _receiver = tracker.begin();
2720 tracker.succeed();
2721 assert_eq!(
2722 handle.request_reconnect(),
2723 ReconnectRequestOutcome::AlreadyReconnecting
2724 );
2725 assert!(tracker.is_authenticated());
2726 assert!(
2727 tokio::time::timeout(Duration::from_millis(10), notify.notified())
2728 .await
2729 .is_err(),
2730 "duplicate request should not notify controller",
2731 );
2732
2733 handle
2734 .connection_mode
2735 .store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
2736 assert_eq!(
2737 handle.request_reconnect(),
2738 ReconnectRequestOutcome::Accepted
2739 );
2740 assert!(!tracker.is_authenticated());
2741 tokio::time::timeout(Duration::from_millis(10), notify.notified())
2742 .await
2743 .expect("a later reconnect cycle should notify the controller");
2744 }
2745
2746 #[rstest]
2747 #[case(
2748 ConnectionMode::Disconnect,
2749 true,
2750 ReconnectRequestOutcome::Disconnected
2751 )]
2752 #[case(
2753 ConnectionMode::Reconnect,
2754 true,
2755 ReconnectRequestOutcome::AlreadyReconnecting
2756 )]
2757 #[case(ConnectionMode::Closed, true, ReconnectRequestOutcome::Closed)]
2758 #[case(ConnectionMode::Active, false, ReconnectRequestOutcome::Unsupported)]
2759 #[tokio::test]
2760 async fn rejected_request_preserves_auth_and_does_not_wake_controller(
2761 #[case] mode: ConnectionMode,
2762 #[case] supported: bool,
2763 #[case] expected: ReconnectRequestOutcome,
2764 ) {
2765 let (handle, tracker, notify) = handle(mode, supported);
2766
2767 assert_eq!(handle.request_reconnect(), expected);
2768 assert!(tracker.is_authenticated());
2769 assert!(
2770 tokio::time::timeout(Duration::from_millis(10), notify.notified())
2771 .await
2772 .is_err(),
2773 "rejected request should not notify controller",
2774 );
2775 }
2776
2777 #[rstest]
2778 #[tokio::test]
2779 async fn reconnect_loss_callback_rejects_nested_request() {
2780 let connection_mode = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
2781 let controller_notify = Arc::new(tokio::sync::Notify::new());
2782 let handle_slot = Arc::new(OnceLock::<WebSocketReconnectHandle>::new());
2783 let handle_slot_callback = Arc::clone(&handle_slot);
2784 let nested_outcomes = Arc::new(parking_lot::Mutex::new(Vec::new()));
2785 let nested_outcomes_callback = Arc::clone(&nested_outcomes);
2786 let states = Arc::new(parking_lot::Mutex::new(Vec::new()));
2787 let states_callback = Arc::clone(&states);
2788 let sink = SocketStateSink::new(move |state| {
2789 states_callback.lock().push(state);
2790 nested_outcomes_callback
2791 .lock()
2792 .push(handle_slot_callback.get().unwrap().request_reconnect());
2793 });
2794
2795 let handle = WebSocketReconnectHandle {
2796 connection_mode,
2797 auth_tracker: Arc::new(OnceLock::new()),
2798 state_sink: Some(sink),
2799 controller_lifecycle: Arc::new(ControllerLifecycle::new()),
2800 controller_notify,
2801 reconnect_published: Arc::new(AtomicBool::new(true)),
2802 supported: true,
2803 };
2804
2805 handle_slot.set(handle.clone()).unwrap();
2806 let (result_tx, result_rx) = tokio::sync::oneshot::channel();
2807 std::thread::spawn(move || {
2808 _ = result_tx.send(handle.request_reconnect());
2809 });
2810
2811 assert_eq!(
2812 tokio::time::timeout(Duration::from_secs(1), result_rx)
2813 .await
2814 .expect("reentrant reconnect callback deadlocked")
2815 .unwrap(),
2816 ReconnectRequestOutcome::Accepted
2817 );
2818 assert_eq!(
2819 *nested_outcomes.lock(),
2820 vec![ReconnectRequestOutcome::AlreadyReconnecting]
2821 );
2822 assert_eq!(*states.lock(), vec![SocketState::Disconnected]);
2823 }
2824
2825 #[rstest]
2826 fn closed_stream_handle_remains_unsupported() {
2827 let (handle, tracker, _notify) = handle(ConnectionMode::Closed, false);
2828 handle.controller_lifecycle.close_and_abort();
2829
2830 assert_eq!(
2831 handle.request_reconnect(),
2832 ReconnectRequestOutcome::Unsupported
2833 );
2834 assert!(tracker.is_authenticated());
2835 }
2836}
2837
2838impl Debug for WebSocketClient {
2839 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2840 f.debug_struct(stringify!(WebSocketClient)).finish()
2841 }
2842}
2843
2844#[bon::bon]
2845impl WebSocketClient {
2846 #[builder(
2865 builder_type = WebSocketClientStreamBuilder,
2866 finish_fn = connect
2867 )]
2868 pub async fn stream_builder(
2869 config: WebSocketConfig,
2870 #[builder(default)] keyed_quotas: Vec<(String, Quota)>,
2871 default_quota: Option<Quota>,
2872 state_sink: Option<SocketStateSink>,
2873 ) -> Result<(MessageReader, Self), TransportError> {
2874 install_cryptographic_provider();
2875 validate_honored_inbound_limits(&config)?;
2876
2877 let connect_timeout = Duration::from_secs(10);
2880
2881 let (writer, reader) = dst::time::timeout(
2882 connect_timeout,
2883 Box::pin(WebSocketClientInner::connect_with_server(
2884 &config.url,
2885 config.headers.clone(),
2886 config.backend,
2887 config.proxy_url.as_deref(),
2888 config.max_message_size_bytes,
2889 config.max_frame_size_bytes,
2890 )),
2891 )
2892 .await
2893 .map_err(|_| {
2894 TransportError::Io(std::io::Error::new(
2895 std::io::ErrorKind::TimedOut,
2896 format!(
2897 "connection timed out after {}s",
2898 connect_timeout.as_secs_f64()
2899 ),
2900 ))
2901 })??;
2902
2903 let inner =
2905 WebSocketClientInner::new_with_writer_and_state_sink(config, writer, state_sink)?;
2906
2907 let connection_mode = Arc::clone(&inner.connection_mode);
2908 let connection_epoch = Arc::clone(&inner.connection_epoch);
2909 let state_notify = Arc::clone(&inner.state_notify);
2910 let controller_notify = Arc::clone(&inner.controller_notify);
2911 let reconnect_published = Arc::clone(&inner.reconnect_published);
2912 let connect_timeout = inner.connect_timeout;
2913 let auth_tracker = Arc::clone(&inner.auth_tracker);
2914 let reconnect_buffer_waits_for_auth = Arc::clone(&inner.reconnect_buffer_waits_for_auth);
2915 let reconnect_headers = inner.reconnect_headers.clone();
2916 let state_sink = inner.state_sink.clone();
2917 let rate_limiter = Self::resolve_rate_limiter(default_quota, keyed_quotas, None)?;
2918 let writer_tx = inner.writer_tx.clone();
2919 let controller_lifecycle = Arc::new(ControllerLifecycle::new());
2920
2921 let controller_task = Self::spawn_controller_task(
2922 inner,
2923 Arc::clone(&connection_mode),
2924 Arc::clone(&state_notify),
2925 Arc::clone(&auth_tracker),
2926 Arc::clone(&controller_lifecycle),
2927 Arc::clone(&controller_notify),
2928 Arc::clone(&reconnect_published),
2929 );
2930 let abort_handle = controller_task.abort_handle();
2931 controller_lifecycle.set_abort(move || abort_handle.abort());
2932
2933 Ok((
2934 reader,
2935 Self {
2936 controller_task,
2937 connection_mode,
2938 connection_epoch,
2939 state_notify,
2940 connect_timeout,
2941 rate_limiter,
2942 writer_tx,
2943 auth_tracker,
2944 reconnect_buffer_waits_for_auth,
2945 reconnect_headers,
2946 state_sink,
2947 controller_lifecycle,
2948 controller_notify,
2949 reconnect_published,
2950 reconnect_supported: false,
2951 },
2952 ))
2953 }
2954
2955 #[builder(finish_fn = connect)]
2999 pub async fn builder(
3000 config: WebSocketConfig,
3001 message_handler: MessageHandler,
3002 ping_handler: Option<PingHandler>,
3003 close_handler: Option<CloseHandler>,
3004 #[builder(default)] keyed_quotas: Vec<(String, Quota)>,
3005 default_quota: Option<Quota>,
3006 rate_limiter: Option<Arc<RateLimiter<Ustr, MonotonicClock>>>,
3007 state_sink: Option<SocketStateSink>,
3008 connection_rate_limiter: Option<Arc<RateLimiter<Ustr, MonotonicClock>>>,
3009 #[builder(default)] connection_rate_keys: Arc<[Ustr]>,
3010 initial_connect_retry_policy: Option<InitialConnectRetryPolicy>,
3011 cancellation_token: Option<CancellationToken>,
3012 ) -> Result<Self, TransportError> {
3013 let rate_limiter = Self::resolve_rate_limiter(default_quota, keyed_quotas, rate_limiter)?;
3014 let connection_rate_limit =
3015 Self::resolve_connection_rate_limit(connection_rate_limiter, connection_rate_keys)?;
3016 Self::connect_with_handler_scoped(
3017 config,
3018 IncomingHandler::Message(message_handler),
3019 ping_handler.map(IncomingPingHandler::Ping),
3020 close_handler,
3021 rate_limiter,
3022 state_sink,
3023 connection_rate_limit,
3024 InitialConnectOptions {
3025 retry_policy: initial_connect_retry_policy,
3026 cancellation_token,
3027 },
3028 )
3029 .await
3030 }
3031
3032 #[builder(
3058 builder_type = WebSocketClientEpochBuilder,
3059 finish_fn = connect
3060 )]
3061 pub async fn epoch_builder(
3062 config: WebSocketConfig,
3063 epoch_handler: EpochMessageHandler,
3064 ping_handler: Option<PingHandler>,
3065 epoch_ping_handler: Option<EpochPingHandler>,
3066 close_handler: Option<CloseHandler>,
3067 #[builder(default)] keyed_quotas: Vec<(String, Quota)>,
3068 default_quota: Option<Quota>,
3069 rate_limiter: Option<Arc<RateLimiter<Ustr, MonotonicClock>>>,
3070 state_sink: Option<SocketStateSink>,
3071 connection_rate_limiter: Option<Arc<RateLimiter<Ustr, MonotonicClock>>>,
3072 #[builder(default)] connection_rate_keys: Arc<[Ustr]>,
3073 initial_connect_retry_policy: Option<InitialConnectRetryPolicy>,
3074 cancellation_token: Option<CancellationToken>,
3075 ) -> Result<Self, TransportError> {
3076 let ping_handler = match (ping_handler, epoch_ping_handler) {
3077 (Some(_), Some(_)) => {
3078 return Err(TransportError::Io(std::io::Error::new(
3079 std::io::ErrorKind::InvalidInput,
3080 "Cannot configure both ping_handler and epoch_ping_handler",
3081 )));
3082 }
3083 (Some(handler), None) => Some(IncomingPingHandler::Ping(handler)),
3084 (None, Some(handler)) => Some(IncomingPingHandler::Epoch(handler)),
3085 (None, None) => None,
3086 };
3087
3088 let rate_limiter = Self::resolve_rate_limiter(default_quota, keyed_quotas, rate_limiter)?;
3089 let connection_rate_limit =
3090 Self::resolve_connection_rate_limit(connection_rate_limiter, connection_rate_keys)?;
3091 Self::connect_with_handler_scoped(
3092 config,
3093 IncomingHandler::Epoch(epoch_handler),
3094 ping_handler,
3095 close_handler,
3096 rate_limiter,
3097 state_sink,
3098 connection_rate_limit,
3099 InitialConnectOptions {
3100 retry_policy: initial_connect_retry_policy,
3101 cancellation_token,
3102 },
3103 )
3104 .await
3105 }
3106
3107 fn resolve_connection_rate_limit(
3108 rate_limiter: Option<Arc<RateLimiter<Ustr, MonotonicClock>>>,
3109 keys: Arc<[Ustr]>,
3110 ) -> Result<Option<ConnectionRateLimit>, TransportError> {
3111 if rate_limiter.is_none() && !keys.is_empty() {
3112 return Err(TransportError::Io(std::io::Error::new(
3113 std::io::ErrorKind::InvalidInput,
3114 "Connection rate keys require a connection rate limiter",
3115 )));
3116 }
3117
3118 if rate_limiter.is_some() && keys.is_empty() {
3119 return Err(TransportError::Io(std::io::Error::new(
3120 std::io::ErrorKind::InvalidInput,
3121 "Connection rate limiter requires at least one connection rate key",
3122 )));
3123 }
3124
3125 Ok(rate_limiter.map(|limiter| ConnectionRateLimit { limiter, keys }))
3126 }
3127
3128 fn resolve_rate_limiter(
3129 default_quota: Option<Quota>,
3130 keyed_quotas: Vec<(String, Quota)>,
3131 rate_limiter: Option<Arc<RateLimiter<Ustr, MonotonicClock>>>,
3132 ) -> Result<Arc<RateLimiter<Ustr, MonotonicClock>>, TransportError> {
3133 if let Some(rate_limiter) = rate_limiter {
3134 if default_quota.is_some() || !keyed_quotas.is_empty() {
3135 return Err(TransportError::Io(std::io::Error::new(
3136 std::io::ErrorKind::InvalidInput,
3137 "Cannot combine a shared rate limiter with quota configuration",
3138 )));
3139 }
3140
3141 return Ok(rate_limiter);
3142 }
3143
3144 let keyed_quotas = keyed_quotas
3145 .into_iter()
3146 .map(|(key, quota)| (Ustr::from(&key), quota))
3147 .collect();
3148 Ok(Arc::new(RateLimiter::new_with_quota(
3149 default_quota,
3150 keyed_quotas,
3151 )))
3152 }
3153
3154 #[expect(
3155 clippy::too_many_arguments,
3156 reason = "handler delivery and connection controls are independent options"
3157 )]
3158 async fn connect_with_handler_scoped(
3159 config: WebSocketConfig,
3160 handler: IncomingHandler,
3161 ping_handler: Option<IncomingPingHandler>,
3162 close_handler: Option<CloseHandler>,
3163 rate_limiter: Arc<RateLimiter<Ustr, MonotonicClock>>,
3164 state_sink: Option<SocketStateSink>,
3165 connection_rate_limit: Option<ConnectionRateLimit>,
3166 initial_connect_options: InitialConnectOptions,
3167 ) -> Result<Self, TransportError> {
3168 log::debug!("Connecting");
3169 let inner = WebSocketClientInner::connect_url_with_handler(
3170 config,
3171 Some(handler),
3172 ping_handler,
3173 close_handler,
3174 state_sink,
3175 connection_rate_limit,
3176 initial_connect_options,
3177 )
3178 .await?;
3179 let connection_mode = Arc::clone(&inner.connection_mode);
3180 let connection_epoch = Arc::clone(&inner.connection_epoch);
3181 let state_notify = Arc::clone(&inner.state_notify);
3182 let controller_notify = Arc::clone(&inner.controller_notify);
3183 let reconnect_published = Arc::clone(&inner.reconnect_published);
3184 let writer_tx = inner.writer_tx.clone();
3185 let connect_timeout = inner.connect_timeout;
3186 let auth_tracker = Arc::clone(&inner.auth_tracker);
3187 let reconnect_buffer_waits_for_auth = Arc::clone(&inner.reconnect_buffer_waits_for_auth);
3188 let reconnect_headers = inner.reconnect_headers.clone();
3189 let state_sink = inner.state_sink.clone();
3190 let controller_lifecycle = Arc::new(ControllerLifecycle::new());
3191
3192 let controller_task = Self::spawn_controller_task(
3193 inner,
3194 Arc::clone(&connection_mode),
3195 Arc::clone(&state_notify),
3196 Arc::clone(&auth_tracker),
3197 Arc::clone(&controller_lifecycle),
3198 Arc::clone(&controller_notify),
3199 Arc::clone(&reconnect_published),
3200 );
3201 let abort_handle = controller_task.abort_handle();
3202 controller_lifecycle.set_abort(move || abort_handle.abort());
3203
3204 Ok(Self {
3205 controller_task,
3206 connection_mode,
3207 connection_epoch,
3208 state_notify,
3209 connect_timeout,
3210 rate_limiter,
3211 writer_tx,
3212 auth_tracker,
3213 reconnect_buffer_waits_for_auth,
3214 reconnect_headers,
3215 state_sink,
3216 controller_lifecycle,
3217 controller_notify,
3218 reconnect_published,
3219 reconnect_supported: true,
3220 })
3221 }
3222
3223 #[must_use]
3225 pub fn reconnect_headers(&self) -> ReconnectHeaders {
3226 self.reconnect_headers.clone()
3227 }
3228
3229 #[must_use]
3231 pub fn reconnect_handle(&self) -> WebSocketReconnectHandle {
3232 WebSocketReconnectHandle {
3233 connection_mode: Arc::clone(&self.connection_mode),
3234 auth_tracker: Arc::clone(&self.auth_tracker),
3235 state_sink: self.state_sink.clone(),
3236 controller_lifecycle: Arc::clone(&self.controller_lifecycle),
3237 controller_notify: Arc::clone(&self.controller_notify),
3238 reconnect_published: Arc::clone(&self.reconnect_published),
3239 supported: self.reconnect_supported,
3240 }
3241 }
3242
3243 #[must_use]
3248 pub fn request_reconnect(&self) -> bool {
3249 self.reconnect_handle().request_reconnect() == ReconnectRequestOutcome::Accepted
3250 }
3251
3252 #[must_use]
3254 pub fn connection_mode(&self) -> ConnectionMode {
3255 ConnectionMode::from_atomic(&self.connection_mode)
3256 }
3257
3258 #[must_use]
3263 pub fn connection_epoch(&self) -> u64 {
3264 self.connection_epoch.load(Ordering::Acquire)
3265 }
3266
3267 #[must_use]
3272 pub fn connection_mode_atomic(&self) -> Arc<AtomicU8> {
3273 Arc::clone(&self.connection_mode)
3274 }
3275
3276 #[must_use]
3281 pub fn connection_epoch_atomic(&self) -> Arc<AtomicU64> {
3282 Arc::clone(&self.connection_epoch)
3283 }
3284
3285 #[inline]
3290 #[must_use]
3291 pub fn is_active(&self) -> bool {
3292 self.connection_mode().is_active()
3293 }
3294
3295 #[must_use]
3297 pub fn is_disconnected(&self) -> bool {
3298 self.controller_task.is_finished()
3299 }
3300
3301 #[inline]
3306 #[must_use]
3307 pub fn is_reconnecting(&self) -> bool {
3308 self.connection_mode().is_reconnect()
3309 }
3310
3311 pub fn set_auth_tracker(&self, tracker: AuthTracker, reconnect_buffer_waits_for_auth: bool) {
3322 let _ = self.auth_tracker.set(tracker);
3323 self.reconnect_buffer_waits_for_auth
3324 .store(reconnect_buffer_waits_for_auth, Ordering::Release);
3325 }
3326
3327 #[inline]
3331 #[must_use]
3332 pub fn is_disconnecting(&self) -> bool {
3333 self.connection_mode().is_disconnect()
3334 }
3335
3336 #[inline]
3342 #[must_use]
3343 pub fn is_closed(&self) -> bool {
3344 self.connection_mode().is_closed()
3345 }
3346
3347 #[inline]
3351 fn check_not_terminal(&self) -> Result<(), SendError> {
3352 match self.connection_mode() {
3353 ConnectionMode::Disconnect | ConnectionMode::Closed => Err(SendError::Closed),
3354 _ => Ok(()),
3355 }
3356 }
3357
3358 async fn await_rate_limit_or_closed(&self, keys: Option<&[Ustr]>) -> Result<(), SendError> {
3360 const CHECK_INTERVAL_MS: u64 = 100;
3361
3362 tokio::select! {
3363 biased;
3364 () = self.rate_limiter.await_keys_ready(keys) => Ok(()),
3365 () = async {
3366 loop {
3367 let mut notified = pin!(self.state_notify.notified());
3369 notified.as_mut().enable();
3370
3371 if matches!(self.connection_mode(), ConnectionMode::Disconnect | ConnectionMode::Closed) {
3372 break;
3373 }
3374 tokio::select! {
3375 biased;
3376 () = notified => {}
3377 () = dst::time::sleep(Duration::from_millis(CHECK_INTERVAL_MS)) => {}
3378 }
3379 }
3380 } => Err(SendError::Closed),
3381 }
3382 }
3383
3384 async fn wait_for_active(&self) -> Result<(), SendError> {
3390 const FALLBACK_INTERVAL_MS: u64 = 100;
3391
3392 let mode = self.connection_mode();
3393 if mode.is_active() {
3394 return Ok(());
3395 }
3396
3397 if matches!(mode, ConnectionMode::Disconnect | ConnectionMode::Closed) {
3398 return Err(SendError::Closed);
3399 }
3400
3401 log::debug!("Waiting for client to become ACTIVE before sending...");
3402
3403 let fallback_interval = Duration::from_millis(FALLBACK_INTERVAL_MS);
3404
3405 dst::time::timeout(self.connect_timeout, async {
3406 loop {
3407 let mut notified = pin!(self.state_notify.notified());
3409 notified.as_mut().enable();
3410
3411 let mode = self.connection_mode();
3412 if mode.is_active() {
3413 return Ok(());
3414 }
3415
3416 if matches!(mode, ConnectionMode::Disconnect | ConnectionMode::Closed) {
3417 return Err(());
3418 }
3419
3420 tokio::select! {
3421 biased;
3422 () = notified => {}
3423 () = dst::time::sleep(fallback_interval) => {}
3424 }
3425 }
3426 })
3427 .await
3428 .map_err(|_| SendError::Timeout)?
3429 .map_err(|()| SendError::Closed)
3430 }
3431
3432 pub fn notify_closed(&self) {
3445 let mode = self.connection_mode();
3446 if mode.is_disconnect() || mode.is_closed() {
3447 return;
3448 }
3449
3450 log::debug!("Stream reader signalled EOF, transitioning to CLOSED");
3451
3452 if ConnectionMode::close_websocket_on_loss(&self.connection_mode, self.state_sink.as_ref())
3453 {
3454 fail_registered_auth(self.auth_tracker.as_ref(), "WebSocket client closed");
3455 self.state_notify.notify_waiters();
3456 }
3457 }
3458
3459 pub async fn disconnect(&self) {
3463 log::debug!("Disconnecting");
3464
3465 if ConnectionMode::request_disconnect(&self.connection_mode)
3467 && let Some(tracker) = self.auth_tracker.get()
3468 {
3469 tracker.fail("WebSocket client disconnected");
3470 }
3471
3472 self.state_notify.notify_waiters();
3473
3474 if dst::time::timeout(Duration::from_secs(GRACEFUL_SHUTDOWN_TIMEOUT_SECS), async {
3475 while !self.is_disconnected() {
3476 dst::time::sleep(Duration::from_millis(CONNECTION_STATE_CHECK_INTERVAL_MS)).await;
3477 }
3478
3479 if !self.controller_task.is_finished() {
3480 self.controller_task.abort();
3481 log_task_aborted("controller");
3482 }
3483 })
3484 .await
3485 == Ok(())
3486 {
3487 log::debug!("Controller task finished");
3488 } else {
3489 log::warn!("Timeout waiting for controller task to finish");
3490
3491 if !self.controller_task.is_finished() {
3492 self.controller_task.abort();
3493 log_task_aborted("controller");
3494 }
3495
3496 self.connection_mode
3497 .store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
3498 }
3499 }
3500
3501 #[allow(unused_variables)]
3511 pub async fn send_text(&self, data: String, keys: Option<&[Ustr]>) -> Result<(), SendError> {
3512 self.check_not_terminal()?;
3513
3514 self.await_rate_limit_or_closed(keys).await?;
3515 self.wait_for_active().await?;
3516
3517 log::trace!("Sending text frame ({} bytes)", data.len());
3518
3519 let msg = Message::Text(data.into());
3520 self.writer_tx.send(WriterCommand::Send(msg))
3521 }
3522
3523 pub async fn send_text_on_connection(
3540 &self,
3541 data: String,
3542 keys: Option<&[Ustr]>,
3543 connection_epoch: u64,
3544 ) -> Result<(), SendError> {
3545 self.check_not_terminal()?;
3546 self.await_rate_limit_or_closed(keys).await?;
3547 self.wait_for_active().await?;
3548
3549 log::trace!(
3550 "Sending text frame once: epoch={connection_epoch} ({} bytes)",
3551 data.len()
3552 );
3553
3554 let (response_tx, response_rx) = tokio::sync::oneshot::channel();
3555 self.writer_tx
3556 .send_control(WriterCommand::SendOnConnection {
3557 message: Message::Text(data.into()),
3558 connection_epoch,
3559 response_tx,
3560 })?;
3561
3562 response_rx
3563 .await
3564 .map_err(|e| SendError::BrokenPipe(e.to_string()))?
3565 }
3566
3567 #[allow(unknown_lints, reason = "Clippy lint is unavailable on Rust 1.97")]
3579 #[expect(
3580 clippy::unused_async,
3581 clippy::unused_async_trait_impl,
3582 reason = "skipping instead of waiting removes the only await; the signature is public API shared with the other send methods"
3583 )]
3584 pub async fn send_pong(&self, data: Vec<u8>) -> Result<(), SendError> {
3585 validate_pong_payload(&data)?;
3586
3587 if !self.connection_mode().is_active() {
3588 log::debug!("Skipping pong: connection not active");
3589 return Ok(());
3590 }
3591
3592 log::trace!("Sending pong frame ({} bytes)", data.len());
3593
3594 self.writer_tx
3595 .send_control(WriterCommand::Send(Message::Pong(data.into())))
3596 }
3597
3598 #[allow(unknown_lints, reason = "Clippy lint is unavailable on Rust 1.97")]
3608 #[expect(
3609 clippy::unused_async,
3610 clippy::unused_async_trait_impl,
3611 reason = "the public send API is async even though this method only enqueues"
3612 )]
3613 pub async fn send_pong_on_connection(
3614 &self,
3615 data: Vec<u8>,
3616 connection_epoch: u64,
3617 ) -> Result<(), SendError> {
3618 validate_pong_payload(&data)?;
3619
3620 if !self.connection_mode().is_active() {
3621 log::debug!("Skipping pong: connection not active");
3622 return Ok(());
3623 }
3624
3625 log::trace!(
3626 "Sending pong frame once: epoch={connection_epoch} ({} bytes)",
3627 data.len()
3628 );
3629
3630 self.writer_tx
3631 .send_control(WriterCommand::SendPongOnConnection {
3632 data,
3633 connection_epoch,
3634 })
3635 }
3636
3637 #[allow(unused_variables)]
3647 pub async fn send_bytes(&self, data: Vec<u8>, keys: Option<&[Ustr]>) -> Result<(), SendError> {
3648 self.check_not_terminal()?;
3649
3650 self.await_rate_limit_or_closed(keys).await?;
3651 self.wait_for_active().await?;
3652
3653 log::trace!("Sending binary frame ({} bytes)", data.len());
3654
3655 let msg = Message::Binary(data.into());
3656 self.writer_tx.send(WriterCommand::Send(msg))
3657 }
3658
3659 pub async fn send_close_message(&self) -> Result<(), SendError> {
3665 self.wait_for_active().await?;
3666
3667 let msg = Message::Close(None);
3668 self.writer_tx.send_control(WriterCommand::Send(msg))
3669 }
3670
3671 fn spawn_controller_task(
3672 mut inner: WebSocketClientInner,
3673 connection_mode: Arc<AtomicU8>,
3674 state_notify: Arc<tokio::sync::Notify>,
3675 auth_tracker: Arc<OnceLock<AuthTracker>>,
3676 controller_lifecycle: Arc<ControllerLifecycle>,
3677 controller_notify: Arc<tokio::sync::Notify>,
3678 reconnect_published: Arc<AtomicBool>,
3679 ) -> dst::task::JoinHandle<()> {
3680 dst::task::spawn(async move {
3681 let _activity = controller_lifecycle.activity();
3682 log_task_started("controller");
3683
3684 let fallback_interval = Duration::from_millis(CONTROLLER_FALLBACK_INTERVAL_MS);
3685 let mut reconnected_at = None;
3686
3687 loop {
3688 tokio::select! {
3689 biased;
3690 () = controller_notify.notified() => {}
3691 () = state_notify.notified() => {}
3692 () = dst::time::sleep(fallback_interval) => {}
3693 }
3694
3695 let mut mode = ConnectionMode::from_atomic(&connection_mode);
3696
3697 if mode.is_disconnect() {
3698 log::debug!("Disconnecting");
3699
3700 let timeout = Duration::from_secs(GRACEFUL_SHUTDOWN_TIMEOUT_SECS);
3701 if dst::time::timeout(timeout, async {
3702 dst::time::sleep(Duration::from_millis(GRACEFUL_SHUTDOWN_DELAY_MS)).await;
3704
3705 if let Some(read_fence) = inner.read_fence.take() {
3706 read_fence.invalidate();
3707 }
3708
3709 if let Some(task) = &inner.read_task
3710 && !task.is_finished()
3711 {
3712 task.abort();
3713 log_task_aborted("read");
3714 }
3715
3716 if let Some(task) = &inner.heartbeat_task
3717 && !task.is_finished()
3718 {
3719 task.abort();
3720 log_task_aborted("heartbeat");
3721 }
3722 })
3723 .await
3724 .is_err()
3725 {
3726 log::warn!("Shutdown timed out after {}s", timeout.as_secs());
3727 }
3728
3729 log::debug!("Closed");
3730 break; }
3732
3733 if mode.is_closed() {
3734 log::debug!("Connection closed");
3735 fail_registered_auth(auth_tracker.as_ref(), "WebSocket client closed");
3736 state_notify.notify_waiters();
3737 break;
3738 }
3739
3740 if mode.is_active() && !inner.is_alive() {
3741 let target = if inner.handler.is_none() {
3742 ConnectionMode::Closed
3743 } else {
3744 ConnectionMode::Reconnect
3745 };
3746
3747 let transitioned = if target.is_closed() {
3748 ConnectionMode::close_websocket_on_loss(
3749 &connection_mode,
3750 inner.state_sink.as_ref(),
3751 )
3752 } else {
3753 request_websocket_reconnect(
3754 &connection_mode,
3755 &reconnect_published,
3756 inner.state_sink.as_ref(),
3757 &auth_tracker,
3758 &controller_notify,
3759 || {},
3760 ) == ReconnectRequestOutcome::Accepted
3761 };
3762
3763 if transitioned {
3764 if target.is_closed() {
3765 fail_registered_auth(auth_tracker.as_ref(), "WebSocket client closed");
3766 }
3767
3768 log::info!("Detected dead connection, transitioning to {target:?}");
3769 }
3770
3771 mode = ConnectionMode::from_atomic(&connection_mode);
3772 }
3773
3774 if mode.is_reconnect() {
3775 if let Some(tracker) = auth_tracker.get() {
3776 tracker.invalidate();
3777 }
3778
3779 let reconnect_uptime = reconnected_at
3780 .take()
3781 .map(|started: dst::time::Instant| started.elapsed());
3782 let previous_reconnect_stable = reconnect_uptime
3783 .is_some_and(|uptime| uptime >= RECONNECT_STABILITY_THRESHOLD);
3784
3785 if previous_reconnect_stable {
3786 inner.backoff.reset();
3787 inner.reconnection_attempt_count = 0;
3788 log::debug!(
3789 "WebSocket remained active for at least {}s, resetting reconnect cycle",
3790 RECONNECT_STABILITY_THRESHOLD.as_secs()
3791 );
3792 }
3793
3794 if let Some(max_attempts) = inner.reconnect_max_attempts
3795 && inner.reconnection_attempt_count >= max_attempts
3796 {
3797 log::error!(
3798 "Max reconnection attempts ({max_attempts}) exceeded, transitioning to CLOSED"
3799 );
3800 connection_mode.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
3801 fail_registered_auth(
3802 auth_tracker.as_ref(),
3803 "WebSocket reconnect attempts exhausted",
3804 );
3805 state_notify.notify_waiters();
3806 break;
3807 }
3808
3809 let backoff_delay = if reconnect_uptime.is_some() && !previous_reconnect_stable
3810 {
3811 inner.backoff.next_duration()
3812 } else {
3813 Duration::ZERO
3814 };
3815
3816 let duration = inner.reconnect_throttle.gated_delay(backoff_delay);
3817 if !duration.is_zero() {
3818 log::warn!("Backing off for {}s...", duration.as_secs_f64());
3819
3820 if !wait_reconnect_delay(
3821 duration,
3822 connection_mode.as_ref(),
3823 state_notify.as_ref(),
3824 )
3825 .await
3826 {
3827 log::debug!("Backoff interrupted by terminal state");
3828 continue;
3829 }
3830 }
3831
3832 inner.reconnection_attempt_count += 1;
3833 inner.reconnect_throttle.record_attempt();
3834 log::debug!(
3835 "Reconnection attempt {} of {}",
3836 inner.reconnection_attempt_count,
3837 inner
3838 .reconnect_max_attempts
3839 .map_or_else(|| "unlimited".to_string(), |m| m.to_string())
3840 );
3841
3842 let reconnect_result = tokio::select! {
3844 biased;
3845 result = inner.reconnect_with_outcome() => Some(result),
3846 () = async {
3847 loop {
3848 let mut notified = pin!(state_notify.notified());
3850 notified.as_mut().enable();
3851
3852 if ConnectionMode::from_atomic(&connection_mode).is_disconnect() {
3853 break;
3854 }
3855 notified.await;
3856 }
3857 } => None,
3858 };
3859
3860 match reconnect_result {
3861 None => {
3862 log::debug!("Reconnect interrupted by disconnect");
3863 }
3864 Some(Ok(ReconnectOutcome::Reconnected)) => {
3865 reconnected_at = Some(dst::time::Instant::now());
3866
3867 state_notify.notify_waiters();
3868
3869 if ConnectionMode::from_atomic(&connection_mode).is_active() {
3873 if let Some(ref handler) = inner.handler {
3874 let connection_epoch =
3875 inner.connection_epoch.load(Ordering::Acquire);
3876 let reconnected_msg =
3877 Message::Text(RECONNECTED.to_string().into());
3878 handler.handle(connection_epoch, reconnected_msg);
3879 match handler {
3880 IncomingHandler::Message(_) => {
3881 log::debug!("Sent reconnected message to handler");
3882 }
3883 IncomingHandler::Epoch(_) => {
3884 log::debug!(
3885 "Sent reconnected message to epoch handler: \
3886 epoch={connection_epoch}",
3887 );
3888 }
3889 }
3890 }
3891
3892 log::debug!("Reconnected successfully");
3893 } else {
3894 log::debug!("Skipping reconnect handlers due to disconnect state");
3895 }
3896 }
3897 Some(Ok(ReconnectOutcome::Aborted)) => {
3898 log::debug!("Reconnect aborted");
3899 }
3900 Some(Err(e)) => {
3901 let duration = inner.backoff.next_duration();
3902 log::warn!(
3903 "Reconnect attempt {} failed: {e}",
3904 inner.reconnection_attempt_count
3905 );
3906
3907 if !duration.is_zero() {
3908 log::warn!("Backing off for {}s...", duration.as_secs_f64());
3909
3910 if !wait_reconnect_delay(
3911 duration,
3912 connection_mode.as_ref(),
3913 state_notify.as_ref(),
3914 )
3915 .await
3916 {
3917 log::debug!("Backoff interrupted by terminal state");
3918 }
3919 }
3920 }
3921 }
3922 }
3923 }
3924
3925 inner
3926 .connection_mode
3927 .store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
3928
3929 log_task_stopped("controller");
3930 })
3931 }
3932}
3933
3934fn fail_registered_auth(auth_tracker: &OnceLock<AuthTracker>, reason: &str) {
3935 if let Some(tracker) = auth_tracker.get() {
3936 tracker.fail(reason);
3937 }
3938}
3939
3940fn validate_pong_payload(data: &[u8]) -> Result<(), SendError> {
3941 if data.len() > MAX_CONTROL_FRAME_PAYLOAD_BYTES {
3942 return Err(SendError::InvalidInput(format!(
3943 "pong payload exceeds {MAX_CONTROL_FRAME_PAYLOAD_BYTES} bytes"
3944 )));
3945 }
3946
3947 Ok(())
3948}
3949
3950impl Drop for WebSocketClient {
3951 fn drop(&mut self) {
3952 self.connection_mode
3953 .store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
3954 fail_registered_auth(self.auth_tracker.as_ref(), "WebSocket client closed");
3955 self.state_notify.notify_waiters();
3956 self.controller_notify.notify_waiters();
3957 self.controller_lifecycle.close_and_abort();
3958 }
3959}
3960
3961#[cfg(test)]
3962#[cfg(not(feature = "turmoil"))]
3963#[cfg(not(all(feature = "simulation", madsim)))] #[cfg(target_os = "linux")] mod tests {
3966 use std::{
3967 collections::HashMap,
3968 num::NonZeroU32,
3969 sync::{
3970 Arc,
3971 atomic::{AtomicU8, AtomicUsize, Ordering},
3972 },
3973 time::Duration,
3974 };
3975
3976 use axum::{Router, routing::post};
3977 use futures_util::{SinkExt, StreamExt};
3978 use log::Level;
3979 use nautilus_common::testing::wait_until_async;
3980 use nautilus_core::string::secret::REDACTED;
3981 use parking_lot::Mutex;
3982 use rstest::rstest;
3983 use tokio::{
3984 io::{AsyncReadExt, AsyncWriteExt},
3985 net::TcpListener,
3986 sync::{mpsc, oneshot},
3987 task::{self, JoinHandle},
3988 };
3989 use tokio_tungstenite::{
3990 accept_async, accept_hdr_async,
3991 tungstenite::{
3992 Message as WsMessage,
3993 handshake::server::{self, Callback},
3994 http::HeaderValue,
3995 protocol::frame::{CloseFrame as WsCloseFrame, coding::CloseCode},
3996 },
3997 };
3998
3999 use crate::{
4000 RECONNECTED, SocketState, SocketStateSink,
4001 error::SendError,
4002 http::{HttpClient, Method},
4003 logging::tests::capture_logs_for,
4004 mode::{ConnectionMode, ReadSessionFence},
4005 ratelimiter::quota::Quota,
4006 transport::{CloseFrame, Message, TransportError},
4007 websocket::{
4008 AuthTracker, CloseHandler, InitialConnectRetryPolicy, ReconnectHeaders,
4009 TransportBackend, WebSocketClient, WebSocketClientInner, WebSocketConfig,
4010 auth::AuthState,
4011 types::{WriterCommand, channel_epoch_message_handler, channel_message_handler},
4012 },
4013 };
4014
4015 const SECRET_MARKER: &str = "OUTBOUND_SECRET_MARKER";
4016 const PING_TRIGGER: &str = "send-test-ping";
4017 const HOSTILE_CLOSE_TRIGGER: &str = "close-hostile";
4018 const CODED_CLOSE_TRIGGER: &str = "close-coded";
4019 const CODED_CLOSE_CODE: u16 = 4003;
4020 const CODED_CLOSE_REASON: &str = "draining";
4021 const HOSTILE_CLOSE_REASON: &str = "injected\r\n\u{1b}[31mforged\u{1b}[0m";
4022 const NETWORK_LOG_TARGETS: &[&str] = &[
4023 "nautilus_network::http::client",
4024 "nautilus_network::websocket::client",
4025 ];
4026
4027 struct TestServer {
4028 task: JoinHandle<()>,
4029 port: u16,
4030 }
4031
4032 #[derive(Debug, Clone)]
4033 struct TestCallback {
4034 key: String,
4035 value: HeaderValue,
4036 }
4037
4038 impl Callback for TestCallback {
4039 #[expect(clippy::panic_in_result_fn)]
4040 fn on_request(
4041 self,
4042 request: &server::Request,
4043 response: server::Response,
4044 ) -> Result<server::Response, server::ErrorResponse> {
4045 let _ = response;
4046 let value = request.headers().get(&self.key);
4047 assert!(value.is_some());
4048
4049 if let Some(value) = request.headers().get(&self.key) {
4050 assert_eq!(value, self.value);
4051 }
4052
4053 Ok(response)
4054 }
4055 }
4056
4057 impl TestServer {
4058 async fn setup() -> Self {
4059 let server = TcpListener::bind("127.0.0.1:0").await.unwrap();
4060 let port = TcpListener::local_addr(&server).unwrap().port();
4061
4062 let header_key = "test".to_string();
4063 let header_value = "test".to_string();
4064
4065 let test_call_back = TestCallback {
4066 key: header_key,
4067 value: HeaderValue::from_str(&header_value).unwrap(),
4068 };
4069
4070 let task = task::spawn(async move {
4071 loop {
4073 let (conn, _) = server.accept().await.unwrap();
4074 let mut websocket = accept_hdr_async(conn, test_call_back.clone())
4075 .await
4076 .unwrap();
4077
4078 task::spawn(async move {
4079 while let Some(Ok(msg)) = websocket.next().await {
4080 match msg {
4081 WsMessage::Text(txt) if txt == "close-now" => {
4082 log::debug!("Forcibly closing from server side");
4083 let _ = websocket.close(None).await;
4085 break;
4086 }
4087 WsMessage::Text(txt) if txt == HOSTILE_CLOSE_TRIGGER => {
4088 let _ = websocket
4092 .send(WsMessage::Close(Some(WsCloseFrame {
4093 code: CloseCode::Error,
4094 reason: HOSTILE_CLOSE_REASON.into(),
4095 })))
4096 .await;
4097
4098 break;
4099 }
4100 WsMessage::Text(txt) if txt == CODED_CLOSE_TRIGGER => {
4101 let _ = websocket
4102 .send(WsMessage::Close(Some(WsCloseFrame {
4103 code: CloseCode::from(CODED_CLOSE_CODE),
4104 reason: CODED_CLOSE_REASON.into(),
4105 })))
4106 .await;
4107
4108 break;
4109 }
4110 WsMessage::Text(txt) if txt == PING_TRIGGER => {
4111 let ping = format!("{SECRET_MARKER}:ping");
4112 if websocket.send(WsMessage::Ping(ping.into())).await.is_err() {
4113 break;
4114 }
4115 }
4116 WsMessage::Text(_) | WsMessage::Binary(_) => {
4118 if websocket.send(msg).await.is_err() {
4119 break;
4120 }
4121 }
4122 WsMessage::Close(_frame) => {
4124 let _ = websocket.close(None).await;
4125 break;
4126 }
4127 _ => {}
4129 }
4130 }
4131 });
4132 }
4133 });
4134
4135 Self { task, port }
4136 }
4137 }
4138
4139 impl Drop for TestServer {
4140 fn drop(&mut self) {
4141 self.task.abort();
4142 }
4143 }
4144
4145 async fn setup_test_client(port: u16) -> WebSocketClient {
4146 let config = WebSocketConfig {
4147 url: format!("ws://127.0.0.1:{port}"),
4148 headers: vec![("test".into(), "test".into())],
4149 heartbeat_interval_secs: None,
4150 heartbeat_payload: None,
4151 connect_timeout_ms: None,
4152 reconnect_delay_initial_ms: None,
4153 reconnect_backoff_factor: None,
4154 reconnect_delay_max_ms: None,
4155 reconnect_jitter_ms: None,
4156 reconnect_max_attempts: None,
4157 heartbeat_timeout_secs: None,
4158 idle_timeout_ms: None,
4159 writer_capacity: None,
4160 backend: TransportBackend::Tungstenite,
4161 proxy_url: None,
4162 max_message_size_bytes: None,
4163 max_frame_size_bytes: None,
4164 };
4165
4166 WebSocketClient::builder()
4167 .config(config)
4168 .message_handler(Arc::new(|_| {}))
4169 .connect()
4170 .await
4171 .expect("Failed to connect")
4172 }
4173
4174 #[rstest]
4175 #[case::handler(false)]
4176 #[case::stream(true)]
4177 #[tokio::test]
4178 async fn test_writer_capacity_from_config(#[case] stream_mode: bool) {
4179 let server = TestServer::setup().await;
4180 let config = WebSocketConfig::builder()
4181 .url(format!("ws://127.0.0.1:{}", server.port))
4182 .headers(vec![("test".to_string(), "test".to_string())])
4183 .backend(TransportBackend::Tungstenite)
4184 .writer_capacity(2)
4185 .build()
4186 .unwrap();
4187
4188 let (client, _reader) = if stream_mode {
4189 let (reader, client) = WebSocketClient::stream_builder()
4190 .config(config)
4191 .connect()
4192 .await
4193 .unwrap();
4194 (client, Some(reader))
4195 } else {
4196 let client = WebSocketClient::builder()
4197 .config(config)
4198 .message_handler(Arc::new(|_| {}))
4199 .connect()
4200 .await
4201 .unwrap();
4202 (client, None)
4203 };
4204
4205 client.send_text("first".to_string(), None).await.unwrap();
4207 client.send_bytes(vec![2], None).await.unwrap();
4208 let overflow = client.send_text("overflow".to_string(), None).await;
4209
4210 assert_eq!(client.connection_mode(), ConnectionMode::Active);
4211 assert!(matches!(overflow, Err(SendError::BufferFull)));
4212 client.disconnect().await;
4213 }
4214
4215 #[rstest]
4216 #[tokio::test]
4217 async fn test_send_reports_writer_buffer_full() {
4218 let server = TestServer::setup().await;
4219 let mut client = setup_test_client(server.port).await;
4220 let (writer_tx, mut writer_rx) = crate::writer::channel(1);
4221 client.writer_tx = writer_tx;
4222
4223 client.send_text("first".to_string(), None).await.unwrap();
4224 let result = client.send_bytes(vec![2], None).await;
4225 assert!(matches!(result, Err(SendError::BufferFull)));
4226 let (command, permit) = writer_rx.recv().await.unwrap();
4227 assert!(
4228 matches!(command, WriterCommand::Send(message) if message == Message::text("first"))
4229 );
4230 assert!(matches!(
4231 client.send_text("overflow".to_string(), None).await,
4232 Err(SendError::BufferFull)
4233 ));
4234 drop(permit);
4235 client.send_bytes(vec![3], None).await.unwrap();
4236 let (command, _) = writer_rx.recv().await.unwrap();
4237 assert!(
4238 matches!(command, WriterCommand::Send(message) if message == Message::Binary(vec![3].into()))
4239 );
4240 client.disconnect().await;
4241 }
4242
4243 async fn setup_reconnecting_client(port: u16) -> WebSocketClient {
4244 let config = WebSocketConfig {
4245 url: format!("ws://127.0.0.1:{port}"),
4246 headers: vec![],
4247 heartbeat_interval_secs: None,
4248 heartbeat_payload: None,
4249 connect_timeout_ms: Some(5_000),
4250 reconnect_delay_initial_ms: Some(1),
4251 reconnect_backoff_factor: Some(1.0),
4252 reconnect_delay_max_ms: Some(1),
4253 reconnect_jitter_ms: Some(0),
4254 reconnect_max_attempts: None,
4255 heartbeat_timeout_secs: None,
4256 idle_timeout_ms: None,
4257 writer_capacity: None,
4258 backend: TransportBackend::Tungstenite,
4259 proxy_url: None,
4260 max_message_size_bytes: None,
4261 max_frame_size_bytes: None,
4262 };
4263
4264 WebSocketClient::builder()
4265 .config(config)
4266 .message_handler(Arc::new(|_| {}))
4267 .connect()
4268 .await
4269 .expect("client should connect")
4270 }
4271
4272 async fn wait_for_mode(client: &WebSocketClient, expected: ConnectionMode) {
4273 crate::dst::time::timeout(Duration::from_secs(5), async {
4274 loop {
4275 if ConnectionMode::from_atomic(&client.connection_mode) == expected {
4276 break;
4277 }
4278
4279 crate::dst::time::sleep(Duration::from_millis(1)).await;
4280 }
4281 })
4282 .await
4283 .expect("client should reach expected connection mode");
4284 }
4285
4286 async fn setup_http_test_server() -> (JoinHandle<()>, u16) {
4287 let server = TcpListener::bind("127.0.0.1:0").await.unwrap();
4288 let port = server.local_addr().unwrap().port();
4289
4290 let app = Router::new().route(
4291 "/logging",
4292 post(|| async {
4293 (
4294 [("x-secret-response", SECRET_MARKER)],
4295 format!("{SECRET_MARKER}:response"),
4296 )
4297 }),
4298 );
4299
4300 let task = task::spawn(async move {
4301 axum::serve(server, app).await.unwrap();
4302 });
4303
4304 (task, port)
4305 }
4306
4307 #[rstest]
4308 #[tokio::test]
4309 async fn test_network_logs_omit_payload_bodies() {
4310 let server = TestServer::setup().await;
4311 let client = setup_test_client(server.port).await;
4312 let capture = capture_logs_for(NETWORK_LOG_TARGETS).await;
4313 let binary = format!("{SECRET_MARKER}:binary").into_bytes();
4314 let binary_marker = format!("{binary:?}");
4315
4316 client
4317 .send_text(format!("{SECRET_MARKER}:café"), None)
4318 .await
4319 .unwrap();
4320 client
4321 .send_text_on_connection(
4322 format!("{SECRET_MARKER}:owned-é"),
4323 None,
4324 client.connection_epoch(),
4325 )
4326 .await
4327 .unwrap();
4328 client.send_bytes(binary, None).await.unwrap();
4329 client
4330 .send_text(PING_TRIGGER.to_string(), None)
4331 .await
4332 .unwrap();
4333
4334 tokio::time::timeout(Duration::from_secs(2), async {
4335 loop {
4336 if capture.messages().iter().any(|(level, message)| {
4337 matches!(level, Level::Trace | Level::Warn)
4338 && message == "Received ping frame (27 bytes)"
4339 }) {
4340 break;
4341 }
4342
4343 tokio::task::yield_now().await;
4344 }
4345 })
4346 .await
4347 .expect("timed out waiting for inbound WebSocket metadata log");
4348
4349 let (http_task, http_port) = setup_http_test_server().await;
4350 let invalid_headers =
4351 HashMap::from([("x-secret-default".to_string(), format!("{SECRET_MARKER}\n"))]);
4352 let invalid_header_error = HttpClient::builder()
4353 .headers(invalid_headers)
4354 .build()
4355 .unwrap_err();
4356 let http_client = HttpClient::builder().build().unwrap();
4357 let params = HashMap::from([("secret".to_string(), vec![SECRET_MARKER.to_string()])]);
4358 let headers = HashMap::from([
4359 (
4360 "X-Secret-Request".to_string(),
4361 format!("{SECRET_MARKER}:first"),
4362 ),
4363 (
4364 "x-secret-request".to_string(),
4365 format!("{SECRET_MARKER}:second"),
4366 ),
4367 ]);
4368 let http_body = format!("{SECRET_MARKER}:http-body").into_bytes();
4369 http_client
4370 .request(
4371 Method::POST,
4372 format!("http://127.0.0.1:{http_port}/logging"),
4373 Some(¶ms),
4374 Some(headers),
4375 Some(http_body),
4376 None,
4377 None,
4378 )
4379 .await
4380 .unwrap();
4381
4382 let rejecting = TcpListener::bind("127.0.0.1:0").await.unwrap();
4385 let rejecting_port = rejecting.local_addr().unwrap().port();
4386
4387 let rejecting_server = task::spawn(async move {
4388 loop {
4389 let (mut stream, _) = rejecting.accept().await.unwrap();
4390 let mut request = [0; 1024];
4391 let _ = stream.read(&mut request).await.unwrap();
4392 stream
4393 .write_all(b"HTTP/1.1 503 Service Unavailable\r\n\r\n")
4394 .await
4395 .unwrap();
4396 }
4397 });
4398
4399 let retry_config = WebSocketConfig {
4400 url: format!("ws://127.0.0.1:{rejecting_port}/?token={SECRET_MARKER}"),
4401 headers: vec![],
4402 heartbeat_interval_secs: None,
4403 heartbeat_payload: None,
4404 connect_timeout_ms: Some(5_000),
4405 reconnect_delay_initial_ms: None,
4406 reconnect_backoff_factor: None,
4407 reconnect_delay_max_ms: None,
4408 reconnect_jitter_ms: None,
4409 reconnect_max_attempts: None,
4410 heartbeat_timeout_secs: None,
4411 idle_timeout_ms: None,
4412 writer_capacity: None,
4413 backend: TransportBackend::Tungstenite,
4414 proxy_url: None,
4415 max_message_size_bytes: None,
4416 max_frame_size_bytes: None,
4417 };
4418
4419 let retry_error = WebSocketClient::builder()
4420 .config(retry_config)
4421 .message_handler(Arc::new(|_| {}))
4422 .initial_connect_retry_policy(InitialConnectRetryPolicy {
4423 max_attempts: NonZeroU32::new(2).unwrap(),
4424 delay_initial: Duration::from_millis(1),
4425 delay_max: Duration::from_millis(1),
4426 backoff_factor: 1.0,
4427 jitter_ms: 0,
4428 })
4429 .connect()
4430 .await
4431 .expect_err("server always rejects the upgrade");
4432
4433 rejecting_server.abort();
4434
4435 assert!(
4436 matches!(retry_error, TransportError::UpgradeRejected(503)),
4437 "expected a 503 upgrade rejection, was: {retry_error:?}"
4438 );
4439
4440 let messages: Vec<_> = capture
4441 .messages()
4442 .into_iter()
4443 .filter(|(level, message)| {
4444 matches!(level, Level::Trace | Level::Warn)
4445 && (message.starts_with("Sending ")
4446 || message.starts_with("Received ")
4447 || message.starts_with("Replaced ")
4448 || message.starts_with("WebSocket connection attempt "))
4449 })
4450 .map(|(_, message)| message)
4451 .collect();
4452
4453 let invalid_header_message = invalid_header_error.to_string();
4454
4455 let retry_warning = messages
4458 .iter()
4459 .find(|message| message.starts_with("WebSocket connection attempt 1/2 "))
4460 .expect("initial-connect retry warning was not captured");
4461 assert!(
4462 retry_warning.contains(REDACTED),
4463 "retry warning omitted the redaction placeholder: {retry_warning}"
4464 );
4465
4466 assert!(
4467 messages.iter().all(|message| {
4468 !message.contains(SECRET_MARKER) && !message.contains(&binary_marker)
4469 }),
4470 "network logs exposed the secret marker: {messages:?}"
4471 );
4472 assert!(
4473 !invalid_header_message.contains(SECRET_MARKER),
4474 "invalid header error exposed the secret marker: {invalid_header_message}"
4475 );
4476 assert!(
4477 invalid_header_message.contains("x-secret-default"),
4478 "invalid header error omitted safe header metadata: {invalid_header_message}"
4479 );
4480 assert!(
4481 messages
4482 .iter()
4483 .any(|message| message == "Sending text frame (28 bytes)"),
4484 "text send metadata missing or inaccurate: {messages:?}"
4485 );
4486 assert!(
4487 messages
4488 .iter()
4489 .any(|message| { message == "Sending text frame once: epoch=0 (31 bytes)" }),
4490 "ownership-bound text metadata missing or inaccurate: {messages:?}"
4491 );
4492 assert!(
4493 messages
4494 .iter()
4495 .any(|message| message == "Sending binary frame (29 bytes)"),
4496 "binary send metadata missing or inaccurate: {messages:?}"
4497 );
4498 assert!(
4499 messages
4500 .iter()
4501 .any(|message| message == "Received text frame (28 bytes)"),
4502 "text receive metadata missing or inaccurate: {messages:?}"
4503 );
4504 assert!(
4505 messages
4506 .iter()
4507 .any(|message| message == "Received text frame (31 bytes)"),
4508 "ownership-bound text receive metadata missing or inaccurate: {messages:?}"
4509 );
4510 assert!(
4511 messages
4512 .iter()
4513 .any(|message| message == "Received message <binary> 29 bytes"),
4514 "binary receive metadata missing or inaccurate: {messages:?}"
4515 );
4516 assert!(
4517 messages
4518 .iter()
4519 .any(|message| message == "Received ping frame (27 bytes)"),
4520 "ping receive metadata missing or inaccurate: {messages:?}"
4521 );
4522 assert!(
4524 messages.iter().any(|message| {
4525 message
4526 == "Sending HTTP request: method=POST extra_headers=3 query_bytes=29 \
4527 body_bytes=32"
4528 }),
4529 "HTTP request metadata missing or inaccurate: {messages:?}"
4530 );
4531 assert!(
4532 messages
4533 .iter()
4534 .any(|message| message == "Replaced duplicate request header 'x-secret-request'"),
4535 "duplicate header metadata missing: {messages:?}"
4536 );
4537 assert!(
4538 messages.iter().any(|message| {
4539 message.starts_with("Received HTTP response: status=200 OK headers=")
4540 && message.ends_with(" body_bytes=31")
4541 }),
4542 "HTTP response metadata missing or inaccurate: {messages:?}"
4543 );
4544
4545 client.disconnect().await;
4546 http_task.abort();
4547 }
4548
4549 #[rstest]
4553 #[tokio::test]
4554 async fn test_send_pong_skips_gated_reconnect() {
4555 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
4556 let port = listener.local_addr().unwrap().port();
4557 let (second_accepted_tx, second_accepted_rx) = oneshot::channel();
4558 let (handshake_gate_tx, handshake_gate_rx) = oneshot::channel();
4559 let (observed_tx, mut observed_rx) = mpsc::unbounded_channel();
4560
4561 let server_task = task::spawn(async move {
4562 let (first, _) = listener.accept().await.unwrap();
4563 let first_websocket = accept_async(first).await.unwrap();
4564 drop(first_websocket);
4565
4566 let (second, _) = listener.accept().await.unwrap();
4567 second_accepted_tx.send(()).unwrap();
4568 handshake_gate_rx.await.unwrap();
4569 let mut replacement = accept_async(second).await.unwrap();
4570
4571 while let Some(message) = replacement.next().await {
4572 match message.unwrap() {
4573 WsMessage::Pong(data) => observed_tx.send(data.to_vec()).unwrap(),
4574 WsMessage::Close(_) => {
4575 let _ = replacement.close(None).await;
4576 break;
4577 }
4578 _ => {}
4579 }
4580 }
4581 });
4582
4583 let client = setup_reconnecting_client(port).await;
4584
4585 crate::dst::time::timeout(Duration::from_secs(5), second_accepted_rx)
4586 .await
4587 .expect("replacement connection should be accepted")
4588 .unwrap();
4589 wait_for_mode(&client, ConnectionMode::Reconnect).await;
4590
4591 let rejected = client.send_pong(vec![2; 126]).await;
4592 assert!(matches!(rejected, Err(SendError::InvalidInput(_))));
4593
4594 crate::dst::time::timeout(
4597 Duration::from_secs(1),
4598 client.send_pong(b"stale-pong".to_vec()),
4599 )
4600 .await
4601 .expect("pong raised during reconnect should not wait for the replacement")
4602 .expect("skipped pong should report success");
4603
4604 handshake_gate_tx.send(()).unwrap();
4605 wait_for_mode(&client, ConnectionMode::Active).await;
4606
4607 let fresh_payload = b"fresh-pong".to_vec();
4608 client.send_pong(fresh_payload.clone()).await.unwrap();
4609 assert_eq!(
4610 crate::dst::time::timeout(Duration::from_secs(5), observed_rx.recv())
4611 .await
4612 .expect("replacement connection should receive the fresh pong"),
4613 Some(fresh_payload)
4614 );
4615
4616 client.disconnect().await;
4617 server_task.await.unwrap();
4618 assert_eq!(
4619 observed_rx.try_recv(),
4620 Err(mpsc::error::TryRecvError::Disconnected),
4621 "replacement connection should receive exactly one pong"
4622 );
4623 }
4624
4625 #[rstest]
4628 #[tokio::test]
4629 async fn test_pong_validation_precedes_inactive_skip() {
4630 let server = TestServer::setup().await;
4631 let client = setup_test_client(server.port).await;
4632 client
4633 .connection_mode
4634 .store(ConnectionMode::Reconnect.as_u8(), Ordering::SeqCst);
4635
4636 let result = client.send_pong(vec![2; 126]).await;
4637 let epoch_result = client.send_pong_on_connection(vec![2; 126], 0).await;
4638
4639 assert!(matches!(result, Err(SendError::InvalidInput(_))));
4640 assert!(matches!(epoch_result, Err(SendError::InvalidInput(_))));
4641 }
4642
4643 #[rstest]
4646 #[tokio::test]
4647 async fn test_pong_payload_limit_preserves_connection() {
4648 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
4649 let port = listener.local_addr().unwrap().port();
4650 let (pong_tx, mut pong_rx) = mpsc::unbounded_channel();
4651
4652 let server_task = task::spawn(async move {
4653 let (stream, _) = listener.accept().await.unwrap();
4654 let mut websocket = accept_async(stream).await.unwrap();
4655
4656 while let Some(message) = websocket.next().await {
4657 match message.unwrap() {
4658 WsMessage::Pong(data) => pong_tx.send(data.to_vec()).unwrap(),
4659 WsMessage::Close(_) => {
4660 let _ = websocket.close(None).await;
4661 break;
4662 }
4663 _ => {}
4664 }
4665 }
4666 });
4667
4668 let client = setup_reconnecting_client(port).await;
4669 let accepted = vec![1; 125];
4670 let follow_up = vec![3; 125];
4671
4672 client.send_pong(accepted.clone()).await.unwrap();
4673 assert_eq!(
4674 crate::dst::time::timeout(Duration::from_secs(5), pong_rx.recv())
4675 .await
4676 .unwrap(),
4677 Some(accepted)
4678 );
4679
4680 let rejected = client.send_pong(vec![2; 126]).await;
4681 assert!(matches!(rejected, Err(SendError::InvalidInput(_))));
4682
4683 client.send_pong(follow_up.clone()).await.unwrap();
4684 assert_eq!(
4685 crate::dst::time::timeout(Duration::from_secs(5), pong_rx.recv())
4686 .await
4687 .unwrap(),
4688 Some(follow_up)
4689 );
4690
4691 client.disconnect().await;
4692 server_task.await.unwrap();
4693 }
4694
4695 #[tokio::test]
4696 async fn test_websocket_basic() {
4697 let server = TestServer::setup().await;
4698 let client = setup_test_client(server.port).await;
4699
4700 assert!(!client.is_disconnected());
4701
4702 client.disconnect().await;
4703 assert!(client.is_disconnected());
4704 }
4705
4706 #[rstest]
4707 #[tokio::test]
4708 async fn test_drop_sets_shared_connection_mode_closed() {
4709 let server = TestServer::setup().await;
4710 let client = setup_test_client(server.port).await;
4711 let connection_mode = client.connection_mode_atomic();
4712
4713 drop(client);
4714
4715 assert_eq!(
4716 ConnectionMode::from_atomic(&connection_mode),
4717 ConnectionMode::Closed
4718 );
4719 }
4720
4721 #[rstest]
4722 #[tokio::test]
4723 async fn test_notify_closed_closes_reconnecting_client() {
4724 let server = TestServer::setup().await;
4725 let client = setup_test_client(server.port).await;
4726 client
4727 .connection_mode
4728 .store(ConnectionMode::Reconnect.as_u8(), Ordering::SeqCst);
4729
4730 client.notify_closed();
4731
4732 assert!(client.is_closed());
4733 }
4734
4735 #[tokio::test]
4736 async fn test_websocket_heartbeat() {
4737 let server = TestServer::setup().await;
4738 let client = setup_test_client(server.port).await;
4739
4740 tokio::time::sleep(std::time::Duration::from_secs(3)).await;
4742
4743 client.disconnect().await;
4745 assert!(client.is_disconnected());
4746 }
4747
4748 #[rstest]
4749 #[tokio::test]
4750 async fn test_websocket_reconnect_exhausted() {
4751 let config = WebSocketConfig {
4752 url: "ws://127.0.0.1:9997".into(), headers: vec![],
4754 heartbeat_interval_secs: None,
4755 heartbeat_payload: None,
4756 connect_timeout_ms: None,
4757 reconnect_delay_initial_ms: None,
4758 reconnect_backoff_factor: None,
4759 reconnect_delay_max_ms: None,
4760 reconnect_jitter_ms: None,
4761 reconnect_max_attempts: None,
4762 heartbeat_timeout_secs: None,
4763 idle_timeout_ms: None,
4764 writer_capacity: None,
4765 backend: TransportBackend::Tungstenite,
4766 proxy_url: None,
4767 max_message_size_bytes: None,
4768 max_frame_size_bytes: None,
4769 };
4770
4771 let states = Arc::new(Mutex::new(Vec::new()));
4772 let states_callback = Arc::clone(&states);
4773 let sink = SocketStateSink::new(move |state| {
4774 states_callback.lock().push(state);
4775 });
4776
4777 let res = WebSocketClient::builder()
4778 .config(config)
4779 .message_handler(Arc::new(|_| {}))
4780 .state_sink(sink)
4781 .connect()
4782 .await;
4783 assert!(res.is_err(), "Should fail quickly with no server");
4784 assert_eq!(*states.lock(), Vec::new());
4785 }
4786
4787 #[tokio::test]
4788 async fn test_websocket_forced_close_reconnect() {
4789 let server = TestServer::setup().await;
4790 let client = setup_test_client(server.port).await;
4791
4792 client.send_text("Hello".into(), None).await.unwrap();
4794
4795 client.send_text("close-now".into(), None).await.unwrap();
4797
4798 tokio::time::sleep(std::time::Duration::from_secs(1)).await;
4800
4801 assert!(!client.is_disconnected());
4803
4804 client.disconnect().await;
4806 assert!(client.is_disconnected());
4807 }
4808
4809 #[rstest]
4810 #[case::plain(false)]
4811 #[case::epoch(true)]
4812 #[tokio::test]
4813 async fn test_peer_close_preserves_forwarded_messages(#[case] epoch_aware: bool) {
4814 let timeout = Duration::from_secs(5);
4815 let subscription = "subscribe-retained";
4816 let server = TestServer::setup().await;
4817 let config = WebSocketConfig::builder()
4818 .url(format!("ws://127.0.0.1:{}", server.port))
4819 .headers(vec![("test".to_string(), "test".to_string())])
4820 .backend(TransportBackend::Tungstenite)
4821 .reconnect_delay_initial_ms(50)
4822 .reconnect_jitter_ms(0)
4823 .build()
4824 .unwrap();
4825 let (handler, mut raw_rx) = channel_epoch_message_handler();
4826
4827 let ws_client = if epoch_aware {
4828 WebSocketClient::epoch_builder()
4829 .config(config)
4830 .epoch_handler(handler)
4831 .connect()
4832 .await
4833 } else {
4834 WebSocketClient::builder()
4835 .config(config)
4836 .message_handler(Arc::new(move |message| handler(0, message)))
4837 .connect()
4838 .await
4839 }
4840 .unwrap();
4841
4842 let ws_client = Arc::new(ws_client);
4843 let (bytes_tx, mut bytes_rx) = tokio::sync::mpsc::unbounded_channel();
4844 let forwarding_client = Arc::clone(&ws_client);
4845
4846 let forwarding_task = tokio::spawn(async move {
4847 while let Some((epoch, message)) = raw_rx.recv().await {
4848 let data = match message {
4849 WsMessage::Text(text) => text.as_bytes().to_vec(),
4850 WsMessage::Binary(data) => data.to_vec(),
4851 WsMessage::Close(_) => break,
4852 _ => continue,
4853 };
4854
4855 if data == RECONNECTED.as_bytes() {
4856 forwarding_client
4857 .send_text(subscription.into(), None)
4858 .await
4859 .unwrap();
4860 }
4861
4862 if bytes_tx.send((epoch, data)).is_err() {
4863 break;
4864 }
4865 }
4866 });
4867
4868 ws_client
4869 .send_text(subscription.into(), None)
4870 .await
4871 .unwrap();
4872 let initial = tokio::time::timeout(timeout, bytes_rx.recv())
4873 .await
4874 .expect("initial subscription response timed out")
4875 .expect("forwarding task ended before peer close");
4876
4877 ws_client.send_text("close-now".into(), None).await.unwrap();
4878 let reconnected = tokio::time::timeout(timeout, bytes_rx.recv())
4879 .await
4880 .expect("reconnect notification timed out")
4881 .expect("peer close terminated the forwarding task");
4882 let replayed = tokio::time::timeout(timeout, bytes_rx.recv())
4883 .await
4884 .expect("subscription replay response timed out")
4885 .expect("forwarding task ended before subscription replay");
4886 let payload = vec![3, 7, 11];
4887 ws_client.send_bytes(payload.clone(), None).await.unwrap();
4888 let live = tokio::time::timeout(timeout, bytes_rx.recv())
4889 .await
4890 .expect("replacement connection data timed out")
4891 .expect("forwarding task ended before replacement data");
4892
4893 let expected_epoch = u64::from(epoch_aware);
4894 assert_eq!(initial, (0, subscription.as_bytes().to_vec()));
4895 assert_eq!(
4896 reconnected,
4897 (expected_epoch, RECONNECTED.as_bytes().to_vec())
4898 );
4899 assert_eq!(replayed, (expected_epoch, subscription.as_bytes().to_vec()));
4900 assert_eq!(live, (expected_epoch, payload));
4901 assert_eq!(ws_client.connection_epoch(), 1);
4902 assert!(!forwarding_task.is_finished());
4903
4904 ws_client.disconnect().await;
4905 forwarding_task.abort();
4906 }
4907
4908 #[rstest]
4909 #[case(true)]
4910 #[case(false)]
4911 #[tokio::test]
4912 async fn test_peer_close_frame_is_forwarded_to_close_handler(#[case] reconnect: bool) {
4913 let server = TestServer::setup().await;
4914 let (handler, mut handler_rx) = channel_message_handler();
4915 let (state_tx, mut state_rx) = tokio::sync::mpsc::unbounded_channel();
4916
4917 let client = WebSocketClient::builder()
4918 .config(WebSocketConfig {
4919 url: format!("ws://127.0.0.1:{}", server.port),
4920 headers: vec![("test".into(), "test".into())],
4921 heartbeat_interval_secs: None,
4922 heartbeat_payload: None,
4923 connect_timeout_ms: Some(1_000),
4924 reconnect_delay_initial_ms: Some(1_000),
4925 reconnect_backoff_factor: None,
4926 reconnect_delay_max_ms: None,
4927 reconnect_jitter_ms: Some(0),
4928 reconnect_max_attempts: None,
4929 heartbeat_timeout_secs: None,
4930 idle_timeout_ms: None,
4931 writer_capacity: None,
4932 backend: TransportBackend::Tungstenite,
4933 proxy_url: None,
4934 max_message_size_bytes: None,
4935 max_frame_size_bytes: None,
4936 })
4937 .message_handler(Arc::new(|_| {}))
4938 .state_sink(SocketStateSink::new(move |state| {
4939 state_tx.send(state).expect("state channel closed");
4940 }))
4941 .close_handler(Arc::new(move |frame| {
4942 handler(Message::Close(frame));
4943 reconnect
4944 }))
4945 .connect()
4946 .await
4947 .expect("Failed to connect");
4948
4949 let tracker = AuthTracker::new();
4950 let pending_auth = tracker.begin();
4951 client.set_auth_tracker(tracker.clone(), false);
4952
4953 client
4954 .send_text(CODED_CLOSE_TRIGGER.to_string(), None)
4955 .await
4956 .unwrap();
4957
4958 let message = tokio::time::timeout(Duration::from_secs(10), handler_rx.recv())
4959 .await
4960 .expect("close frame was not forwarded")
4961 .expect("handler channel closed");
4962
4963 let WsMessage::Close(Some(frame)) = message else {
4964 panic!("expected a close frame, was {message:?}");
4965 };
4966
4967 let connected = state_rx.recv().await.expect("initial socket state");
4968 let disconnected = tokio::time::timeout(Duration::from_secs(5), state_rx.recv())
4969 .await
4970 .expect("peer close must publish socket loss")
4971 .expect("state channel closed");
4972
4973 assert_eq!(u16::from(frame.code), CODED_CLOSE_CODE);
4974 assert_eq!(frame.reason.as_str(), CODED_CLOSE_REASON);
4975 assert_eq!(connected, SocketState::Connected);
4976 assert_eq!(disconnected, SocketState::Disconnected);
4977
4978 if !reconnect {
4979 let auth_result = tokio::time::timeout(Duration::from_secs(1), pending_auth)
4980 .await
4981 .expect("terminal close should resolve pending authentication")
4982 .expect("authentication sender should report its terminal result");
4983
4984 assert_eq!(auth_result, Err("WebSocket client closed".to_string()));
4985 assert_eq!(tracker.auth_state(), AuthState::Failed);
4986 }
4987
4988 client.disconnect().await;
4989 }
4990
4991 #[rstest]
4992 #[case(true, ConnectionMode::Active)]
4993 #[case(false, ConnectionMode::Closed)]
4994 fn test_peer_close_callback_controls_reconnect(
4995 #[case] reconnect: bool,
4996 #[case] expected_mode: ConnectionMode,
4997 ) {
4998 let mode = AtomicU8::new(ConnectionMode::Active.as_u8());
4999 let fence = ReadSessionFence::new();
5000 let received = Arc::new(Mutex::new(Vec::new()));
5001 let states = Arc::new(Mutex::new(Vec::new()));
5002 let sink_states = Arc::clone(&states);
5003 let sink = SocketStateSink::new(move |state| sink_states.lock().push(state));
5004 let callback_received = Arc::clone(&received);
5005 let callback: CloseHandler = Arc::new(move |frame| {
5006 callback_received.lock().push(frame);
5007 reconnect
5008 });
5009
5010 let frame = CloseFrame::new(CODED_CLOSE_CODE, CODED_CLOSE_REASON);
5011
5012 WebSocketClientInner::dispatch_close_frame(
5013 &mode,
5014 &fence,
5015 Some(&callback),
5016 Some(&sink),
5017 Some(frame.clone()),
5018 );
5019
5020 assert_eq!(*received.lock(), vec![Some(frame)]);
5021 assert_eq!(ConnectionMode::from_atomic(&mode), expected_mode);
5022 assert_eq!(
5023 *states.lock(),
5024 if reconnect {
5025 Vec::new()
5026 } else {
5027 vec![SocketState::Disconnected]
5028 },
5029 );
5030 }
5031
5032 #[rstest]
5033 fn test_peer_close_callback_ignores_retired_reader() {
5034 let mode = AtomicU8::new(ConnectionMode::Active.as_u8());
5035 let fence = ReadSessionFence::new();
5036 fence.invalidate();
5037 let calls = Arc::new(AtomicUsize::new(0));
5038 let callback_calls = Arc::clone(&calls);
5039 let callback: CloseHandler = Arc::new(move |_| {
5040 callback_calls.fetch_add(1, Ordering::SeqCst);
5041 false
5042 });
5043
5044 WebSocketClientInner::dispatch_close_frame(&mode, &fence, Some(&callback), None, None);
5045
5046 assert_eq!(calls.load(Ordering::SeqCst), 0);
5047 assert_eq!(ConnectionMode::from_atomic(&mode), ConnectionMode::Active);
5048 }
5049
5050 #[rstest]
5052 #[tokio::test]
5053 async fn test_close_frame_reason_is_redacted_in_logs() {
5054 let server = TestServer::setup().await;
5055 let client = setup_test_client(server.port).await;
5056 let capture = capture_logs_for(&["nautilus_network::websocket::client"]).await;
5057
5058 client
5059 .send_text(HOSTILE_CLOSE_TRIGGER.to_string(), None)
5060 .await
5061 .unwrap();
5062
5063 tokio::time::timeout(Duration::from_secs(5), async {
5064 loop {
5065 if capture.messages().iter().any(|(_, message)| {
5066 message.starts_with("Received close frame, terminating: code=")
5067 }) {
5068 break;
5069 }
5070
5071 tokio::task::yield_now().await;
5072 }
5073 })
5074 .await
5075 .expect("timed out waiting for the close frame log");
5076
5077 let close_frame_logs: Vec<String> = capture
5078 .messages()
5079 .iter()
5080 .map(|(_, message)| message.clone())
5081 .filter(|message| message.starts_with("Received close frame, terminating: code="))
5082 .collect();
5083
5084 let expected = "Received close frame, terminating: code=1011, reason=<redacted>";
5085 assert!(
5086 close_frame_logs.iter().any(|message| message == expected),
5087 "close frame reason was not redacted: {close_frame_logs:?}"
5088 );
5089 assert!(
5090 close_frame_logs.iter().all(|message| {
5091 !message.contains('\r') && !message.contains('\n') && !message.contains('\u{1b}')
5092 }),
5093 "close frame log forged line breaks or escape sequences: {close_frame_logs:?}"
5094 );
5095
5096 assert!(
5097 close_frame_logs
5098 .iter()
5099 .all(|message| !message.contains("injected") && !message.contains("forged"))
5100 );
5101
5102 client.disconnect().await;
5103 }
5104
5105 #[rstest]
5106 #[tokio::test]
5107 async fn test_state_sink_reports_initial_loss_and_recovery() {
5108 let server = TestServer::setup().await;
5109
5110 let config = WebSocketConfig {
5111 url: format!("ws://127.0.0.1:{}", server.port),
5112 headers: vec![("test".into(), "test".into())],
5113 heartbeat_interval_secs: None,
5114 heartbeat_payload: None,
5115 connect_timeout_ms: Some(1_000),
5116 reconnect_delay_initial_ms: Some(1),
5117 reconnect_backoff_factor: Some(1.0),
5118 reconnect_delay_max_ms: Some(1),
5119 reconnect_jitter_ms: Some(0),
5120 reconnect_max_attempts: Some(3),
5121 heartbeat_timeout_secs: None,
5122 idle_timeout_ms: None,
5123 writer_capacity: None,
5124 backend: TransportBackend::Tungstenite,
5125 proxy_url: None,
5126 max_message_size_bytes: None,
5127 max_frame_size_bytes: None,
5128 };
5129
5130 let states = Arc::new(Mutex::new(Vec::new()));
5131 let states_callback = Arc::clone(&states);
5132 let sink = SocketStateSink::new(move |state| {
5133 states_callback.lock().push(state);
5134 });
5135
5136 let client = WebSocketClient::builder()
5137 .config(config)
5138 .message_handler(Arc::new(|_| {}))
5139 .state_sink(sink)
5140 .connect()
5141 .await
5142 .unwrap();
5143
5144 assert_eq!(*states.lock(), vec![SocketState::Connected]);
5145
5146 client.send_text("close-now".into(), None).await.unwrap();
5147 wait_until_async(
5148 || {
5149 let states = Arc::clone(&states);
5150 async move { states.lock().len() == 3 }
5151 },
5152 Duration::from_secs(5),
5153 )
5154 .await;
5155
5156 assert_eq!(
5157 *states.lock(),
5158 vec![
5159 SocketState::Connected,
5160 SocketState::Disconnected,
5161 SocketState::Connected,
5162 ]
5163 );
5164
5165 client.disconnect().await;
5166 assert_eq!(states.lock().len(), 3);
5167 }
5168
5169 #[rstest]
5170 #[tokio::test]
5171 async fn test_drop_suppresses_socket_state_event() {
5172 let server = TestServer::setup().await;
5173
5174 let config = WebSocketConfig {
5175 url: format!("ws://127.0.0.1:{}", server.port),
5176 headers: vec![("test".into(), "test".into())],
5177 heartbeat_interval_secs: None,
5178 heartbeat_payload: None,
5179 connect_timeout_ms: Some(1_000),
5180 reconnect_delay_initial_ms: Some(1),
5181 reconnect_backoff_factor: Some(1.0),
5182 reconnect_delay_max_ms: Some(1),
5183 reconnect_jitter_ms: Some(0),
5184 reconnect_max_attempts: Some(3),
5185 heartbeat_timeout_secs: None,
5186 idle_timeout_ms: None,
5187 writer_capacity: None,
5188 backend: TransportBackend::Tungstenite,
5189 proxy_url: None,
5190 max_message_size_bytes: None,
5191 max_frame_size_bytes: None,
5192 };
5193
5194 let states = Arc::new(Mutex::new(Vec::new()));
5195 let states_callback = Arc::clone(&states);
5196 let sink = SocketStateSink::new(move |state| {
5197 states_callback.lock().push(state);
5198 });
5199
5200 let client = WebSocketClient::builder()
5201 .config(config)
5202 .message_handler(Arc::new(|_| {}))
5203 .state_sink(sink)
5204 .connect()
5205 .await
5206 .unwrap();
5207
5208 drop(client);
5209 crate::dst::time::sleep(Duration::from_millis(25)).await;
5210
5211 assert_eq!(*states.lock(), vec![SocketState::Connected]);
5212 }
5213
5214 #[rstest]
5215 #[tokio::test]
5216 async fn test_stream_state_sink_reports_reader_loss() {
5217 let server = TestServer::setup().await;
5218
5219 let config = WebSocketConfig {
5220 url: format!("ws://127.0.0.1:{}", server.port),
5221 headers: vec![("test".into(), "test".into())],
5222 heartbeat_interval_secs: None,
5223 heartbeat_payload: None,
5224 connect_timeout_ms: None,
5225 reconnect_delay_initial_ms: None,
5226 reconnect_backoff_factor: None,
5227 reconnect_delay_max_ms: None,
5228 reconnect_jitter_ms: None,
5229 reconnect_max_attempts: None,
5230 heartbeat_timeout_secs: None,
5231 idle_timeout_ms: None,
5232 writer_capacity: None,
5233 backend: TransportBackend::Tungstenite,
5234 proxy_url: None,
5235 max_message_size_bytes: None,
5236 max_frame_size_bytes: None,
5237 };
5238
5239 let states = Arc::new(Mutex::new(Vec::new()));
5240 let states_callback = Arc::clone(&states);
5241 let sink = SocketStateSink::new(move |state| {
5242 states_callback.lock().push(state);
5243 });
5244
5245 let (_reader, client) = WebSocketClient::stream_builder()
5246 .config(config)
5247 .state_sink(sink)
5248 .connect()
5249 .await
5250 .unwrap();
5251
5252 client.notify_closed();
5253
5254 assert!(client.is_closed());
5255 assert_eq!(
5256 *states.lock(),
5257 vec![SocketState::Connected, SocketState::Disconnected]
5258 );
5259 }
5260
5261 #[rstest]
5262 #[tokio::test]
5263 async fn test_state_sink_emits_no_retry_or_exhaustion_events() {
5264 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
5265 let port = listener.local_addr().unwrap().port();
5266 let server_task = task::spawn(async move {
5267 let (connection, _) = listener.accept().await.unwrap();
5268 let mut websocket = accept_async(connection).await.unwrap();
5269 while let Some(Ok(message)) = websocket.next().await {
5270 if matches!(&message, WsMessage::Text(text) if text.as_str() == "close-now") {
5271 websocket.close(None).await.unwrap();
5272 break;
5273 }
5274 }
5275 });
5276
5277 let config = WebSocketConfig {
5278 url: format!("ws://127.0.0.1:{port}"),
5279 headers: vec![],
5280 heartbeat_interval_secs: None,
5281 heartbeat_payload: None,
5282 connect_timeout_ms: Some(100),
5283 reconnect_delay_initial_ms: Some(1),
5284 reconnect_backoff_factor: Some(1.0),
5285 reconnect_delay_max_ms: Some(1),
5286 reconnect_jitter_ms: Some(0),
5287 reconnect_max_attempts: Some(2),
5288 heartbeat_timeout_secs: None,
5289 idle_timeout_ms: None,
5290 writer_capacity: None,
5291 backend: TransportBackend::Tungstenite,
5292 proxy_url: None,
5293 max_message_size_bytes: None,
5294 max_frame_size_bytes: None,
5295 };
5296
5297 let states = Arc::new(Mutex::new(Vec::new()));
5298 let states_callback = Arc::clone(&states);
5299 let sink = SocketStateSink::new(move |state| {
5300 states_callback.lock().push(state);
5301 });
5302
5303 let client = WebSocketClient::builder()
5304 .config(config)
5305 .message_handler(Arc::new(|_| {}))
5306 .state_sink(sink)
5307 .connect()
5308 .await
5309 .unwrap();
5310
5311 client.send_text("close-now".into(), None).await.unwrap();
5312 wait_until_async(
5313 || async { client.is_disconnected() },
5314 Duration::from_secs(5),
5315 )
5316 .await;
5317 assert_eq!(
5318 *states.lock(),
5319 vec![SocketState::Connected, SocketState::Disconnected]
5320 );
5321
5322 server_task.await.unwrap();
5323 }
5324
5325 #[rstest]
5326 #[case(
5327 "invalid header",
5328 "value",
5329 "Invalid WebSocket reconnect header name: invalid HTTP header name"
5330 )]
5331 #[case(
5332 "x-test",
5333 "invalid\nvalue",
5334 "Invalid WebSocket reconnect header value: failed to parse header value"
5335 )]
5336 fn test_reconnect_headers_reject_invalid_update(
5337 #[case] name: &str,
5338 #[case] value: &str,
5339 #[case] expected_message: &str,
5340 ) {
5341 let initial = vec![("x-existing".to_string(), "initial".to_string())];
5342 let headers = ReconnectHeaders::new(initial.clone());
5343
5344 let error = headers.update(name, value).unwrap_err();
5345
5346 let TransportError::Io(error) = error else {
5347 panic!("expected I/O error, was {error:?}");
5348 };
5349
5350 assert_eq!(error.kind(), std::io::ErrorKind::InvalidInput);
5351 assert_eq!(error.to_string(), expected_message);
5352 assert_eq!(headers.snapshot(), initial);
5353 }
5354
5355 #[tokio::test]
5356 #[allow(clippy::result_large_err)]
5357 async fn test_reconnect_uses_updated_headers_without_interrupting_active_connection() {
5358 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
5359 let port = listener.local_addr().unwrap().port();
5360 let (header_tx, mut header_rx) = tokio::sync::mpsc::unbounded_channel();
5361
5362 let server_task = task::spawn(async move {
5363 loop {
5364 let (conn, _) = listener.accept().await.unwrap();
5365 let header_tx = header_tx.clone();
5366
5367 let mut websocket = accept_hdr_async(
5368 conn,
5369 move |request: &server::Request, response: server::Response| {
5370 let values = request
5371 .headers()
5372 .get_all("authorization")
5373 .iter()
5374 .map(|value| value.to_str().unwrap().to_string())
5375 .collect::<Vec<_>>();
5376 header_tx.send(values).unwrap();
5377 Ok(response)
5378 },
5379 )
5380 .await
5381 .unwrap();
5382
5383 task::spawn(async move {
5384 while let Some(Ok(msg)) = websocket.next().await {
5385 if matches!(&msg, WsMessage::Text(text) if text.as_str() == "close-now") {
5386 let _ = websocket.close(None).await;
5387 break;
5388 }
5389 }
5390 });
5391 }
5392 });
5393
5394 let config = WebSocketConfig {
5395 url: format!("ws://127.0.0.1:{port}"),
5396 headers: vec![("Authorization".into(), "Bearer initial".into())],
5397 heartbeat_interval_secs: None,
5398 heartbeat_payload: None,
5399 connect_timeout_ms: Some(1_000),
5400 reconnect_delay_initial_ms: Some(50),
5401 reconnect_delay_max_ms: Some(50),
5402 reconnect_backoff_factor: Some(1.0),
5403 reconnect_jitter_ms: Some(0),
5404 reconnect_max_attempts: None,
5405 heartbeat_timeout_secs: None,
5406 idle_timeout_ms: None,
5407 writer_capacity: None,
5408 backend: TransportBackend::Tungstenite,
5409 proxy_url: None,
5410 max_message_size_bytes: None,
5411 max_frame_size_bytes: None,
5412 };
5413
5414 let client = WebSocketClient::builder()
5415 .config(config)
5416 .message_handler(Arc::new(|_| {}))
5417 .connect()
5418 .await
5419 .unwrap();
5420
5421 let initial = header_rx.recv().await.unwrap();
5422 let reconnect_headers = client.reconnect_headers();
5423 reconnect_headers
5424 .update("authorization", "Bearer refreshed")
5425 .unwrap();
5426
5427 tokio::time::sleep(Duration::from_millis(100)).await;
5428 assert_eq!(initial, vec!["Bearer initial"]);
5429 assert!(client.is_active());
5430 assert!(header_rx.try_recv().is_err());
5431 assert!(!format!("{reconnect_headers:?}").contains("refreshed"));
5432
5433 client.send_text("close-now".into(), None).await.unwrap();
5434 let refreshed = tokio::time::timeout(Duration::from_secs(3), header_rx.recv())
5435 .await
5436 .unwrap()
5437 .unwrap();
5438
5439 assert_eq!(refreshed, vec!["Bearer refreshed"]);
5440
5441 client.disconnect().await;
5442 server_task.abort();
5443 }
5444
5445 #[tokio::test]
5446 async fn test_rate_limiter() {
5447 let server = TestServer::setup().await;
5448 let quota = Quota::per_second(NonZeroU32::new(2).unwrap()).unwrap();
5449
5450 let config = WebSocketConfig {
5451 url: format!("ws://127.0.0.1:{}", server.port),
5452 headers: vec![("test".into(), "test".into())],
5453 heartbeat_interval_secs: None,
5454 heartbeat_payload: None,
5455 connect_timeout_ms: None,
5456 reconnect_delay_initial_ms: None,
5457 reconnect_backoff_factor: None,
5458 reconnect_delay_max_ms: None,
5459 reconnect_jitter_ms: None,
5460 reconnect_max_attempts: None,
5461 heartbeat_timeout_secs: None,
5462 idle_timeout_ms: None,
5463 writer_capacity: None,
5464 backend: TransportBackend::Tungstenite,
5465 proxy_url: None,
5466 max_message_size_bytes: None,
5467 max_frame_size_bytes: None,
5468 };
5469
5470 let client = WebSocketClient::builder()
5471 .config(config)
5472 .message_handler(Arc::new(|_| {}))
5473 .keyed_quotas(vec![("default".into(), quota)])
5474 .connect()
5475 .await
5476 .unwrap();
5477
5478 let keys: [ustr::Ustr; 1] = [ustr::Ustr::from("default")];
5481 let start = std::time::Instant::now();
5482 client
5483 .send_text("test1".into(), Some(keys.as_slice()))
5484 .await
5485 .unwrap();
5486 client
5487 .send_text("test2".into(), Some(keys.as_slice()))
5488 .await
5489 .unwrap();
5490 let after_burst = start.elapsed();
5491 client
5492 .send_text("test3".into(), Some(keys.as_slice()))
5493 .await
5494 .unwrap();
5495 let after_third = start.elapsed();
5496
5497 assert!(
5498 after_burst < std::time::Duration::from_millis(300),
5499 "Burst sends should not be rate limited, took {after_burst:?}"
5500 );
5501 assert!(
5502 after_third >= std::time::Duration::from_millis(400),
5503 "Third send should wait for quota replenishment, took {after_third:?}"
5504 );
5505
5506 client.disconnect().await;
5508 assert!(client.is_disconnected());
5509 }
5510
5511 #[tokio::test]
5512 async fn test_concurrent_writers() {
5513 let server = TestServer::setup().await;
5514 let client = Arc::new(setup_test_client(server.port).await);
5515
5516 let mut handles = vec![];
5517
5518 for i in 0..10 {
5519 let client = Arc::clone(&client);
5520 handles.push(task::spawn(async move {
5521 client.send_text(format!("test{i}"), None).await.unwrap();
5522 }));
5523 }
5524
5525 for handle in handles {
5526 handle.await.unwrap();
5527 }
5528
5529 client.disconnect().await;
5531 assert!(client.is_disconnected());
5532 }
5533}
5534
5535#[cfg(test)]
5536#[cfg(not(feature = "turmoil"))]
5537#[cfg(not(all(feature = "simulation", madsim)))] mod rust_tests {
5539 use std::{
5540 pin::Pin,
5541 sync::{
5542 Arc, OnceLock,
5543 atomic::{AtomicBool, AtomicU8, AtomicUsize, Ordering},
5544 },
5545 task::{Context, Poll},
5546 };
5547
5548 use futures_util::{SinkExt, StreamExt};
5549 #[cfg(target_os = "linux")]
5550 use log::Level;
5551 use nautilus_common::testing::wait_until_async;
5552 use parking_lot::{Condvar, Mutex};
5553 use rstest::rstest;
5554 #[cfg(feature = "transport-sockudo")]
5555 use sockudo_ws::handshake as sockudo_handshake;
5556 #[cfg(feature = "transport-sockudo")]
5557 use tokio::io::AsyncRead;
5558 use tokio::{
5559 io::{AsyncReadExt, AsyncWriteExt},
5560 net::TcpListener,
5561 sync::oneshot,
5562 task::{self, JoinHandle},
5563 time::{Duration, sleep},
5564 };
5565 use tokio_tungstenite::{accept_async, tungstenite::Message as WsMessage};
5566 #[cfg(feature = "transport-sockudo")]
5567 use tokio_tungstenite::{
5568 accept_hdr_async,
5569 tungstenite::{
5570 handshake::server::{self, Callback},
5571 http::HeaderValue,
5572 },
5573 };
5574
5575 use super::*;
5576 #[cfg(target_os = "linux")]
5577 use crate::logging::tests::capture_logs_for;
5578 use crate::{
5579 SocketState,
5580 websocket::types::{channel_epoch_message_handler, channel_message_handler},
5581 };
5582
5583 const TEST_TIMEOUT: Duration = Duration::from_secs(10);
5584
5585 #[cfg(all(feature = "transport-sockudo", target_os = "linux"))]
5588 const RECONNECT_LOG_TARGETS: &[&str] = &[
5589 "nautilus_network::transport::sockudo",
5590 "nautilus_network::websocket::client",
5591 ];
5592
5593 struct CondvarReleaseGuard<'a> {
5594 release: &'a (Mutex<bool>, Condvar),
5595 }
5596
5597 impl<'a> CondvarReleaseGuard<'a> {
5598 fn new(release: &'a (Mutex<bool>, Condvar)) -> Self {
5599 Self { release }
5600 }
5601
5602 fn release(&self) {
5603 let (lock, condvar) = self.release;
5604 let mut released = lock.lock();
5605 *released = true;
5606 condvar.notify_all();
5607 }
5608 }
5609
5610 impl Drop for CondvarReleaseGuard<'_> {
5611 fn drop(&mut self) {
5612 self.release();
5613 }
5614 }
5615
5616 async fn recv_rendezvous<T: Send + 'static>(
5617 receiver: std::sync::mpsc::Receiver<T>,
5618 name: &'static str,
5619 ) -> T {
5620 let receive_task = tokio::task::spawn_blocking(move || receiver.recv_timeout(TEST_TIMEOUT));
5621
5622 match tokio::time::timeout(TEST_TIMEOUT * 2, receive_task).await {
5623 Ok(Ok(Ok(value))) => value,
5624 Ok(Ok(Err(e))) => {
5625 panic!("{name} did not arrive within the test timeout: {e}")
5626 }
5627 Ok(Err(e)) => panic!("{name} receive task failed: {e}"),
5628 Err(e) => panic!("{name} receive task did not finish: {e}"),
5629 }
5630 }
5631
5632 async fn await_task_termination(task: tokio::task::JoinHandle<()>, name: &'static str) {
5633 match tokio::time::timeout(TEST_TIMEOUT, task).await {
5634 Ok(Ok(())) => {}
5635 Ok(Err(e)) if e.is_cancelled() => {}
5636 Ok(Err(e)) => panic!("{name} failed: {e}"),
5637 Err(e) => panic!("{name} did not terminate within the test timeout: {e}"),
5638 }
5639 }
5640
5641 fn reconnect_test_config(port: u16) -> WebSocketConfig {
5642 WebSocketConfig {
5643 url: format!("ws://127.0.0.1:{port}"),
5644 headers: vec![],
5645 heartbeat_interval_secs: None,
5646 heartbeat_payload: None,
5647 connect_timeout_ms: Some(1_000),
5648 reconnect_delay_initial_ms: None,
5649 reconnect_delay_max_ms: None,
5650 reconnect_backoff_factor: None,
5651 reconnect_jitter_ms: None,
5652 reconnect_max_attempts: None,
5653 heartbeat_timeout_secs: None,
5654 idle_timeout_ms: None,
5655 writer_capacity: None,
5656 backend: TransportBackend::Tungstenite,
5657 proxy_url: None,
5658 max_message_size_bytes: None,
5659 max_frame_size_bytes: None,
5660 }
5661 }
5662
5663 fn initial_connect_retry_policy(
5664 max_attempts: u32,
5665 initial_delay_ms: u64,
5666 ) -> InitialConnectRetryPolicy {
5667 InitialConnectRetryPolicy {
5668 max_attempts: std::num::NonZeroU32::new(max_attempts).unwrap(),
5669 delay_initial: Duration::from_millis(initial_delay_ms),
5670 delay_max: Duration::from_millis(initial_delay_ms),
5671 backoff_factor: 2.0,
5672 jitter_ms: 0,
5673 }
5674 }
5675
5676 #[rstest]
5677 #[case(TransportError::UpgradeRejected(408), true)]
5678 #[case(TransportError::UpgradeRejected(425), true)]
5679 #[case(TransportError::UpgradeRejected(429), true)]
5680 #[case(TransportError::UpgradeRejected(500), true)]
5681 #[case(TransportError::UpgradeRejected(599), true)]
5682 #[case(TransportError::UpgradeRejected(404), false)]
5683 #[case(TransportError::UpgradeRejected(600), false)]
5684 #[case(TransportError::ProxyConnectRejected(429), true)]
5685 #[case(TransportError::ProxyConnectRejected(407), false)]
5686 #[case(TransportError::ConnectionClosed, true)]
5687 #[case(
5688 TransportError::Io(std::io::Error::from(std::io::ErrorKind::TimedOut)),
5689 true
5690 )]
5691 #[case(
5692 TransportError::Io(std::io::Error::from(std::io::ErrorKind::Interrupted)),
5693 true
5694 )]
5695 #[case(
5696 TransportError::Io(std::io::Error::from(std::io::ErrorKind::InvalidInput)),
5697 false
5698 )]
5699 #[case(
5700 TransportError::Io(std::io::Error::from(std::io::ErrorKind::InvalidData)),
5701 false
5702 )]
5703 #[case(
5704 TransportError::Io(std::io::Error::from(std::io::ErrorKind::Unsupported)),
5705 false
5706 )]
5707 #[case(
5708 TransportError::Io(std::io::Error::from(std::io::ErrorKind::PermissionDenied)),
5709 false
5710 )]
5711 #[case(TransportError::Handshake("malformed".to_string()), false)]
5712 fn initial_connect_error_classification(#[case] error: TransportError, #[case] expected: bool) {
5713 assert_eq!(is_retryable_initial_connect_error(&error), expected);
5714 }
5715
5716 async fn rejected_upgrade_attempts(
5717 status: u16,
5718 backend: TransportBackend,
5719 retry_policy: Option<InitialConnectRetryPolicy>,
5720 ) -> u64 {
5721 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
5722 let port = listener.local_addr().unwrap().port();
5723 let accepted = Arc::new(AtomicU64::new(0));
5724 let accepted_server = Arc::clone(&accepted);
5725
5726 let server = tokio::spawn(async move {
5727 loop {
5728 let (mut stream, _) = listener.accept().await.unwrap();
5729 accepted_server.fetch_add(1, Ordering::SeqCst);
5730 let mut request = Vec::new();
5731
5732 loop {
5733 let mut chunk = [0; 1024];
5734 let n = stream.read(&mut chunk).await.unwrap();
5735 if n == 0 {
5736 break;
5737 }
5738
5739 request.extend_from_slice(&chunk[..n]);
5740 if request.windows(4).any(|window| window == b"\r\n\r\n") {
5741 break;
5742 }
5743 }
5744
5745 stream
5746 .write_all(format!("HTTP/1.1 {status}\r\n\r\n").as_bytes())
5747 .await
5748 .unwrap();
5749 }
5750 });
5751
5752 let (handler, _rx) = channel_message_handler();
5753 let mut config = reconnect_test_config(port);
5754 config.backend = backend;
5755 WebSocketClient::builder()
5756 .config(config)
5757 .message_handler(handler)
5758 .maybe_initial_connect_retry_policy(retry_policy)
5759 .connect()
5760 .await
5761 .expect_err("server always rejects the upgrade");
5762
5763 let attempts = accepted.load(Ordering::SeqCst);
5764 server.abort();
5765 attempts
5766 }
5767
5768 async fn closed_proxy_attempts() -> u64 {
5769 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
5770 let proxy_addr = listener.local_addr().unwrap();
5771 let accepted = Arc::new(AtomicU64::new(0));
5772 let accepted_server = Arc::clone(&accepted);
5773
5774 let server = tokio::spawn(async move {
5775 loop {
5776 let (mut stream, _) = listener.accept().await.unwrap();
5777 accepted_server.fetch_add(1, Ordering::SeqCst);
5778 let mut request = [0; 1024];
5779 let _ = stream.read(&mut request).await.unwrap();
5780 }
5781 });
5782
5783 let (handler, _rx) = channel_message_handler();
5784 let mut config = reconnect_test_config(9);
5785 config.proxy_url = Some(format!("http://{proxy_addr}"));
5786 WebSocketClient::builder()
5787 .config(config)
5788 .message_handler(handler)
5789 .initial_connect_retry_policy(initial_connect_retry_policy(3, 1))
5790 .connect()
5791 .await
5792 .expect_err("proxy always closes before its CONNECT response");
5793
5794 let attempts = accepted.load(Ordering::SeqCst);
5795 server.abort();
5796 attempts
5797 }
5798
5799 #[rstest]
5800 #[tokio::test]
5801 async fn initial_connect_cancellation_interrupts_in_flight_handshake() {
5802 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
5803 let port = listener.local_addr().unwrap().port();
5804 let (accepted_tx, accepted_rx) = tokio::sync::oneshot::channel();
5805
5806 let server = tokio::spawn(async move {
5807 let (_stream, _) = listener.accept().await.unwrap();
5808 let _ = accepted_tx.send(());
5809 std::future::pending::<()>().await;
5810 });
5811
5812 let token = CancellationToken::new();
5813 let (handler, _rx) = channel_message_handler();
5814 let mut config = reconnect_test_config(port);
5815 config.connect_timeout_ms = Some(10_000);
5816 let connect = WebSocketClient::builder()
5817 .config(config)
5818 .message_handler(handler)
5819 .initial_connect_retry_policy(initial_connect_retry_policy(5, 30_000))
5820 .cancellation_token(token.clone())
5821 .connect();
5822 tokio::pin!(connect);
5823
5824 tokio::select! {
5825 biased;
5826 result = &mut connect => panic!("connect completed before cancellation: {result:?}"),
5827 result = accepted_rx => result.unwrap(),
5828 }
5829 token.cancel();
5830
5831 let err = tokio::time::timeout(Duration::from_millis(250), connect)
5832 .await
5833 .expect("cancellation should interrupt the handshake")
5834 .expect_err("cancelled connect should fail");
5835 assert!(
5836 matches!(err, TransportError::Io(ref e) if e.kind() == std::io::ErrorKind::Interrupted)
5837 );
5838 server.abort();
5839 }
5840
5841 #[rstest]
5842 #[tokio::test]
5843 async fn timed_out_initial_connect_attempt_is_retried() {
5844 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
5845 let port = listener.local_addr().unwrap().port();
5846 let accepted = Arc::new(AtomicU64::new(0));
5847 let accepted_server = Arc::clone(&accepted);
5848
5849 let server = tokio::spawn(async move {
5852 let mut held = Vec::new();
5853
5854 loop {
5855 let (stream, _) = listener.accept().await.unwrap();
5856 accepted_server.fetch_add(1, Ordering::SeqCst);
5857 held.push(stream);
5858 }
5859 });
5860
5861 let (handler, _rx) = channel_message_handler();
5862 let mut config = reconnect_test_config(port);
5863 config.connect_timeout_ms = Some(50);
5864 let err = WebSocketClient::builder()
5865 .config(config)
5866 .message_handler(handler)
5867 .initial_connect_retry_policy(initial_connect_retry_policy(3, 1))
5868 .connect()
5869 .await
5870 .expect_err("every attempt times out");
5871
5872 assert!(
5873 matches!(err, TransportError::Io(ref e) if e.kind() == std::io::ErrorKind::TimedOut)
5874 );
5875
5876 let counted = tokio::time::timeout(Duration::from_secs(1), async {
5881 while accepted.load(Ordering::SeqCst) < 3 {
5882 tokio::time::sleep(Duration::from_millis(5)).await;
5883 }
5884 })
5885 .await;
5886
5887 assert!(
5888 counted.is_ok(),
5889 "expected 3 connection attempts, observed {}",
5890 accepted.load(Ordering::SeqCst)
5891 );
5892 assert_eq!(accepted.load(Ordering::SeqCst), 3);
5895 server.abort();
5896 }
5897
5898 #[rstest]
5899 #[case::too_many_requests(429)]
5900 #[case::server_error(503)]
5901 #[tokio::test]
5902 async fn transient_upgrade_rejection_is_retried(#[case] status: u16) {
5903 assert_eq!(
5904 rejected_upgrade_attempts(
5905 status,
5906 TransportBackend::Tungstenite,
5907 Some(initial_connect_retry_policy(3, 1)),
5908 )
5909 .await,
5910 3
5911 );
5912 }
5913
5914 #[rstest]
5915 #[tokio::test]
5916 async fn initial_connect_without_retry_policy_makes_one_attempt() {
5917 assert_eq!(
5918 rejected_upgrade_attempts(503, TransportBackend::Tungstenite, None).await,
5919 1
5920 );
5921 }
5922
5923 #[cfg(feature = "transport-sockudo")]
5924 #[rstest]
5925 #[tokio::test]
5926 async fn sockudo_too_many_requests_upgrade_rejection_is_retried() {
5927 assert_eq!(
5928 rejected_upgrade_attempts(
5929 429,
5930 TransportBackend::Sockudo,
5931 Some(initial_connect_retry_policy(3, 1)),
5932 )
5933 .await,
5934 3
5935 );
5936 }
5937
5938 #[rstest]
5939 #[tokio::test]
5940 async fn permanent_upgrade_rejection_is_not_retried() {
5941 assert_eq!(
5942 rejected_upgrade_attempts(
5943 404,
5944 TransportBackend::Tungstenite,
5945 Some(initial_connect_retry_policy(3, 1)),
5946 )
5947 .await,
5948 1
5949 );
5950 }
5951
5952 #[rstest]
5953 #[tokio::test]
5954 async fn proxy_close_before_connect_response_is_retried() {
5955 assert_eq!(closed_proxy_attempts().await, 3);
5956 }
5957
5958 #[rstest]
5959 #[tokio::test]
5960 async fn initial_connect_cancellation_interrupts_retry_backoff() {
5961 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
5962 let port = listener.local_addr().unwrap().port();
5963 let accepted = Arc::new(AtomicU64::new(0));
5964 let accepted_server = Arc::clone(&accepted);
5965 let (rejected_tx, rejected_rx) = oneshot::channel();
5966 let server = tokio::spawn(async move {
5967 let mut rejected_tx = Some(rejected_tx);
5968
5969 loop {
5970 let (mut stream, _) = listener.accept().await.unwrap();
5971 accepted_server.fetch_add(1, Ordering::SeqCst);
5972 let mut request = [0; 1024];
5973 let _ = stream.read(&mut request).await.unwrap();
5974 stream
5975 .write_all(b"HTTP/1.1 503 Service Unavailable\r\n\r\n")
5976 .await
5977 .unwrap();
5978
5979 if let Some(rejected_tx) = rejected_tx.take() {
5980 rejected_tx.send(()).unwrap();
5981 }
5982 }
5983 });
5984
5985 let token = CancellationToken::new();
5986 let (handler, _rx) = channel_message_handler();
5987 let connect = WebSocketClient::builder()
5988 .config(reconnect_test_config(port))
5989 .message_handler(handler)
5990 .initial_connect_retry_policy(initial_connect_retry_policy(5, 30_000))
5991 .cancellation_token(token.clone())
5992 .connect();
5993 tokio::pin!(connect);
5994
5995 tokio::select! {
5996 biased;
5997 result = &mut connect => panic!("connect completed before backoff: {result:?}"),
5998 result = rejected_rx => result.unwrap(),
5999 }
6000 assert!(futures_util::poll!(&mut connect).is_pending());
6001 token.cancel();
6002
6003 let error = tokio::time::timeout(Duration::from_millis(250), connect)
6004 .await
6005 .expect("cancellation should interrupt retry backoff")
6006 .expect_err("cancelled connect should fail");
6007
6008 assert!(
6009 matches!(error, TransportError::Io(ref error) if error.kind() == std::io::ErrorKind::Interrupted)
6010 );
6011 assert_eq!(accepted.load(Ordering::SeqCst), 1);
6012 server.abort();
6013 }
6014
6015 #[rstest]
6016 #[tokio::test]
6017 async fn permanent_initial_connect_error_does_not_wait_for_retry_ladder() {
6018 let (handler, _rx) = channel_message_handler();
6019 let mut config = reconnect_test_config(1);
6020 config.url = "not a websocket URL".to_string();
6021 let connect = WebSocketClient::builder()
6022 .config(config)
6023 .message_handler(handler)
6024 .initial_connect_retry_policy(initial_connect_retry_policy(5, 30_000))
6025 .connect();
6026
6027 let err = tokio::time::timeout(Duration::from_millis(250), connect)
6028 .await
6029 .expect("permanent error should not enter retry backoff")
6030 .expect_err("invalid URL should fail");
6031 assert!(matches!(
6032 err,
6033 TransportError::InvalidUrl(_) | TransportError::Handshake(_)
6034 ));
6035 }
6036
6037 #[rstest]
6038 #[tokio::test(start_paused = true)]
6039 async fn connection_rate_limit_gates_initial_connect_and_reconnect() {
6040 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
6041 let port = listener.local_addr().unwrap().port();
6042 let server = tokio::spawn(async move {
6043 let (stream, _) = listener.accept().await.unwrap();
6044 let _first = accept_async(stream).await.unwrap();
6045 let (stream, _) = listener.accept().await.unwrap();
6046 let _second = accept_async(stream).await.unwrap();
6047 std::future::pending::<()>().await;
6048 });
6049
6050 let key = Ustr::from("connection-attempt");
6051
6052 let limiter = Arc::new(RateLimiter::new_with_quota(
6053 None,
6054 vec![(key, Quota::with_period(Duration::from_secs(1)).unwrap())],
6055 ));
6056 limiter.await_keys_ready(Some(&[key])).await;
6057
6058 let rate_limit = ConnectionRateLimit {
6059 limiter,
6060 keys: Arc::from([key]),
6061 };
6062
6063 let (handler, _rx) = channel_message_handler();
6064 let mut config = reconnect_test_config(port);
6065 config.connect_timeout_ms = Some(10_000);
6066 let initial = WebSocketClientInner::connect_url_with_handler(
6067 config,
6068 Some(IncomingHandler::Message(handler)),
6069 None,
6070 None,
6071 None,
6072 Some(rate_limit),
6073 InitialConnectOptions::default(),
6074 );
6075 tokio::pin!(initial);
6076 assert!(futures_util::poll!(&mut initial).is_pending());
6077 tokio::time::advance(Duration::from_secs(1)).await;
6078 tokio::task::yield_now().await;
6079
6080 tokio::time::resume();
6083 let mut inner = initial.await.unwrap();
6084 tokio::time::pause();
6085
6086 inner
6087 .connection_mode
6088 .store(ConnectionMode::Reconnect.as_u8(), Ordering::SeqCst);
6089 let reconnect = inner.reconnect_with_outcome();
6090 tokio::pin!(reconnect);
6091 assert!(futures_util::poll!(&mut reconnect).is_pending());
6092 tokio::time::advance(Duration::from_secs(1)).await;
6093 tokio::task::yield_now().await;
6094 tokio::time::resume();
6095 assert_eq!(reconnect.await.unwrap(), ReconnectOutcome::Reconnected);
6096
6097 server.abort();
6098 }
6099
6100 #[rstest]
6101 #[tokio::test(start_paused = true)]
6102 async fn initial_connect_cancellation_interrupts_connection_rate_limit_wait() {
6103 let key = Ustr::from("connection-attempt");
6104
6105 let limiter = Arc::new(RateLimiter::new_with_quota(
6106 None,
6107 vec![(key, Quota::with_period(Duration::from_secs(1)).unwrap())],
6108 ));
6109 limiter.await_keys_ready(Some(&[key])).await;
6110
6111 let rate_limit = ConnectionRateLimit {
6112 limiter,
6113 keys: Arc::from([key]),
6114 };
6115
6116 let token = CancellationToken::new();
6117 let (handler, _rx) = channel_message_handler();
6118
6119 let connect = WebSocketClientInner::connect_url_with_handler(
6120 reconnect_test_config(1),
6121 Some(IncomingHandler::Message(handler)),
6122 None,
6123 None,
6124 None,
6125 Some(rate_limit),
6126 InitialConnectOptions {
6127 retry_policy: None,
6128 cancellation_token: Some(token.clone()),
6129 },
6130 );
6131
6132 tokio::pin!(connect);
6133 assert!(futures_util::poll!(&mut connect).is_pending());
6134
6135 token.cancel();
6136 let error = connect.await.expect_err("cancelled connect should fail");
6137
6138 assert!(
6139 matches!(error, TransportError::Io(ref error) if error.kind() == std::io::ErrorKind::Interrupted)
6140 );
6141 }
6142
6143 #[rstest]
6144 #[tokio::test]
6145 async fn test_reconnect_outcome_is_aborted_before_connect() {
6146 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
6147 let port = listener.local_addr().unwrap().port();
6148 let server = tokio::spawn(async move {
6149 let (stream, _) = listener.accept().await.unwrap();
6150 let _websocket = accept_async(stream).await.unwrap();
6151 std::future::pending::<()>().await;
6152 });
6153
6154 let (handler, _rx) = channel_message_handler();
6155 let mut inner =
6156 WebSocketClientInner::connect_url(reconnect_test_config(port), Some(handler), None)
6157 .await
6158 .unwrap();
6159 inner
6160 .connection_mode
6161 .store(ConnectionMode::Disconnect.as_u8(), Ordering::SeqCst);
6162
6163 let outcome = inner.reconnect_with_outcome().await.unwrap();
6164
6165 assert_eq!(outcome, ReconnectOutcome::Aborted);
6166 server.abort();
6167 }
6168
6169 #[rstest]
6170 #[tokio::test]
6171 async fn test_stream_reconnect_outcome_is_aborted_and_notifies_closed() {
6172 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
6173 let port = listener.local_addr().unwrap().port();
6174 let server = tokio::spawn(async move {
6175 let (stream, _) = listener.accept().await.unwrap();
6176 let _websocket = accept_async(stream).await.unwrap();
6177 std::future::pending::<()>().await;
6178 });
6179
6180 let mut inner = WebSocketClientInner::connect_url(reconnect_test_config(port), None, None)
6181 .await
6182 .unwrap();
6183 let state_notify = Arc::clone(&inner.state_notify);
6184 let mut notified = std::pin::pin!(state_notify.notified());
6185 notified.as_mut().enable();
6186
6187 let outcome = inner.reconnect_with_outcome().await.unwrap();
6188
6189 assert_eq!(outcome, ReconnectOutcome::Aborted);
6190 assert_eq!(
6191 ConnectionMode::from_atomic(&inner.connection_mode),
6192 ConnectionMode::Closed
6193 );
6194 tokio::time::timeout(TEST_TIMEOUT, notified)
6195 .await
6196 .expect("stream close notification was not published");
6197 server.abort();
6198 }
6199
6200 #[rstest]
6201 #[tokio::test]
6202 async fn test_reconnect_outcome_is_reconnected_with_handler() {
6203 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
6204 let port = listener.local_addr().unwrap().port();
6205 let server = tokio::spawn(async move {
6206 let (stream, _) = listener.accept().await.unwrap();
6207 let _first = accept_async(stream).await.unwrap();
6208 let (stream, _) = listener.accept().await.unwrap();
6209 let mut second = accept_async(stream).await.unwrap();
6210 second
6211 .send(WsMessage::Text("replacement".into()))
6212 .await
6213 .unwrap();
6214 std::future::pending::<()>().await;
6215 });
6216
6217 let (epoch_handler, mut epoch_rx) = channel_epoch_message_handler();
6218 let mut inner = WebSocketClientInner::connect_url_with_handler(
6219 reconnect_test_config(port),
6220 Some(IncomingHandler::Epoch(epoch_handler)),
6221 None,
6222 None,
6223 None,
6224 None,
6225 InitialConnectOptions::default(),
6226 )
6227 .await
6228 .unwrap();
6229 inner
6230 .connection_mode
6231 .store(ConnectionMode::Reconnect.as_u8(), Ordering::SeqCst);
6232
6233 let outcome = inner.reconnect_with_outcome().await.unwrap();
6234
6235 assert_eq!(outcome, ReconnectOutcome::Reconnected);
6236 assert_eq!(
6237 ConnectionMode::from_atomic(&inner.connection_mode),
6238 ConnectionMode::Active
6239 );
6240 let (epoch, message) = tokio::time::timeout(TEST_TIMEOUT, epoch_rx.recv())
6241 .await
6242 .expect("replacement epoch message was not delivered")
6243 .expect("epoch handler channel closed");
6244 assert_eq!(epoch, 1);
6245 assert_eq!(message, WsMessage::Text("replacement".into()));
6246 server.abort();
6247 }
6248
6249 #[rstest]
6250 #[tokio::test]
6251 async fn test_inner_drop_invalidates_read_fence_and_aborts_tasks() {
6252 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
6253 let port = listener.local_addr().unwrap().port();
6254 let server = tokio::spawn(async move {
6255 let (stream, _) = listener.accept().await.unwrap();
6256 let _websocket = accept_async(stream).await.unwrap();
6257 std::future::pending::<()>().await;
6258 });
6259
6260 let mut config = reconnect_test_config(port);
6261 config.heartbeat_interval_secs = Some(60);
6262 let (handler, _handler_rx) = channel_message_handler();
6263 let inner = WebSocketClientInner::connect_url(config, Some(handler), None)
6264 .await
6265 .unwrap();
6266 let read_fence = inner
6267 .read_fence
6268 .clone()
6269 .expect("read fence should exist in handler mode");
6270 let read_abort = inner
6271 .read_task
6272 .as_ref()
6273 .expect("read task should be spawned in handler mode")
6274 .abort_handle();
6275 let write_abort = inner.write_task.abort_handle();
6276 let heartbeat_abort = inner
6277 .heartbeat_task
6278 .as_ref()
6279 .expect("heartbeat task should be spawned for a configured heartbeat")
6280 .abort_handle();
6281
6282 assert!(read_fence.is_valid(), "read fence should start valid");
6283 assert!(
6284 !read_abort.is_finished(),
6285 "read task should be running before drop"
6286 );
6287 assert!(
6288 !write_abort.is_finished(),
6289 "write task should be running before drop"
6290 );
6291 assert!(
6292 !heartbeat_abort.is_finished(),
6293 "heartbeat task should be running before drop"
6294 );
6295
6296 drop(inner);
6297 wait_until_async(
6298 || async {
6299 read_abort.is_finished()
6300 && write_abort.is_finished()
6301 && heartbeat_abort.is_finished()
6302 },
6303 TEST_TIMEOUT,
6304 )
6305 .await;
6306
6307 assert!(!read_fence.is_valid(), "read fence was not invalidated");
6308 assert!(read_abort.is_finished(), "read task was not aborted");
6309 assert!(write_abort.is_finished(), "write task was not aborted");
6310 assert!(
6311 heartbeat_abort.is_finished(),
6312 "heartbeat task was not aborted"
6313 );
6314 server.abort();
6315 }
6316
6317 struct RecordingServer {
6318 task: JoinHandle<()>,
6319 port: u16,
6320 messages: Arc<tokio::sync::Mutex<Vec<String>>>,
6321 control_frames: Arc<tokio::sync::Mutex<Vec<String>>>,
6322 connections: Arc<AtomicUsize>,
6323 }
6324
6325 #[cfg(feature = "transport-sockudo")]
6326 async fn read_http_request<S>(stream: &mut S) -> Vec<u8>
6327 where
6328 S: AsyncRead + Unpin,
6329 {
6330 let mut buf = Vec::new();
6331 let mut chunk = [0u8; 256];
6332 loop {
6333 let n = stream.read(&mut chunk).await.unwrap();
6334 assert!(n > 0, "HTTP request closed before headers completed");
6335 buf.extend_from_slice(&chunk[..n]);
6336 if buf.windows(4).any(|window| window == b"\r\n\r\n") {
6337 return buf;
6338 }
6339 }
6340 }
6341
6342 #[cfg(feature = "transport-sockudo")]
6343 fn extract_header<'a>(request: &'a str, name: &str) -> Option<&'a str> {
6344 request.lines().find_map(|line| {
6345 let (header_name, header_value) = line.split_once(':')?;
6346 if header_name.eq_ignore_ascii_case(name) {
6347 Some(header_value.trim())
6348 } else {
6349 None
6350 }
6351 })
6352 }
6353
6354 #[cfg(feature = "transport-sockudo")]
6355 #[derive(Debug, Clone)]
6356 struct HeaderAssertCallback {
6357 key: String,
6358 value: HeaderValue,
6359 }
6360
6361 #[cfg(feature = "transport-sockudo")]
6362 impl Callback for HeaderAssertCallback {
6363 #[expect(
6364 clippy::panic_in_result_fn,
6365 reason = "assertion failures should fail the test"
6366 )]
6367 fn on_request(
6368 self,
6369 request: &server::Request,
6370 response: server::Response,
6371 ) -> Result<server::Response, server::ErrorResponse> {
6372 assert_eq!(request.headers().get(&self.key), Some(&self.value));
6373 Ok(response)
6374 }
6375 }
6376
6377 impl RecordingServer {
6378 async fn setup() -> Self {
6379 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
6380 let port = listener.local_addr().unwrap().port();
6381 let messages = Arc::new(tokio::sync::Mutex::new(Vec::new()));
6382 let messages_clone = Arc::clone(&messages);
6383 let control_frames = Arc::new(tokio::sync::Mutex::new(Vec::new()));
6384 let control_frames_clone = Arc::clone(&control_frames);
6385 let connections = Arc::new(AtomicUsize::new(0));
6386 let connections_clone = Arc::clone(&connections);
6387
6388 let task = task::spawn(async move {
6389 loop {
6390 let (stream, _) = listener.accept().await.unwrap();
6391 let mut websocket = accept_async(stream).await.unwrap();
6392 connections_clone.fetch_add(1, Ordering::SeqCst);
6393 let messages = Arc::clone(&messages_clone);
6394 let control_frames = Arc::clone(&control_frames_clone);
6395
6396 task::spawn(async move {
6397 while let Some(Ok(msg)) = websocket.next().await {
6398 match msg {
6399 WsMessage::Text(text) => {
6400 messages.lock().await.push(text.to_string());
6401 }
6402 WsMessage::Pong(data) => {
6403 control_frames
6404 .lock()
6405 .await
6406 .push(format!("pong:{}", String::from_utf8_lossy(&data)));
6407 }
6408 WsMessage::Close(_) => {
6409 let _ = websocket.close(None).await;
6410 break;
6411 }
6412 _ => {}
6413 }
6414 }
6415 });
6416 }
6417 });
6418
6419 Self {
6420 task,
6421 port,
6422 messages,
6423 control_frames,
6424 connections,
6425 }
6426 }
6427
6428 async fn messages(&self) -> Vec<String> {
6429 self.messages.lock().await.clone()
6430 }
6431
6432 async fn control_frames(&self) -> Vec<String> {
6433 self.control_frames.lock().await.clone()
6434 }
6435
6436 async fn wait_for_connections(&self, expected: usize) {
6437 wait_until_async(
6438 || async { self.connections.load(Ordering::SeqCst) == expected },
6439 TEST_TIMEOUT,
6440 )
6441 .await;
6442 }
6443 }
6444
6445 impl Drop for RecordingServer {
6446 fn drop(&mut self) {
6447 self.task.abort();
6448 }
6449 }
6450
6451 #[rstest]
6452 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
6453 async fn test_manual_reconnect_waits_for_slow_loss_callback_and_new_auth() {
6454 let server = RecordingServer::setup().await;
6455 let tracker = AuthTracker::new();
6456 let _initial_auth = tracker.begin();
6457 tracker.succeed();
6458 let states = Arc::new(Mutex::new(Vec::new()));
6459 let states_callback = Arc::clone(&states);
6460 let auth_at_loss = Arc::new(Mutex::new(Vec::new()));
6461 let auth_at_loss_callback = Arc::clone(&auth_at_loss);
6462 let tracker_callback = tracker.clone();
6463 let callback_release = Arc::new((Mutex::new(false), Condvar::new()));
6464 let callback_release_guard = CondvarReleaseGuard::new(callback_release.as_ref());
6465 let callback_release_clone = Arc::clone(&callback_release);
6466 let (callback_entered_tx, callback_entered_rx) = std::sync::mpsc::channel();
6467
6468 let sink = SocketStateSink::new(move |state| {
6469 states_callback.lock().push(state);
6470 if state == SocketState::Disconnected {
6471 auth_at_loss_callback
6472 .lock()
6473 .push(tracker_callback.auth_state());
6474 callback_entered_tx.send(()).unwrap();
6475 let (lock, condvar) = callback_release_clone.as_ref();
6476 let mut released = lock.lock();
6477 while !*released {
6478 condvar.wait(&mut released);
6479 }
6480 }
6481 });
6482
6483 let (handler, mut handler_rx) = channel_message_handler();
6484 let client = WebSocketClient::builder()
6485 .config(reconnect_test_config(server.port))
6486 .message_handler(handler)
6487 .state_sink(sink)
6488 .connect()
6489 .await
6490 .unwrap();
6491 client.set_auth_tracker(tracker.clone(), true);
6492 server.wait_for_connections(1).await;
6493
6494 let handle = client.reconnect_handle();
6495 let (request_tx, request_rx) = std::sync::mpsc::channel();
6496 let request_thread = std::thread::spawn(move || {
6497 request_tx.send(handle.request_reconnect()).unwrap();
6498 });
6499
6500 recv_rendezvous(callback_entered_rx, "slow reconnect callback entry").await;
6501 client
6502 .writer_tx
6503 .send(WriterCommand::Send(Message::text("buffered")))
6504 .unwrap();
6505 server.wait_for_connections(2).await;
6506 tokio::time::sleep(Duration::from_millis(250)).await;
6507
6508 assert_eq!(client.connection_mode(), ConnectionMode::Reconnect);
6509 assert!(!client.reconnect_published.load(Ordering::SeqCst));
6510 assert_eq!(tracker.auth_state(), AuthState::Unauthenticated);
6511 assert_eq!(*auth_at_loss.lock(), vec![AuthState::Unauthenticated]);
6512 assert_eq!(
6513 *states.lock(),
6514 vec![SocketState::Connected, SocketState::Disconnected]
6515 );
6516 assert!(server.messages().await.is_empty());
6517
6518 callback_release_guard.release();
6519 assert_eq!(
6520 recv_rendezvous(request_rx, "manual reconnect result").await,
6521 ReconnectRequestOutcome::Accepted
6522 );
6523 request_thread.join().unwrap();
6524 wait_until_async(|| async { client.is_active() }, TEST_TIMEOUT).await;
6525 wait_until_async(
6526 || {
6527 let states = Arc::clone(&states);
6528 async move { states.lock().len() == 3 }
6529 },
6530 TEST_TIMEOUT,
6531 )
6532 .await;
6533
6534 let notification = tokio::time::timeout(TEST_TIMEOUT, handler_rx.recv())
6535 .await
6536 .expect("reconnect notification was not delivered")
6537 .expect("handler channel closed");
6538 assert_eq!(notification, WsMessage::Text(RECONNECTED.into()));
6539 assert!(handler_rx.try_recv().is_err());
6540 tokio::time::sleep(Duration::from_millis(200)).await;
6541 assert!(server.messages().await.is_empty());
6542
6543 let _replacement_auth = tracker.begin();
6544 tracker.succeed();
6545 wait_until_async(
6546 || {
6547 let messages = Arc::clone(&server.messages);
6548 async move { messages.lock().await.len() == 1 }
6549 },
6550 TEST_TIMEOUT,
6551 )
6552 .await;
6553
6554 client.send_text("live".into(), None).await.unwrap();
6555 wait_until_async(
6556 || {
6557 let messages = Arc::clone(&server.messages);
6558 async move { messages.lock().await.len() == 2 }
6559 },
6560 TEST_TIMEOUT,
6561 )
6562 .await;
6563
6564 assert_eq!(server.messages().await, vec!["buffered", "live"]);
6565 assert_eq!(server.connections.load(Ordering::SeqCst), 2);
6566 assert_eq!(
6567 *states.lock(),
6568 vec![
6569 SocketState::Connected,
6570 SocketState::Disconnected,
6571 SocketState::Connected,
6572 ]
6573 );
6574
6575 client.disconnect().await;
6576 }
6577
6578 #[rstest]
6579 #[tokio::test]
6580 async fn test_reconnect_handle_is_closed_after_client_drop() {
6581 let server = RecordingServer::setup().await;
6582 let (handler, _handler_rx) = channel_message_handler();
6583 let client = WebSocketClient::builder()
6584 .config(reconnect_test_config(server.port))
6585 .message_handler(handler)
6586 .connect()
6587 .await
6588 .unwrap();
6589 let tracker = AuthTracker::new();
6590 let pending_auth = tracker.begin();
6591 client.set_auth_tracker(tracker.clone(), true);
6592 let controller_abort = client.controller_task.abort_handle();
6593 let handle = client.reconnect_handle();
6594
6595 drop(client);
6596 wait_until_async(|| async { controller_abort.is_finished() }, TEST_TIMEOUT).await;
6597
6598 assert_eq!(handle.request_reconnect(), ReconnectRequestOutcome::Closed);
6599 assert_eq!(tracker.auth_state(), AuthState::Failed);
6600 assert_eq!(
6601 tokio::time::timeout(TEST_TIMEOUT, pending_auth)
6602 .await
6603 .expect("client drop should resolve pending authentication")
6604 .expect("authentication sender should report its terminal result"),
6605 Err("WebSocket client closed".to_string())
6606 );
6607 }
6608
6609 #[rstest]
6610 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
6611 async fn test_concurrent_drop_closes_accepted_reconnect() {
6612 let server = RecordingServer::setup().await;
6613 let tracker = AuthTracker::new();
6614 let _initial_auth = tracker.begin();
6615 tracker.succeed();
6616 let callback_release = Arc::new((Mutex::new(false), Condvar::new()));
6617 let callback_release_guard = CondvarReleaseGuard::new(callback_release.as_ref());
6618 let callback_release_clone = Arc::clone(&callback_release);
6619 let (callback_entered_tx, callback_entered_rx) = std::sync::mpsc::channel();
6620 let states = Arc::new(Mutex::new(Vec::new()));
6621 let states_callback = Arc::clone(&states);
6622 let sink = SocketStateSink::new(move |state| {
6623 states_callback.lock().push(state);
6624 if state == SocketState::Disconnected {
6625 callback_entered_tx.send(()).unwrap();
6626 let (lock, condvar) = callback_release_clone.as_ref();
6627 let mut released = lock.lock();
6628 while !*released {
6629 condvar.wait(&mut released);
6630 }
6631 }
6632 });
6633
6634 let (handler, _handler_rx) = channel_message_handler();
6635 let client = WebSocketClient::builder()
6636 .config(reconnect_test_config(server.port))
6637 .message_handler(handler)
6638 .state_sink(sink)
6639 .connect()
6640 .await
6641 .unwrap();
6642 client.set_auth_tracker(tracker.clone(), true);
6643 let controller_abort = client.controller_task.abort_handle();
6644 let connection_mode = Arc::clone(&client.connection_mode);
6645 let handle = client.reconnect_handle();
6646 let surviving_handle = handle.clone();
6647 let (request_tx, request_rx) = std::sync::mpsc::channel();
6648 let request_thread = std::thread::spawn(move || {
6649 request_tx.send(handle.request_reconnect()).unwrap();
6650 });
6651
6652 recv_rendezvous(callback_entered_rx, "concurrent drop callback entry").await;
6653 drop(client);
6654 wait_until_async(|| async { controller_abort.is_finished() }, TEST_TIMEOUT).await;
6655 assert_eq!(
6656 ConnectionMode::from_atomic(&connection_mode),
6657 ConnectionMode::Closed
6658 );
6659 assert_eq!(tracker.auth_state(), AuthState::Failed);
6660
6661 callback_release_guard.release();
6662 assert_eq!(
6663 recv_rendezvous(request_rx, "concurrent drop reconnect result").await,
6664 ReconnectRequestOutcome::Accepted
6665 );
6666 request_thread.join().unwrap();
6667 assert_eq!(
6668 surviving_handle.request_reconnect(),
6669 ReconnectRequestOutcome::Closed
6670 );
6671 assert_eq!(
6672 *states.lock(),
6673 vec![SocketState::Connected, SocketState::Disconnected]
6674 );
6675 }
6676
6677 #[rstest]
6678 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
6679 async fn test_reconnect_callback_can_drop_client() {
6680 let server = RecordingServer::setup().await;
6681 let client_slot = Arc::new(Mutex::new(None::<WebSocketClient>));
6682 let client_slot_callback = Arc::clone(&client_slot);
6683 let states = Arc::new(Mutex::new(Vec::new()));
6684 let states_callback = Arc::clone(&states);
6685 let sink = SocketStateSink::new(move |state| {
6686 states_callback.lock().push(state);
6687 if state == SocketState::Disconnected {
6688 drop(client_slot_callback.lock().take());
6689 }
6690 });
6691
6692 let (handler, _handler_rx) = channel_message_handler();
6693 let client = WebSocketClient::builder()
6694 .config(reconnect_test_config(server.port))
6695 .message_handler(handler)
6696 .state_sink(sink)
6697 .connect()
6698 .await
6699 .unwrap();
6700 let tracker = AuthTracker::new();
6701 let _initial_auth = tracker.begin();
6702 tracker.succeed();
6703 client.set_auth_tracker(tracker.clone(), true);
6704 let controller_abort = client.controller_task.abort_handle();
6705 let connection_mode = Arc::clone(&client.connection_mode);
6706 let handle = client.reconnect_handle();
6707 let surviving_handle = handle.clone();
6708 *client_slot.lock() = Some(client);
6709 let (result_tx, result_rx) = std::sync::mpsc::channel();
6710 std::thread::spawn(move || {
6711 result_tx.send(handle.request_reconnect()).unwrap();
6712 });
6713
6714 assert_eq!(
6715 recv_rendezvous(result_rx, "callback drop reconnect result").await,
6716 ReconnectRequestOutcome::Accepted
6717 );
6718 wait_until_async(|| async { controller_abort.is_finished() }, TEST_TIMEOUT).await;
6719
6720 assert!(client_slot.lock().is_none());
6721 assert_eq!(
6722 ConnectionMode::from_atomic(&connection_mode),
6723 ConnectionMode::Closed
6724 );
6725 assert_eq!(tracker.auth_state(), AuthState::Failed);
6726 assert_eq!(
6727 surviving_handle.request_reconnect(),
6728 ReconnectRequestOutcome::Closed
6729 );
6730 assert_eq!(
6731 *states.lock(),
6732 vec![SocketState::Connected, SocketState::Disconnected]
6733 );
6734 }
6735
6736 #[rstest]
6737 #[tokio::test]
6738 async fn test_reconnect_then_disconnect() {
6739 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
6741 let port = listener.local_addr().unwrap().port();
6742
6743 let server = task::spawn(async move {
6745 let (stream, _) = listener.accept().await.unwrap();
6746 let ws = accept_async(stream).await.unwrap();
6747 drop(ws);
6748 sleep(Duration::from_secs(1)).await;
6750 });
6751
6752 let (handler, _rx) = channel_message_handler();
6754
6755 let config = WebSocketConfig {
6757 url: format!("ws://127.0.0.1:{port}"),
6758 headers: vec![],
6759 heartbeat_interval_secs: None,
6760 heartbeat_payload: None,
6761 connect_timeout_ms: Some(1_000),
6762 reconnect_delay_initial_ms: Some(50),
6763 reconnect_delay_max_ms: Some(100),
6764 reconnect_backoff_factor: Some(1.0),
6765 reconnect_jitter_ms: Some(0),
6766 reconnect_max_attempts: None,
6767 heartbeat_timeout_secs: None,
6768 idle_timeout_ms: None,
6769 writer_capacity: None,
6770 backend: TransportBackend::Tungstenite,
6771 proxy_url: None,
6772 max_message_size_bytes: None,
6773 max_frame_size_bytes: None,
6774 };
6775
6776 let client = WebSocketClient::builder()
6778 .config(config)
6779 .message_handler(handler)
6780 .connect()
6781 .await
6782 .unwrap();
6783
6784 sleep(Duration::from_millis(100)).await;
6786 client.disconnect().await;
6788 assert!(client.is_disconnected());
6789 server.abort();
6790 }
6791
6792 #[rstest]
6793 #[tokio::test]
6794 async fn test_reconnect_state_flips_when_reader_stops() {
6795 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
6797 let port = listener.local_addr().unwrap().port();
6798 let server = task::spawn(async move {
6799 if let Ok((stream, _)) = listener.accept().await
6800 && let Ok(ws) = accept_async(stream).await
6801 {
6802 drop(ws);
6803 }
6804
6805 sleep(Duration::from_millis(50)).await;
6806 });
6807
6808 let (handler, _rx) = channel_message_handler();
6809
6810 let config = WebSocketConfig {
6811 url: format!("ws://127.0.0.1:{port}"),
6812 headers: vec![],
6813 heartbeat_interval_secs: None,
6814 heartbeat_payload: None,
6815 connect_timeout_ms: Some(1_000),
6816 reconnect_delay_initial_ms: Some(50),
6817 reconnect_delay_max_ms: Some(100),
6818 reconnect_backoff_factor: Some(1.0),
6819 reconnect_jitter_ms: Some(0),
6820 reconnect_max_attempts: None,
6821 heartbeat_timeout_secs: None,
6822 idle_timeout_ms: None,
6823 writer_capacity: None,
6824 backend: TransportBackend::Tungstenite,
6825 proxy_url: None,
6826 max_message_size_bytes: None,
6827 max_frame_size_bytes: None,
6828 };
6829
6830 let client = WebSocketClient::builder()
6831 .config(config)
6832 .message_handler(handler)
6833 .connect()
6834 .await
6835 .unwrap();
6836
6837 tokio::time::timeout(Duration::from_secs(2), async {
6838 loop {
6839 if client.is_reconnecting() {
6840 break;
6841 }
6842
6843 tokio::time::sleep(Duration::from_millis(10)).await;
6844 }
6845 })
6846 .await
6847 .expect("client did not enter RECONNECT state");
6848
6849 client.disconnect().await;
6850 server.abort();
6851 }
6852
6853 #[rstest]
6854 #[tokio::test]
6855 async fn test_stream_mode_disables_auto_reconnect() {
6856 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
6859 let port = listener.local_addr().unwrap().port();
6860 let server = task::spawn(async move {
6861 if let Ok((stream, _)) = listener.accept().await
6862 && let Ok(_ws) = accept_async(stream).await
6863 {
6864 sleep(Duration::from_millis(100)).await;
6866 }
6867 });
6868
6869 let config = WebSocketConfig {
6870 url: format!("ws://127.0.0.1:{port}"),
6871 headers: vec![],
6872 heartbeat_interval_secs: None,
6873 heartbeat_payload: None,
6874 connect_timeout_ms: Some(1_000),
6875 reconnect_delay_initial_ms: Some(50),
6876 reconnect_delay_max_ms: Some(100),
6877 reconnect_backoff_factor: Some(1.0),
6878 reconnect_jitter_ms: Some(0),
6879 reconnect_max_attempts: None,
6880 heartbeat_timeout_secs: None,
6881 idle_timeout_ms: None,
6882 writer_capacity: None,
6883 backend: TransportBackend::Tungstenite,
6884 proxy_url: None,
6885 max_message_size_bytes: None,
6886 max_frame_size_bytes: None,
6887 };
6888
6889 let (_reader, _client) = WebSocketClient::stream_builder()
6890 .config(config)
6891 .connect()
6892 .await
6893 .unwrap();
6894
6895 server.abort();
6903 }
6904
6905 #[rstest]
6906 #[tokio::test]
6907 async fn test_message_handler_mode_allows_auto_reconnect() {
6908 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
6910 let port = listener.local_addr().unwrap().port();
6911
6912 let server = task::spawn(async move {
6913 if let Ok((stream, _)) = listener.accept().await
6915 && let Ok(ws) = accept_async(stream).await
6916 {
6917 drop(ws);
6918 }
6919
6920 sleep(Duration::from_millis(50)).await;
6921 });
6922
6923 let (handler, _rx) = channel_message_handler();
6924
6925 let config = WebSocketConfig {
6926 url: format!("ws://127.0.0.1:{port}"),
6927 headers: vec![],
6928 heartbeat_interval_secs: None,
6929 heartbeat_payload: None,
6930 connect_timeout_ms: Some(1_000),
6931 reconnect_delay_initial_ms: Some(50),
6932 reconnect_delay_max_ms: Some(100),
6933 reconnect_backoff_factor: Some(1.0),
6934 reconnect_jitter_ms: Some(0),
6935 reconnect_max_attempts: None,
6936 heartbeat_timeout_secs: None,
6937 idle_timeout_ms: None,
6938 writer_capacity: None,
6939 backend: TransportBackend::Tungstenite,
6940 proxy_url: None,
6941 max_message_size_bytes: None,
6942 max_frame_size_bytes: None,
6943 };
6944
6945 let client = WebSocketClient::builder()
6946 .config(config)
6947 .message_handler(handler)
6948 .connect()
6949 .await
6950 .unwrap();
6951
6952 tokio::time::timeout(Duration::from_secs(2), async {
6954 loop {
6955 if client.is_reconnecting() || client.is_closed() {
6956 break;
6957 }
6958
6959 tokio::time::sleep(Duration::from_millis(10)).await;
6960 }
6961 })
6962 .await
6963 .expect("client should attempt reconnection or close");
6964
6965 assert!(
6968 client.is_reconnecting() || client.is_closed(),
6969 "Client with message handler should attempt reconnection"
6970 );
6971
6972 client.disconnect().await;
6973 server.abort();
6974 }
6975
6976 #[rstest]
6977 #[tokio::test]
6978 async fn test_handler_mode_reconnect_with_new_connection() {
6979 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
6981 let port = listener.local_addr().unwrap().port();
6982
6983 let server = task::spawn(async move {
6984 if let Ok((stream, _)) = listener.accept().await
6986 && let Ok(ws) = accept_async(stream).await
6987 {
6988 drop(ws);
6989 }
6990
6991 sleep(Duration::from_millis(100)).await;
6993
6994 if let Ok((stream, _)) = listener.accept().await
6996 && let Ok(mut ws) = accept_async(stream).await
6997 {
6998 use futures_util::SinkExt;
6999 let _ = ws
7000 .send(WsMessage::Text("reconnected".to_string().into()))
7001 .await;
7002 sleep(Duration::from_secs(1)).await;
7003 }
7004 });
7005
7006 let (handler, mut rx) = channel_message_handler();
7007
7008 let config = WebSocketConfig {
7009 url: format!("ws://127.0.0.1:{port}"),
7010 headers: vec![],
7011 heartbeat_interval_secs: None,
7012 heartbeat_payload: None,
7013 connect_timeout_ms: Some(2_000),
7014 reconnect_delay_initial_ms: Some(50),
7015 reconnect_delay_max_ms: Some(200),
7016 reconnect_backoff_factor: Some(1.5),
7017 reconnect_jitter_ms: Some(10),
7018 reconnect_max_attempts: None,
7019 heartbeat_timeout_secs: None,
7020 idle_timeout_ms: None,
7021 writer_capacity: None,
7022 backend: TransportBackend::Tungstenite,
7023 proxy_url: None,
7024 max_message_size_bytes: None,
7025 max_frame_size_bytes: None,
7026 };
7027
7028 let client = WebSocketClient::builder()
7029 .config(config)
7030 .message_handler(handler)
7031 .connect()
7032 .await
7033 .unwrap();
7034
7035 let result = tokio::time::timeout(Duration::from_secs(5), async {
7037 loop {
7038 if let Ok(msg) = rx.try_recv()
7039 && matches!(msg, WsMessage::Text(ref text) if AsRef::<str>::as_ref(text) == "reconnected")
7040 {
7041 return true;
7042 }
7043
7044 tokio::time::sleep(Duration::from_millis(10)).await;
7045 }
7046 })
7047 .await;
7048
7049 assert!(
7050 result.is_ok(),
7051 "Should receive message after reconnection within timeout"
7052 );
7053
7054 client.disconnect().await;
7055 server.abort();
7056 }
7057
7058 #[rstest]
7059 #[tokio::test]
7060 async fn test_stream_mode_no_auto_reconnect() {
7061 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
7064 let port = listener.local_addr().unwrap().port();
7065
7066 let server = task::spawn(async move {
7067 if let Ok((stream, _)) = listener.accept().await
7069 && let Ok(mut ws) = accept_async(stream).await
7070 {
7071 use futures_util::SinkExt;
7072 let _ = ws.send(WsMessage::Text("hello".to_string().into())).await;
7073 sleep(Duration::from_millis(50)).await;
7074 }
7076 });
7077
7078 let config = WebSocketConfig {
7079 url: format!("ws://127.0.0.1:{port}"),
7080 headers: vec![],
7081 heartbeat_interval_secs: None,
7082 heartbeat_payload: None,
7083 connect_timeout_ms: Some(1_000),
7084 reconnect_delay_initial_ms: Some(50),
7085 reconnect_delay_max_ms: Some(100),
7086 reconnect_backoff_factor: Some(1.0),
7087 reconnect_jitter_ms: Some(0),
7088 reconnect_max_attempts: None,
7089 heartbeat_timeout_secs: None,
7090 idle_timeout_ms: None,
7091 writer_capacity: None,
7092 backend: TransportBackend::Tungstenite,
7093 proxy_url: None,
7094 max_message_size_bytes: None,
7095 max_frame_size_bytes: None,
7096 };
7097
7098 let (mut reader, client) = WebSocketClient::stream_builder()
7099 .config(config)
7100 .connect()
7101 .await
7102 .unwrap();
7103
7104 assert!(client.is_active(), "Client should start as active");
7106
7107 let msg = reader.next().await;
7109 assert!(
7110 matches!(&msg, Some(Ok(Message::Text(bytes))) if bytes.as_ref() == b"hello"),
7111 "Should receive initial message"
7112 );
7113
7114 while let Some(msg) = reader.next().await {
7116 if msg.is_err() || matches!(msg, Ok(Message::Close(_))) {
7117 break;
7118 }
7119 }
7120
7121 sleep(Duration::from_millis(200)).await;
7124 assert!(
7125 client.is_active(),
7126 "Stream mode client stays ACTIVE before notify_closed()"
7127 );
7128
7129 client.notify_closed();
7131
7132 assert!(
7133 client.is_closed(),
7134 "Stream mode client should be CLOSED after notify_closed()"
7135 );
7136 assert!(
7137 !client.is_reconnecting(),
7138 "Stream mode client should never attempt reconnection"
7139 );
7140
7141 client.disconnect().await;
7142 server.abort();
7143 }
7144
7145 #[rstest]
7146 #[tokio::test]
7147 async fn test_send_timeout_uses_configured_connect_timeout() {
7148 use nautilus_common::testing::wait_until_async;
7151
7152 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
7153 let port = listener.local_addr().unwrap().port();
7154
7155 let server = task::spawn(async move {
7156 if let Ok((stream, _)) = listener.accept().await
7158 && let Ok(ws) = accept_async(stream).await
7159 {
7160 drop(ws);
7161 }
7162
7163 sleep(Duration::from_mins(1)).await;
7165 });
7166
7167 let (handler, _rx) = channel_message_handler();
7168
7169 let config = WebSocketConfig {
7171 url: format!("ws://127.0.0.1:{port}"),
7172 headers: vec![],
7173 heartbeat_interval_secs: None,
7174 heartbeat_payload: None,
7175 connect_timeout_ms: Some(2_000), reconnect_delay_initial_ms: Some(50),
7177 reconnect_delay_max_ms: Some(100),
7178 reconnect_backoff_factor: Some(1.0),
7179 reconnect_jitter_ms: Some(0),
7180 reconnect_max_attempts: None,
7181 heartbeat_timeout_secs: None,
7182 idle_timeout_ms: None,
7183 writer_capacity: None,
7184 backend: TransportBackend::Tungstenite,
7185 proxy_url: None,
7186 max_message_size_bytes: None,
7187 max_frame_size_bytes: None,
7188 };
7189
7190 let client = WebSocketClient::builder()
7191 .config(config)
7192 .message_handler(handler)
7193 .connect()
7194 .await
7195 .unwrap();
7196
7197 wait_until_async(
7199 || async { client.is_reconnecting() },
7200 Duration::from_secs(3),
7201 )
7202 .await;
7203
7204 let start = std::time::Instant::now();
7206 let send_result = client.send_text("test".to_string(), None).await;
7207 let elapsed = start.elapsed();
7208
7209 assert!(
7210 send_result.is_err(),
7211 "Send should fail when client stuck in RECONNECT"
7212 );
7213 assert!(
7214 matches!(send_result, Err(crate::error::SendError::Timeout)),
7215 "Send should return Timeout error, was: {send_result:?}"
7216 );
7217 assert!(
7220 elapsed >= Duration::from_millis(1800),
7221 "Send should timeout after at least 2s (configured timeout), took {elapsed:?}"
7222 );
7223
7224 client.disconnect().await;
7225 server.abort();
7226 }
7227
7228 #[rstest]
7229 #[tokio::test]
7230 async fn test_send_waits_during_reconnection() {
7231 use nautilus_common::testing::wait_until_async;
7233
7234 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
7235 let port = listener.local_addr().unwrap().port();
7236
7237 let server = task::spawn(async move {
7238 if let Ok((stream, _)) = listener.accept().await
7240 && let Ok(ws) = accept_async(stream).await
7241 {
7242 drop(ws);
7243 }
7244
7245 sleep(Duration::from_millis(500)).await;
7247
7248 if let Ok((stream, _)) = listener.accept().await
7250 && let Ok(mut ws) = accept_async(stream).await
7251 {
7252 while let Some(Ok(msg)) = ws.next().await {
7254 if ws.send(msg).await.is_err() {
7255 break;
7256 }
7257 }
7258 }
7259 });
7260
7261 let (handler, _rx) = channel_message_handler();
7262
7263 let config = WebSocketConfig {
7264 url: format!("ws://127.0.0.1:{port}"),
7265 headers: vec![],
7266 heartbeat_interval_secs: None,
7267 heartbeat_payload: None,
7268 connect_timeout_ms: Some(5_000), reconnect_delay_initial_ms: Some(100),
7270 reconnect_delay_max_ms: Some(200),
7271 reconnect_backoff_factor: Some(1.0),
7272 reconnect_jitter_ms: Some(0),
7273 reconnect_max_attempts: None,
7274 heartbeat_timeout_secs: None,
7275 idle_timeout_ms: None,
7276 writer_capacity: None,
7277 backend: TransportBackend::Tungstenite,
7278 proxy_url: None,
7279 max_message_size_bytes: None,
7280 max_frame_size_bytes: None,
7281 };
7282
7283 let client = WebSocketClient::builder()
7284 .config(config)
7285 .message_handler(handler)
7286 .connect()
7287 .await
7288 .unwrap();
7289
7290 wait_until_async(
7292 || async { client.is_reconnecting() },
7293 Duration::from_secs(2),
7294 )
7295 .await;
7296
7297 let send_result = tokio::time::timeout(
7299 Duration::from_secs(3),
7300 client.send_text("test_message".to_string(), None),
7301 )
7302 .await;
7303
7304 assert!(
7305 send_result.is_ok() && send_result.unwrap().is_ok(),
7306 "Send should succeed after waiting for reconnection"
7307 );
7308
7309 client.disconnect().await;
7310 server.abort();
7311 }
7312
7313 #[rstest]
7314 #[tokio::test]
7315 async fn test_rate_limiter_before_active_wait() {
7316 use std::{num::NonZeroU32, sync::Arc};
7321
7322 use nautilus_common::testing::wait_until_async;
7323
7324 use crate::ratelimiter::quota::Quota;
7325
7326 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
7327 let port = listener.local_addr().unwrap().port();
7328
7329 let server = task::spawn(async move {
7330 if let Ok((stream, _)) = listener.accept().await
7332 && let Ok(mut ws) = accept_async(stream).await
7333 {
7334 if let Some(Ok(_)) = ws.next().await {
7336 drop(ws);
7337 }
7338 }
7339
7340 sleep(Duration::from_millis(500)).await;
7342
7343 if let Ok((stream, _)) = listener.accept().await
7345 && let Ok(mut ws) = accept_async(stream).await
7346 {
7347 while let Some(Ok(msg)) = ws.next().await {
7348 if ws.send(msg).await.is_err() {
7349 break;
7350 }
7351 }
7352 }
7353 });
7354
7355 let (handler, _rx) = channel_message_handler();
7356
7357 let config = WebSocketConfig {
7358 url: format!("ws://127.0.0.1:{port}"),
7359 headers: vec![],
7360 heartbeat_interval_secs: None,
7361 heartbeat_payload: None,
7362 connect_timeout_ms: Some(5_000),
7363 reconnect_delay_initial_ms: Some(50),
7364 reconnect_delay_max_ms: Some(100),
7365 reconnect_backoff_factor: Some(1.0),
7366 reconnect_jitter_ms: Some(0),
7367 reconnect_max_attempts: None,
7368 heartbeat_timeout_secs: None,
7369 idle_timeout_ms: None,
7370 writer_capacity: None,
7371 backend: TransportBackend::Tungstenite,
7372 proxy_url: None,
7373 max_message_size_bytes: None,
7374 max_frame_size_bytes: None,
7375 };
7376
7377 let quota = Quota::per_second(NonZeroU32::new(1).unwrap())
7379 .unwrap()
7380 .allow_burst(NonZeroU32::new(1).unwrap());
7381
7382 let client = Arc::new(
7383 WebSocketClient::builder()
7384 .config(config)
7385 .message_handler(handler)
7386 .keyed_quotas(vec![("test_key".to_string(), quota)])
7387 .connect()
7388 .await
7389 .unwrap(),
7390 );
7391
7392 let test_key: [Ustr; 1] = [Ustr::from("test_key")];
7394 client
7395 .send_text("msg1".to_string(), Some(test_key.as_slice()))
7396 .await
7397 .unwrap();
7398
7399 wait_until_async(
7401 || async { client.is_reconnecting() },
7402 Duration::from_secs(2),
7403 )
7404 .await;
7405
7406 let start = std::time::Instant::now();
7408 let send_result = client
7409 .send_text("msg2".to_string(), Some(test_key.as_slice()))
7410 .await;
7411 let elapsed = start.elapsed();
7412
7413 assert!(
7415 send_result.is_ok(),
7416 "Send should succeed after rate limit + reconnection, was: {send_result:?}"
7417 );
7418 assert!(
7422 elapsed >= Duration::from_millis(850),
7423 "Should wait for rate limit (~1s), waited {elapsed:?}"
7424 );
7425
7426 client.disconnect().await;
7427 server.abort();
7428 }
7429
7430 #[rstest]
7431 #[tokio::test]
7432 async fn test_disconnect_during_reconnect_exits_cleanly() {
7433 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
7436 let port = listener.local_addr().unwrap().port();
7437
7438 let server = task::spawn(async move {
7439 if let Ok((stream, _)) = listener.accept().await
7441 && let Ok(ws) = accept_async(stream).await
7442 {
7443 drop(ws);
7444 }
7445
7446 sleep(Duration::from_mins(1)).await;
7448 });
7449
7450 let (handler, _rx) = channel_message_handler();
7451
7452 let config = WebSocketConfig {
7453 url: format!("ws://127.0.0.1:{port}"),
7454 headers: vec![],
7455 heartbeat_interval_secs: None,
7456 heartbeat_payload: None,
7457 connect_timeout_ms: Some(2_000), reconnect_delay_initial_ms: Some(100),
7459 reconnect_delay_max_ms: Some(200),
7460 reconnect_backoff_factor: Some(1.0),
7461 reconnect_jitter_ms: Some(0),
7462 reconnect_max_attempts: None,
7463 heartbeat_timeout_secs: None,
7464 idle_timeout_ms: None,
7465 writer_capacity: None,
7466 backend: TransportBackend::Tungstenite,
7467 proxy_url: None,
7468 max_message_size_bytes: None,
7469 max_frame_size_bytes: None,
7470 };
7471
7472 let client = WebSocketClient::builder()
7473 .config(config)
7474 .message_handler(handler)
7475 .connect()
7476 .await
7477 .unwrap();
7478
7479 tokio::time::timeout(Duration::from_secs(2), async {
7481 while !client.is_reconnecting() {
7482 sleep(Duration::from_millis(10)).await;
7483 }
7484 })
7485 .await
7486 .expect("Client should enter RECONNECT state");
7487
7488 client.disconnect().await;
7490
7491 assert!(
7493 client.is_disconnected(),
7494 "Client should be cleanly disconnected"
7495 );
7496
7497 server.abort();
7498 }
7499
7500 #[rstest]
7501 #[tokio::test]
7502 async fn test_send_fails_fast_when_closed_before_rate_limit() {
7503 use std::{num::NonZeroU32, sync::Arc};
7506
7507 use nautilus_common::testing::wait_until_async;
7508
7509 use crate::ratelimiter::quota::Quota;
7510
7511 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
7512 let port = listener.local_addr().unwrap().port();
7513
7514 let server = task::spawn(async move {
7515 if let Ok((stream, _)) = listener.accept().await
7517 && let Ok(ws) = accept_async(stream).await
7518 {
7519 drop(ws);
7520 }
7521
7522 sleep(Duration::from_mins(1)).await;
7523 });
7524
7525 let (handler, _rx) = channel_message_handler();
7526
7527 let config = WebSocketConfig {
7528 url: format!("ws://127.0.0.1:{port}"),
7529 headers: vec![],
7530 heartbeat_interval_secs: None,
7531 heartbeat_payload: None,
7532 connect_timeout_ms: Some(5_000),
7533 reconnect_delay_initial_ms: Some(50),
7534 reconnect_delay_max_ms: Some(100),
7535 reconnect_backoff_factor: Some(1.0),
7536 reconnect_jitter_ms: Some(0),
7537 reconnect_max_attempts: None,
7538 heartbeat_timeout_secs: None,
7539 idle_timeout_ms: None,
7540 writer_capacity: None,
7541 backend: TransportBackend::Tungstenite,
7542 proxy_url: None,
7543 max_message_size_bytes: None,
7544 max_frame_size_bytes: None,
7545 };
7546
7547 let quota = Quota::with_period(Duration::from_secs(10))
7550 .unwrap()
7551 .allow_burst(NonZeroU32::new(1).unwrap());
7552
7553 let client = Arc::new(
7554 WebSocketClient::builder()
7555 .config(config)
7556 .message_handler(handler)
7557 .keyed_quotas(vec![("test_key".to_string(), quota)])
7558 .connect()
7559 .await
7560 .unwrap(),
7561 );
7562
7563 wait_until_async(
7565 || async { client.is_reconnecting() || client.is_closed() },
7566 Duration::from_secs(2),
7567 )
7568 .await;
7569
7570 client.disconnect().await;
7572 assert!(
7573 !client.is_active(),
7574 "Client should not be active after disconnect"
7575 );
7576
7577 let start = std::time::Instant::now();
7579 let test_key: [Ustr; 1] = [Ustr::from("test_key")];
7580 let result = client
7581 .send_text("test".to_string(), Some(test_key.as_slice()))
7582 .await;
7583 let elapsed = start.elapsed();
7584
7585 assert!(result.is_err(), "Send should fail when client is closed");
7587 assert!(
7588 matches!(result, Err(crate::error::SendError::Closed)),
7589 "Send should return Closed error, was: {result:?}"
7590 );
7591
7592 assert!(
7594 elapsed < Duration::from_millis(100),
7595 "Send should fail fast without rate limiting, took {elapsed:?}"
7596 );
7597
7598 server.abort();
7599 }
7600
7601 #[rstest]
7602 #[tokio::test]
7603 async fn test_builder_rejects_shared_rate_limiter_with_quotas() {
7604 let (handler, _rx) = channel_message_handler();
7605 let rate_limiter = Arc::new(RateLimiter::new_with_quota(None, Vec::new()));
7606 let result = WebSocketClient::builder()
7607 .config(reconnect_test_config(1))
7608 .message_handler(handler)
7609 .default_quota(Quota::with_period(Duration::from_secs(1)).unwrap())
7610 .rate_limiter(rate_limiter)
7611 .connect()
7612 .await;
7613
7614 assert_eq!(
7615 result.unwrap_err().to_string(),
7616 "I/O error: Cannot combine a shared rate limiter with quota configuration"
7617 );
7618 }
7619
7620 #[rstest]
7621 #[tokio::test]
7622 async fn test_builder_rejects_connection_keys_without_limiter() {
7623 let (handler, _rx) = channel_message_handler();
7624 let result = WebSocketClient::builder()
7625 .config(reconnect_test_config(1))
7626 .message_handler(handler)
7627 .connection_rate_keys(Arc::from([Ustr::from("connect")]))
7628 .connect()
7629 .await;
7630
7631 assert_eq!(
7632 result.unwrap_err().to_string(),
7633 "I/O error: Connection rate keys require a connection rate limiter"
7634 );
7635 }
7636
7637 #[rstest]
7638 #[tokio::test]
7639 async fn test_builder_rejects_connection_limiter_without_keys() {
7640 let (handler, _rx) = channel_message_handler();
7641 let rate_limiter = Arc::new(RateLimiter::new_with_quota(None, Vec::new()));
7642 let result = WebSocketClient::builder()
7643 .config(reconnect_test_config(1))
7644 .message_handler(handler)
7645 .connection_rate_limiter(rate_limiter)
7646 .connect()
7647 .await;
7648
7649 assert_eq!(
7650 result.unwrap_err().to_string(),
7651 "I/O error: Connection rate limiter requires at least one connection rate key"
7652 );
7653 }
7654
7655 #[rstest]
7656 #[tokio::test]
7657 async fn test_epoch_builder_rejects_both_ping_handler_types() {
7658 let (handler, _rx) = channel_epoch_message_handler();
7659 let ping_handler: PingHandler = Arc::new(|_| {});
7660 let epoch_ping_handler: EpochPingHandler = Arc::new(|_, _| {});
7661 let result = WebSocketClient::epoch_builder()
7662 .config(reconnect_test_config(1))
7663 .epoch_handler(handler)
7664 .ping_handler(ping_handler)
7665 .epoch_ping_handler(epoch_ping_handler)
7666 .connect()
7667 .await;
7668
7669 assert_eq!(
7670 result.unwrap_err().to_string(),
7671 "I/O error: Cannot configure both ping_handler and epoch_ping_handler"
7672 );
7673 }
7674
7675 #[rstest]
7676 #[tokio::test]
7677 async fn test_connect_url_rejects_invalid_reconnect_timing_before_connect() {
7678 let (handler, _rx) = channel_message_handler();
7679
7680 let config = WebSocketConfig {
7681 url: "ws://127.0.0.1:1".to_string(),
7682 headers: vec![],
7683 heartbeat_interval_secs: None,
7684 heartbeat_payload: None,
7685 connect_timeout_ms: Some(0),
7686 reconnect_delay_initial_ms: Some(100),
7687 reconnect_delay_max_ms: Some(500),
7688 reconnect_backoff_factor: Some(1.5),
7689 reconnect_jitter_ms: Some(0),
7690 reconnect_max_attempts: None,
7691 heartbeat_timeout_secs: None,
7692 idle_timeout_ms: None,
7693 writer_capacity: None,
7694 backend: TransportBackend::Tungstenite,
7695 proxy_url: None,
7696 max_message_size_bytes: None,
7697 max_frame_size_bytes: None,
7698 };
7699
7700 let err = WebSocketClientInner::connect_url(config, Some(handler), None)
7701 .await
7702 .expect_err("invalid reconnect timing should be rejected");
7703
7704 match err {
7705 TransportError::Io(error) => {
7706 assert_eq!(error.kind(), std::io::ErrorKind::InvalidInput);
7707 assert!(
7708 error.to_string().contains("connect_timeout_ms"),
7709 "error should mention zero reconnect timeout, was: {error}"
7710 );
7711 }
7712 other => panic!("expected InvalidInput IO error, was: {other:?}"),
7713 }
7714 }
7715
7716 #[rstest]
7717 #[tokio::test]
7718 async fn test_connect_url_rejects_invalid_reconnect_backoff_before_connect() {
7719 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
7720 let port = listener.local_addr().unwrap().port();
7721 let accepted = Arc::new(std::sync::atomic::AtomicBool::new(false));
7722 let accepted_clone = Arc::clone(&accepted);
7723
7724 let server = task::spawn(async move {
7725 let (stream, _) = listener.accept().await.unwrap();
7726 accepted_clone.store(true, Ordering::SeqCst);
7727 accept_async(stream).await.unwrap();
7728 });
7729
7730 let (handler, _rx) = channel_message_handler();
7731
7732 let config = WebSocketConfig {
7733 url: format!("ws://127.0.0.1:{port}"),
7734 headers: vec![],
7735 heartbeat_interval_secs: None,
7736 heartbeat_payload: None,
7737 connect_timeout_ms: Some(1_000),
7738 reconnect_delay_initial_ms: Some(50),
7739 reconnect_delay_max_ms: Some(100),
7740 reconnect_backoff_factor: Some(100.1),
7741 reconnect_jitter_ms: Some(0),
7742 reconnect_max_attempts: None,
7743 heartbeat_timeout_secs: None,
7744 idle_timeout_ms: None,
7745 writer_capacity: None,
7746 backend: TransportBackend::Tungstenite,
7747 proxy_url: None,
7748 max_message_size_bytes: None,
7749 max_frame_size_bytes: None,
7750 };
7751
7752 let error = WebSocketClientInner::connect_url(config, Some(handler), None)
7753 .await
7754 .expect_err("invalid reconnect backoff should be rejected");
7755
7756 match error {
7757 TransportError::Io(error) => {
7758 assert_eq!(error.kind(), std::io::ErrorKind::InvalidInput);
7759 assert!(
7760 error.to_string().contains("factor"),
7761 "error should mention the invalid factor, was: {error}"
7762 );
7763 }
7764 other => panic!("expected InvalidInput IO error, was: {other:?}"),
7765 }
7766
7767 assert!(
7768 !accepted.load(Ordering::SeqCst),
7769 "invalid reconnect backoff must be rejected before connecting"
7770 );
7771 server.abort();
7772 }
7773
7774 #[rstest]
7775 #[tokio::test]
7776 async fn test_client_without_handler_sets_stream_mode() {
7777 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
7781 let port = listener.local_addr().unwrap().port();
7782
7783 let server = task::spawn(async move {
7784 if let Ok((stream, _)) = listener.accept().await
7786 && let Ok(ws) = accept_async(stream).await
7787 {
7788 drop(ws); }
7790 });
7791
7792 let config = WebSocketConfig {
7793 url: format!("ws://127.0.0.1:{port}"),
7794 headers: vec![],
7795 heartbeat_interval_secs: None,
7796 heartbeat_payload: None,
7797 connect_timeout_ms: Some(1_000),
7798 reconnect_delay_initial_ms: Some(100),
7799 reconnect_delay_max_ms: Some(500),
7800 reconnect_backoff_factor: Some(1.5),
7801 reconnect_jitter_ms: Some(0),
7802 reconnect_max_attempts: None,
7803 heartbeat_timeout_secs: None,
7804 idle_timeout_ms: None,
7805 writer_capacity: None,
7806 backend: TransportBackend::Tungstenite,
7807 proxy_url: None,
7808 max_message_size_bytes: None,
7809 max_frame_size_bytes: None,
7810 };
7811
7812 let inner = WebSocketClientInner::connect_url(config, None, None)
7814 .await
7815 .unwrap();
7816
7817 assert!(
7819 inner.handler.is_none(),
7820 "Client without handler should not retain an internal handler"
7821 );
7822
7823 server.abort();
7827 }
7828
7829 #[rstest]
7830 #[tokio::test]
7831 async fn test_idle_timeout_triggers_reconnect() {
7832 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
7833 let port = listener.local_addr().unwrap().port();
7834
7835 let server = task::spawn(async move {
7837 let (stream, _) = listener.accept().await.unwrap();
7838 let _ws = accept_async(stream).await.unwrap();
7839 sleep(Duration::from_secs(5)).await;
7841 });
7842
7843 let (handler, _rx) = channel_message_handler();
7844
7845 let config = WebSocketConfig {
7846 url: format!("ws://127.0.0.1:{port}"),
7847 headers: vec![],
7848 heartbeat_interval_secs: None,
7849 heartbeat_payload: None,
7850 connect_timeout_ms: Some(2_000),
7851 reconnect_delay_initial_ms: Some(50),
7852 reconnect_delay_max_ms: Some(100),
7853 reconnect_backoff_factor: Some(1.0),
7854 reconnect_jitter_ms: Some(0),
7855 reconnect_max_attempts: Some(1),
7856 heartbeat_timeout_secs: None,
7857 idle_timeout_ms: Some(500),
7858 writer_capacity: None,
7859 backend: TransportBackend::Tungstenite,
7860 proxy_url: None,
7861 max_message_size_bytes: None,
7862 max_frame_size_bytes: None,
7863 };
7864
7865 let client = WebSocketClient::builder()
7866 .config(config)
7867 .message_handler(handler)
7868 .connect()
7869 .await
7870 .unwrap();
7871
7872 assert!(client.is_active());
7873
7874 wait_until_async(
7876 || async { client.is_reconnecting() || client.is_disconnected() },
7877 Duration::from_secs(3),
7878 )
7879 .await;
7880
7881 assert!(
7882 !client.is_active(),
7883 "Client should not be active after idle timeout"
7884 );
7885
7886 client.disconnect().await;
7887 server.abort();
7888 }
7889
7890 #[rstest]
7891 #[tokio::test]
7892 async fn test_idle_timeout_resets_on_data() {
7893 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
7894 let port = listener.local_addr().unwrap().port();
7895
7896 let server = task::spawn(async move {
7898 let (stream, _) = listener.accept().await.unwrap();
7899 let mut ws = accept_async(stream).await.unwrap();
7900
7901 for _ in 0..10 {
7902 sleep(Duration::from_millis(200)).await;
7903
7904 if ws.send(WsMessage::Text("ping".into())).await.is_err() {
7905 break;
7906 }
7907 }
7908 });
7909
7910 let (handler, _rx) = channel_message_handler();
7911
7912 let config = WebSocketConfig {
7913 url: format!("ws://127.0.0.1:{port}"),
7914 headers: vec![],
7915 heartbeat_interval_secs: None,
7916 heartbeat_payload: None,
7917 connect_timeout_ms: Some(2_000),
7918 reconnect_delay_initial_ms: Some(50),
7919 reconnect_delay_max_ms: Some(100),
7920 reconnect_backoff_factor: Some(1.0),
7921 reconnect_jitter_ms: Some(0),
7922 reconnect_max_attempts: Some(1),
7923 heartbeat_timeout_secs: None,
7924 idle_timeout_ms: Some(1_000),
7925 writer_capacity: None,
7926 backend: TransportBackend::Tungstenite,
7927 proxy_url: None,
7928 max_message_size_bytes: None,
7929 max_frame_size_bytes: None,
7930 };
7931
7932 let client = WebSocketClient::builder()
7933 .config(config)
7934 .message_handler(handler)
7935 .connect()
7936 .await
7937 .unwrap();
7938
7939 assert!(client.is_active());
7940
7941 sleep(Duration::from_millis(1_500)).await;
7943
7944 assert!(
7945 client.is_active(),
7946 "Client should remain active when data is flowing"
7947 );
7948
7949 client.disconnect().await;
7950 server.abort();
7951 }
7952
7953 #[rstest]
7954 #[tokio::test]
7955 async fn test_idle_timeout_fires_when_only_pings_received() {
7956 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
7962 let port = listener.local_addr().unwrap().port();
7963 let server = task::spawn(async move {
7964 let (stream, _) = listener.accept().await.unwrap();
7965 let mut ws = accept_async(stream).await.unwrap();
7966
7967 for _ in 0..60 {
7968 sleep(Duration::from_millis(100)).await;
7969
7970 if ws.send(WsMessage::Ping(Vec::new().into())).await.is_err() {
7971 break;
7972 }
7973 }
7974 });
7975
7976 let (handler, _rx) = channel_message_handler();
7977
7978 let config = WebSocketConfig {
7979 url: format!("ws://127.0.0.1:{port}"),
7980 headers: vec![],
7981 heartbeat_interval_secs: None,
7982 heartbeat_payload: None,
7983 connect_timeout_ms: Some(2_000),
7984 reconnect_delay_initial_ms: Some(50),
7985 reconnect_delay_max_ms: Some(100),
7986 reconnect_backoff_factor: Some(1.0),
7987 reconnect_jitter_ms: Some(0),
7988 reconnect_max_attempts: Some(1),
7989 heartbeat_timeout_secs: None,
7990 idle_timeout_ms: Some(500),
7991 writer_capacity: None,
7992 backend: TransportBackend::Tungstenite,
7993 proxy_url: None,
7994 max_message_size_bytes: None,
7995 max_frame_size_bytes: None,
7996 };
7997
7998 let client = WebSocketClient::builder()
7999 .config(config)
8000 .message_handler(handler)
8001 .connect()
8002 .await
8003 .unwrap();
8004
8005 assert!(client.is_active());
8006
8007 wait_until_async(
8011 || async { client.is_reconnecting() || client.is_disconnected() },
8012 Duration::from_millis(1_500),
8013 )
8014 .await;
8015
8016 assert!(
8017 !client.is_active(),
8018 "Client should not be active after idle timeout when only pings/pongs flow"
8019 );
8020
8021 client.disconnect().await;
8022 server.abort();
8023 }
8024
8025 #[rstest]
8026 #[tokio::test]
8027 async fn test_idle_timeout_fires_when_only_pongs_received() {
8028 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
8033 let port = listener.local_addr().unwrap().port();
8034 let server = task::spawn(async move {
8035 let (stream, _) = listener.accept().await.unwrap();
8036 let mut ws = accept_async(stream).await.unwrap();
8037
8038 let deadline = tokio::time::Instant::now() + Duration::from_secs(6);
8042 while tokio::time::Instant::now() < deadline {
8043 if let Ok(Some(Err(_)) | None) =
8044 tokio::time::timeout(Duration::from_millis(100), ws.next()).await
8045 {
8046 break;
8047 }
8048 }
8049 });
8050
8051 let (handler, _rx) = channel_message_handler();
8052
8053 let config = WebSocketConfig {
8054 url: format!("ws://127.0.0.1:{port}"),
8055 headers: vec![],
8056 heartbeat_interval_secs: Some(1),
8057 heartbeat_payload: None,
8058 connect_timeout_ms: Some(2_000),
8059 reconnect_delay_initial_ms: Some(50),
8060 reconnect_delay_max_ms: Some(100),
8061 reconnect_backoff_factor: Some(1.0),
8062 reconnect_jitter_ms: Some(0),
8063 reconnect_max_attempts: Some(1),
8064 heartbeat_timeout_secs: None,
8065 idle_timeout_ms: Some(1_500),
8066 writer_capacity: None,
8067 backend: TransportBackend::Tungstenite,
8068 proxy_url: None,
8069 max_message_size_bytes: None,
8070 max_frame_size_bytes: None,
8071 };
8072
8073 let client = WebSocketClient::builder()
8074 .config(config)
8075 .message_handler(handler)
8076 .connect()
8077 .await
8078 .unwrap();
8079
8080 assert!(client.is_active());
8081
8082 wait_until_async(
8086 || async { client.is_reconnecting() || client.is_disconnected() },
8087 Duration::from_millis(2_500),
8088 )
8089 .await;
8090
8091 assert!(
8092 !client.is_active(),
8093 "Client should not be active after idle timeout when only pongs flow"
8094 );
8095
8096 client.disconnect().await;
8097 server.abort();
8098 }
8099
8100 #[rstest]
8101 #[tokio::test]
8102 async fn test_disconnect_during_backoff_exits_promptly() {
8103 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
8107 let port = listener.local_addr().unwrap().port();
8108
8109 let server = task::spawn(async move {
8110 if let Ok((stream, _)) = listener.accept().await {
8112 let _ = accept_async(stream).await;
8113 }
8114
8115 sleep(Duration::from_mins(1)).await;
8117 });
8118
8119 let (handler, _rx) = channel_message_handler();
8120
8121 let config = WebSocketConfig {
8122 url: format!("ws://127.0.0.1:{port}"),
8123 headers: vec![],
8124 heartbeat_interval_secs: None,
8125 heartbeat_payload: None,
8126 connect_timeout_ms: Some(1_000),
8127 reconnect_delay_initial_ms: Some(10_000), reconnect_delay_max_ms: Some(10_000),
8129 reconnect_backoff_factor: Some(1.0),
8130 reconnect_jitter_ms: Some(0),
8131 reconnect_max_attempts: None,
8132 heartbeat_timeout_secs: None,
8133 idle_timeout_ms: None,
8134 writer_capacity: None,
8135 backend: TransportBackend::Tungstenite,
8136 proxy_url: None,
8137 max_message_size_bytes: None,
8138 max_frame_size_bytes: None,
8139 };
8140
8141 let client = WebSocketClient::builder()
8142 .config(config)
8143 .message_handler(handler)
8144 .connect()
8145 .await
8146 .unwrap();
8147
8148 wait_until_async(
8150 || async { client.is_reconnecting() },
8151 Duration::from_secs(3),
8152 )
8153 .await;
8154
8155 sleep(Duration::from_millis(1_500)).await;
8157
8158 let start = std::time::Instant::now();
8160 client.disconnect().await;
8161 let elapsed = start.elapsed();
8162
8163 assert!(client.is_disconnected(), "Client should be disconnected");
8164 assert!(
8166 elapsed < Duration::from_secs(2),
8167 "Disconnect should interrupt backoff sleep, took {elapsed:?}"
8168 );
8169
8170 server.abort();
8171 }
8172
8173 #[rstest]
8174 #[tokio::test]
8175 async fn test_rate_limit_cancelled_on_disconnect() {
8176 use std::{num::NonZeroU32, sync::Arc};
8179
8180 use crate::ratelimiter::quota::Quota;
8181
8182 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
8183 let port = listener.local_addr().unwrap().port();
8184 let server = task::spawn(async move {
8185 if let Ok((stream, _)) = listener.accept().await {
8186 let mut ws = accept_async(stream).await.unwrap();
8187 while let Some(Ok(msg)) = ws.next().await {
8189 if ws.send(msg).await.is_err() {
8190 break;
8191 }
8192 }
8193 }
8194 });
8195
8196 let (handler, _rx) = channel_message_handler();
8197
8198 let config = WebSocketConfig {
8199 url: format!("ws://127.0.0.1:{port}"),
8200 headers: vec![],
8201 heartbeat_interval_secs: None,
8202 heartbeat_payload: None,
8203 connect_timeout_ms: Some(5_000),
8204 reconnect_delay_initial_ms: Some(100),
8205 reconnect_delay_max_ms: Some(500),
8206 reconnect_backoff_factor: Some(1.5),
8207 reconnect_jitter_ms: Some(0),
8208 reconnect_max_attempts: None,
8209 heartbeat_timeout_secs: None,
8210 idle_timeout_ms: None,
8211 writer_capacity: None,
8212 backend: TransportBackend::Tungstenite,
8213 proxy_url: None,
8214 max_message_size_bytes: None,
8215 max_frame_size_bytes: None,
8216 };
8217
8218 let quota = Quota::with_period(Duration::from_mins(1))
8220 .unwrap()
8221 .allow_burst(NonZeroU32::new(1).unwrap());
8222
8223 let client = Arc::new(
8224 WebSocketClient::builder()
8225 .config(config)
8226 .message_handler(handler)
8227 .keyed_quotas(vec![("rate_key".to_string(), quota)])
8228 .connect()
8229 .await
8230 .unwrap(),
8231 );
8232
8233 let test_key: [Ustr; 1] = [Ustr::from("rate_key")];
8234
8235 client
8237 .send_text("exhaust".to_string(), Some(test_key.as_slice()))
8238 .await
8239 .unwrap();
8240
8241 let client_clone = Arc::clone(&client);
8243 let send_handle = task::spawn(async move {
8244 client_clone
8245 .send_text("blocked".to_string(), Some(&[Ustr::from("rate_key")]))
8246 .await
8247 });
8248
8249 sleep(Duration::from_millis(200)).await;
8251
8252 let start = std::time::Instant::now();
8254 client.disconnect().await;
8255 let elapsed_disconnect = start.elapsed();
8256
8257 let result = tokio::time::timeout(Duration::from_secs(2), send_handle)
8259 .await
8260 .expect("Send task should complete quickly")
8261 .expect("Send task should not panic");
8262
8263 assert!(
8264 matches!(result, Err(crate::error::SendError::Closed)),
8265 "Blocked send should return Closed, was: {result:?}"
8266 );
8267
8268 assert!(
8270 elapsed_disconnect < Duration::from_secs(3),
8271 "Disconnect should not wait for rate limiter, took {elapsed_disconnect:?}"
8272 );
8273
8274 server.abort();
8275 }
8276
8277 #[rstest]
8278 #[tokio::test]
8279 async fn test_stream_mode_transitions_to_closed_on_dead_write_task() {
8280 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
8284 let port = listener.local_addr().unwrap().port();
8285 let server = task::spawn(async move {
8286 if let Ok((stream, _)) = listener.accept().await
8287 && let Ok(ws) = accept_async(stream).await
8288 {
8289 drop(ws);
8291 }
8292 });
8293
8294 let config = WebSocketConfig {
8295 url: format!("ws://127.0.0.1:{port}"),
8296 headers: vec![],
8297 heartbeat_interval_secs: None,
8298 heartbeat_payload: None,
8299 connect_timeout_ms: Some(1_000),
8300 reconnect_delay_initial_ms: Some(50),
8301 reconnect_delay_max_ms: Some(100),
8302 reconnect_backoff_factor: Some(1.0),
8303 reconnect_jitter_ms: Some(0),
8304 reconnect_max_attempts: None,
8305 heartbeat_timeout_secs: None,
8306 idle_timeout_ms: None,
8307 writer_capacity: None,
8308 backend: TransportBackend::Tungstenite,
8309 proxy_url: None,
8310 max_message_size_bytes: None,
8311 max_frame_size_bytes: None,
8312 };
8313
8314 let (_reader, client) = WebSocketClient::stream_builder()
8315 .config(config)
8316 .connect()
8317 .await
8318 .unwrap();
8319
8320 assert!(client.is_active(), "Client should start active");
8321
8322 sleep(Duration::from_millis(100)).await;
8324
8325 for _ in 0..20 {
8327 let _ = client.send_text("ping".to_string(), None).await;
8328 sleep(Duration::from_millis(50)).await;
8329
8330 if !client.is_active() {
8331 break;
8332 }
8333 }
8334
8335 wait_until_async(|| async { !client.is_active() }, Duration::from_secs(5)).await;
8337
8338 assert!(
8340 client.is_closed() || client.is_disconnected(),
8341 "Stream mode should transition to CLOSED, not RECONNECT. \
8342 is_reconnecting={}, is_closed={}, is_disconnected={}",
8343 client.is_reconnecting(),
8344 client.is_closed(),
8345 client.is_disconnected(),
8346 );
8347 assert!(
8348 !client.is_reconnecting(),
8349 "Stream mode should never attempt reconnection"
8350 );
8351
8352 server.abort();
8353 }
8354
8355 #[derive(Default)]
8356 struct BlockingFailState {
8357 send_entered: AtomicBool,
8358 send_entered_notify: tokio::sync::Notify,
8359 released: AtomicBool,
8360 fail: AtomicBool,
8361 waker: parking_lot::Mutex<Option<std::task::Waker>>,
8362 }
8363
8364 impl BlockingFailState {
8365 fn trigger_failure(&self) {
8366 self.fail.store(true, Ordering::SeqCst);
8367 self.release_send();
8368 }
8369
8370 fn release_send(&self) {
8371 self.released.store(true, Ordering::SeqCst);
8372
8373 if let Some(waker) = self.waker.lock().take() {
8374 waker.wake();
8375 }
8376 }
8377 }
8378
8379 struct BlockingFailTransport {
8381 state: Arc<BlockingFailState>,
8382 }
8383
8384 impl futures_util::Stream for BlockingFailTransport {
8385 type Item = Result<Message, TransportError>;
8386
8387 fn poll_next(
8388 self: std::pin::Pin<&mut Self>,
8389 _cx: &mut std::task::Context<'_>,
8390 ) -> std::task::Poll<Option<Self::Item>> {
8391 std::task::Poll::Pending
8392 }
8393 }
8394
8395 impl futures_util::Sink<Message> for BlockingFailTransport {
8396 type Error = TransportError;
8397
8398 fn poll_ready(
8399 self: std::pin::Pin<&mut Self>,
8400 _cx: &mut std::task::Context<'_>,
8401 ) -> std::task::Poll<Result<(), Self::Error>> {
8402 std::task::Poll::Ready(Ok(()))
8403 }
8404
8405 fn start_send(self: std::pin::Pin<&mut Self>, _item: Message) -> Result<(), Self::Error> {
8406 Ok(())
8407 }
8408
8409 fn poll_flush(
8410 self: std::pin::Pin<&mut Self>,
8411 cx: &mut std::task::Context<'_>,
8412 ) -> std::task::Poll<Result<(), Self::Error>> {
8413 *self.state.waker.lock() = Some(cx.waker().clone());
8416 self.state.send_entered.store(true, Ordering::SeqCst);
8417 self.state.send_entered_notify.notify_one();
8418
8419 if !self.state.released.load(Ordering::SeqCst) {
8420 std::task::Poll::Pending
8421 } else if self.state.fail.load(Ordering::SeqCst) {
8422 std::task::Poll::Ready(Err(TransportError::ConnectionReset))
8423 } else {
8424 std::task::Poll::Ready(Ok(()))
8425 }
8426 }
8427
8428 fn poll_close(
8429 self: std::pin::Pin<&mut Self>,
8430 _cx: &mut std::task::Context<'_>,
8431 ) -> std::task::Poll<Result<(), Self::Error>> {
8432 std::task::Poll::Ready(Ok(()))
8433 }
8434 }
8435
8436 struct BlockingMessageState {
8437 polled_tx: Mutex<Option<std::sync::mpsc::Sender<()>>>,
8438 release: (Mutex<bool>, parking_lot::Condvar),
8439 message: Mutex<Option<Message>>,
8440 }
8441
8442 struct BlockingMessageTransport {
8443 state: Arc<BlockingMessageState>,
8444 }
8445
8446 impl futures_util::Stream for BlockingMessageTransport {
8447 type Item = Result<Message, TransportError>;
8448
8449 fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
8450 if let Some(polled_tx) = self.state.polled_tx.lock().take() {
8451 polled_tx.send(()).unwrap();
8452 }
8453
8454 let (lock, condvar) = &self.state.release;
8455 let mut released = lock.lock();
8456 while !*released {
8457 condvar.wait(&mut released);
8458 }
8459
8460 Poll::Ready(self.state.message.lock().take().map(Ok))
8461 }
8462 }
8463
8464 impl futures_util::Sink<Message> for BlockingMessageTransport {
8465 type Error = TransportError;
8466
8467 fn poll_ready(
8468 self: Pin<&mut Self>,
8469 _cx: &mut Context<'_>,
8470 ) -> Poll<Result<(), Self::Error>> {
8471 Poll::Ready(Ok(()))
8472 }
8473
8474 fn start_send(self: Pin<&mut Self>, _item: Message) -> Result<(), Self::Error> {
8475 Ok(())
8476 }
8477
8478 fn poll_flush(
8479 self: Pin<&mut Self>,
8480 _cx: &mut Context<'_>,
8481 ) -> Poll<Result<(), Self::Error>> {
8482 Poll::Ready(Ok(()))
8483 }
8484
8485 fn poll_close(
8486 self: Pin<&mut Self>,
8487 _cx: &mut Context<'_>,
8488 ) -> Poll<Result<(), Self::Error>> {
8489 Poll::Ready(Ok(()))
8490 }
8491 }
8492
8493 struct RecordingState {
8494 messages: Arc<Mutex<Vec<Message>>>,
8495 recorded_notify: tokio::sync::Notify,
8496 }
8497
8498 struct RecordingTransport {
8499 state: Arc<RecordingState>,
8500 }
8501
8502 impl futures_util::Stream for RecordingTransport {
8503 type Item = Result<Message, TransportError>;
8504
8505 fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
8506 Poll::Pending
8507 }
8508 }
8509
8510 impl futures_util::Sink<Message> for RecordingTransport {
8511 type Error = TransportError;
8512
8513 fn poll_ready(
8514 self: Pin<&mut Self>,
8515 _cx: &mut Context<'_>,
8516 ) -> Poll<Result<(), Self::Error>> {
8517 Poll::Ready(Ok(()))
8518 }
8519
8520 fn start_send(self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> {
8521 self.state.messages.lock().push(item);
8522 self.state.recorded_notify.notify_one();
8523 Ok(())
8524 }
8525
8526 fn poll_flush(
8527 self: Pin<&mut Self>,
8528 _cx: &mut Context<'_>,
8529 ) -> Poll<Result<(), Self::Error>> {
8530 Poll::Ready(Ok(()))
8531 }
8532
8533 fn poll_close(
8534 self: Pin<&mut Self>,
8535 _cx: &mut Context<'_>,
8536 ) -> Poll<Result<(), Self::Error>> {
8537 Poll::Ready(Ok(()))
8538 }
8539 }
8540
8541 #[rstest]
8542 #[tokio::test(start_paused = true)]
8543 async fn test_pong_is_bound_to_connection_epoch() {
8544 let initial_state = Arc::new(RecordingState {
8545 messages: Arc::new(Mutex::new(Vec::new())),
8546 recorded_notify: tokio::sync::Notify::new(),
8547 });
8548
8549 let initial_transport: BoxedWsTransport = Box::pin(RecordingTransport {
8550 state: initial_state,
8551 });
8552
8553 let (writer, _reader) = initial_transport.split();
8554 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
8555 let state_notify = Arc::new(tokio::sync::Notify::new());
8556 let auth_tracker = Arc::new(OnceLock::new());
8557 let reconnect_buffer_waits_for_auth = Arc::new(AtomicBool::new(false));
8558 let (writer_tx, writer_rx) = crate::writer::channel(crate::writer::DEFAULT_WRITER_CAPACITY);
8559 let write_task = WebSocketClientInner::spawn_write_task(
8560 Arc::clone(&connection_state),
8561 Arc::clone(&state_notify),
8562 Arc::new(AtomicBool::new(true)),
8563 writer,
8564 writer_rx,
8565 Arc::new(AtomicU64::new(0)),
8566 auth_tracker,
8567 reconnect_buffer_waits_for_auth,
8568 None,
8569 );
8570
8571 let recorded = Arc::new(Mutex::new(Vec::new()));
8572
8573 let replacement_state = Arc::new(RecordingState {
8574 messages: Arc::clone(&recorded),
8575 recorded_notify: tokio::sync::Notify::new(),
8576 });
8577
8578 let replacement_transport: BoxedWsTransport = Box::pin(RecordingTransport {
8579 state: replacement_state,
8580 });
8581
8582 let (replacement_writer, _reader) = replacement_transport.split();
8583 let (update_tx, update_rx) = tokio::sync::oneshot::channel();
8584 writer_tx
8585 .send(WriterCommand::Update(replacement_writer, update_tx))
8586 .unwrap();
8587 writer_tx
8588 .send_control(WriterCommand::SendPongOnConnection {
8589 data: b"stale-pong".to_vec(),
8590 connection_epoch: 0,
8591 })
8592 .unwrap();
8593
8594 tokio::time::advance(Duration::from_millis(100)).await;
8595 assert_eq!(update_rx.await.unwrap(), 1);
8596
8597 let (sentinel_tx, sentinel_rx) = tokio::sync::oneshot::channel();
8598 writer_tx
8599 .send_control(WriterCommand::SendOnConnection {
8600 message: Message::text("sentinel-1"),
8601 connection_epoch: 1,
8602 response_tx: sentinel_tx,
8603 })
8604 .unwrap();
8605
8606 sentinel_rx.await.unwrap().unwrap();
8607 assert_eq!(recorded.lock().as_slice(), &[Message::text("sentinel-1")]);
8608
8609 writer_tx
8610 .send_control(WriterCommand::SendPongOnConnection {
8611 data: b"fresh-pong".to_vec(),
8612 connection_epoch: 1,
8613 })
8614 .unwrap();
8615
8616 let (sentinel_tx, sentinel_rx) = tokio::sync::oneshot::channel();
8617 writer_tx
8618 .send_control(WriterCommand::SendOnConnection {
8619 message: Message::text("sentinel-2"),
8620 connection_epoch: 1,
8621 response_tx: sentinel_tx,
8622 })
8623 .unwrap();
8624
8625 sentinel_rx.await.unwrap().unwrap();
8626 assert_eq!(
8627 recorded.lock().as_slice(),
8628 &[
8629 Message::text("sentinel-1"),
8630 Message::Pong(b"fresh-pong".to_vec().into()),
8631 Message::text("sentinel-2"),
8632 ]
8633 );
8634
8635 connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
8636 state_notify.notify_waiters();
8637 drop(writer_tx);
8638 write_task.await.unwrap();
8639 }
8640
8641 #[rstest]
8642 #[case(Message::text("stale"))]
8643 #[case(Message::Binary(vec![1, 2, 3].into()))]
8644 #[case(Message::Ping(vec![1, 2, 3].into()))]
8645 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
8646 async fn test_message_handler_drops_old_session_message(#[case] message: Message) {
8647 let (polled_tx, polled_rx) = std::sync::mpsc::channel();
8648
8649 let state = Arc::new(BlockingMessageState {
8650 polled_tx: Mutex::new(Some(polled_tx)),
8651 release: (Mutex::new(false), Condvar::new()),
8652 message: Mutex::new(Some(message)),
8653 });
8654
8655 let release_guard = CondvarReleaseGuard::new(&state.release);
8656
8657 let transport: BoxedWsTransport = Box::pin(BlockingMessageTransport {
8658 state: Arc::clone(&state),
8659 });
8660
8661 let (_writer, reader) = transport.split();
8662 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
8663 let state_notify = Arc::new(tokio::sync::Notify::new());
8664 let read_fence = ReadSessionFence::new();
8665 let message_count = Arc::new(AtomicUsize::new(0));
8666 let ping_count = Arc::new(AtomicUsize::new(0));
8667 let message_count_clone = Arc::clone(&message_count);
8668 let ping_count_clone = Arc::clone(&ping_count);
8669 let message_handler: MessageHandler =
8670 Arc::new(move |_| _ = message_count_clone.fetch_add(1, Ordering::SeqCst));
8671 let message_handler = IncomingHandler::Message(message_handler);
8672 let ping_handler: PingHandler =
8673 Arc::new(move |_| _ = ping_count_clone.fetch_add(1, Ordering::SeqCst));
8674 let ping_handler = IncomingPingHandler::Ping(ping_handler);
8675
8676 let read_task = WebSocketClientInner::spawn_message_handler_task(
8677 Arc::clone(&connection_state),
8678 state_notify,
8679 read_fence.clone(),
8680 reader,
8681 0,
8682 Some(&message_handler),
8683 Some(&ping_handler),
8684 None,
8685 None,
8686 None,
8687 None,
8688 );
8689
8690 recv_rendezvous(polled_rx, "WebSocket reader poll entry").await;
8691 connection_state.store(ConnectionMode::Reconnect.as_u8(), Ordering::SeqCst);
8692 read_fence.invalidate();
8693 read_task.abort();
8694 connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
8695 release_guard.release();
8696 await_task_termination(read_task, "old WebSocket read task").await;
8697
8698 assert_eq!(message_count.load(Ordering::SeqCst), 0);
8699 assert_eq!(ping_count.load(Ordering::SeqCst), 0);
8700 }
8701
8702 #[rstest]
8703 #[tokio::test]
8704 async fn test_reconnect_buffer_drain_stops_after_reconnect_request() {
8705 let state = Arc::new(BlockingFailState::default());
8706
8707 let transport: BoxedWsTransport = Box::pin(BlockingFailTransport {
8708 state: Arc::clone(&state),
8709 });
8710
8711 let (mut writer, _reader) = transport.split();
8712 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
8713 let task_connection_state = Arc::clone(&connection_state);
8714
8715 let drain_task = tokio::spawn(async move {
8716 let auth_tracker = Arc::new(OnceLock::new());
8717 let reconnect_buffer_waits_for_auth = AtomicBool::new(false);
8718 let (sender, mut receiver) = crate::writer::channel(2);
8719 sender.send(Message::text("admitted")).unwrap();
8720 sender.send(Message::text("held-for-reconnect")).unwrap();
8721 let mut buffer = VecDeque::from([
8722 receiver.recv().await.unwrap(),
8723 receiver.recv().await.unwrap(),
8724 ]);
8725 let send_error = WebSocketClientInner::drain_reconnect_buffer(
8726 &mut buffer,
8727 &mut writer,
8728 &task_connection_state,
8729 &auth_tracker,
8730 &reconnect_buffer_waits_for_auth,
8731 )
8732 .await;
8733 (buffer, send_error)
8734 });
8735
8736 state.send_entered_notify.notified().await;
8737 connection_state.store(ConnectionMode::Reconnect.as_u8(), Ordering::SeqCst);
8738 state.release_send();
8739
8740 let (buffer, send_error) = tokio::time::timeout(TEST_TIMEOUT, drain_task)
8741 .await
8742 .expect("buffer drain should stop after reconnect acceptance")
8743 .unwrap();
8744 assert!(!send_error);
8745 assert_eq!(
8746 buffer
8747 .into_iter()
8748 .map(|(message, _)| message)
8749 .collect::<Vec<_>>(),
8750 vec![Message::text("held-for-reconnect")]
8751 );
8752 }
8753
8754 #[rstest]
8755 #[case(0)]
8756 #[case(usize::MAX)]
8757 #[tokio::test]
8758 async fn test_stream_writer_rejects_invalid_capacity(#[case] capacity: usize) {
8759 let transport: BoxedWsTransport = Box::pin(BlockingFailTransport {
8760 state: Arc::new(BlockingFailState::default()),
8761 });
8762
8763 let (writer, _reader) = transport.split();
8764 let mut config = WebSocketConfig::builder()
8765 .url("ws://localhost".to_string())
8766 .build()
8767 .unwrap();
8768 config.writer_capacity = Some(capacity);
8769
8770 let result = WebSocketClientInner::new_with_writer(config, writer).await;
8771
8772 assert!(
8773 matches!(result, Err(TransportError::Io(e)) if e.kind() == std::io::ErrorKind::InvalidInput && e.to_string().contains("writer_capacity"))
8774 );
8775 }
8776
8777 #[rstest]
8778 #[case(0)]
8779 #[case(usize::MAX)]
8780 #[tokio::test]
8781 async fn test_handlerless_connect_rejects_invalid_capacity(#[case] capacity: usize) {
8782 let mut config = WebSocketConfig::builder()
8783 .url("ws://127.0.0.1:1".to_string())
8784 .build()
8785 .unwrap();
8786 config.writer_capacity = Some(capacity);
8787
8788 let result = WebSocketClientInner::connect_url(config, None, None).await;
8789
8790 assert!(
8791 matches!(result, Err(TransportError::Io(e)) if e.kind() == std::io::ErrorKind::InvalidInput && e.to_string().contains("writer_capacity"))
8792 );
8793 }
8794
8795 #[rstest]
8796 #[tokio::test(start_paused = true)]
8797 async fn test_stalled_websocket_send_reconnects_and_replays() {
8798 let state = Arc::new(BlockingFailState::default());
8799
8800 let transport: BoxedWsTransport = Box::pin(BlockingFailTransport {
8801 state: Arc::clone(&state),
8802 });
8803
8804 let (writer, _reader) = transport.split();
8805 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
8806 let state_notify = Arc::new(tokio::sync::Notify::new());
8807 let auth_tracker = Arc::new(OnceLock::new());
8808 let reconnect_buffer_waits_for_auth = Arc::new(AtomicBool::new(false));
8809 let (writer_tx, writer_rx) = crate::writer::channel(2);
8810 let states = Arc::new(Mutex::new(Vec::new()));
8811 let states_callback = Arc::clone(&states);
8812 let sink = SocketStateSink::new(move |state| {
8813 states_callback.lock().push(state);
8814 });
8815
8816 let write_task = WebSocketClientInner::spawn_write_task(
8817 Arc::clone(&connection_state),
8818 Arc::clone(&state_notify),
8819 Arc::new(AtomicBool::new(true)),
8820 writer,
8821 writer_rx,
8822 Arc::new(AtomicU64::new(0)),
8823 Arc::clone(&auth_tracker),
8824 reconnect_buffer_waits_for_auth,
8825 Some(sink),
8826 );
8827
8828 writer_tx
8829 .send(WriterCommand::Send(Message::text("complete-message")))
8830 .unwrap();
8831 state.send_entered_notify.notified().await;
8832 writer_tx
8833 .send(WriterCommand::Send(Message::text("second-message")))
8834 .unwrap();
8835 assert_eq!(
8836 ConnectionMode::from_atomic(&connection_state),
8837 ConnectionMode::Active
8838 );
8839
8840 for _ in 0..100 {
8841 assert!(matches!(
8842 writer_tx.send(WriterCommand::Send(Message::text("overflow"))),
8843 Err(SendError::BufferFull)
8844 ));
8845 }
8846
8847 tokio::time::advance(Duration::from_secs(WRITE_TIMEOUT_SECS)).await;
8848 tokio::task::yield_now().await;
8849 assert_eq!(
8850 ConnectionMode::from_atomic(&connection_state),
8851 ConnectionMode::Reconnect
8852 );
8853 assert!(matches!(
8854 writer_tx.send(WriterCommand::Send(Message::text("overflow"))),
8855 Err(SendError::BufferFull)
8856 ));
8857
8858 let recorded = Arc::new(Mutex::new(Vec::new()));
8859
8860 let recording_state = Arc::new(RecordingState {
8861 messages: Arc::clone(&recorded),
8862 recorded_notify: tokio::sync::Notify::new(),
8863 });
8864
8865 let transport: BoxedWsTransport = Box::pin(RecordingTransport {
8866 state: Arc::clone(&recording_state),
8867 });
8868
8869 let (new_writer, _reader) = transport.split();
8870 let (update_tx, update_rx) = tokio::sync::oneshot::channel();
8871 writer_tx
8872 .send_update(WriterCommand::Update(new_writer, update_tx))
8873 .unwrap();
8874
8875 tokio::time::advance(Duration::from_secs(GRACEFUL_SHUTDOWN_TIMEOUT_SECS)).await;
8876 assert_eq!(
8877 tokio::time::timeout(Duration::from_secs(1), update_rx)
8878 .await
8879 .expect("writer update should not remain queued behind a stalled send")
8880 .unwrap(),
8881 1,
8882 "the replacement sink should install as connection epoch 1"
8883 );
8884 assert_eq!(
8885 ConnectionMode::from_atomic(&connection_state),
8886 ConnectionMode::Reconnect
8887 );
8888
8889 connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
8890 state_notify.notify_waiters();
8891 tokio::time::advance(Duration::from_millis(CONNECTION_STATE_CHECK_INTERVAL_MS)).await;
8892 recording_state.recorded_notify.notified().await;
8893 assert_eq!(
8894 recorded.lock().as_slice(),
8895 &[
8896 Message::text("complete-message"),
8897 Message::text("second-message")
8898 ]
8899 );
8900
8901 writer_tx
8902 .send(WriterCommand::Send(Message::text("second-message")))
8903 .unwrap();
8904
8905 connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
8906 state_notify.notify_waiters();
8907 drop(writer_tx);
8908 write_task.await.unwrap();
8909
8910 assert_eq!(*states.lock(), vec![SocketState::Disconnected]);
8911 }
8912
8913 #[rstest]
8914 #[case(Message::Ping(vec![1, 2, 3].into()))]
8915 #[case(Message::Pong(vec![4, 5, 6].into()))]
8916 #[case(Message::Close(None))]
8917 #[tokio::test(start_paused = true)]
8918 async fn test_stalled_control_frame_is_not_replayed(#[case] control: Message) {
8919 let state = Arc::new(BlockingFailState::default());
8922
8923 let transport: BoxedWsTransport = Box::pin(BlockingFailTransport {
8924 state: Arc::clone(&state),
8925 });
8926
8927 let (writer, _reader) = transport.split();
8928 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
8929 let state_notify = Arc::new(tokio::sync::Notify::new());
8930 let (writer_tx, writer_rx) = crate::writer::channel(crate::writer::DEFAULT_WRITER_CAPACITY);
8931 let write_task = WebSocketClientInner::spawn_write_task(
8932 Arc::clone(&connection_state),
8933 Arc::clone(&state_notify),
8934 Arc::new(AtomicBool::new(true)),
8935 writer,
8936 writer_rx,
8937 Arc::new(AtomicU64::new(0)),
8938 Arc::new(OnceLock::new()),
8939 Arc::new(AtomicBool::new(false)),
8940 None,
8941 );
8942
8943 writer_tx.send(WriterCommand::Send(control)).unwrap();
8944 state.send_entered_notify.notified().await;
8945
8946 let recorded = Arc::new(Mutex::new(Vec::new()));
8947
8948 let recording_state = Arc::new(RecordingState {
8949 messages: Arc::clone(&recorded),
8950 recorded_notify: tokio::sync::Notify::new(),
8951 });
8952
8953 let transport: BoxedWsTransport = Box::pin(RecordingTransport {
8954 state: recording_state,
8955 });
8956
8957 let (new_writer, _reader) = transport.split();
8958 let (update_tx, update_rx) = tokio::sync::oneshot::channel();
8959 writer_tx
8960 .send(WriterCommand::Update(new_writer, update_tx))
8961 .unwrap();
8962
8963 tokio::time::advance(Duration::from_secs(WRITE_TIMEOUT_SECS)).await;
8964 assert_eq!(
8965 tokio::time::timeout(Duration::from_secs(1), update_rx)
8966 .await
8967 .expect("writer update should not remain queued behind a stalled send")
8968 .unwrap(),
8969 1,
8970 "the replacement sink should install as connection epoch 1"
8971 );
8972 assert_eq!(
8973 ConnectionMode::from_atomic(&connection_state),
8974 ConnectionMode::Reconnect,
8975 "a failed control-frame write should still trigger a reconnect"
8976 );
8977
8978 connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
8979 state_notify.notify_waiters();
8980
8981 for _ in 0..5 {
8983 tokio::time::advance(Duration::from_millis(CONNECTION_STATE_CHECK_INTERVAL_MS)).await;
8984 }
8985
8986 let replayed = recorded.lock().clone();
8987 assert!(
8988 replayed.is_empty(),
8989 "a failed control frame must not reach the replacement connection, was {replayed:?}"
8990 );
8991
8992 connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
8993 state_notify.notify_waiters();
8994 drop(writer_tx);
8995 write_task.await.unwrap();
8996 }
8997
8998 #[rstest]
8999 #[case(Message::Ping(vec![1, 2, 3].into()))]
9000 #[case(Message::Pong(vec![4, 5, 6].into()))]
9001 #[case(Message::Close(None))]
9002 #[tokio::test(start_paused = true)]
9003 async fn test_control_frame_enqueued_during_reconnect_is_not_replayed(
9004 #[case] control: Message,
9005 ) {
9006 let recorded = Arc::new(Mutex::new(Vec::new()));
9009
9010 let recording_state = Arc::new(RecordingState {
9011 messages: Arc::clone(&recorded),
9012 recorded_notify: tokio::sync::Notify::new(),
9013 });
9014
9015 let transport: BoxedWsTransport = Box::pin(RecordingTransport {
9016 state: recording_state,
9017 });
9018
9019 let (writer, _reader) = transport.split();
9020 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Reconnect.as_u8()));
9021 let state_notify = Arc::new(tokio::sync::Notify::new());
9022 let (writer_tx, writer_rx) = crate::writer::channel(crate::writer::DEFAULT_WRITER_CAPACITY);
9023 let write_task = WebSocketClientInner::spawn_write_task(
9024 Arc::clone(&connection_state),
9025 Arc::clone(&state_notify),
9026 Arc::new(AtomicBool::new(true)),
9027 writer,
9028 writer_rx,
9029 Arc::new(AtomicU64::new(0)),
9030 Arc::new(OnceLock::new()),
9031 Arc::new(AtomicBool::new(false)),
9032 None,
9033 );
9034
9035 writer_tx.send(WriterCommand::Send(control)).unwrap();
9036 tokio::time::advance(Duration::from_millis(CONNECTION_STATE_CHECK_INTERVAL_MS)).await;
9037
9038 connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
9039 state_notify.notify_waiters();
9040
9041 for _ in 0..5 {
9043 tokio::time::advance(Duration::from_millis(CONNECTION_STATE_CHECK_INTERVAL_MS)).await;
9044 }
9045
9046 let replayed = recorded.lock().clone();
9047 assert!(
9048 replayed.is_empty(),
9049 "a control frame enqueued during reconnect must not reach the replacement connection, was {replayed:?}"
9050 );
9051
9052 connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
9053 state_notify.notify_waiters();
9054 drop(writer_tx);
9055 write_task.await.unwrap();
9056 }
9057
9058 #[rstest]
9059 #[tokio::test(start_paused = true)]
9060 async fn test_stalled_text_heartbeat_is_not_replayed() {
9061 let state = Arc::new(BlockingFailState::default());
9062
9063 let transport: BoxedWsTransport = Box::pin(BlockingFailTransport {
9064 state: Arc::clone(&state),
9065 });
9066
9067 let (writer, _reader) = transport.split();
9068 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
9069 let state_notify = Arc::new(tokio::sync::Notify::new());
9070 let (writer_tx, writer_rx) = crate::writer::channel(crate::writer::DEFAULT_WRITER_CAPACITY);
9071 let write_task = WebSocketClientInner::spawn_write_task(
9072 Arc::clone(&connection_state),
9073 Arc::clone(&state_notify),
9074 Arc::new(AtomicBool::new(true)),
9075 writer,
9076 writer_rx,
9077 Arc::new(AtomicU64::new(0)),
9078 Arc::new(OnceLock::new()),
9079 Arc::new(AtomicBool::new(false)),
9080 None,
9081 );
9082
9083 writer_tx
9084 .send(WriterCommand::Heartbeat(Message::text(
9085 "{\"op\":\"heartbeat\"}",
9086 )))
9087 .unwrap();
9088 state.send_entered_notify.notified().await;
9089
9090 let recorded = Arc::new(Mutex::new(Vec::new()));
9091
9092 let recording_state = Arc::new(RecordingState {
9093 messages: Arc::clone(&recorded),
9094 recorded_notify: tokio::sync::Notify::new(),
9095 });
9096
9097 let transport: BoxedWsTransport = Box::pin(RecordingTransport {
9098 state: recording_state,
9099 });
9100
9101 let (new_writer, _reader) = transport.split();
9102 let (update_tx, update_rx) = tokio::sync::oneshot::channel();
9103 writer_tx
9104 .send(WriterCommand::Update(new_writer, update_tx))
9105 .unwrap();
9106
9107 tokio::time::advance(Duration::from_secs(WRITE_TIMEOUT_SECS)).await;
9108 assert_eq!(
9109 tokio::time::timeout(Duration::from_secs(1), update_rx)
9110 .await
9111 .expect("writer update should not remain queued behind a stalled send")
9112 .unwrap(),
9113 1,
9114 "the replacement sink should install as connection epoch 1"
9115 );
9116 assert_eq!(
9117 ConnectionMode::from_atomic(&connection_state),
9118 ConnectionMode::Reconnect,
9119 "a failed heartbeat write should still trigger a reconnect"
9120 );
9121
9122 connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
9123 state_notify.notify_waiters();
9124
9125 for _ in 0..5 {
9126 tokio::time::advance(Duration::from_millis(CONNECTION_STATE_CHECK_INTERVAL_MS)).await;
9127 }
9128
9129 let replayed = recorded.lock().clone();
9130 assert!(
9131 replayed.is_empty(),
9132 "a failed text heartbeat must not reach the replacement connection, was {replayed:?}"
9133 );
9134
9135 connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
9136 state_notify.notify_waiters();
9137 drop(writer_tx);
9138 write_task.await.unwrap();
9139 }
9140
9141 #[rstest]
9142 #[tokio::test(start_paused = true)]
9143 async fn test_text_heartbeat_enqueued_during_reconnect_is_not_replayed() {
9144 let recorded = Arc::new(Mutex::new(Vec::new()));
9145
9146 let recording_state = Arc::new(RecordingState {
9147 messages: Arc::clone(&recorded),
9148 recorded_notify: tokio::sync::Notify::new(),
9149 });
9150
9151 let transport: BoxedWsTransport = Box::pin(RecordingTransport {
9152 state: recording_state,
9153 });
9154
9155 let (writer, _reader) = transport.split();
9156 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Reconnect.as_u8()));
9157 let state_notify = Arc::new(tokio::sync::Notify::new());
9158 let (writer_tx, writer_rx) = crate::writer::channel(crate::writer::DEFAULT_WRITER_CAPACITY);
9159 let write_task = WebSocketClientInner::spawn_write_task(
9160 Arc::clone(&connection_state),
9161 Arc::clone(&state_notify),
9162 Arc::new(AtomicBool::new(true)),
9163 writer,
9164 writer_rx,
9165 Arc::new(AtomicU64::new(0)),
9166 Arc::new(OnceLock::new()),
9167 Arc::new(AtomicBool::new(false)),
9168 None,
9169 );
9170
9171 writer_tx
9172 .send(WriterCommand::Heartbeat(Message::text(
9173 "{\"op\":\"heartbeat\"}",
9174 )))
9175 .unwrap();
9176 tokio::time::advance(Duration::from_millis(CONNECTION_STATE_CHECK_INTERVAL_MS)).await;
9177
9178 connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
9179 state_notify.notify_waiters();
9180
9181 for _ in 0..5 {
9182 tokio::time::advance(Duration::from_millis(CONNECTION_STATE_CHECK_INTERVAL_MS)).await;
9183 }
9184
9185 let replayed = recorded.lock().clone();
9186 assert!(
9187 replayed.is_empty(),
9188 "a text heartbeat enqueued during reconnect must not reach the replacement connection, was {replayed:?}"
9189 );
9190
9191 connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
9192 state_notify.notify_waiters();
9193 drop(writer_tx);
9194 write_task.await.unwrap();
9195 }
9196
9197 #[rstest]
9198 #[case::text(
9199 Some("{\"op\":\"heartbeat\"}"),
9200 Message::text("{\"op\":\"heartbeat\"}")
9201 )]
9202 #[case::ping(None, Message::Ping(vec![].into()))]
9203 #[tokio::test(start_paused = true)]
9204 async fn test_heartbeat_task_enqueues_writer_heartbeat_command(
9205 #[case] payload: Option<&str>,
9206 #[case] expected: Message,
9207 ) {
9208 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
9209 let (writer_tx, mut writer_rx) =
9210 crate::writer::channel(crate::writer::DEFAULT_WRITER_CAPACITY);
9211 let task = WebSocketClientInner::spawn_heartbeat_task(
9212 Arc::clone(&connection_state),
9213 1,
9214 payload.map(ToString::to_string),
9215 writer_tx,
9216 );
9217
9218 tokio::time::advance(Duration::from_secs(1)).await;
9219 let cmd = writer_rx
9220 .recv()
9221 .await
9222 .expect("heartbeat task should enqueue");
9223
9224 match cmd.0 {
9225 WriterCommand::Heartbeat(msg) => assert_eq!(msg, expected),
9226 other => panic!("expected Heartbeat, was {other:?}"),
9227 }
9228
9229 connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
9230 tokio::time::advance(Duration::from_secs(1)).await;
9231 task.await.unwrap();
9232 }
9233
9234 #[cfg(target_os = "linux")]
9235 #[rstest]
9236 #[tokio::test]
9237 async fn test_heartbeat_task_warns_when_writer_channel_closed() {
9238 let capture = capture_logs_for(&["nautilus_network::websocket::client"]).await;
9239 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
9240 let (writer_tx, mut writer_rx) =
9241 crate::writer::channel(crate::writer::DEFAULT_WRITER_CAPACITY);
9242 let task = WebSocketClientInner::spawn_heartbeat_task(
9243 Arc::clone(&connection_state),
9244 1,
9245 None,
9246 writer_tx,
9247 );
9248
9249 let _ = tokio::time::timeout(Duration::from_secs(2), writer_rx.recv())
9250 .await
9251 .expect("timed out waiting for the first heartbeat")
9252 .expect("heartbeat channel closed before the first heartbeat");
9253 drop(writer_rx);
9254 tokio::time::sleep(Duration::from_millis(1_200)).await;
9255
9256 let messages = capture.messages();
9257 assert!(
9258 messages.iter().any(|(level, message)| {
9259 *level == Level::Warn && message.contains("Failed to send heartbeat to writer task")
9260 }),
9261 "expected WARN when the writer channel is closed, was {messages:?}"
9262 );
9263 assert!(
9264 !messages.iter().any(|(level, _)| *level == Level::Error),
9265 "a closed writer channel must not log ERROR, was {messages:?}"
9266 );
9267
9268 connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
9269 tokio::time::timeout(Duration::from_secs(2), task)
9270 .await
9271 .expect("heartbeat task should stop after close")
9272 .unwrap();
9273 }
9274
9275 #[rstest]
9276 #[tokio::test(start_paused = true)]
9277 async fn test_stalled_ownership_bound_send_times_out_without_replay() {
9278 let state = Arc::new(BlockingFailState::default());
9279
9280 let transport: BoxedWsTransport = Box::pin(BlockingFailTransport {
9281 state: Arc::clone(&state),
9282 });
9283
9284 let (writer, _reader) = transport.split();
9285 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
9286 let state_notify = Arc::new(tokio::sync::Notify::new());
9287 let auth_tracker = Arc::new(OnceLock::new());
9288 let reconnect_buffer_waits_for_auth = Arc::new(AtomicBool::new(false));
9289 let (writer_tx, writer_rx) = crate::writer::channel(crate::writer::DEFAULT_WRITER_CAPACITY);
9290 let write_task = WebSocketClientInner::spawn_write_task(
9291 Arc::clone(&connection_state),
9292 Arc::clone(&state_notify),
9293 Arc::new(AtomicBool::new(true)),
9294 writer,
9295 writer_rx,
9296 Arc::new(AtomicU64::new(0)),
9297 Arc::clone(&auth_tracker),
9298 reconnect_buffer_waits_for_auth,
9299 None,
9300 );
9301
9302 let (response_tx, response_rx) = tokio::sync::oneshot::channel();
9303 writer_tx
9304 .send_control(WriterCommand::SendOnConnection {
9305 message: Message::text("ownership-bound"),
9306 connection_epoch: 0,
9307 response_tx,
9308 })
9309 .unwrap();
9310
9311 state.send_entered_notify.notified().await;
9312
9313 tokio::time::advance(Duration::from_secs(WRITE_TIMEOUT_SECS)).await;
9314 let outcome = tokio::time::timeout(Duration::from_secs(1), response_rx)
9315 .await
9316 .expect("a stalled ownership-bound send must not wedge the writer task")
9317 .unwrap();
9318 assert!(
9319 matches!(outcome, Err(SendError::WriteTimeout)),
9320 "expected the write deadline to be reported, was {outcome:?}"
9321 );
9322 assert_eq!(
9323 ConnectionMode::from_atomic(&connection_state),
9324 ConnectionMode::Reconnect
9325 );
9326
9327 let recorded = Arc::new(Mutex::new(Vec::new()));
9330
9331 let recording_state = Arc::new(RecordingState {
9332 messages: Arc::clone(&recorded),
9333 recorded_notify: tokio::sync::Notify::new(),
9334 });
9335
9336 let transport: BoxedWsTransport = Box::pin(RecordingTransport {
9337 state: Arc::clone(&recording_state),
9338 });
9339
9340 let (new_writer, _reader) = transport.split();
9341 let (update_tx, update_rx) = tokio::sync::oneshot::channel();
9342 writer_tx
9343 .send(WriterCommand::Update(new_writer, update_tx))
9344 .unwrap();
9345
9346 tokio::time::advance(Duration::from_secs(GRACEFUL_SHUTDOWN_TIMEOUT_SECS)).await;
9347 assert_eq!(
9348 tokio::time::timeout(Duration::from_secs(1), update_rx)
9349 .await
9350 .expect("writer update should not remain queued behind a stalled send")
9351 .unwrap(),
9352 1
9353 );
9354
9355 connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
9356 state_notify.notify_waiters();
9357 tokio::time::advance(Duration::from_millis(CONNECTION_STATE_CHECK_INTERVAL_MS)).await;
9358
9359 for name in ["sentinel-1", "sentinel-2"] {
9365 let (sentinel_tx, sentinel_rx) = tokio::sync::oneshot::channel();
9366 writer_tx
9367 .send_control(WriterCommand::SendOnConnection {
9368 message: Message::text(name),
9369 connection_epoch: 1,
9370 response_tx: sentinel_tx,
9371 })
9372 .unwrap();
9373
9374 recording_state.recorded_notify.notified().await;
9375 sentinel_rx
9376 .await
9377 .unwrap()
9378 .expect("the sentinel should send on the replacement connection");
9379 }
9380
9381 assert_eq!(
9382 recorded.lock().as_slice(),
9383 &[Message::text("sentinel-1"), Message::text("sentinel-2")],
9384 "an ownership-bound message must never be replayed after its deadline expires"
9385 );
9386
9387 connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
9388 state_notify.notify_waiters();
9389 drop(writer_tx);
9390 write_task.await.unwrap();
9391 }
9392
9393 #[rstest]
9394 #[tokio::test(start_paused = true)]
9395 async fn test_stalled_websocket_replay_reconnects_and_retries_buffer() {
9396 let initial_messages = Arc::new(Mutex::new(Vec::new()));
9397
9398 let initial_recording_state = Arc::new(RecordingState {
9399 messages: initial_messages,
9400 recorded_notify: tokio::sync::Notify::new(),
9401 });
9402
9403 let transport: BoxedWsTransport = Box::pin(RecordingTransport {
9404 state: initial_recording_state,
9405 });
9406
9407 let (writer, _reader) = transport.split();
9408 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Reconnect.as_u8()));
9409 let state_notify = Arc::new(tokio::sync::Notify::new());
9410 let auth_tracker = Arc::new(OnceLock::new());
9411 let reconnect_buffer_waits_for_auth = Arc::new(AtomicBool::new(false));
9412 let (writer_tx, writer_rx) = crate::writer::channel(crate::writer::DEFAULT_WRITER_CAPACITY);
9413 let write_task = WebSocketClientInner::spawn_write_task(
9414 Arc::clone(&connection_state),
9415 Arc::clone(&state_notify),
9416 Arc::new(AtomicBool::new(true)),
9417 writer,
9418 writer_rx,
9419 Arc::new(AtomicU64::new(0)),
9420 Arc::clone(&auth_tracker),
9421 reconnect_buffer_waits_for_auth,
9422 None,
9423 );
9424
9425 writer_tx
9426 .send(WriterCommand::Send(Message::text("buffered-message")))
9427 .unwrap();
9428
9429 let blocking_state = Arc::new(BlockingFailState::default());
9430
9431 let transport: BoxedWsTransport = Box::pin(BlockingFailTransport {
9432 state: Arc::clone(&blocking_state),
9433 });
9434
9435 let (blocking_writer, _reader) = transport.split();
9436 let (blocking_tx, blocking_rx) = tokio::sync::oneshot::channel();
9437 writer_tx
9438 .send(WriterCommand::Update(blocking_writer, blocking_tx))
9439 .unwrap();
9440 assert_eq!(blocking_rx.await.unwrap(), 1);
9441
9442 connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
9443 state_notify.notify_waiters();
9444 tokio::time::advance(Duration::from_millis(CONNECTION_STATE_CHECK_INTERVAL_MS)).await;
9445 blocking_state.send_entered_notify.notified().await;
9446
9447 let recorded = Arc::new(Mutex::new(Vec::new()));
9448
9449 let recording_state = Arc::new(RecordingState {
9450 messages: Arc::clone(&recorded),
9451 recorded_notify: tokio::sync::Notify::new(),
9452 });
9453
9454 let transport: BoxedWsTransport = Box::pin(RecordingTransport {
9455 state: Arc::clone(&recording_state),
9456 });
9457
9458 let (new_writer, _reader) = transport.split();
9459 let (update_tx, update_rx) = tokio::sync::oneshot::channel();
9460 writer_tx
9461 .send(WriterCommand::Update(new_writer, update_tx))
9462 .unwrap();
9463
9464 tokio::time::advance(Duration::from_secs(GRACEFUL_SHUTDOWN_TIMEOUT_SECS)).await;
9465 assert_eq!(
9466 tokio::time::timeout(Duration::from_secs(1), update_rx)
9467 .await
9468 .expect("writer update should not remain queued behind stalled replay")
9469 .unwrap(),
9470 2,
9471 "the second replacement sink should install as connection epoch 2"
9472 );
9473 assert_eq!(
9474 ConnectionMode::from_atomic(&connection_state),
9475 ConnectionMode::Reconnect
9476 );
9477
9478 connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
9479 state_notify.notify_waiters();
9480 tokio::time::advance(Duration::from_millis(CONNECTION_STATE_CHECK_INTERVAL_MS)).await;
9481 recording_state.recorded_notify.notified().await;
9482 assert_eq!(
9483 recorded.lock().as_slice(),
9484 &[Message::text("buffered-message")]
9485 );
9486
9487 connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
9488 state_notify.notify_waiters();
9489 drop(writer_tx);
9490 write_task.await.unwrap();
9491 }
9492
9493 #[rstest]
9494 #[tokio::test]
9495 async fn test_new_with_writer_rejects_zero_heartbeat() {
9496 let transport: BoxedWsTransport = Box::pin(BlockingFailTransport {
9499 state: Arc::new(BlockingFailState::default()),
9500 });
9501
9502 let (writer, _reader) = transport.split();
9503
9504 let config = WebSocketConfig {
9505 url: "ws://127.0.0.1:1".to_string(),
9506 headers: vec![],
9507 heartbeat_interval_secs: Some(0),
9508 heartbeat_payload: None,
9509 connect_timeout_ms: None,
9510 reconnect_delay_initial_ms: None,
9511 reconnect_delay_max_ms: None,
9512 reconnect_backoff_factor: None,
9513 reconnect_jitter_ms: None,
9514 reconnect_max_attempts: None,
9515 heartbeat_timeout_secs: None,
9516 idle_timeout_ms: None,
9517 writer_capacity: None,
9518 backend: TransportBackend::Tungstenite,
9519 proxy_url: None,
9520 max_message_size_bytes: None,
9521 max_frame_size_bytes: None,
9522 };
9523
9524 let err = WebSocketClientInner::new_with_writer(config, writer)
9525 .await
9526 .expect_err("zero heartbeat should be rejected in stream mode");
9527 assert!(
9528 err.to_string()
9529 .contains("Heartbeat interval cannot be zero"),
9530 "error should mention zero heartbeat, was: {err}"
9531 );
9532 }
9533
9534 #[rstest]
9535 #[tokio::test]
9536 async fn test_connect_times_out_on_silent_server() {
9537 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
9540 let port = listener.local_addr().unwrap().port();
9541
9542 let server = task::spawn(async move {
9543 if let Ok((_stream, _)) = listener.accept().await {
9545 sleep(Duration::from_secs(30)).await;
9546 }
9547 });
9548
9549 let (handler, _rx) = channel_message_handler();
9550
9551 let config = WebSocketConfig {
9552 url: format!("ws://127.0.0.1:{port}"),
9553 headers: vec![],
9554 heartbeat_interval_secs: None,
9555 heartbeat_payload: None,
9556 connect_timeout_ms: Some(500),
9557 reconnect_delay_initial_ms: Some(50),
9558 reconnect_delay_max_ms: Some(100),
9559 reconnect_backoff_factor: Some(1.0),
9560 reconnect_jitter_ms: Some(0),
9561 reconnect_max_attempts: None,
9562 heartbeat_timeout_secs: None,
9563 idle_timeout_ms: None,
9564 writer_capacity: None,
9565 backend: TransportBackend::Tungstenite,
9566 proxy_url: None,
9567 max_message_size_bytes: None,
9568 max_frame_size_bytes: None,
9569 };
9570
9571 let result = tokio::time::timeout(
9572 Duration::from_secs(5),
9573 WebSocketClient::builder()
9574 .config(config)
9575 .message_handler(handler)
9576 .connect(),
9577 )
9578 .await
9579 .expect("connect should not hang on a silent server");
9580
9581 assert!(result.is_err(), "connect should fail with a timeout error");
9582 let err_msg = result.unwrap_err().to_string();
9583 assert!(
9584 err_msg.contains("timed out"),
9585 "error should mention the timeout, was: {err_msg}"
9586 );
9587
9588 server.abort();
9589 }
9590
9591 #[rstest]
9592 #[tokio::test]
9593 async fn test_reconnect_succeeds_with_timeout_shorter_than_swap_ceremony() {
9594 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
9600 let port = listener.local_addr().unwrap().port();
9601 let (release_first_tx, release_first_rx) = tokio::sync::oneshot::channel();
9602
9603 let server = task::spawn(async move {
9604 if let Ok((stream, _)) = listener.accept().await
9607 && let Ok(ws) = accept_async(stream).await
9608 {
9609 let _ = release_first_rx.await;
9610 drop(ws);
9611 }
9612
9613 loop {
9616 let Ok((stream, _)) = listener.accept().await else {
9617 break;
9618 };
9619
9620 if let Ok(mut ws) = accept_async(stream).await {
9621 loop {
9622 if ws
9623 .send(WsMessage::Text("reconnected-msg".to_string().into()))
9624 .await
9625 .is_err()
9626 {
9627 break;
9628 }
9629
9630 sleep(Duration::from_millis(50)).await;
9631 }
9632 }
9633 }
9634 });
9635
9636 let (handler, mut rx) = channel_message_handler();
9637
9638 let config = WebSocketConfig {
9639 url: format!("ws://127.0.0.1:{port}"),
9640 headers: vec![],
9641 heartbeat_interval_secs: None,
9642 heartbeat_payload: None,
9643 connect_timeout_ms: Some(150), reconnect_delay_initial_ms: Some(25),
9645 reconnect_delay_max_ms: Some(50),
9646 reconnect_backoff_factor: Some(1.0),
9647 reconnect_jitter_ms: Some(0),
9648 reconnect_max_attempts: None,
9649 heartbeat_timeout_secs: None,
9650 idle_timeout_ms: None,
9651 writer_capacity: None,
9652 backend: TransportBackend::Tungstenite,
9653 proxy_url: None,
9654 max_message_size_bytes: None,
9655 max_frame_size_bytes: None,
9656 };
9657
9658 let client = WebSocketClient::builder()
9659 .config(config)
9660 .message_handler(handler)
9661 .connect()
9662 .await
9663 .expect("Initial WebSocket connection should succeed before triggering reconnect");
9664 wait_until_async(|| async { client.is_active() }, Duration::from_secs(2)).await;
9665 release_first_tx
9666 .send(())
9667 .expect("server should still be holding the first connection");
9668
9669 let received = tokio::time::timeout(Duration::from_secs(5), async {
9670 loop {
9671 match rx.recv().await {
9672 Some(WsMessage::Text(text)) if text.as_str() == "reconnected-msg" => {
9673 return true;
9674 }
9675 Some(_) => {}
9676 None => return false,
9677 }
9678 }
9679 })
9680 .await
9681 .expect("Timed out waiting for a message from the replacement WebSocket connection");
9682
9683 assert!(
9684 received,
9685 "Replacement WebSocket message channel closed before delivering the reconnect message"
9686 );
9687 assert!(
9688 client.is_active(),
9689 "WebSocket client should be active after delivering the reconnect message"
9690 );
9691
9692 client.disconnect().await;
9693 server.abort();
9694 }
9695
9696 #[rstest]
9697 #[tokio::test]
9698 async fn test_idle_timeout_fires_under_ping_flood() {
9699 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
9703 let port = listener.local_addr().unwrap().port();
9704 let server = task::spawn(async move {
9705 let (stream, _) = listener.accept().await.unwrap();
9706 let mut ws = accept_async(stream).await.unwrap();
9707
9708 for _ in 0..600 {
9709 sleep(Duration::from_millis(5)).await;
9710
9711 if ws.send(WsMessage::Ping(Vec::new().into())).await.is_err() {
9712 break;
9713 }
9714 }
9715 });
9716
9717 let (handler, _rx) = channel_message_handler();
9718
9719 let config = WebSocketConfig {
9720 url: format!("ws://127.0.0.1:{port}"),
9721 headers: vec![],
9722 heartbeat_interval_secs: None,
9723 heartbeat_payload: None,
9724 connect_timeout_ms: Some(2_000),
9725 reconnect_delay_initial_ms: Some(50),
9726 reconnect_delay_max_ms: Some(100),
9727 reconnect_backoff_factor: Some(1.0),
9728 reconnect_jitter_ms: Some(0),
9729 reconnect_max_attempts: Some(1),
9730 heartbeat_timeout_secs: None,
9731 idle_timeout_ms: Some(500),
9732 writer_capacity: None,
9733 backend: TransportBackend::Tungstenite,
9734 proxy_url: None,
9735 max_message_size_bytes: None,
9736 max_frame_size_bytes: None,
9737 };
9738
9739 let client = WebSocketClient::builder()
9740 .config(config)
9741 .message_handler(handler)
9742 .connect()
9743 .await
9744 .unwrap();
9745
9746 assert!(client.is_active());
9747
9748 wait_until_async(
9749 || async { client.is_reconnecting() || client.is_disconnected() },
9750 Duration::from_millis(1_500),
9751 )
9752 .await;
9753
9754 assert!(
9755 !client.is_active(),
9756 "Client should not be active after idle timeout under a ping flood"
9757 );
9758
9759 client.disconnect().await;
9760 server.abort();
9761 }
9762
9763 #[rstest]
9764 #[tokio::test]
9765 async fn test_send_failure_does_not_overwrite_disconnect() {
9766 let state = Arc::new(BlockingFailState::default());
9767
9768 let transport: BoxedWsTransport = Box::pin(BlockingFailTransport {
9769 state: Arc::clone(&state),
9770 });
9771
9772 let (writer, _reader) = transport.split();
9773
9774 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
9775 let state_notify = Arc::new(tokio::sync::Notify::new());
9776 let auth_tracker = Arc::new(OnceLock::new());
9777 let reconnect_buffer_waits_for_auth = Arc::new(AtomicBool::new(false));
9778 let connection_epoch = Arc::new(AtomicU64::new(0));
9779
9780 let (writer_tx, writer_rx) = crate::writer::channel(crate::writer::DEFAULT_WRITER_CAPACITY);
9781 let write_task = WebSocketClientInner::spawn_write_task(
9782 Arc::clone(&connection_state),
9783 Arc::clone(&state_notify),
9784 Arc::new(AtomicBool::new(true)),
9785 writer,
9786 writer_rx,
9787 connection_epoch,
9788 Arc::clone(&auth_tracker),
9789 Arc::clone(&reconnect_buffer_waits_for_auth),
9790 None,
9791 );
9792
9793 writer_tx
9794 .send(WriterCommand::Send(Message::text("doomed")))
9795 .unwrap();
9796
9797 wait_until_async(
9799 || {
9800 let state = Arc::clone(&state);
9801 async move { state.send_entered.load(Ordering::SeqCst) }
9802 },
9803 Duration::from_secs(2),
9804 )
9805 .await;
9806
9807 connection_state.store(ConnectionMode::Disconnect.as_u8(), Ordering::SeqCst);
9810 state.trigger_failure();
9811
9812 tokio::time::timeout(Duration::from_secs(2), write_task)
9813 .await
9814 .expect("write task should exit after disconnect")
9815 .unwrap();
9816
9817 assert_eq!(
9818 ConnectionMode::from_atomic(&connection_state),
9819 ConnectionMode::Disconnect,
9820 "Send failure must not resurrect a disconnecting client into RECONNECT"
9821 );
9822 }
9823
9824 #[tokio::test]
9825 async fn send_on_connection_write_failure_reports_broken_pipe_and_reconnects() {
9826 let state = Arc::new(BlockingFailState::default());
9827
9828 let transport: BoxedWsTransport = Box::pin(BlockingFailTransport {
9829 state: Arc::clone(&state),
9830 });
9831
9832 let (writer, _reader) = transport.split();
9833 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
9834 let state_notify = Arc::new(tokio::sync::Notify::new());
9835 let connection_epoch = Arc::new(AtomicU64::new(0));
9836 let (writer_tx, writer_rx) = crate::writer::channel(crate::writer::DEFAULT_WRITER_CAPACITY);
9837 let write_task = WebSocketClientInner::spawn_write_task(
9838 Arc::clone(&connection_state),
9839 Arc::clone(&state_notify),
9840 Arc::new(AtomicBool::new(true)),
9841 writer,
9842 writer_rx,
9843 Arc::clone(&connection_epoch),
9844 Arc::new(OnceLock::new()),
9845 Arc::new(AtomicBool::new(false)),
9846 None,
9847 );
9848
9849 let (response_tx, response_rx) = tokio::sync::oneshot::channel();
9850 writer_tx
9851 .send_control(WriterCommand::SendOnConnection {
9852 message: Message::text("doomed"),
9853 connection_epoch: 0,
9854 response_tx,
9855 })
9856 .unwrap();
9857
9858 wait_until_async(
9859 || {
9860 let state = Arc::clone(&state);
9861 async move { state.send_entered.load(Ordering::SeqCst) }
9862 },
9863 Duration::from_secs(2),
9864 )
9865 .await;
9866
9867 state.trigger_failure();
9868
9869 match response_rx.await.unwrap().unwrap_err() {
9870 SendError::BrokenPipe(message) => assert_eq!(message, "connection reset"),
9871 other => panic!("expected broken-pipe send error, was {other:?}"),
9872 }
9873
9874 wait_until_async(
9875 || async {
9876 ConnectionMode::from_atomic(&connection_state) == ConnectionMode::Reconnect
9877 },
9878 Duration::from_secs(2),
9879 )
9880 .await;
9881
9882 assert_eq!(
9883 ConnectionMode::from_atomic(&connection_state),
9884 ConnectionMode::Reconnect,
9885 );
9886 assert_eq!(connection_epoch.load(Ordering::Acquire), 0);
9887
9888 connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
9889 state_notify.notify_waiters();
9890 drop(writer_tx);
9891 write_task.await.unwrap();
9892 }
9893
9894 #[tokio::test]
9895 async fn send_on_connection_rejects_stale_epoch_without_replay() {
9896 let server = RecordingServer::setup().await;
9897 let url = format!("ws://127.0.0.1:{}", server.port);
9898 let (writer, _reader) = WebSocketClientInner::connect_with_server(
9899 &url,
9900 vec![],
9901 TransportBackend::Tungstenite,
9902 None,
9903 None,
9904 None,
9905 )
9906 .await
9907 .unwrap();
9908 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
9909 let state_notify = Arc::new(tokio::sync::Notify::new());
9910 let auth_tracker = Arc::new(OnceLock::new());
9911 let connection_epoch = Arc::new(AtomicU64::new(0));
9912 let (writer_tx, writer_rx) = crate::writer::channel(crate::writer::DEFAULT_WRITER_CAPACITY);
9913 let write_task = WebSocketClientInner::spawn_write_task(
9914 Arc::clone(&connection_state),
9915 Arc::clone(&state_notify),
9916 Arc::new(AtomicBool::new(true)),
9917 writer,
9918 writer_rx,
9919 Arc::clone(&connection_epoch),
9920 auth_tracker,
9921 Arc::new(AtomicBool::new(false)),
9922 None,
9923 );
9924
9925 let (response_tx, response_rx) = tokio::sync::oneshot::channel();
9926 writer_tx
9927 .send_control(WriterCommand::SendOnConnection {
9928 message: Message::text("epoch-0"),
9929 connection_epoch: 0,
9930 response_tx,
9931 })
9932 .unwrap();
9933
9934 response_rx.await.unwrap().unwrap();
9935
9936 connection_state.store(ConnectionMode::Reconnect.as_u8(), Ordering::SeqCst);
9937 let (response_tx, response_rx) = tokio::sync::oneshot::channel();
9938 writer_tx
9939 .send_control(WriterCommand::SendOnConnection {
9940 message: Message::text("during-reconnect"),
9941 connection_epoch: 0,
9942 response_tx,
9943 })
9944 .unwrap();
9945
9946 assert!(matches!(
9947 response_rx.await.unwrap(),
9948 Err(SendError::ConnectionChanged),
9949 ));
9950
9951 let (replacement, _reader) = WebSocketClientInner::connect_with_server(
9952 &url,
9953 vec![],
9954 TransportBackend::Tungstenite,
9955 None,
9956 None,
9957 None,
9958 )
9959 .await
9960 .unwrap();
9961 let (update_tx, update_rx) = tokio::sync::oneshot::channel();
9962 writer_tx
9963 .send(WriterCommand::Update(replacement, update_tx))
9964 .unwrap();
9965 assert_eq!(update_rx.await.unwrap(), 1);
9966 connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
9967
9968 let (response_tx, response_rx) = tokio::sync::oneshot::channel();
9969 writer_tx
9970 .send_control(WriterCommand::SendOnConnection {
9971 message: Message::text("stale"),
9972 connection_epoch: 0,
9973 response_tx,
9974 })
9975 .unwrap();
9976
9977 assert!(matches!(
9978 response_rx.await.unwrap(),
9979 Err(SendError::ConnectionChanged),
9980 ));
9981
9982 let (response_tx, response_rx) = tokio::sync::oneshot::channel();
9983 writer_tx
9984 .send_control(WriterCommand::SendOnConnection {
9985 message: Message::text("epoch-1"),
9986 connection_epoch: 1,
9987 response_tx,
9988 })
9989 .unwrap();
9990
9991 response_rx.await.unwrap().unwrap();
9992
9993 connection_state.store(ConnectionMode::Reconnect.as_u8(), Ordering::SeqCst);
9994 let (second_replacement, _reader) = WebSocketClientInner::connect_with_server(
9995 &url,
9996 vec![],
9997 TransportBackend::Tungstenite,
9998 None,
9999 None,
10000 None,
10001 )
10002 .await
10003 .unwrap();
10004 let (update_tx, update_rx) = tokio::sync::oneshot::channel();
10005 writer_tx
10006 .send(WriterCommand::Update(second_replacement, update_tx))
10007 .unwrap();
10008 assert_eq!(update_rx.await.unwrap(), 2);
10009 connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
10010
10011 let (response_tx, response_rx) = tokio::sync::oneshot::channel();
10012 writer_tx
10013 .send_control(WriterCommand::SendOnConnection {
10014 message: Message::text("stale-after-second-reconnect"),
10015 connection_epoch: 1,
10016 response_tx,
10017 })
10018 .unwrap();
10019
10020 assert!(matches!(
10021 response_rx.await.unwrap(),
10022 Err(SendError::ConnectionChanged),
10023 ));
10024
10025 let (response_tx, response_rx) = tokio::sync::oneshot::channel();
10026 writer_tx
10027 .send_control(WriterCommand::SendOnConnection {
10028 message: Message::text("epoch-2"),
10029 connection_epoch: 2,
10030 response_tx,
10031 })
10032 .unwrap();
10033
10034 response_rx.await.unwrap().unwrap();
10035
10036 wait_until_async(
10037 || {
10038 let messages = Arc::clone(&server.messages);
10039 async move { messages.lock().await.len() == 3 }
10040 },
10041 Duration::from_secs(2),
10042 )
10043 .await;
10044
10045 assert_eq!(connection_epoch.load(Ordering::Acquire), 2);
10046 assert_eq!(
10047 server.messages().await,
10048 vec![
10049 "epoch-0".to_string(),
10050 "epoch-1".to_string(),
10051 "epoch-2".to_string(),
10052 ],
10053 );
10054
10055 connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
10056 state_notify.notify_waiters();
10057 drop(writer_tx);
10058 write_task.abort();
10059 }
10060
10061 #[rstest]
10062 fn test_reconnect_buffer_action_requires_active_mode() {
10063 let connection_state = AtomicU8::new(ConnectionMode::Active.as_u8());
10064 let reconnect_buffer_waits_for_auth = AtomicBool::new(true);
10065 let auth_tracker = Arc::new(OnceLock::new());
10066 let tracker = AuthTracker::new();
10067 tracker.succeed();
10068 auth_tracker.set(tracker).unwrap();
10069
10070 assert_eq!(
10071 WebSocketClientInner::reconnect_buffer_action(
10072 &reconnect_buffer_waits_for_auth,
10073 &auth_tracker,
10074 &connection_state,
10075 ),
10076 ReconnectBufferAction::Drain,
10077 );
10078
10079 connection_state.store(ConnectionMode::Reconnect.as_u8(), Ordering::SeqCst);
10080 assert_eq!(
10081 WebSocketClientInner::reconnect_buffer_action(
10082 &reconnect_buffer_waits_for_auth,
10083 &auth_tracker,
10084 &connection_state,
10085 ),
10086 ReconnectBufferAction::Wait,
10087 );
10088 }
10089
10090 #[tokio::test]
10091 async fn test_write_task_waits_for_auth_before_replaying_buffer() {
10092 use nautilus_common::testing::wait_until_async;
10093
10094 let server = RecordingServer::setup().await;
10095 let url = format!("ws://127.0.0.1:{}", server.port);
10096 let (writer, _reader) = WebSocketClientInner::connect_with_server(
10097 &url,
10098 vec![],
10099 TransportBackend::Tungstenite,
10100 None,
10101 None,
10102 None,
10103 )
10104 .await
10105 .unwrap();
10106
10107 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Reconnect.as_u8()));
10108 let state_notify = Arc::new(tokio::sync::Notify::new());
10109 let auth_tracker = Arc::new(OnceLock::new());
10110 let reconnect_buffer_waits_for_auth = Arc::new(AtomicBool::new(true));
10111 let connection_epoch = Arc::new(AtomicU64::new(0));
10112 let tracker = AuthTracker::new();
10113 auth_tracker.set(tracker.clone()).unwrap();
10114
10115 let (writer_tx, writer_rx) = crate::writer::channel(1);
10116 let write_task = WebSocketClientInner::spawn_write_task(
10117 Arc::clone(&connection_state),
10118 Arc::clone(&state_notify),
10119 Arc::new(AtomicBool::new(true)),
10120 writer,
10121 writer_rx,
10122 Arc::clone(&connection_epoch),
10123 Arc::clone(&auth_tracker),
10124 Arc::clone(&reconnect_buffer_waits_for_auth),
10125 None,
10126 );
10127
10128 writer_tx
10129 .send(WriterCommand::Send(Message::Text("stale".into())))
10130 .unwrap();
10131
10132 let (new_writer, _reader) = WebSocketClientInner::connect_with_server(
10133 &url,
10134 vec![],
10135 TransportBackend::Tungstenite,
10136 None,
10137 None,
10138 None,
10139 )
10140 .await
10141 .unwrap();
10142 let (tx, rx) = tokio::sync::oneshot::channel();
10143 writer_tx
10144 .send_update(WriterCommand::Update(new_writer, tx))
10145 .unwrap();
10146 assert_eq!(rx.await.unwrap(), 1);
10147
10148 connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
10149
10150 tokio::time::sleep(Duration::from_millis(300)).await;
10151 assert!(
10152 server.messages().await.is_empty(),
10153 "buffered messages should wait for re-authentication"
10154 );
10155
10156 assert!(matches!(
10157 writer_tx.send(WriterCommand::Send(Message::text("overflow"))),
10158 Err(SendError::BufferFull)
10159 ));
10160 let (response_tx, response_rx) = tokio::sync::oneshot::channel();
10161 writer_tx
10162 .send_control(WriterCommand::SendOnConnection {
10163 message: Message::text("authenticate"),
10164 connection_epoch: 1,
10165 response_tx,
10166 })
10167 .unwrap();
10168
10169 response_rx.await.unwrap().unwrap();
10170 tracker.succeed();
10171
10172 wait_until_async(
10173 || {
10174 let messages = Arc::clone(&server.messages);
10175 async move { messages.lock().await.len() == 2 }
10176 },
10177 Duration::from_secs(3),
10178 )
10179 .await;
10180
10181 assert_eq!(
10182 server.messages().await,
10183 vec!["authenticate".to_string(), "stale".to_string()]
10184 );
10185
10186 connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
10187 state_notify.notify_waiters();
10188 drop(writer_tx);
10189 write_task.abort();
10190 }
10191
10192 #[rstest]
10193 #[tokio::test]
10194 async fn test_write_task_buffers_send_behind_pending_replay() {
10195 use nautilus_common::testing::wait_until_async;
10196
10197 let server = RecordingServer::setup().await;
10198 let url = format!("ws://127.0.0.1:{}", server.port);
10199 let (writer, _reader) = WebSocketClientInner::connect_with_server(
10200 &url,
10201 vec![],
10202 TransportBackend::Tungstenite,
10203 None,
10204 None,
10205 None,
10206 )
10207 .await
10208 .unwrap();
10209
10210 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Reconnect.as_u8()));
10211 let state_notify = Arc::new(tokio::sync::Notify::new());
10212 let auth_tracker = Arc::new(OnceLock::new());
10213 let reconnect_buffer_waits_for_auth = Arc::new(AtomicBool::new(true));
10214 let connection_epoch = Arc::new(AtomicU64::new(0));
10215 let tracker = AuthTracker::new();
10216 auth_tracker.set(tracker.clone()).unwrap();
10217
10218 let (writer_tx, writer_rx) = crate::writer::channel(2);
10221 let write_task = WebSocketClientInner::spawn_write_task(
10222 Arc::clone(&connection_state),
10223 Arc::clone(&state_notify),
10224 Arc::new(AtomicBool::new(true)),
10225 writer,
10226 writer_rx,
10227 Arc::clone(&connection_epoch),
10228 Arc::clone(&auth_tracker),
10229 Arc::clone(&reconnect_buffer_waits_for_auth),
10230 None,
10231 );
10232
10233 writer_tx
10234 .send(WriterCommand::Send(Message::Text("stale".into())))
10235 .unwrap();
10236
10237 let (new_writer, _reader) = WebSocketClientInner::connect_with_server(
10238 &url,
10239 vec![],
10240 TransportBackend::Tungstenite,
10241 None,
10242 None,
10243 None,
10244 )
10245 .await
10246 .unwrap();
10247 let (tx, rx) = tokio::sync::oneshot::channel();
10248 writer_tx
10249 .send_update(WriterCommand::Update(new_writer, tx))
10250 .unwrap();
10251 assert_eq!(rx.await.unwrap(), 1);
10252
10253 connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
10255 tokio::time::sleep(Duration::from_millis(300)).await;
10256 assert!(
10257 server.messages().await.is_empty(),
10258 "buffered message should wait for re-authentication"
10259 );
10260
10261 writer_tx
10264 .send(WriterCommand::Send(Message::Text("fresh".into())))
10265 .unwrap();
10266 tokio::time::sleep(Duration::from_millis(300)).await;
10267 assert!(
10268 server.messages().await.is_empty(),
10269 "a send arriving before the gate opens must not reach the connection, was {:?}",
10270 server.messages().await
10271 );
10272
10273 let (response_tx, response_rx) = tokio::sync::oneshot::channel();
10274 writer_tx
10275 .send_control(WriterCommand::SendOnConnection {
10276 message: Message::text("authenticate"),
10277 connection_epoch: 1,
10278 response_tx,
10279 })
10280 .unwrap();
10281
10282 response_rx.await.unwrap().unwrap();
10283 tracker.succeed();
10284
10285 wait_until_async(
10286 || {
10287 let messages = Arc::clone(&server.messages);
10288 async move { messages.lock().await.len() == 3 }
10289 },
10290 Duration::from_secs(3),
10291 )
10292 .await;
10293
10294 assert_eq!(
10295 server.messages().await,
10296 vec![
10297 "authenticate".to_string(),
10298 "stale".to_string(),
10299 "fresh".to_string(),
10300 ],
10301 "replay must preserve FIFO order behind the authentication handshake"
10302 );
10303
10304 connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
10305 state_notify.notify_waiters();
10306 drop(writer_tx);
10307 write_task.abort();
10308 }
10309
10310 #[rstest]
10311 #[tokio::test]
10312 async fn test_write_task_control_send_reaches_connection_behind_pending_replay() {
10313 use nautilus_common::testing::wait_until_async;
10314
10315 let server = RecordingServer::setup().await;
10316 let url = format!("ws://127.0.0.1:{}", server.port);
10317 let (writer, _reader) = WebSocketClientInner::connect_with_server(
10318 &url,
10319 vec![],
10320 TransportBackend::Tungstenite,
10321 None,
10322 None,
10323 None,
10324 )
10325 .await
10326 .unwrap();
10327
10328 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Reconnect.as_u8()));
10329 let state_notify = Arc::new(tokio::sync::Notify::new());
10330 let auth_tracker = Arc::new(OnceLock::new());
10331 let reconnect_buffer_waits_for_auth = Arc::new(AtomicBool::new(true));
10332 let connection_epoch = Arc::new(AtomicU64::new(0));
10333 let tracker = AuthTracker::new();
10334 auth_tracker.set(tracker.clone()).unwrap();
10335
10336 let (writer_tx, writer_rx) = crate::writer::channel(2);
10337 let write_task = WebSocketClientInner::spawn_write_task(
10338 Arc::clone(&connection_state),
10339 Arc::clone(&state_notify),
10340 Arc::new(AtomicBool::new(true)),
10341 writer,
10342 writer_rx,
10343 Arc::clone(&connection_epoch),
10344 Arc::clone(&auth_tracker),
10345 Arc::clone(&reconnect_buffer_waits_for_auth),
10346 None,
10347 );
10348
10349 writer_tx
10350 .send(WriterCommand::Send(Message::Text("stale".into())))
10351 .unwrap();
10352
10353 let (new_writer, _reader) = WebSocketClientInner::connect_with_server(
10354 &url,
10355 vec![],
10356 TransportBackend::Tungstenite,
10357 None,
10358 None,
10359 None,
10360 )
10361 .await
10362 .unwrap();
10363 let (tx, rx) = tokio::sync::oneshot::channel();
10364 writer_tx
10365 .send_update(WriterCommand::Update(new_writer, tx))
10366 .unwrap();
10367 assert_eq!(rx.await.unwrap(), 1);
10368
10369 connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
10371 tokio::time::sleep(Duration::from_millis(300)).await;
10372 assert!(
10373 server.messages().await.is_empty(),
10374 "buffered message should wait for re-authentication"
10375 );
10376
10377 writer_tx
10381 .send_control(WriterCommand::Send(Message::Pong(
10382 b"keepalive".to_vec().into(),
10383 )))
10384 .unwrap();
10385
10386 wait_until_async(
10387 || {
10388 let control_frames = Arc::clone(&server.control_frames);
10389 async move { !control_frames.lock().await.is_empty() }
10390 },
10391 Duration::from_secs(3),
10392 )
10393 .await;
10394
10395 assert_eq!(
10396 server.control_frames().await,
10397 vec!["pong:keepalive".to_string()],
10398 "a control frame sent during the wait must reach the new connection without replay"
10399 );
10400 assert!(
10401 server.messages().await.is_empty(),
10402 "ordinary sends must still wait for re-authentication, was {:?}",
10403 server.messages().await
10404 );
10405
10406 let (response_tx, response_rx) = tokio::sync::oneshot::channel();
10407 writer_tx
10408 .send_control(WriterCommand::SendOnConnection {
10409 message: Message::text("authenticate"),
10410 connection_epoch: 1,
10411 response_tx,
10412 })
10413 .unwrap();
10414
10415 response_rx.await.unwrap().unwrap();
10416 tracker.succeed();
10417
10418 wait_until_async(
10419 || {
10420 let messages = Arc::clone(&server.messages);
10421 async move { messages.lock().await.len() == 2 }
10422 },
10423 Duration::from_secs(3),
10424 )
10425 .await;
10426
10427 assert_eq!(
10428 server.messages().await,
10429 vec!["authenticate".to_string(), "stale".to_string()],
10430 "the control frame must not disturb ordinary replay ordering"
10431 );
10432
10433 connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
10434 state_notify.notify_waiters();
10435 drop(writer_tx);
10436 write_task.abort();
10437 }
10438
10439 #[tokio::test]
10440 async fn test_write_task_discards_buffer_after_auth_failure() {
10441 let server = RecordingServer::setup().await;
10442 let url = format!("ws://127.0.0.1:{}", server.port);
10443 let (writer, _reader) = WebSocketClientInner::connect_with_server(
10444 &url,
10445 vec![],
10446 TransportBackend::Tungstenite,
10447 None,
10448 None,
10449 None,
10450 )
10451 .await
10452 .unwrap();
10453
10454 let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Reconnect.as_u8()));
10455 let state_notify = Arc::new(tokio::sync::Notify::new());
10456 let auth_tracker = Arc::new(OnceLock::new());
10457 let reconnect_buffer_waits_for_auth = Arc::new(AtomicBool::new(true));
10458 let connection_epoch = Arc::new(AtomicU64::new(0));
10459 let tracker = AuthTracker::new();
10460 auth_tracker.set(tracker.clone()).unwrap();
10461
10462 let (writer_tx, writer_rx) = crate::writer::channel(crate::writer::DEFAULT_WRITER_CAPACITY);
10463 let write_task = WebSocketClientInner::spawn_write_task(
10464 Arc::clone(&connection_state),
10465 Arc::clone(&state_notify),
10466 Arc::new(AtomicBool::new(true)),
10467 writer,
10468 writer_rx,
10469 Arc::clone(&connection_epoch),
10470 Arc::clone(&auth_tracker),
10471 Arc::clone(&reconnect_buffer_waits_for_auth),
10472 None,
10473 );
10474
10475 writer_tx
10476 .send(WriterCommand::Send(Message::Text("stale".into())))
10477 .unwrap();
10478
10479 let (new_writer, _reader) = WebSocketClientInner::connect_with_server(
10480 &url,
10481 vec![],
10482 TransportBackend::Tungstenite,
10483 None,
10484 None,
10485 None,
10486 )
10487 .await
10488 .unwrap();
10489 let (tx, rx) = tokio::sync::oneshot::channel();
10490 writer_tx
10491 .send(WriterCommand::Update(new_writer, tx))
10492 .unwrap();
10493 assert_eq!(rx.await.unwrap(), 1);
10494
10495 connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
10496 tracker.fail("rejected");
10497 tracker.invalidate();
10498 assert_eq!(tracker.auth_state(), AuthState::Failed);
10499 tokio::time::sleep(Duration::from_millis(300)).await;
10500 assert!(
10501 server.messages().await.is_empty(),
10502 "buffered messages should be discarded after authentication failure"
10503 );
10504
10505 let _auth_receiver = tracker.begin();
10506 tracker.succeed();
10507 tokio::time::sleep(Duration::from_millis(300)).await;
10508 assert!(
10509 server.messages().await.is_empty(),
10510 "discarded buffered messages should not replay on a later auth success"
10511 );
10512
10513 connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
10514 state_notify.notify_waiters();
10515 drop(writer_tx);
10516 write_task.abort();
10517 }
10518
10519 #[rstest]
10520 #[tokio::test]
10521 async fn test_zero_idle_timeout_rejected() {
10522 let (handler, _rx) = channel_message_handler();
10523
10524 let config = WebSocketConfig {
10525 url: "ws://127.0.0.1:9999".to_string(),
10526 headers: vec![],
10527 heartbeat_interval_secs: None,
10528 heartbeat_payload: None,
10529 connect_timeout_ms: None,
10530 reconnect_delay_initial_ms: None,
10531 reconnect_delay_max_ms: None,
10532 reconnect_backoff_factor: None,
10533 reconnect_jitter_ms: None,
10534 reconnect_max_attempts: None,
10535 heartbeat_timeout_secs: None,
10536 idle_timeout_ms: Some(0),
10537 writer_capacity: None,
10538 backend: TransportBackend::Tungstenite,
10539 proxy_url: None,
10540 max_message_size_bytes: None,
10541 max_frame_size_bytes: None,
10542 };
10543
10544 let result = WebSocketClient::builder()
10545 .config(config)
10546 .message_handler(handler)
10547 .connect()
10548 .await;
10549
10550 assert!(result.is_err(), "Zero idle timeout should be rejected");
10551 let err_msg = result.unwrap_err().to_string();
10552 assert!(
10553 err_msg.contains("idle_timeout_ms"),
10554 "Error should name the offending field, was: {err_msg}"
10555 );
10556 }
10557
10558 #[rstest]
10559 #[tokio::test]
10560 async fn test_zero_heartbeat_timeout_rejected() {
10561 let (handler, _rx) = channel_message_handler();
10562
10563 let config = WebSocketConfig {
10564 url: "ws://127.0.0.1:9999".to_string(),
10565 headers: vec![],
10566 heartbeat_interval_secs: Some(30),
10567 heartbeat_payload: None,
10568 connect_timeout_ms: None,
10569 reconnect_delay_initial_ms: None,
10570 reconnect_delay_max_ms: None,
10571 reconnect_backoff_factor: None,
10572 reconnect_jitter_ms: None,
10573 reconnect_max_attempts: None,
10574 heartbeat_timeout_secs: Some(0),
10575 idle_timeout_ms: None,
10576 writer_capacity: None,
10577 backend: TransportBackend::Tungstenite,
10578 proxy_url: None,
10579 max_message_size_bytes: None,
10580 max_frame_size_bytes: None,
10581 };
10582
10583 let result = WebSocketClient::builder()
10584 .config(config)
10585 .message_handler(handler)
10586 .connect()
10587 .await;
10588
10589 assert!(result.is_err(), "Zero heartbeat timeout should be rejected");
10590 let err_msg = result.unwrap_err().to_string();
10591 assert!(
10592 err_msg.contains("heartbeat_timeout_secs"),
10593 "Error should name the offending field, was: {err_msg}"
10594 );
10595 }
10596
10597 #[cfg(all(feature = "transport-sockudo", not(feature = "turmoil")))]
10598 #[rstest]
10599 #[tokio::test]
10600 async fn test_sockudo_backend_rejects_reserved_headers_before_connect() {
10601 let (handler, _rx) = channel_message_handler();
10602
10603 let config = WebSocketConfig {
10604 url: "ws://127.0.0.1:1".to_string(),
10605 headers: vec![("Host".to_string(), "example.com".to_string())],
10606 heartbeat_interval_secs: None,
10607 heartbeat_payload: None,
10608 connect_timeout_ms: None,
10609 reconnect_delay_initial_ms: None,
10610 reconnect_delay_max_ms: None,
10611 reconnect_backoff_factor: None,
10612 reconnect_jitter_ms: None,
10613 reconnect_max_attempts: None,
10614 heartbeat_timeout_secs: None,
10615 idle_timeout_ms: None,
10616 writer_capacity: None,
10617 backend: TransportBackend::Sockudo,
10618 proxy_url: None,
10619 max_message_size_bytes: None,
10620 max_frame_size_bytes: None,
10621 };
10622
10623 let err = WebSocketClient::builder()
10624 .config(config)
10625 .message_handler(handler)
10626 .connect()
10627 .await
10628 .expect_err("reserved header should fail before TCP connect");
10629
10630 assert!(
10631 err.to_string()
10632 .contains("reserved upgrade header not allowed in extra_headers"),
10633 "expected reserved-header failure, was: {err}"
10634 );
10635 }
10636
10637 #[cfg(all(feature = "transport-sockudo", target_os = "linux"))]
10639 #[rstest]
10640 #[tokio::test]
10641 async fn test_sockudo_reconnect_recovers_from_retryable_upgrade_rejection() {
10642 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
10643 let port = listener.local_addr().unwrap().port();
10644
10645 let server = task::spawn(async move {
10646 let (stream, _) = listener.accept().await.unwrap();
10648 drop(accept_async(stream).await.unwrap());
10649
10650 let (mut stream, _) = listener.accept().await.unwrap();
10652 let _request = read_http_request(&mut stream).await;
10653 stream
10654 .write_all(b"HTTP/1.1 502 Bad Gateway\r\n\r\n")
10655 .await
10656 .unwrap();
10657 stream.flush().await.unwrap();
10658 sleep(Duration::from_millis(50)).await;
10659 drop(stream);
10660
10661 let (stream, _) = listener.accept().await.unwrap();
10663 let mut ws = accept_async(stream).await.unwrap();
10664 ws.send(WsMessage::Text("recovered".to_string().into()))
10665 .await
10666 .unwrap();
10667 sleep(Duration::from_secs(1)).await;
10668 });
10669
10670 let (handler, mut rx) = channel_message_handler();
10671
10672 let config = WebSocketConfig {
10673 url: format!(
10674 "ws://user-secret:password-secret@127.0.0.1:{port}/path-secret?token=query-secret"
10675 ),
10676 headers: vec![],
10677 heartbeat_interval_secs: None,
10678 heartbeat_payload: None,
10679 connect_timeout_ms: Some(2_000),
10680 reconnect_delay_initial_ms: Some(50),
10681 reconnect_delay_max_ms: Some(100),
10682 reconnect_backoff_factor: Some(1.0),
10683 reconnect_jitter_ms: Some(0),
10684 reconnect_max_attempts: None,
10685 heartbeat_timeout_secs: None,
10686 idle_timeout_ms: None,
10687 writer_capacity: None,
10688 backend: TransportBackend::Sockudo,
10689 proxy_url: None,
10690 max_message_size_bytes: None,
10691 max_frame_size_bytes: None,
10692 };
10693
10694 let capture = capture_logs_for(RECONNECT_LOG_TARGETS).await;
10698
10699 let client = WebSocketClient::builder()
10700 .config(config)
10701 .message_handler(handler)
10702 .connect()
10703 .await
10704 .expect("sockudo connect should succeed");
10705
10706 let recovered = tokio::time::timeout(Duration::from_secs(10), async {
10707 loop {
10708 if let Ok(WsMessage::Text(text)) = rx.try_recv()
10709 && text.as_str() == "recovered"
10710 {
10711 return;
10712 }
10713
10714 tokio::time::sleep(Duration::from_millis(10)).await;
10715 }
10716 })
10717 .await;
10718
10719 let messages = capture.messages();
10720 client.disconnect().await;
10721 server.abort();
10722
10723 recovered.expect("client should recover after the retryable upgrade rejection");
10724
10725 let errors: Vec<_> = messages
10726 .iter()
10727 .filter(|(level, _)| *level == Level::Error)
10728 .collect();
10729 assert!(
10730 errors.is_empty(),
10731 "a retryable rejection must not log at ERROR, was: {errors:?}"
10732 );
10733
10734 assert!(
10736 messages.iter().any(|(level, message)| *level == Level::Warn
10737 && message.contains(&format!(
10738 "Sockudo handshake rejected by 127.0.0.1:{port} with retryable status 502"
10739 ))),
10740 "expected a warning naming the retryable rejection, was: {messages:?}"
10741 );
10742
10743 for secret in [
10744 "user-secret",
10745 "password-secret",
10746 "path-secret",
10747 "query-secret",
10748 ] {
10749 assert!(
10750 messages
10751 .iter()
10752 .all(|(_, message)| !message.contains(secret)),
10753 "logs must not contain {secret}, was: {messages:?}"
10754 );
10755 }
10756 }
10757
10758 #[cfg(all(feature = "transport-sockudo", not(feature = "turmoil")))]
10759 #[rstest]
10760 #[tokio::test]
10761 async fn test_sockudo_backend_replays_leftover_without_custom_headers() {
10762 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
10763 let port = listener.local_addr().unwrap().port();
10764 let server = task::spawn(async move {
10765 if let Ok((mut stream, _)) = listener.accept().await {
10766 let request = read_http_request(&mut stream).await;
10767 let request = String::from_utf8(request).unwrap();
10768 let sec_websocket_key = extract_header(&request, "Sec-WebSocket-Key").unwrap();
10769 let accept = sockudo_handshake::generate_accept_key(sec_websocket_key);
10770 let mut response = format!(
10771 concat!(
10772 "HTTP/1.1 101 Switching Protocols\r\n",
10773 "Upgrade: websocket\r\n",
10774 "Connection: Upgrade\r\n",
10775 "Sec-WebSocket-Accept: {}\r\n",
10776 "\r\n",
10777 ),
10778 accept
10779 )
10780 .into_bytes();
10781 response.extend_from_slice(b"\x81\x05hello");
10782 stream.write_all(&response).await.unwrap();
10783 }
10784 });
10785
10786 let (handler, mut rx) = channel_message_handler();
10787
10788 let config = WebSocketConfig {
10789 url: format!("ws://127.0.0.1:{port}/ws"),
10790 headers: vec![],
10791 heartbeat_interval_secs: None,
10792 heartbeat_payload: None,
10793 connect_timeout_ms: Some(2_000),
10794 reconnect_delay_initial_ms: Some(50),
10795 reconnect_delay_max_ms: Some(100),
10796 reconnect_backoff_factor: Some(1.0),
10797 reconnect_jitter_ms: Some(0),
10798 reconnect_max_attempts: None,
10799 heartbeat_timeout_secs: None,
10800 idle_timeout_ms: None,
10801 writer_capacity: None,
10802 backend: TransportBackend::Sockudo,
10803 proxy_url: None,
10804 max_message_size_bytes: None,
10805 max_frame_size_bytes: None,
10806 };
10807
10808 let client = WebSocketClient::builder()
10809 .config(config)
10810 .message_handler(handler)
10811 .connect()
10812 .await
10813 .expect("sockudo connect without custom headers");
10814
10815 let received = tokio::time::timeout(Duration::from_secs(3), async {
10816 loop {
10817 if let Ok(msg) = rx.try_recv() {
10818 return msg;
10819 }
10820
10821 tokio::time::sleep(Duration::from_millis(10)).await;
10822 }
10823 })
10824 .await
10825 .expect("did not receive leftover frame before timeout");
10826
10827 match received {
10828 WsMessage::Text(t) => assert_eq!(t.as_str(), "hello"),
10829 other => panic!("expected text, was {other:?}"),
10830 }
10831
10832 client.disconnect().await;
10833 tokio::time::timeout(Duration::from_secs(3), server)
10834 .await
10835 .expect("server did not close before timeout")
10836 .unwrap();
10837 }
10838
10839 #[cfg(all(feature = "transport-sockudo", not(feature = "turmoil")))]
10840 #[rstest]
10841 #[tokio::test]
10842 async fn test_sockudo_backend_sends_custom_headers() {
10843 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
10844 let port = listener.local_addr().unwrap().port();
10845 let server = task::spawn(async move {
10846 if let Ok((stream, _)) = listener.accept().await {
10847 let callback = HeaderAssertCallback {
10848 key: "X-Test".to_string(),
10849 value: HeaderValue::from_static("value"),
10850 };
10851
10852 if let Ok(mut ws) = accept_hdr_async(stream, callback).await {
10853 while let Some(Ok(msg)) = ws.next().await {
10854 if msg.is_text() || msg.is_binary() {
10855 if ws.send(msg).await.is_err() {
10856 break;
10857 }
10858
10859 continue;
10860 }
10861
10862 if msg.is_close() {
10863 let _ = ws.close(None).await;
10864 break;
10865 }
10866 }
10867 }
10868 }
10869 });
10870
10871 let (handler, mut rx) = channel_message_handler();
10872
10873 let config = WebSocketConfig {
10874 url: format!("ws://127.0.0.1:{port}"),
10875 headers: vec![("X-Test".to_string(), "value".to_string())],
10876 heartbeat_interval_secs: None,
10877 heartbeat_payload: None,
10878 connect_timeout_ms: Some(2_000),
10879 reconnect_delay_initial_ms: Some(50),
10880 reconnect_delay_max_ms: Some(100),
10881 reconnect_backoff_factor: Some(1.0),
10882 reconnect_jitter_ms: Some(0),
10883 reconnect_max_attempts: None,
10884 heartbeat_timeout_secs: None,
10885 idle_timeout_ms: None,
10886 writer_capacity: None,
10887 backend: TransportBackend::Sockudo,
10888 proxy_url: None,
10889 max_message_size_bytes: None,
10890 max_frame_size_bytes: None,
10891 };
10892
10893 let client = WebSocketClient::builder()
10894 .config(config)
10895 .message_handler(handler)
10896 .connect()
10897 .await
10898 .expect("sockudo connect with custom headers");
10899
10900 client.send_text("ping".to_string(), None).await.unwrap();
10901
10902 let received = tokio::time::timeout(Duration::from_secs(3), async {
10903 loop {
10904 if let Ok(msg) = rx.try_recv() {
10905 return msg;
10906 }
10907
10908 tokio::time::sleep(Duration::from_millis(10)).await;
10909 }
10910 })
10911 .await
10912 .expect("did not receive echo before timeout");
10913
10914 match received {
10915 WsMessage::Text(t) => assert_eq!(t.as_str(), "ping"),
10916 other => panic!("expected text, was {other:?}"),
10917 }
10918
10919 client.disconnect().await;
10920 tokio::time::timeout(Duration::from_secs(3), server)
10921 .await
10922 .expect("server did not close before timeout")
10923 .unwrap();
10924 }
10925
10926 #[cfg(all(feature = "transport-sockudo", not(feature = "turmoil")))]
10927 #[rstest]
10928 #[tokio::test]
10929 async fn test_sockudo_backend_round_trip_text() {
10930 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
10932 let port = listener.local_addr().unwrap().port();
10933 let server = task::spawn(async move {
10934 if let Ok((stream, _)) = listener.accept().await
10935 && let Ok(mut ws) = accept_async(stream).await
10936 {
10937 while let Some(Ok(msg)) = ws.next().await {
10938 match msg {
10939 WsMessage::Text(_) | WsMessage::Binary(_) => {
10940 if ws.send(msg).await.is_err() {
10941 break;
10942 }
10943 }
10944 WsMessage::Close(_) => {
10945 let _ = ws.close(None).await;
10946 break;
10947 }
10948 _ => {}
10949 }
10950 }
10951 }
10952 });
10953
10954 let (handler, mut rx) = channel_message_handler();
10955
10956 let config = WebSocketConfig {
10957 url: format!("ws://127.0.0.1:{port}"),
10958 headers: vec![],
10959 heartbeat_interval_secs: None,
10960 heartbeat_payload: None,
10961 connect_timeout_ms: Some(2_000),
10962 reconnect_delay_initial_ms: Some(50),
10963 reconnect_delay_max_ms: Some(100),
10964 reconnect_backoff_factor: Some(1.0),
10965 reconnect_jitter_ms: Some(0),
10966 reconnect_max_attempts: None,
10967 heartbeat_timeout_secs: None,
10968 idle_timeout_ms: None,
10969 writer_capacity: None,
10970 backend: TransportBackend::Sockudo,
10971 proxy_url: None,
10972 max_message_size_bytes: None,
10973 max_frame_size_bytes: None,
10974 };
10975
10976 let client = WebSocketClient::builder()
10977 .config(config)
10978 .message_handler(handler)
10979 .connect()
10980 .await
10981 .expect("sockudo connect");
10982
10983 client.send_text("ping".to_string(), None).await.unwrap();
10984
10985 let received = tokio::time::timeout(Duration::from_secs(3), async {
10986 loop {
10987 if let Ok(msg) = rx.try_recv() {
10988 return msg;
10989 }
10990
10991 tokio::time::sleep(Duration::from_millis(10)).await;
10992 }
10993 })
10994 .await
10995 .expect("did not receive echo before timeout");
10996
10997 match received {
10998 WsMessage::Text(t) => assert_eq!(t.as_str(), "ping"),
10999 other => panic!("expected text, was {other:?}"),
11000 }
11001
11002 client.disconnect().await;
11003 server.abort();
11004 }
11005
11006 #[cfg(all(feature = "transport-sockudo", not(feature = "turmoil")))]
11007 #[rstest]
11008 #[case::ws_default_port("ws://example.com/ws", "example.com", "example.com", 80, "/ws", false)]
11009 #[case::wss_default_port(
11010 "wss://example.com/ws",
11011 "example.com",
11012 "example.com",
11013 443,
11014 "/ws",
11015 true
11016 )]
11017 #[case::ws_explicit_default(
11020 "ws://example.com:80/ws",
11021 "example.com",
11022 "example.com",
11023 80,
11024 "/ws",
11025 false
11026 )]
11027 #[case::ws_non_default(
11028 "ws://example.com:8443/feed",
11029 "example.com",
11030 "example.com:8443",
11031 8443,
11032 "/feed",
11033 false
11034 )]
11035 #[case::wss_non_default(
11036 "wss://example.com:9443/feed",
11037 "example.com",
11038 "example.com:9443",
11039 9443,
11040 "/feed",
11041 true
11042 )]
11043 #[case::root_path(
11044 "ws://example.com:9000/",
11045 "example.com",
11046 "example.com:9000",
11047 9000,
11048 "/",
11049 false
11050 )]
11051 #[case::query_string(
11052 "ws://example.com/feed?token=abc&channel=trades",
11053 "example.com",
11054 "example.com",
11055 80,
11056 "/feed?token=abc&channel=trades",
11057 false
11058 )]
11059 #[case::ipv6_default("ws://[::1]/feed", "::1", "[::1]", 80, "/feed", false)]
11061 #[case::ipv6_explicit_port("ws://[::1]:9000/feed", "::1", "[::1]:9000", 9000, "/feed", false)]
11062 #[case::ipv6_wss(
11063 "wss://[2001:db8::1]:8443/",
11064 "2001:db8::1",
11065 "[2001:db8::1]:8443",
11066 8443,
11067 "/",
11068 true
11069 )]
11070 fn sockudo_target_parses_url(
11071 #[case] url: &str,
11072 #[case] host: &str,
11073 #[case] host_header: &str,
11074 #[case] port: u16,
11075 #[case] path: &str,
11076 #[case] is_tls: bool,
11077 ) {
11078 let target = super::SockudoTarget::parse(url).expect("parse should succeed");
11079 assert_eq!(target.host, host);
11080 assert_eq!(target.host_header, host_header);
11081 assert_eq!(target.port, port);
11082 assert_eq!(target.path, path);
11083 assert_eq!(target.is_tls, is_tls);
11084 }
11085
11086 #[cfg(all(feature = "transport-sockudo", not(feature = "turmoil")))]
11087 #[rstest]
11088 fn sockudo_target_rejects_unsupported_scheme() {
11089 let err = super::SockudoTarget::parse("http://example.com/feed").expect_err("not a ws URL");
11090 let msg = err.to_string();
11091 assert!(
11092 msg.contains("expected ws:// or wss://"),
11093 "unexpected error: {msg}"
11094 );
11095 }
11096
11097 #[cfg(all(feature = "transport-sockudo", not(feature = "turmoil")))]
11098 #[rstest]
11099 fn sockudo_target_rejects_malformed_url() {
11100 const SECRET: &str = "malformed-websocket-url-secret";
11101 let url = format!("not a url {SECRET}");
11102 let err = super::SockudoTarget::parse(&url).expect_err("malformed URL");
11103 let message = err.to_string();
11104
11105 assert!(
11106 matches!(err, super::TransportError::InvalidUrl(_)),
11107 "expected InvalidUrl, was: {err:?}"
11108 );
11109 assert!(!message.contains(SECRET));
11110 assert!(!message.contains(&url));
11111 }
11112}
11113
11114#[cfg(test)]
11115mod property_tests {
11116 use std::{
11117 collections::{HashSet, VecDeque},
11118 sync::{Arc, OnceLock, atomic::AtomicBool},
11119 };
11120
11121 use proptest::prelude::*;
11122 use rstest::rstest;
11123
11124 use super::{super::auth::AuthResultReceiver, *};
11125
11126 const AUTH_FAILED: &str = "model auth failed";
11127
11128 #[derive(Debug, Clone)]
11129 enum ReconnectBufferTraceOp {
11130 BeginAuth,
11131 AuthSucceeds,
11132 AuthFails,
11133 AuthInvalidates,
11134 ReconnectStarts,
11135 ReconnectCompletes,
11136 BufferedMessage(u8),
11137 }
11138
11139 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
11140 enum ModelConnectionMode {
11141 Active,
11142 Reconnect,
11143 }
11144
11145 #[derive(Debug, Clone, Copy)]
11146 enum ExpectedReconnectBufferAction {
11147 Drain,
11148 Wait,
11149 Discard,
11150 }
11151
11152 #[derive(Debug)]
11153 struct ReconnectBufferModel {
11154 mode: ModelConnectionMode,
11155 auth_state: AuthState,
11156 buffer: VecDeque<String>,
11157 released: Vec<String>,
11158 discarded: Vec<String>,
11159 live_sent: Vec<String>,
11160 handler_controls: Vec<&'static str>,
11161 next_message_index: usize,
11162 }
11163
11164 impl ReconnectBufferModel {
11165 fn new() -> Self {
11166 Self {
11167 mode: ModelConnectionMode::Active,
11168 auth_state: AuthState::Unauthenticated,
11169 buffer: VecDeque::new(),
11170 released: Vec::new(),
11171 discarded: Vec::new(),
11172 live_sent: Vec::new(),
11173 handler_controls: Vec::new(),
11174 next_message_index: 0,
11175 }
11176 }
11177
11178 fn next_payload(&mut self, raw: u8) -> String {
11179 let payload = format!("message-{}-{raw}", self.next_message_index);
11180 self.next_message_index += 1;
11181 payload
11182 }
11183
11184 fn expected_action(&self, waits_for_auth: bool) -> ExpectedReconnectBufferAction {
11185 if !waits_for_auth {
11186 return ExpectedReconnectBufferAction::Drain;
11187 }
11188
11189 match self.auth_state {
11190 AuthState::Authenticated => ExpectedReconnectBufferAction::Drain,
11191 AuthState::Failed => ExpectedReconnectBufferAction::Discard,
11192 AuthState::Unauthenticated => ExpectedReconnectBufferAction::Wait,
11193 }
11194 }
11195 }
11196
11197 fn reconnect_buffer_trace_op_strategy() -> impl Strategy<Value = ReconnectBufferTraceOp> {
11198 prop_oneof![
11199 Just(ReconnectBufferTraceOp::BeginAuth),
11200 Just(ReconnectBufferTraceOp::AuthSucceeds),
11201 Just(ReconnectBufferTraceOp::AuthFails),
11202 Just(ReconnectBufferTraceOp::AuthInvalidates),
11203 Just(ReconnectBufferTraceOp::ReconnectStarts),
11204 Just(ReconnectBufferTraceOp::ReconnectCompletes),
11205 any::<u8>().prop_map(ReconnectBufferTraceOp::BufferedMessage),
11206 ]
11207 }
11208
11209 fn reconnect_buffer_actions_match(
11210 actual: ReconnectBufferAction,
11211 expected: ExpectedReconnectBufferAction,
11212 ) -> bool {
11213 matches!(
11214 (actual, expected),
11215 (
11216 ReconnectBufferAction::Drain,
11217 ExpectedReconnectBufferAction::Drain
11218 ) | (
11219 ReconnectBufferAction::Wait,
11220 ExpectedReconnectBufferAction::Wait
11221 ) | (
11222 ReconnectBufferAction::Discard,
11223 ExpectedReconnectBufferAction::Discard
11224 )
11225 )
11226 }
11227
11228 fn apply_ready_reconnect_buffer_action(
11229 model: &mut ReconnectBufferModel,
11230 reconnect_buffer_waits_for_auth: &AtomicBool,
11231 auth_tracker: &Arc<OnceLock<AuthTracker>>,
11232 waits_for_auth: bool,
11233 step: usize,
11234 op: &ReconnectBufferTraceOp,
11235 ) -> Result<(), TestCaseError> {
11236 if model.mode != ModelConnectionMode::Active || model.buffer.is_empty() {
11237 return Ok(());
11238 }
11239
11240 let expected = model.expected_action(waits_for_auth);
11241 let actual = WebSocketClientInner::can_drain_reconnect_buffer(
11242 reconnect_buffer_waits_for_auth,
11243 auth_tracker,
11244 );
11245
11246 prop_assert!(
11247 reconnect_buffer_actions_match(actual, expected),
11248 "reconnect buffer action mismatch at step {}, op {:?}, waits_for_auth={}, auth_state={:?}",
11249 step,
11250 op,
11251 waits_for_auth,
11252 model.auth_state
11253 );
11254
11255 match expected {
11256 ExpectedReconnectBufferAction::Drain => {
11257 model.released.extend(model.buffer.drain(..));
11258 }
11259 ExpectedReconnectBufferAction::Wait => {}
11260 ExpectedReconnectBufferAction::Discard => {
11261 model.discarded.extend(model.buffer.drain(..));
11262 }
11263 }
11264
11265 Ok(())
11266 }
11267
11268 fn assert_reconnected_control_stays_separate(
11269 model: &ReconnectBufferModel,
11270 step: usize,
11271 ) -> Result<(), TestCaseError> {
11272 prop_assert!(
11273 model
11274 .handler_controls
11275 .iter()
11276 .all(|message| *message == RECONNECTED),
11277 "handler control stream contained a non-RECONNECTED message at step {}",
11278 step
11279 );
11280 prop_assert!(
11281 !model.buffer.iter().any(|message| message == RECONNECTED),
11282 "RECONNECTED control message entered reconnect buffer at step {}",
11283 step
11284 );
11285 prop_assert!(
11286 !model.released.iter().any(|message| message == RECONNECTED),
11287 "RECONNECTED control message entered replayed messages at step {}",
11288 step
11289 );
11290 prop_assert!(
11291 !model.discarded.iter().any(|message| message == RECONNECTED),
11292 "RECONNECTED control message entered discarded messages at step {}",
11293 step
11294 );
11295 prop_assert!(
11296 !model.live_sent.iter().any(|message| message == RECONNECTED),
11297 "RECONNECTED control message entered application sends at step {}",
11298 step
11299 );
11300
11301 Ok(())
11302 }
11303
11304 fn assert_messages_accounted_once(
11305 model: &ReconnectBufferModel,
11306 step: usize,
11307 ) -> Result<(), TestCaseError> {
11308 let mut seen = HashSet::new();
11309
11310 for message in model
11311 .released
11312 .iter()
11313 .chain(model.discarded.iter())
11314 .chain(model.buffer.iter())
11315 .chain(model.live_sent.iter())
11316 {
11317 prop_assert!(
11318 seen.insert(message.as_str()),
11319 "message {} appeared more than once at step {}",
11320 message,
11321 step
11322 );
11323 }
11324
11325 Ok(())
11326 }
11327
11328 fn apply_reconnect_buffer_trace_op(
11329 model: &mut ReconnectBufferModel,
11330 tracker: &AuthTracker,
11331 auth_receivers: &mut Vec<AuthResultReceiver>,
11332 op: &ReconnectBufferTraceOp,
11333 ) -> Result<(), TestCaseError> {
11334 match op {
11335 ReconnectBufferTraceOp::BeginAuth => {
11336 auth_receivers.push(tracker.begin());
11337 model.auth_state = AuthState::Unauthenticated;
11338 }
11339 ReconnectBufferTraceOp::AuthSucceeds => {
11340 tracker.succeed();
11341 model.auth_state = AuthState::Authenticated;
11342 }
11343 ReconnectBufferTraceOp::AuthFails => {
11344 tracker.fail(AUTH_FAILED);
11345 model.auth_state = AuthState::Failed;
11346 }
11347 ReconnectBufferTraceOp::AuthInvalidates => {
11348 tracker.invalidate();
11349
11350 if model.auth_state == AuthState::Authenticated {
11351 model.auth_state = AuthState::Unauthenticated;
11352 }
11353 }
11354 ReconnectBufferTraceOp::ReconnectStarts => {
11355 tracker.invalidate();
11356
11357 if model.auth_state == AuthState::Authenticated {
11358 model.auth_state = AuthState::Unauthenticated;
11359 }
11360
11361 model.mode = ModelConnectionMode::Reconnect;
11362 }
11363 ReconnectBufferTraceOp::ReconnectCompletes => {
11364 model.mode = ModelConnectionMode::Active;
11365 model.handler_controls.push(RECONNECTED);
11366 }
11367 ReconnectBufferTraceOp::BufferedMessage(raw) => {
11368 let payload = model.next_payload(*raw);
11369 prop_assert_ne!(payload.as_str(), RECONNECTED);
11370
11371 if model.mode == ModelConnectionMode::Reconnect {
11372 model.buffer.push_back(payload);
11373 } else {
11374 model.live_sent.push(payload);
11375 }
11376 }
11377 }
11378
11379 Ok(())
11380 }
11381
11382 proptest! {
11383 #![proptest_config(ProptestConfig::with_cases(256))]
11384
11385 #[rstest]
11388 fn test_reconnect_buffer_trace_matches_auth_gate_model(
11389 waits_for_auth in any::<bool>(),
11390 ops in proptest::collection::vec(reconnect_buffer_trace_op_strategy(), 1..100)
11391 ) {
11392 let auth_tracker = Arc::new(OnceLock::new());
11393 let reconnect_buffer_waits_for_auth = AtomicBool::new(waits_for_auth);
11394 let tracker = AuthTracker::new();
11395 auth_tracker.set(tracker.clone()).unwrap();
11396 let mut auth_receivers = Vec::new();
11397 let mut model = ReconnectBufferModel::new();
11398
11399 for (step, op) in ops.iter().enumerate() {
11400 apply_reconnect_buffer_trace_op(
11401 &mut model,
11402 &tracker,
11403 &mut auth_receivers,
11404 op,
11405 )?;
11406
11407 prop_assert_eq!(
11408 tracker.auth_state(),
11409 model.auth_state,
11410 "auth state mismatch at step {}, op {:?}",
11411 step,
11412 op
11413 );
11414
11415 apply_ready_reconnect_buffer_action(
11416 &mut model,
11417 &reconnect_buffer_waits_for_auth,
11418 &auth_tracker,
11419 waits_for_auth,
11420 step,
11421 op,
11422 )?;
11423 assert_reconnected_control_stays_separate(&model, step)?;
11424 prop_assert_eq!(
11425 model.handler_controls.len(),
11426 ops[..=step]
11427 .iter()
11428 .filter(|op| matches!(op, ReconnectBufferTraceOp::ReconnectCompletes))
11429 .count(),
11430 "handler control count mismatch at step {}",
11431 step
11432 );
11433 assert_messages_accounted_once(&model, step)?;
11434 }
11435 }
11436
11437 #[rstest]
11440 fn test_reconnect_buffer_releases_after_auth_success_once(
11441 payloads in proptest::collection::vec(any::<u8>(), 1..32),
11442 extra_success_ticks in 0usize..16
11443 ) {
11444 let auth_tracker = Arc::new(OnceLock::new());
11445 let reconnect_buffer_waits_for_auth = AtomicBool::new(true);
11446 let tracker = AuthTracker::new();
11447 auth_tracker.set(tracker.clone()).unwrap();
11448 let mut auth_receivers = Vec::new();
11449 let mut model = ReconnectBufferModel::new();
11450
11451 apply_reconnect_buffer_trace_op(
11452 &mut model,
11453 &tracker,
11454 &mut auth_receivers,
11455 &ReconnectBufferTraceOp::ReconnectStarts,
11456 )?;
11457 apply_reconnect_buffer_trace_op(
11458 &mut model,
11459 &tracker,
11460 &mut auth_receivers,
11461 &ReconnectBufferTraceOp::BeginAuth,
11462 )?;
11463
11464 for payload in payloads {
11465 apply_reconnect_buffer_trace_op(
11466 &mut model,
11467 &tracker,
11468 &mut auth_receivers,
11469 &ReconnectBufferTraceOp::BufferedMessage(payload),
11470 )?;
11471 }
11472
11473 let buffered_len = model.buffer.len();
11474 apply_reconnect_buffer_trace_op(
11475 &mut model,
11476 &tracker,
11477 &mut auth_receivers,
11478 &ReconnectBufferTraceOp::ReconnectCompletes,
11479 )?;
11480 apply_ready_reconnect_buffer_action(
11481 &mut model,
11482 &reconnect_buffer_waits_for_auth,
11483 &auth_tracker,
11484 true,
11485 0,
11486 &ReconnectBufferTraceOp::ReconnectCompletes,
11487 )?;
11488
11489 prop_assert_eq!(model.released.len(), 0);
11490 prop_assert_eq!(model.buffer.len(), buffered_len);
11491
11492 apply_reconnect_buffer_trace_op(
11493 &mut model,
11494 &tracker,
11495 &mut auth_receivers,
11496 &ReconnectBufferTraceOp::AuthSucceeds,
11497 )?;
11498 apply_ready_reconnect_buffer_action(
11499 &mut model,
11500 &reconnect_buffer_waits_for_auth,
11501 &auth_tracker,
11502 true,
11503 1,
11504 &ReconnectBufferTraceOp::AuthSucceeds,
11505 )?;
11506
11507 prop_assert_eq!(model.released.len(), buffered_len);
11508 prop_assert!(model.buffer.is_empty());
11509 assert_messages_accounted_once(&model, 1)?;
11510
11511 for tick in 0..extra_success_ticks {
11512 apply_reconnect_buffer_trace_op(
11513 &mut model,
11514 &tracker,
11515 &mut auth_receivers,
11516 &ReconnectBufferTraceOp::AuthSucceeds,
11517 )?;
11518 apply_ready_reconnect_buffer_action(
11519 &mut model,
11520 &reconnect_buffer_waits_for_auth,
11521 &auth_tracker,
11522 true,
11523 tick + 2,
11524 &ReconnectBufferTraceOp::AuthSucceeds,
11525 )?;
11526 prop_assert_eq!(
11527 model.released.len(),
11528 buffered_len,
11529 "buffered messages replayed more than once at tick {}",
11530 tick
11531 );
11532 }
11533 }
11534
11535 #[rstest]
11538 fn test_reconnect_buffer_discards_after_auth_failure(
11539 before_failure_payloads in proptest::collection::vec(any::<u8>(), 0..16),
11540 after_failure_payloads in proptest::collection::vec(any::<u8>(), 1..16),
11541 later_success_ticks in 0usize..16
11542 ) {
11543 let auth_tracker = Arc::new(OnceLock::new());
11544 let reconnect_buffer_waits_for_auth = AtomicBool::new(true);
11545 let tracker = AuthTracker::new();
11546 auth_tracker.set(tracker.clone()).unwrap();
11547 let mut auth_receivers = Vec::new();
11548 let mut model = ReconnectBufferModel::new();
11549
11550 apply_reconnect_buffer_trace_op(
11551 &mut model,
11552 &tracker,
11553 &mut auth_receivers,
11554 &ReconnectBufferTraceOp::ReconnectStarts,
11555 )?;
11556 apply_reconnect_buffer_trace_op(
11557 &mut model,
11558 &tracker,
11559 &mut auth_receivers,
11560 &ReconnectBufferTraceOp::BeginAuth,
11561 )?;
11562
11563 for payload in before_failure_payloads {
11564 apply_reconnect_buffer_trace_op(
11565 &mut model,
11566 &tracker,
11567 &mut auth_receivers,
11568 &ReconnectBufferTraceOp::BufferedMessage(payload),
11569 )?;
11570 }
11571
11572 apply_reconnect_buffer_trace_op(
11573 &mut model,
11574 &tracker,
11575 &mut auth_receivers,
11576 &ReconnectBufferTraceOp::AuthFails,
11577 )?;
11578
11579 for payload in after_failure_payloads {
11580 apply_reconnect_buffer_trace_op(
11581 &mut model,
11582 &tracker,
11583 &mut auth_receivers,
11584 &ReconnectBufferTraceOp::BufferedMessage(payload),
11585 )?;
11586 }
11587
11588 let buffered_len = model.buffer.len();
11589 apply_reconnect_buffer_trace_op(
11590 &mut model,
11591 &tracker,
11592 &mut auth_receivers,
11593 &ReconnectBufferTraceOp::ReconnectCompletes,
11594 )?;
11595 apply_ready_reconnect_buffer_action(
11596 &mut model,
11597 &reconnect_buffer_waits_for_auth,
11598 &auth_tracker,
11599 true,
11600 0,
11601 &ReconnectBufferTraceOp::ReconnectCompletes,
11602 )?;
11603
11604 prop_assert_eq!(model.discarded.len(), buffered_len);
11605 prop_assert!(model.released.is_empty());
11606 prop_assert!(model.buffer.is_empty());
11607 assert_messages_accounted_once(&model, 0)?;
11608
11609 for tick in 0..later_success_ticks {
11610 apply_reconnect_buffer_trace_op(
11611 &mut model,
11612 &tracker,
11613 &mut auth_receivers,
11614 &ReconnectBufferTraceOp::BeginAuth,
11615 )?;
11616 apply_reconnect_buffer_trace_op(
11617 &mut model,
11618 &tracker,
11619 &mut auth_receivers,
11620 &ReconnectBufferTraceOp::AuthSucceeds,
11621 )?;
11622 apply_ready_reconnect_buffer_action(
11623 &mut model,
11624 &reconnect_buffer_waits_for_auth,
11625 &auth_tracker,
11626 true,
11627 tick + 1,
11628 &ReconnectBufferTraceOp::AuthSucceeds,
11629 )?;
11630 prop_assert!(
11631 model.released.is_empty(),
11632 "discarded messages replayed after later auth success at tick {}",
11633 tick
11634 );
11635 }
11636 }
11637 }
11638}
11639
11640#[cfg(test)]
11641#[cfg(feature = "turmoil")]
11642mod turmoil_tests {
11643 use std::{sync::Arc, time::Duration};
11644
11645 use futures_util::{SinkExt, StreamExt};
11646 use nautilus_common::testing::wait_until_async;
11647 use rstest::rstest;
11648 use tokio_tungstenite::{accept_async, tungstenite::Message as WsMessage};
11649 use turmoil::{Builder, net};
11650
11651 use super::*;
11652 use crate::websocket::types::channel_message_handler;
11653
11654 const AUTH_BUFFER_WAIT_SEED: u64 = 0xA17B_0001;
11655 const AUTH_BUFFER_DISCARD_SEED: u64 = 0xA17B_0002;
11656
11657 fn seeded_turmoil_builder(seed: u64) -> Builder {
11658 let mut builder = Builder::new();
11659 builder.rng_seed(seed);
11660 builder
11661 }
11662
11663 #[rstest]
11664 fn test_turmoil_reconnect_buffer_waits_for_auth() {
11665 let mut sim = seeded_turmoil_builder(AUTH_BUFFER_WAIT_SEED).build();
11666 let messages = Arc::new(tokio::sync::Mutex::new(Vec::new()));
11667 let server_messages = Arc::clone(&messages);
11668 sim.host("server", move || {
11669 let messages = Arc::clone(&server_messages);
11670 auth_buffer_server(messages)
11671 });
11672
11673 sim.client("client", async move {
11674 let tracker = AuthTracker::new();
11675 let (handler, _rx) = channel_message_handler();
11676 let client = WebSocketClient::builder()
11677 .config(turmoil_websocket_config())
11678 .message_handler(handler)
11679 .connect()
11680 .await
11681 .expect("Should connect");
11682
11683 client.set_auth_tracker(tracker.clone(), true);
11684 assert!(client.is_active(), "Client should start active");
11685
11686 wait_until_async(
11687 || async { client.is_reconnecting() },
11688 Duration::from_secs(3),
11689 )
11690 .await;
11691
11692 client
11693 .writer_tx
11694 .send(WriterCommand::Send(Message::Text("stale".into())))
11695 .unwrap();
11696
11697 wait_until_async(|| async { client.is_active() }, Duration::from_secs(3)).await;
11698
11699 let _auth_receiver = tracker.begin();
11700
11701 tokio::time::sleep(Duration::from_millis(300)).await;
11702 assert!(
11703 messages.lock().await.is_empty(),
11704 "buffered messages should wait for auth after reconnect"
11705 );
11706
11707 tracker.succeed();
11708
11709 wait_until_async(
11710 || {
11711 let messages = Arc::clone(&messages);
11712 async move { messages.lock().await.as_slice() == ["stale"] }
11713 },
11714 Duration::from_secs(3),
11715 )
11716 .await;
11717
11718 assert_eq!(messages.lock().await.as_slice(), ["stale"]);
11719
11720 client.disconnect().await;
11721 assert!(client.is_disconnected());
11722
11723 Ok(())
11724 });
11725
11726 sim.run().unwrap();
11727 }
11728
11729 #[rstest]
11730 fn test_turmoil_reconnect_buffer_discards_after_auth_failure() {
11731 let mut sim = seeded_turmoil_builder(AUTH_BUFFER_DISCARD_SEED).build();
11732 let messages = Arc::new(tokio::sync::Mutex::new(Vec::new()));
11733 let server_messages = Arc::clone(&messages);
11734 sim.host("server", move || {
11735 let messages = Arc::clone(&server_messages);
11736 auth_buffer_server(messages)
11737 });
11738
11739 sim.client("client", async move {
11740 let tracker = AuthTracker::new();
11741 let (handler, _rx) = channel_message_handler();
11742 let client = WebSocketClient::builder()
11743 .config(turmoil_websocket_config())
11744 .message_handler(handler)
11745 .connect()
11746 .await
11747 .expect("Should connect");
11748
11749 client.set_auth_tracker(tracker.clone(), true);
11750 assert!(client.is_active(), "Client should start active");
11751
11752 wait_until_async(
11753 || async { client.is_reconnecting() },
11754 Duration::from_secs(3),
11755 )
11756 .await;
11757
11758 client
11759 .writer_tx
11760 .send(WriterCommand::Send(Message::Text("stale".into())))
11761 .unwrap();
11762
11763 wait_until_async(|| async { client.is_active() }, Duration::from_secs(3)).await;
11764
11765 let _auth_receiver = tracker.begin();
11766 tracker.fail("rejected");
11767
11768 tokio::time::sleep(Duration::from_millis(300)).await;
11769 assert!(
11770 messages.lock().await.is_empty(),
11771 "buffered messages should be discarded after auth failure"
11772 );
11773
11774 let _retry_auth_receiver = tracker.begin();
11775 tracker.succeed();
11776
11777 tokio::time::sleep(Duration::from_millis(300)).await;
11778 assert!(
11779 messages.lock().await.is_empty(),
11780 "discarded messages should not replay on a later auth success"
11781 );
11782
11783 client.disconnect().await;
11784 assert!(client.is_disconnected());
11785
11786 Ok(())
11787 });
11788
11789 sim.run().unwrap();
11790 }
11791
11792 fn turmoil_websocket_config() -> WebSocketConfig {
11793 WebSocketConfig {
11794 url: "ws://server:8080".to_string(),
11795 headers: vec![],
11796 heartbeat_interval_secs: None,
11797 heartbeat_payload: None,
11798 connect_timeout_ms: Some(5_000),
11799 reconnect_delay_initial_ms: Some(50),
11800 reconnect_delay_max_ms: Some(200),
11801 reconnect_backoff_factor: Some(1.0),
11802 reconnect_jitter_ms: Some(0),
11803 reconnect_max_attempts: None,
11804 heartbeat_timeout_secs: None,
11805 idle_timeout_ms: None,
11806 writer_capacity: None,
11807 backend: TransportBackend::Tungstenite,
11808 proxy_url: None,
11809 max_message_size_bytes: None,
11810 max_frame_size_bytes: None,
11811 }
11812 }
11813
11814 async fn auth_buffer_server(
11815 messages: Arc<tokio::sync::Mutex<Vec<String>>>,
11816 ) -> Result<(), Box<dyn std::error::Error>> {
11817 let listener = net::TcpListener::bind("0.0.0.0:8080").await?;
11818
11819 let (stream, _) = listener.accept().await?;
11820 let mut websocket = accept_async(stream).await?;
11821 let _ = websocket.send(WsMessage::Text("first".into())).await;
11822 drop(websocket);
11823
11824 tokio::time::sleep(Duration::from_millis(200)).await;
11825
11826 let (stream, _) = listener.accept().await?;
11827 let mut websocket = accept_async(stream).await?;
11828
11829 while let Some(msg) = websocket.next().await {
11830 match msg {
11831 Ok(WsMessage::Text(text)) => {
11832 messages.lock().await.push(text.to_string());
11833 }
11834 Ok(WsMessage::Close(_)) => {
11835 let _ = websocket.close(None).await;
11836 break;
11837 }
11838 Ok(_) => {}
11839 Err(_) => break,
11840 }
11841 }
11842
11843 Ok(())
11844 }
11845}