nautilus_network/transport/
sockudo.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16//! `sockudo-ws` backend for the transport abstraction.
17//!
18//! Mirrors the layout of the [`tungstenite`](super::tungstenite) module: provides
19//! `From`/`TryFrom` conversions between the neutral [`Message`] / [`TransportError`]
20//! and sockudo's native types, plus a [`SockudoTransport<S>`] adapter that lifts a
21//! sockudo [`WebSocketStream<S>`] into the backend-agnostic [`WsTransport`] trait.
22//!
23//! The `Message` enums are structurally identical: both carry payloads as `bytes::Bytes`
24//! across all five variants, so conversions are zero-copy and infallible.
25//!
26//! Request bytes come from sockudo's header-aware builder; the response loop stays
27//! local so upgrade rejections map to [`TransportError::UpgradeRejected`] with
28//! host-context logging instead of collapsing into a handshake failure.
29
30use std::{
31    pin::Pin,
32    task::{Context, Poll},
33};
34
35use bytes::{Bytes, BytesMut};
36use futures_util::{Sink, Stream, StreamExt};
37use sockudo_ws::{
38    HandshakeResult,
39    error::{CloseReason as SockudoCloseReason, Error as SockudoError},
40    handshake,
41    protocol::Message as SockudoMessage,
42    stream::{SplitWriter, WebSocketStream},
43};
44use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
45use tokio_util::sync::ReusableBoxFuture;
46
47use super::{
48    error::{TransportError, retryable_status},
49    message::{CloseFrame, Message},
50    stream::WsTransport,
51};
52
53const MAX_HTTP_HEADER_SIZE: usize = 8192;
54
55// WebSocket upgrade headers we always set, plus body-framing headers that have
56// no place on a GET upgrade.
57const RESERVED_UPGRADE_HEADERS: &[&str] = &[
58    "host",
59    "upgrade",
60    "connection",
61    "sec-websocket-key",
62    "sec-websocket-version",
63    "sec-websocket-protocol",
64    "sec-websocket-extensions",
65    "content-length",
66    "transfer-encoding",
67    "te",
68    "trailer",
69];
70
71/// Performs the client handshake, building the upgrade request with sockudo's
72/// header-aware builder and reading the response locally.
73///
74/// The local response loop preserves the rejection status for
75/// [`TransportError::UpgradeRejected`] and the retry-aware log severity; caller
76/// pre-validates `extra_headers` via [`validate_extra_headers`].
77pub(crate) async fn client_handshake_with_headers<S>(
78    stream: &mut S,
79    host: &str,
80    path: &str,
81    extra_headers: &[(String, String)],
82) -> Result<HandshakeResult, TransportError>
83where
84    S: AsyncRead + AsyncWrite + Unpin,
85{
86    use tokio::io::{AsyncReadExt, AsyncWriteExt};
87
88    let key = handshake::generate_key();
89    let request =
90        handshake::build_request_with_headers(host, path, &key, None, None, Some(extra_headers))?;
91
92    stream
93        .write_all(&request)
94        .await
95        .map_err(TransportError::from)?;
96    stream.flush().await.map_err(TransportError::from)?;
97
98    let mut buf = BytesMut::with_capacity(4096);
99    loop {
100        if buf.len() > MAX_HTTP_HEADER_SIZE {
101            return Err(SockudoError::InvalidHttp("response too large").into());
102        }
103
104        let n = stream
105            .read_buf(&mut buf)
106            .await
107            .map_err(TransportError::from)?;
108
109        if n == 0 {
110            return Err(TransportError::ConnectionClosed);
111        }
112
113        let parsed = match handshake::parse_response(&buf) {
114            Ok(parsed) => parsed,
115            Err(e) => {
116                let status = rejected_upgrade_status(&e, &buf);
117                log_handshake_response(host, &e, &buf, status);
118                return Err(
119                    status.map_or_else(|| TransportError::from(e), TransportError::UpgradeRejected)
120                );
121            }
122        };
123
124        if let Some((res, consumed)) = parsed {
125            let accept = res.accept.ok_or_else(|| {
126                let e = SockudoError::HandshakeFailed("missing Sec-WebSocket-Accept");
127                log_handshake_response(host, &e, &buf, None);
128                TransportError::from(e)
129            })?;
130
131            if !handshake::validate_accept_key(&key, accept) {
132                let e = SockudoError::HandshakeFailed("invalid Sec-WebSocket-Accept");
133                log_handshake_response(host, &e, &buf, None);
134                return Err(e.into());
135            }
136
137            let res_protocol = res.protocol.map(String::from);
138            let res_extensions = res.extensions.map(String::from);
139
140            let leftover = if consumed < buf.len() {
141                Some(buf.split_off(consumed).freeze())
142            } else {
143                None
144            };
145
146            return Ok(HandshakeResult {
147                path: path.to_string(),
148                protocol: res_protocol,
149                extensions: res_extensions,
150                leftover,
151            });
152        }
153    }
154}
155
156// Gated on the error variant rather than its message. `HandshakeFailed` means the response parsed
157// as HTTP but failed handshake validation, which is the only case where a rejection status is
158// recoverable; `InvalidHttp` stays permanent. Matching the diagnostic text would put an upstream
159// string into transport semantics, and recovering after any error would let a plausible status
160// line with malformed headers pass as a rejection.
161fn rejected_upgrade_status(err: &SockudoError, buf: &[u8]) -> Option<u16> {
162    if !matches!(err, SockudoError::HandshakeFailed(_)) {
163        return None;
164    }
165
166    let status_line_end = buf.windows(2).position(|window| window == b"\r\n")?;
167    let status_line = std::str::from_utf8(&buf[..status_line_end]).ok()?;
168    let mut parts = status_line.split_whitespace();
169    let version = parts.next()?;
170    let status = parts.next()?;
171    if !version.starts_with("HTTP/1.") || status.len() != 3 {
172        return None;
173    }
174
175    // A 101 that still failed validation is a protocol fault, not a rejection.
176    status.parse().ok().filter(|status| *status != 101)
177}
178
179// A rejection the client goes on to retry is reported at WARN: an ERROR arms the
180// `shutdown_on_error` trigger, which stops the node while the reconnect loop is still recovering.
181// Permanent rejections, malformed responses, and failed accept-key checks retain ERROR handling.
182//
183// `host` names the endpoint so a node running several connections shows which one failed. It is
184// the `Host` header built from the URL authority, so it carries no userinfo, path, or signed
185// query; the request path stays out of the message for that reason.
186fn log_handshake_response(host: &str, err: &SockudoError, buf: &BytesMut, status: Option<u16>) {
187    let response_bytes = buf.len();
188
189    match status {
190        Some(status) if retryable_status(status) => log::warn!(
191            "Sockudo handshake rejected by {host} with retryable status {status}; response bytes={response_bytes}"
192        ),
193        Some(status) => log::error!(
194            "Sockudo handshake rejected by {host} with permanent status {status}; response bytes={response_bytes}"
195        ),
196        None => log::error!(
197            "Sockudo handshake failed for {host}: {err}; response bytes={response_bytes}"
198        ),
199    }
200}
201
202pub(crate) fn validate_extra_headers(headers: &[(String, String)]) -> Result<(), SockudoError> {
203    for (name, value) in headers {
204        validate_extra_header(name, value)?;
205    }
206
207    Ok(())
208}
209
210fn validate_extra_header(name: &str, value: &str) -> Result<(), SockudoError> {
211    let parsed_name = name
212        .parse::<http::HeaderName>()
213        .map_err(|_| SockudoError::InvalidHttp("invalid header name"))?;
214
215    if RESERVED_UPGRADE_HEADERS.contains(&parsed_name.as_str()) {
216        return Err(SockudoError::InvalidHttp(
217            "reserved upgrade header not allowed in extra_headers",
218        ));
219    }
220
221    http::HeaderValue::from_str(value)
222        .map_err(|_| SockudoError::InvalidHttp("invalid header value"))?;
223    Ok(())
224}
225
226/// Replay bytes read during the handshake before forwarding to the inner IO.
227pub(crate) struct PrefixedIo<S> {
228    inner: S,
229    prefix: Bytes,
230}
231
232impl<S> PrefixedIo<S> {
233    pub(crate) const fn new(inner: S, prefix: Bytes) -> Self {
234        Self { inner, prefix }
235    }
236}
237
238impl<S> AsyncRead for PrefixedIo<S>
239where
240    S: AsyncRead + Unpin,
241{
242    fn poll_read(
243        mut self: Pin<&mut Self>,
244        cx: &mut Context<'_>,
245        buf: &mut ReadBuf<'_>,
246    ) -> Poll<std::io::Result<()>> {
247        if !self.prefix.is_empty() {
248            let n = self.prefix.len().min(buf.remaining());
249            let chunk = self.prefix.split_to(n);
250            buf.put_slice(&chunk);
251            return Poll::Ready(Ok(()));
252        }
253
254        Pin::new(&mut self.inner).poll_read(cx, buf)
255    }
256}
257
258impl<S> AsyncWrite for PrefixedIo<S>
259where
260    S: AsyncWrite + Unpin,
261{
262    fn poll_write(
263        mut self: Pin<&mut Self>,
264        cx: &mut Context<'_>,
265        buf: &[u8],
266    ) -> Poll<std::io::Result<usize>> {
267        Pin::new(&mut self.inner).poll_write(cx, buf)
268    }
269
270    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
271        Pin::new(&mut self.inner).poll_flush(cx)
272    }
273
274    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
275        Pin::new(&mut self.inner).poll_shutdown(cx)
276    }
277}
278
279impl From<SockudoMessage> for Message {
280    fn from(value: SockudoMessage) -> Self {
281        match value {
282            SockudoMessage::Text(b) => Self::Text(b),
283            SockudoMessage::Binary(b) => Self::Binary(b),
284            SockudoMessage::Ping(b) => Self::Ping(b),
285            SockudoMessage::Pong(b) => Self::Pong(b),
286            SockudoMessage::Close(reason) => Self::Close(reason.map(Into::into)),
287        }
288    }
289}
290
291impl From<Message> for SockudoMessage {
292    /// Converts a neutral [`Message`] into a Sockudo [`SockudoMessage`].
293    ///
294    /// Conversion is infallible: both enums carry payloads as `bytes::Bytes` across
295    /// all variants. Sockudo validates UTF-8 on Text frames at parse time, not at
296    /// send time, so feeding it non-UTF-8 bytes via [`Self::Text`] is the caller's
297    /// responsibility.
298    fn from(value: Message) -> Self {
299        match value {
300            Message::Text(b) => Self::Text(b),
301            Message::Binary(b) => Self::Binary(b),
302            Message::Ping(b) => Self::Ping(b),
303            Message::Pong(b) => Self::Pong(b),
304            Message::Close(frame) => Self::Close(frame.map(Into::into)),
305        }
306    }
307}
308
309impl From<SockudoCloseReason> for CloseFrame {
310    fn from(value: SockudoCloseReason) -> Self {
311        Self {
312            code: value.code,
313            reason: value.reason,
314        }
315    }
316}
317
318impl From<CloseFrame> for SockudoCloseReason {
319    fn from(value: CloseFrame) -> Self {
320        Self {
321            code: value.code,
322            reason: value.reason,
323        }
324    }
325}
326
327impl From<SockudoError> for TransportError {
328    fn from(value: SockudoError) -> Self {
329        match value {
330            SockudoError::Io(e) => Self::Io(e),
331            SockudoError::ConnectionClosed => Self::ConnectionClosed,
332            SockudoError::ConnectionReset => Self::ConnectionReset,
333            SockudoError::Closed(reason) => Self::ClosedByPeer(reason.map(Into::into)),
334            SockudoError::MessageTooLarge => Self::MessageTooLarge,
335            SockudoError::FrameTooLarge => Self::FrameTooLarge,
336            SockudoError::InvalidUtf8 => Self::InvalidUtf8,
337            SockudoError::InvalidFrame(msg) | SockudoError::Protocol(msg) => {
338                Self::Protocol(msg.to_string())
339            }
340            SockudoError::InvalidHttp(msg) | SockudoError::HandshakeFailed(msg) => {
341                Self::Handshake(msg.to_string())
342            }
343
344            // Keepalive and idle deadlines are dead connections, TimedOut takes
345            // the connection-drop warn path.
346            timeout @ (SockudoError::HeartbeatTimeout | SockudoError::IdleTimeout) => Self::Io(
347                std::io::Error::new(std::io::ErrorKind::TimedOut, timeout.to_string()),
348            ),
349            other => Self::Other(other.to_string()),
350        }
351    }
352}
353
354/// Adapter that lifts a `sockudo-ws` [`WebSocketStream<S>`] into a
355/// backend-agnostic [`WsTransport`].
356///
357/// Translates messages and errors to the neutral types on the way through
358/// `Stream::poll_next` and `Sink<Message>::start_send` / `poll_*`. The
359/// underlying stream is owned and forwarded to via pin projection.
360///
361/// The underlying stream drives automatic control responses while reading.
362pub struct SockudoTransport<S> {
363    inner: WebSocketStream<S>,
364}
365
366impl<S> SockudoTransport<S> {
367    /// Wraps an established Sockudo WebSocket stream.
368    #[inline]
369    #[must_use]
370    pub const fn new(inner: WebSocketStream<S>) -> Self {
371        Self { inner }
372    }
373
374    /// Consumes the adapter and returns the underlying stream.
375    #[inline]
376    pub fn into_inner(self) -> WebSocketStream<S> {
377        self.inner
378    }
379
380    /// Borrows the underlying stream.
381    #[inline]
382    pub const fn get_ref(&self) -> &WebSocketStream<S> {
383        &self.inner
384    }
385}
386
387impl<S> std::fmt::Debug for SockudoTransport<S> {
388    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
389        f.debug_struct(stringify!(SockudoTransport))
390            .finish_non_exhaustive()
391    }
392}
393
394impl<S> Stream for SockudoTransport<S>
395where
396    S: AsyncRead + AsyncWrite + Unpin,
397{
398    type Item = Result<Message, TransportError>;
399
400    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
401        match Pin::new(&mut self.inner).poll_next(cx) {
402            Poll::Ready(Some(Ok(msg))) => Poll::Ready(Some(Ok(Message::from(msg)))),
403            Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(TransportError::from(e)))),
404            Poll::Ready(None) => Poll::Ready(None),
405            Poll::Pending => Poll::Pending,
406        }
407    }
408}
409
410impl<S> Sink<Message> for SockudoTransport<S>
411where
412    S: AsyncRead + AsyncWrite + Unpin,
413{
414    type Error = TransportError;
415
416    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
417        Pin::new(&mut self.inner)
418            .poll_ready(cx)
419            .map_err(TransportError::from)
420    }
421
422    fn start_send(mut self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> {
423        Pin::new(&mut self.inner)
424            .start_send(SockudoMessage::from(item))
425            .map_err(TransportError::from)
426    }
427
428    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
429        Pin::new(&mut self.inner)
430            .poll_flush(cx)
431            .map_err(TransportError::from)
432    }
433
434    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
435        Pin::new(&mut self.inner)
436            .poll_close(cx)
437            .map_err(TransportError::from)
438    }
439}
440
441impl<S> SockudoTransport<S>
442where
443    S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
444{
445    /// Returns a neutral transport backed by Sockudo's native split handles.
446    ///
447    /// Inbound data can progress independently of blocked application writes. Sockudo's
448    /// connection task drives automatic control frames and enforces connection deadlines.
449    /// Call this before sending or receiving messages. A Tokio runtime is required.
450    pub fn into_split(self) -> impl WsTransport {
451        let (reader, writer) = self.inner.split();
452
453        // Keep the decoder stationary: unfold otherwise moves the entire reader for every message
454        let read = futures_util::stream::unfold(Box::new(reader), |mut reader| async move {
455            reader.next().await.map(|message| {
456                (
457                    message.map(Message::from).map_err(TransportError::from),
458                    reader,
459                )
460            })
461        })
462        .fuse();
463
464        SockudoSplitTransport {
465            read: Box::pin(read),
466            writer: Some(writer),
467            write: ReusableBoxFuture::new(std::future::pending()),
468        }
469    }
470}
471
472struct SockudoSplitTransport<R, S> {
473    read: Pin<Box<R>>,
474    writer: Option<SplitWriter<S>>,
475    write: ReusableBoxFuture<'static, (SplitWriter<S>, Result<(), TransportError>)>,
476}
477
478impl<R, S> Stream for SockudoSplitTransport<R, S>
479where
480    S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
481    R: Stream<Item = Result<Message, TransportError>>,
482{
483    type Item = Result<Message, TransportError>;
484
485    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
486        self.read.as_mut().poll_next(cx)
487    }
488}
489
490impl<R, S> Sink<Message> for SockudoSplitTransport<R, S>
491where
492    S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
493    R: Stream<Item = Result<Message, TransportError>>,
494{
495    type Error = TransportError;
496
497    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
498        std::task::ready!(self.poll_pending_write(cx))?;
499
500        if self.writer.as_ref().is_none_or(SplitWriter::is_closed) {
501            return Poll::Ready(Err(TransportError::ConnectionClosed));
502        }
503
504        Poll::Ready(Ok(()))
505    }
506
507    fn start_send(mut self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> {
508        let writer = self.writer.take().ok_or(TransportError::ConnectionClosed)?;
509        self.write.set(write_split(writer, Some(item)));
510        Ok(())
511    }
512
513    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
514        if let Some(writer) = self.writer.take() {
515            self.write.set(write_split(writer, None));
516        }
517
518        self.poll_pending_write(cx)
519    }
520
521    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
522        std::task::ready!(self.poll_pending_write(cx))?;
523
524        if self.writer.as_ref().is_none_or(SplitWriter::is_closed) {
525            return Poll::Ready(Ok(()));
526        }
527
528        self.as_mut()
529            .start_send(Message::Close(Some(CloseFrame::new(
530                CloseFrame::NORMAL,
531                "",
532            ))))?;
533        self.poll_pending_write(cx)
534    }
535}
536
537impl<R, S> SockudoSplitTransport<R, S> {
538    fn poll_pending_write(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), TransportError>> {
539        if self.writer.is_some() {
540            return Poll::Ready(Ok(()));
541        }
542
543        let (writer, result) = std::task::ready!(self.write.poll(cx));
544        self.writer = Some(writer);
545        Poll::Ready(result)
546    }
547}
548
549async fn write_split<S>(
550    mut writer: SplitWriter<S>,
551    message: Option<Message>,
552) -> (SplitWriter<S>, Result<(), TransportError>)
553where
554    S: AsyncWrite + Unpin,
555{
556    // Retain the native send future across polls: cancelling it can terminate a partial frame
557    let result = match message {
558        Some(message) => writer.send(message.into()).await,
559        None => writer.flush().await,
560    }
561    .map_err(TransportError::from);
562
563    (writer, result)
564}
565
566const _: fn() = || {
567    fn assert_ws_transport<T: WsTransport>() {}
568    assert_ws_transport::<SockudoTransport<tokio::net::TcpStream>>();
569};
570
571#[cfg(test)]
572mod tests {
573    use std::future::Future;
574
575    use bytes::Bytes;
576    use rstest::rstest;
577    #[cfg(not(feature = "turmoil"))]
578    use sockudo_ws::handshake::generate_accept_key;
579    #[cfg(not(feature = "turmoil"))]
580    use tokio::io::{AsyncRead, AsyncReadExt, AsyncWriteExt, duplex};
581
582    use super::*;
583
584    #[cfg(not(feature = "turmoil"))]
585    #[tokio::test]
586    async fn test_unified_read_flush_preserves_send_completion_wakeup() {
587        use futures_util::{SinkExt, StreamExt};
588
589        let (client_io, peer_io) = duplex(32);
590        let config = sockudo_ws::Config::builder()
591            .auto_ping(false)
592            .idle_timeout(0)
593            .build();
594        let transport = SockudoTransport::new(WebSocketStream::client(client_io, config.clone()));
595        let (mut writer, mut reader) = transport.split();
596        let mut peer = WebSocketStream::server(peer_io, config);
597
598        let send = tokio::spawn(async move { writer.send(Message::text("x".repeat(4096))).await });
599        tokio::task::yield_now().await;
600        peer.send(SockudoMessage::Ping(Bytes::from_static(b"probe")))
601            .await
602            .unwrap();
603        assert_eq!(
604            reader.next().await.unwrap().unwrap(),
605            Message::Ping(Bytes::from_static(b"probe"))
606        );
607
608        let read = tokio::spawn(async move { reader.next().await });
609
610        tokio::time::timeout(std::time::Duration::from_secs(1), async {
611            assert!(matches!(peer.next().await, Some(Ok(SockudoMessage::Text(data))) if data.len() == 4096));
612            assert!(matches!(peer.next().await, Some(Ok(SockudoMessage::Pong(data))) if data.as_ref() == b"probe"));
613            send.await.unwrap().unwrap();
614        })
615        .await
616        .expect("a completed write must wake the sender without its own timeout");
617
618        drop(peer);
619        read.await.unwrap();
620    }
621
622    #[cfg(not(feature = "turmoil"))]
623    #[tokio::test]
624    async fn test_split_reads_while_application_write_is_blocked() {
625        use futures_util::{SinkExt, StreamExt};
626
627        let (client_io, mut peer) = duplex(32);
628        let config = sockudo_ws::Config::builder()
629            .auto_ping(false)
630            .idle_timeout(0)
631            .build();
632        let transport =
633            SockudoTransport::new(WebSocketStream::client(client_io, config.clone())).into_split();
634        let (mut writer, mut reader) = transport.split();
635        let mut send = Box::pin(writer.send(Message::text("x".repeat(4096))));
636        futures_util::future::poll_fn(|cx| {
637            assert!(send.as_mut().poll(cx).is_pending());
638            Poll::Ready(())
639        })
640        .await;
641
642        peer.write_all(b"\x81\x04peer").await.unwrap();
643
644        let message = tokio::time::timeout(std::time::Duration::from_secs(1), reader.next())
645            .await
646            .unwrap()
647            .unwrap()
648            .unwrap();
649
650        assert_eq!(message, Message::text("peer"));
651    }
652
653    #[cfg(not(feature = "turmoil"))]
654    #[rstest]
655    #[case(false)]
656    #[case(true)]
657    #[tokio::test]
658    async fn test_transport_retains_send_after_outer_future_is_cancelled(#[case] native: bool) {
659        use futures_util::{SinkExt, StreamExt};
660
661        let (client_io, peer_io) = duplex(32);
662        let config = sockudo_ws::Config::builder()
663            .auto_ping(false)
664            .idle_timeout(0)
665            .build();
666        let transport = SockudoTransport::new(WebSocketStream::client(client_io, config.clone()));
667        let mut transport: Pin<Box<dyn WsTransport>> = if native {
668            Box::pin(transport.into_split())
669        } else {
670            Box::pin(transport)
671        };
672
673        let payload = Message::text("x".repeat(4096));
674        {
675            let mut send = Box::pin(transport.send(payload.clone()));
676            futures_util::future::poll_fn(|cx| {
677                assert!(send.as_mut().poll(cx).is_pending());
678                Poll::Ready(())
679            })
680            .await;
681        }
682
683        let peer_task = tokio::spawn(async move {
684            let mut peer = WebSocketStream::server(peer_io, config);
685            let first = peer.next().await.unwrap().unwrap();
686            let second = peer.next().await.unwrap().unwrap();
687            (Message::from(first), Message::from(second))
688        });
689
690        transport.flush().await.unwrap();
691        transport.send(Message::text("after")).await.unwrap();
692        let (first, second) = peer_task.await.unwrap();
693
694        assert_eq!(first, payload);
695        assert_eq!(second, Message::text("after"));
696    }
697
698    #[cfg(not(feature = "turmoil"))]
699    #[tokio::test]
700    async fn test_split_retains_partial_read_after_outer_future_is_cancelled() {
701        use futures_util::StreamExt;
702
703        let (client_io, mut peer) = duplex(64);
704        let config = sockudo_ws::Config::builder()
705            .auto_ping(false)
706            .idle_timeout(0)
707            .build();
708        let mut transport =
709            SockudoTransport::new(WebSocketStream::client(client_io, config)).into_split();
710        peer.write_all(b"\x81\x08pa").await.unwrap();
711        {
712            let mut receive = Box::pin(transport.next());
713            futures_util::future::poll_fn(|cx| {
714                assert!(receive.as_mut().poll(cx).is_pending());
715                Poll::Ready(())
716            })
717            .await;
718        }
719
720        peer.write_all(b"rtial!").await.unwrap();
721        let message = transport.next().await.unwrap().unwrap();
722
723        assert_eq!(message, Message::text("partial!"));
724    }
725
726    #[cfg(not(feature = "turmoil"))]
727    #[tokio::test]
728    async fn test_split_read_stays_ended() {
729        let (client_io, mut peer) = duplex(64);
730        let config = sockudo_ws::Config::builder()
731            .auto_ping(false)
732            .idle_timeout(0)
733            .build();
734        let mut transport =
735            SockudoTransport::new(WebSocketStream::client(client_io, config)).into_split();
736        peer.write_all(b"\x88\x02\x03\xe8").await.unwrap();
737        let close = transport.next().await.unwrap().unwrap();
738        let ended = transport.next().await;
739        let still_ended = transport.next().await;
740
741        assert_eq!(close, Message::Close(Some(CloseFrame::new(1000, ""))));
742        assert!(ended.is_none());
743        assert!(still_ended.is_none());
744    }
745
746    #[cfg(not(feature = "turmoil"))]
747    #[tokio::test]
748    async fn test_split_drives_pong_without_another_read() {
749        use futures_util::StreamExt;
750
751        let (client_io, mut peer) = duplex(64);
752        let config = sockudo_ws::Config::builder()
753            .auto_ping(false)
754            .idle_timeout(0)
755            .build();
756        let mut transport =
757            SockudoTransport::new(WebSocketStream::client(client_io, config.clone())).into_split();
758        peer.write_all(b"\x89\x04ping").await.unwrap();
759        let received = transport.next().await.unwrap().unwrap();
760        let mut peer = WebSocketStream::server(peer, config);
761        let response = tokio::time::timeout(std::time::Duration::from_secs(1), peer.next())
762            .await
763            .unwrap()
764            .unwrap()
765            .unwrap();
766
767        assert_eq!(received, Message::Ping(Bytes::from_static(b"ping")));
768        assert_eq!(
769            Message::from(response),
770            Message::Pong(Bytes::from_static(b"ping"))
771        );
772    }
773
774    #[cfg(not(feature = "turmoil"))]
775    #[tokio::test]
776    async fn test_split_close_sends_normal_status() {
777        use futures_util::{SinkExt, StreamExt};
778
779        let (client_io, peer_io) = duplex(64);
780        let config = sockudo_ws::Config::builder()
781            .auto_ping(false)
782            .idle_timeout(0)
783            .build();
784        let mut transport =
785            SockudoTransport::new(WebSocketStream::client(client_io, config.clone())).into_split();
786        transport.close().await.unwrap();
787        transport.close().await.unwrap();
788        let mut peer = WebSocketStream::server(peer_io, config);
789        let received = peer.next().await.unwrap().unwrap();
790
791        assert_eq!(
792            Message::from(received),
793            Message::Close(Some(CloseFrame::new(1000, "")))
794        );
795    }
796
797    #[cfg(not(feature = "turmoil"))]
798    #[tokio::test]
799    async fn test_split_close_after_explicit_close_preserves_reason() {
800        use futures_util::{SinkExt, StreamExt};
801
802        let (client_io, peer_io) = duplex(64);
803        let config = sockudo_ws::Config::builder()
804            .auto_ping(false)
805            .idle_timeout(0)
806            .build();
807        let mut transport =
808            SockudoTransport::new(WebSocketStream::client(client_io, config.clone())).into_split();
809
810        let close = Message::Close(Some(CloseFrame {
811            code: 3001,
812            reason: "finished".into(),
813        }));
814
815        transport.send(close.clone()).await.unwrap();
816        transport.close().await.unwrap();
817        transport.close().await.unwrap();
818        let mut peer = WebSocketStream::server(peer_io, config);
819        let received = peer.next().await.unwrap().unwrap();
820        let flush = transport.flush().await;
821
822        assert_eq!(Message::from(received), close);
823        assert!(matches!(flush, Err(TransportError::ConnectionClosed)));
824    }
825
826    #[cfg(feature = "turmoil")]
827    #[tokio::test(start_paused = true)]
828    async fn test_sockudo_idle_deadline_uses_simulated_clock() {
829        use futures_util::StreamExt;
830
831        let (client_io, _peer) = tokio::io::duplex(64);
832        let config = sockudo_ws::Config::builder()
833            .auto_ping(false)
834            .idle_timeout(1)
835            .build();
836        let mut transport =
837            SockudoTransport::new(WebSocketStream::client(client_io, config.clone())).into_split();
838        let mut receive = Box::pin(transport.next());
839        futures_util::future::poll_fn(|cx| {
840            assert!(receive.as_mut().poll(cx).is_pending());
841            Poll::Ready(())
842        })
843        .await;
844
845        tokio::time::advance(std::time::Duration::from_millis(1001)).await;
846        tokio::task::yield_now().await;
847        let result =
848            futures_util::future::poll_fn(|cx| Poll::Ready(receive.as_mut().poll(cx))).await;
849
850        let Poll::Ready(Some(Err(TransportError::Io(e)))) = result else {
851            panic!("expected simulated idle timeout, was {result:?}");
852        };
853
854        assert_eq!(e.kind(), std::io::ErrorKind::TimedOut);
855        assert_eq!(e.to_string(), SockudoError::IdleTimeout.to_string());
856    }
857
858    #[cfg(not(feature = "turmoil"))]
859    async fn read_http_request<S>(stream: &mut S) -> Vec<u8>
860    where
861        S: AsyncRead + Unpin,
862    {
863        let mut buf = Vec::new();
864        let mut chunk = [0u8; 256];
865        loop {
866            let n = stream.read(&mut chunk).await.unwrap();
867            assert!(n > 0, "HTTP request closed before headers completed");
868            buf.extend_from_slice(&chunk[..n]);
869            if buf.windows(4).any(|window| window == b"\r\n\r\n") {
870                return buf;
871            }
872        }
873    }
874
875    #[cfg(not(feature = "turmoil"))]
876    fn build_test_response(sec_websocket_key: &str, extra_bytes: &[u8]) -> Vec<u8> {
877        let accept = generate_accept_key(sec_websocket_key);
878        let mut response = format!(
879            concat!(
880                "HTTP/1.1 101 Switching Protocols\r\n",
881                "Upgrade: websocket\r\n",
882                "Connection: Upgrade\r\n",
883                "Sec-WebSocket-Accept: {}\r\n",
884                "\r\n",
885            ),
886            accept
887        )
888        .into_bytes();
889        response.extend_from_slice(extra_bytes);
890        response
891    }
892
893    #[cfg(not(feature = "turmoil"))]
894    fn extract_header<'a>(request: &'a str, name: &str) -> Option<&'a str> {
895        request.lines().find_map(|line| {
896            let (header_name, header_value) = line.split_once(':')?;
897            if header_name.eq_ignore_ascii_case(name) {
898                Some(header_value.trim())
899            } else {
900                None
901            }
902        })
903    }
904
905    #[rstest]
906    #[tokio::test]
907    #[cfg(not(feature = "turmoil"))]
908    async fn handshake_rejects_eof_before_response() {
909        let (mut client, mut server) = duplex(4096);
910        let peer = tokio::spawn(async move {
911            read_http_request(&mut server).await;
912        });
913
914        let error = client_handshake_with_headers(&mut client, "localhost", "/", &[])
915            .await
916            .unwrap_err();
917        peer.await.unwrap();
918
919        assert!(matches!(error, TransportError::ConnectionClosed));
920    }
921
922    #[rstest]
923    #[case::at_limit(MAX_HTTP_HEADER_SIZE, "connection closed")]
924    #[case::above_limit(MAX_HTTP_HEADER_SIZE + 1, "handshake failed: response too large")]
925    #[tokio::test]
926    #[cfg(not(feature = "turmoil"))]
927    async fn handshake_rejects_incomplete_headers_at_size_boundary(
928        #[case] response_size: usize,
929        #[case] expected: &str,
930    ) {
931        let (mut client, mut server) = duplex(MAX_HTTP_HEADER_SIZE * 2);
932        let peer = tokio::spawn(async move {
933            read_http_request(&mut server).await;
934            let mut response = b"HTTP/1.1 101 Switching Protocols\r\nX-Padding: ".to_vec();
935            response.resize(response_size, b'x');
936            server.write_all(&response).await.unwrap();
937        });
938
939        let error = client_handshake_with_headers(&mut client, "localhost", "/", &[])
940            .await
941            .unwrap_err();
942        peer.await.unwrap();
943
944        assert_eq!(error.to_string(), expected);
945    }
946
947    #[tokio::test]
948    #[cfg(not(feature = "turmoil"))]
949    async fn client_handshake_with_headers_sends_custom_headers() {
950        let (mut client, mut server) = duplex(4096);
951        let headers = vec![
952            ("ok-access-key".to_string(), "key-1".to_string()),
953            ("ok-access-passphrase".to_string(), "pass-1".to_string()),
954        ];
955
956        let server_task = tokio::spawn(async move {
957            let request = read_http_request(&mut server).await;
958            let request = String::from_utf8(request).unwrap();
959
960            assert!(request.starts_with("GET /ws/v5/public-sbe?instId=BTC-USDT HTTP/1.1\r\n"));
961            assert_eq!(extract_header(&request, "Host"), Some("ws.okx.com:8443"));
962            assert_eq!(extract_header(&request, "ok-access-key"), Some("key-1"));
963            assert_eq!(
964                extract_header(&request, "ok-access-passphrase"),
965                Some("pass-1")
966            );
967
968            let sec_websocket_key = extract_header(&request, "Sec-WebSocket-Key").unwrap();
969            let response = build_test_response(sec_websocket_key, &[]);
970            server.write_all(&response).await.unwrap();
971        });
972
973        let handshake = client_handshake_with_headers(
974            &mut client,
975            "ws.okx.com:8443",
976            "/ws/v5/public-sbe?instId=BTC-USDT",
977            &headers,
978        )
979        .await
980        .unwrap();
981
982        assert_eq!(handshake.path, "/ws/v5/public-sbe?instId=BTC-USDT");
983        assert!(handshake.leftover.is_none());
984        server_task.await.unwrap();
985    }
986
987    #[rstest]
988    #[cfg(not(feature = "turmoil"))]
989    #[case::host("Host")]
990    #[case::upgrade("Upgrade")]
991    #[case::connection("Connection")]
992    #[case::sec_websocket_key("Sec-WebSocket-Key")]
993    #[case::sec_websocket_version("Sec-WebSocket-Version")]
994    #[case::sec_websocket_protocol("Sec-WebSocket-Protocol")]
995    #[case::sec_websocket_extensions("Sec-WebSocket-Extensions")]
996    #[case::content_length("Content-Length")]
997    #[case::transfer_encoding("Transfer-Encoding")]
998    #[case::te("TE")]
999    #[case::trailer("Trailer")]
1000    fn validate_extra_header_rejects_reserved_upgrade_headers(#[case] name: &str) {
1001        let err = validate_extra_header(name, "value").unwrap_err();
1002
1003        assert!(matches!(
1004            err,
1005            SockudoError::InvalidHttp("reserved upgrade header not allowed in extra_headers")
1006        ));
1007    }
1008
1009    #[tokio::test]
1010    #[cfg(not(feature = "turmoil"))]
1011    async fn client_handshake_with_headers_rejects_missing_accept() {
1012        let (mut client, mut server) = duplex(4096);
1013        let server_task = tokio::spawn(async move {
1014            let _request = read_http_request(&mut server).await;
1015            server
1016                .write_all(
1017                    b"HTTP/1.1 101 Switching Protocols\r\n\
1018                      Upgrade: websocket\r\n\
1019                      Connection: Upgrade\r\n\
1020                      \r\n",
1021                )
1022                .await
1023                .unwrap();
1024        });
1025
1026        let err = client_handshake_with_headers(&mut client, "example.com", "/ws", &[])
1027            .await
1028            .unwrap_err();
1029
1030        assert!(matches!(
1031            err,
1032            TransportError::Handshake(ref msg) if msg == "missing Sec-WebSocket-Accept"
1033        ));
1034        server_task.await.unwrap();
1035    }
1036
1037    #[tokio::test]
1038    #[cfg(not(feature = "turmoil"))]
1039    async fn client_handshake_with_headers_preserves_rejected_status() {
1040        let (mut client, mut server) = duplex(4096);
1041        let server_task = tokio::spawn(async move {
1042            let _request = read_http_request(&mut server).await;
1043            server.write_all(b"HTTP/1.1 429\r\n\r\n").await.unwrap();
1044        });
1045
1046        let err = client_handshake_with_headers(&mut client, "example.com", "/ws", &[])
1047            .await
1048            .unwrap_err();
1049
1050        assert!(matches!(err, TransportError::UpgradeRejected(429)));
1051        server_task.await.unwrap();
1052    }
1053
1054    #[tokio::test]
1055    #[cfg(not(feature = "turmoil"))]
1056    async fn client_handshake_with_headers_returns_leftover_bytes() {
1057        let (mut client, mut server) = duplex(4096);
1058        let extra = b"\x81\x05hello";
1059
1060        let server_task = tokio::spawn(async move {
1061            let request = read_http_request(&mut server).await;
1062            let request = String::from_utf8(request).unwrap();
1063            let sec_websocket_key = extract_header(&request, "Sec-WebSocket-Key").unwrap();
1064            let response = build_test_response(sec_websocket_key, extra);
1065            server.write_all(&response).await.unwrap();
1066        });
1067
1068        let handshake = client_handshake_with_headers(&mut client, "example.com", "/ws", &[])
1069            .await
1070            .unwrap();
1071
1072        assert_eq!(handshake.leftover.as_deref(), Some(extra.as_slice()));
1073        server_task.await.unwrap();
1074    }
1075
1076    #[tokio::test]
1077    #[cfg(not(feature = "turmoil"))]
1078    async fn prefixed_io_replays_leftover_before_socket() {
1079        let (client, mut server) = duplex(4096);
1080        let mut prefixed = PrefixedIo::new(client, Bytes::from_static(b"abc"));
1081
1082        let server_task = tokio::spawn(async move {
1083            server.write_all(b"def").await.unwrap();
1084        });
1085
1086        let mut buf = [0u8; 6];
1087        prefixed.read_exact(&mut buf).await.unwrap();
1088
1089        assert_eq!(&buf, b"abcdef");
1090        server_task.await.unwrap();
1091    }
1092
1093    #[rstest]
1094    fn round_trip_text() {
1095        let original = SockudoMessage::Text(Bytes::from_static(b"hello"));
1096        let neutral: Message = original.into();
1097        assert!(neutral.is_text());
1098        assert_eq!(neutral.as_bytes(), b"hello");
1099
1100        let back: SockudoMessage = neutral.into();
1101        match back {
1102            SockudoMessage::Text(b) => assert_eq!(&b[..], b"hello"),
1103            other => panic!("expected text, was {other:?}"),
1104        }
1105    }
1106
1107    #[rstest]
1108    fn round_trip_binary() {
1109        let original = SockudoMessage::Binary(Bytes::from_static(&[1, 2, 3]));
1110        let neutral: Message = original.into();
1111        assert_eq!(neutral.as_bytes(), &[1, 2, 3]);
1112
1113        let back: SockudoMessage = neutral.into();
1114        match back {
1115            SockudoMessage::Binary(b) => assert_eq!(&b[..], &[1, 2, 3]),
1116            other => panic!("expected binary, was {other:?}"),
1117        }
1118    }
1119
1120    #[rstest]
1121    fn round_trip_ping_pong() {
1122        let neutral: Message = SockudoMessage::Ping(Bytes::from_static(b"p")).into();
1123        assert!(neutral.is_ping());
1124
1125        let neutral: Message = SockudoMessage::Pong(Bytes::from_static(b"q")).into();
1126        assert!(neutral.is_pong());
1127    }
1128
1129    #[rstest]
1130    fn close_frame_round_trip() {
1131        let original = SockudoMessage::Close(Some(SockudoCloseReason {
1132            code: 1000,
1133            reason: "bye".into(),
1134        }));
1135
1136        let neutral: Message = original.into();
1137
1138        let Message::Close(Some(frame)) = &neutral else {
1139            panic!("expected close frame");
1140        };
1141
1142        assert_eq!(frame.code, 1000);
1143        assert_eq!(frame.reason, "bye");
1144
1145        let back: SockudoMessage = neutral.into();
1146
1147        let SockudoMessage::Close(Some(reason)) = back else {
1148            panic!("expected close frame");
1149        };
1150
1151        assert_eq!(reason.code, 1000);
1152        assert_eq!(reason.reason, "bye");
1153    }
1154
1155    #[rstest]
1156    fn error_translation_closed() {
1157        let err: TransportError = SockudoError::ConnectionClosed.into();
1158        assert!(matches!(err, TransportError::ConnectionClosed));
1159    }
1160
1161    #[rstest]
1162    fn error_translation_utf8() {
1163        let err: TransportError = SockudoError::InvalidUtf8.into();
1164        assert!(matches!(err, TransportError::InvalidUtf8));
1165    }
1166
1167    #[rstest]
1168    fn error_translation_handshake() {
1169        let err: TransportError = SockudoError::HandshakeFailed("bad").into();
1170        assert!(matches!(err, TransportError::Handshake(_)));
1171    }
1172
1173    #[rstest]
1174    #[case(SockudoError::HeartbeatTimeout, "WebSocket Pong deadline expired")]
1175    #[case(SockudoError::IdleTimeout, "WebSocket inbound idle deadline expired")]
1176    fn error_translation_timeouts_are_timed_out_io(
1177        #[case] sockudo: SockudoError,
1178        #[case] message: &str,
1179    ) {
1180        let err: TransportError = sockudo.into();
1181
1182        let TransportError::Io(io_err) = &err else {
1183            panic!("expected I/O timeout, was: {err:?}");
1184        };
1185
1186        assert_eq!(io_err.kind(), std::io::ErrorKind::TimedOut);
1187        assert_eq!(io_err.to_string(), message);
1188        assert_eq!(err.to_string(), format!("I/O error: {message}"));
1189    }
1190
1191    // The log-capture harness is Linux-only for CI stability.
1192    #[cfg(not(feature = "turmoil"))]
1193    #[cfg(target_os = "linux")]
1194    #[cfg(not(all(feature = "simulation", madsim)))]
1195    mod handshake_logging {
1196        use log::Level;
1197        use rstest::rstest;
1198        use tokio::io::{AsyncWriteExt, duplex};
1199
1200        use super::read_http_request;
1201        use crate::{
1202            logging::tests::capture_logs_for,
1203            transport::{error::TransportError, sockudo::client_handshake_with_headers},
1204        };
1205
1206        const LOG_TARGETS: &[&str] = &["nautilus_network::transport::sockudo"];
1207        const HOST: &str = "ws.example.com:8443";
1208        const PATH_SECRET: &str = "handshake-path-secret";
1209
1210        #[rstest]
1211        #[case::retryable_bad_gateway("HTTP/1.1 502 Bad Gateway\r\n\r\n", Level::Warn, Some(502))]
1212        #[case::retryable_rate_limited(
1213            "HTTP/1.1 429 Too Many Requests\r\n\r\n",
1214            Level::Warn,
1215            Some(429)
1216        )]
1217        #[case::permanent_unauthorized(
1218            "HTTP/1.1 401 Unauthorized\r\n\r\n",
1219            Level::Error,
1220            Some(401)
1221        )]
1222        #[case::permanent_not_found("HTTP/1.1 404 Not Found\r\n\r\n", Level::Error, Some(404))]
1223        #[case::malformed_status_line("NOT-HTTP 502 Bad Gateway\r\n\r\n", Level::Error, None)]
1224        #[case::missing_accept_key(
1225            "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\r\n",
1226            Level::Error,
1227            None
1228        )]
1229        #[case::invalid_accept_key(
1230            "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\
1231             Sec-WebSocket-Accept: AAAAAAAAAAAAAAAAAAAAAAAAAAA=\r\n\r\n",
1232            Level::Error,
1233            None
1234        )]
1235        #[tokio::test]
1236        async fn handshake_failure_log_level_follows_retry_policy(
1237            #[case] response: &'static str,
1238            #[case] expected_level: Level,
1239            #[case] expected_status: Option<u16>,
1240        ) {
1241            let capture = capture_logs_for(LOG_TARGETS).await;
1242            let (mut client, mut server) = duplex(4096);
1243            let server_task = tokio::spawn(async move {
1244                let _request = read_http_request(&mut server).await;
1245                server.write_all(response.as_bytes()).await.unwrap();
1246                server.flush().await.unwrap();
1247            });
1248
1249            let err = client_handshake_with_headers(
1250                &mut client,
1251                HOST,
1252                &format!("/ws?token={PATH_SECRET}"),
1253                &[],
1254            )
1255            .await
1256            .expect_err("handshake should fail");
1257            server_task.await.unwrap();
1258
1259            match expected_status {
1260                Some(status) => assert!(
1261                    matches!(err, TransportError::UpgradeRejected(actual) if actual == status),
1262                    "expected upgrade rejection {status}, was: {err:?}"
1263                ),
1264                None => assert!(
1265                    matches!(err, TransportError::Handshake(_)),
1266                    "expected a permanent handshake failure, was: {err:?}"
1267                ),
1268            }
1269
1270            let messages = capture.messages();
1271            assert_eq!(
1272                messages.len(),
1273                1,
1274                "expected exactly one handshake log, was: {messages:?}"
1275            );
1276
1277            let (level, message) = &messages[0];
1278            assert_eq!(*level, expected_level, "unexpected level for: {message}");
1279            assert!(
1280                message.contains(HOST),
1281                "log should name the endpoint, was: {message}"
1282            );
1283
1284            if let Some(status) = expected_status {
1285                assert!(
1286                    message.contains(&status.to_string()),
1287                    "log should name the rejection status, was: {message}"
1288                );
1289            }
1290
1291            assert!(
1292                !message.contains(PATH_SECRET),
1293                "log must not carry the request path, was: {message}"
1294            );
1295        }
1296    }
1297}
nautilus_network/transport/sockudo.rs

nautilus_network/transport/
sockudo.rs
