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