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};
37use sockudo_ws::{
38    HandshakeResult,
39    error::{CloseReason as SockudoCloseReason, Error as SockudoError},
40    handshake,
41    protocol::Message as SockudoMessage,
42    stream::WebSocketStream,
43};
44use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
45
46use super::{
47    error::{TransportError, retryable_status},
48    message::{CloseFrame, Message},
49    stream::WsTransport,
50};
51
52const MAX_HTTP_HEADER_SIZE: usize = 8192;
53
54// WebSocket upgrade headers we always set, plus body-framing headers that have
55// no place on a GET upgrade.
56const RESERVED_UPGRADE_HEADERS: &[&str] = &[
57    "host",
58    "upgrade",
59    "connection",
60    "sec-websocket-key",
61    "sec-websocket-version",
62    "sec-websocket-protocol",
63    "sec-websocket-extensions",
64    "content-length",
65    "transfer-encoding",
66    "te",
67    "trailer",
68];
69
70/// Performs the client handshake, building the upgrade request with sockudo's
71/// header-aware builder and reading the response locally.
72///
73/// The local response loop preserves the rejection status for
74/// [`TransportError::UpgradeRejected`] and the retry-aware log severity; caller
75/// pre-validates `extra_headers` via [`validate_extra_headers`].
76pub(crate) async fn client_handshake_with_headers<S>(
77    stream: &mut S,
78    host: &str,
79    path: &str,
80    extra_headers: &[(String, String)],
81) -> Result<HandshakeResult, TransportError>
82where
83    S: AsyncRead + AsyncWrite + Unpin,
84{
85    use tokio::io::{AsyncReadExt, AsyncWriteExt};
86
87    let key = handshake::generate_key();
88    let request =
89        handshake::build_request_with_headers(host, path, &key, None, None, Some(extra_headers))?;
90
91    stream
92        .write_all(&request)
93        .await
94        .map_err(TransportError::from)?;
95    stream.flush().await.map_err(TransportError::from)?;
96
97    let mut buf = BytesMut::with_capacity(4096);
98
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        if n == 0 {
109            return Err(TransportError::ConnectionClosed);
110        }
111
112        let parsed = match handshake::parse_response(&buf) {
113            Ok(parsed) => parsed,
114            Err(e) => {
115                let status = rejected_upgrade_status(&e, &buf);
116                log_handshake_response(host, &e, &buf, status);
117                return Err(
118                    status.map_or_else(|| TransportError::from(e), TransportError::UpgradeRejected)
119                );
120            }
121        };
122
123        if let Some((res, consumed)) = parsed {
124            let accept = res.accept.ok_or_else(|| {
125                let e = SockudoError::HandshakeFailed("missing Sec-WebSocket-Accept");
126                log_handshake_response(host, &e, &buf, None);
127                TransportError::from(e)
128            })?;
129
130            if !handshake::validate_accept_key(&key, accept) {
131                let e = SockudoError::HandshakeFailed("invalid Sec-WebSocket-Accept");
132                log_handshake_response(host, &e, &buf, None);
133                return Err(e.into());
134            }
135
136            let res_protocol = res.protocol.map(String::from);
137            let res_extensions = res.extensions.map(String::from);
138            let leftover = if consumed < buf.len() {
139                Some(buf.split_off(consumed).freeze())
140            } else {
141                None
142            };
143
144            return Ok(HandshakeResult {
145                path: path.to_string(),
146                protocol: res_protocol,
147                extensions: res_extensions,
148                leftover,
149            });
150        }
151    }
152}
153
154// Gated on the error variant rather than its message. `HandshakeFailed` means the response parsed
155// as HTTP but failed handshake validation, which is the only case where a rejection status is
156// recoverable; `InvalidHttp` stays permanent. Matching the diagnostic text would put an upstream
157// string into transport semantics, and recovering after any error would let a plausible status
158// line with malformed headers pass as a rejection.
159fn rejected_upgrade_status(err: &SockudoError, buf: &[u8]) -> Option<u16> {
160    if !matches!(err, SockudoError::HandshakeFailed(_)) {
161        return None;
162    }
163
164    let status_line_end = buf.windows(2).position(|window| window == b"\r\n")?;
165    let status_line = std::str::from_utf8(&buf[..status_line_end]).ok()?;
166    let mut parts = status_line.split_whitespace();
167    let version = parts.next()?;
168    let status = parts.next()?;
169    if !version.starts_with("HTTP/1.") || status.len() != 3 {
170        return None;
171    }
172
173    // A 101 that still failed validation is a protocol fault, not a rejection.
174    status.parse().ok().filter(|status| *status != 101)
175}
176
177// A rejection the client goes on to retry is reported at WARN: an ERROR arms the
178// `shutdown_on_error` trigger, which stops the node while the reconnect loop is still recovering.
179// Permanent rejections, malformed responses, and failed accept-key checks retain ERROR handling.
180//
181// `host` names the endpoint so a node running several connections shows which one failed. It is
182// the `Host` header built from the URL authority, so it carries no userinfo, path, or signed
183// query; the request path stays out of the message for that reason.
184fn log_handshake_response(host: &str, err: &SockudoError, buf: &BytesMut, status: Option<u16>) {
185    let response_bytes = buf.len();
186
187    match status {
188        Some(status) if retryable_status(status) => log::warn!(
189            "Sockudo handshake rejected by {host} with retryable status {status}; response bytes={response_bytes}"
190        ),
191        Some(status) => log::error!(
192            "Sockudo handshake rejected by {host} with permanent status {status}; response bytes={response_bytes}"
193        ),
194        None => log::error!(
195            "Sockudo handshake failed for {host}: {err}; response bytes={response_bytes}"
196        ),
197    }
198}
199
200pub(crate) fn validate_extra_headers(headers: &[(String, String)]) -> Result<(), SockudoError> {
201    for (name, value) in headers {
202        validate_extra_header(name, value)?;
203    }
204    Ok(())
205}
206
207fn validate_extra_header(name: &str, value: &str) -> Result<(), SockudoError> {
208    let parsed_name = name
209        .parse::<http::HeaderName>()
210        .map_err(|_| SockudoError::InvalidHttp("invalid header name"))?;
211
212    if RESERVED_UPGRADE_HEADERS.contains(&parsed_name.as_str()) {
213        return Err(SockudoError::InvalidHttp(
214            "reserved upgrade header not allowed in extra_headers",
215        ));
216    }
217
218    http::HeaderValue::from_str(value)
219        .map_err(|_| SockudoError::InvalidHttp("invalid header value"))?;
220    Ok(())
221}
222
223/// Replay bytes read during the handshake before forwarding to the inner IO.
224pub(crate) struct PrefixedIo<S> {
225    inner: S,
226    prefix: Bytes,
227}
228
229impl<S> PrefixedIo<S> {
230    pub(crate) const fn new(inner: S, prefix: Bytes) -> Self {
231        Self { inner, prefix }
232    }
233}
234
235impl<S> AsyncRead for PrefixedIo<S>
236where
237    S: AsyncRead + Unpin,
238{
239    fn poll_read(
240        mut self: Pin<&mut Self>,
241        cx: &mut Context<'_>,
242        buf: &mut ReadBuf<'_>,
243    ) -> Poll<std::io::Result<()>> {
244        if !self.prefix.is_empty() {
245            let n = self.prefix.len().min(buf.remaining());
246            let chunk = self.prefix.split_to(n);
247            buf.put_slice(&chunk);
248            return Poll::Ready(Ok(()));
249        }
250
251        Pin::new(&mut self.inner).poll_read(cx, buf)
252    }
253}
254
255impl<S> AsyncWrite for PrefixedIo<S>
256where
257    S: AsyncWrite + Unpin,
258{
259    fn poll_write(
260        mut self: Pin<&mut Self>,
261        cx: &mut Context<'_>,
262        buf: &[u8],
263    ) -> Poll<std::io::Result<usize>> {
264        Pin::new(&mut self.inner).poll_write(cx, buf)
265    }
266
267    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
268        Pin::new(&mut self.inner).poll_flush(cx)
269    }
270
271    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
272        Pin::new(&mut self.inner).poll_shutdown(cx)
273    }
274}
275
276impl From<SockudoMessage> for Message {
277    fn from(value: SockudoMessage) -> Self {
278        match value {
279            SockudoMessage::Text(b) => Self::Text(b),
280            SockudoMessage::Binary(b) => Self::Binary(b),
281            SockudoMessage::Ping(b) => Self::Ping(b),
282            SockudoMessage::Pong(b) => Self::Pong(b),
283            SockudoMessage::Close(reason) => Self::Close(reason.map(Into::into)),
284        }
285    }
286}
287
288impl From<Message> for SockudoMessage {
289    /// Converts a neutral [`Message`] into a Sockudo [`SockudoMessage`].
290    ///
291    /// Conversion is infallible: both enums carry payloads as `bytes::Bytes` across
292    /// all variants. Sockudo validates UTF-8 on Text frames at parse time, not at
293    /// send time, so feeding it non-UTF-8 bytes via [`Self::Text`] is the caller's
294    /// responsibility.
295    fn from(value: Message) -> Self {
296        match value {
297            Message::Text(b) => Self::Text(b),
298            Message::Binary(b) => Self::Binary(b),
299            Message::Ping(b) => Self::Ping(b),
300            Message::Pong(b) => Self::Pong(b),
301            Message::Close(frame) => Self::Close(frame.map(Into::into)),
302        }
303    }
304}
305
306impl From<SockudoCloseReason> for CloseFrame {
307    fn from(value: SockudoCloseReason) -> Self {
308        Self {
309            code: value.code,
310            reason: value.reason,
311        }
312    }
313}
314
315impl From<CloseFrame> for SockudoCloseReason {
316    fn from(value: CloseFrame) -> Self {
317        Self {
318            code: value.code,
319            reason: value.reason,
320        }
321    }
322}
323
324impl From<SockudoError> for TransportError {
325    fn from(value: SockudoError) -> Self {
326        match value {
327            SockudoError::Io(e) => Self::Io(e),
328            SockudoError::ConnectionClosed => Self::ConnectionClosed,
329            SockudoError::ConnectionReset => Self::ConnectionReset,
330            SockudoError::Closed(reason) => Self::ClosedByPeer(reason.map(Into::into)),
331            SockudoError::MessageTooLarge => Self::MessageTooLarge,
332            SockudoError::FrameTooLarge => Self::FrameTooLarge,
333            SockudoError::InvalidUtf8 => Self::InvalidUtf8,
334            SockudoError::InvalidFrame(msg) | SockudoError::Protocol(msg) => {
335                Self::Protocol(msg.to_string())
336            }
337            SockudoError::InvalidHttp(msg) | SockudoError::HandshakeFailed(msg) => {
338                Self::Handshake(msg.to_string())
339            }
340
341            // Keepalive and idle deadlines are dead connections, TimedOut takes
342            // the connection-drop warn path.
343            timeout @ (SockudoError::HeartbeatTimeout | SockudoError::IdleTimeout) => Self::Io(
344                std::io::Error::new(std::io::ErrorKind::TimedOut, timeout.to_string()),
345            ),
346            other => Self::Other(other.to_string()),
347        }
348    }
349}
350
351/// Adapter that lifts a `sockudo-ws` [`WebSocketStream<S>`] into a
352/// backend-agnostic [`WsTransport`].
353///
354/// Translates messages and errors to the neutral types on the way through
355/// `Stream::poll_next` and `Sink<Message>::start_send` / `poll_*`. The
356/// underlying stream is owned and forwarded to via pin projection.
357///
358/// If flushing an outbound frame returns `Pending`, the next [`Stream::poll_next`] retries the
359/// flush before reading. This prevents queued control responses from being stranded when write
360/// backpressure coincides with a quiet reader.
361pub struct SockudoTransport<S> {
362    inner: WebSocketStream<S>,
363    pending_flush: bool,
364    max_message_size_bytes: Option<usize>,
365}
366
367impl<S> SockudoTransport<S> {
368    /// Wraps an established Sockudo WebSocket stream.
369    #[inline]
370    #[must_use]
371    pub const fn new(inner: WebSocketStream<S>) -> Self {
372        Self {
373            inner,
374            pending_flush: false,
375            max_message_size_bytes: None,
376        }
377    }
378
379    pub(crate) const fn with_max_message_size(mut self, limit: Option<usize>) -> Self {
380        self.max_message_size_bytes = limit;
381        self
382    }
383
384    /// Consumes the adapter and returns the underlying stream.
385    #[inline]
386    pub fn into_inner(self) -> WebSocketStream<S> {
387        self.inner
388    }
389
390    /// Borrows the underlying stream.
391    #[inline]
392    pub const fn get_ref(&self) -> &WebSocketStream<S> {
393        &self.inner
394    }
395}
396
397impl<S> std::fmt::Debug for SockudoTransport<S> {
398    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
399        f.debug_struct(stringify!(SockudoTransport))
400            .finish_non_exhaustive()
401    }
402}
403
404impl<S> Stream for SockudoTransport<S>
405where
406    S: AsyncRead + AsyncWrite + Unpin,
407{
408    type Item = Result<Message, TransportError>;
409
410    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
411        // Drain any flush that returned Pending on a prior poll so queued
412        // control responses (Pong, close reply) reach the peer before the
413        // next read. Errors are dropped here; subsequent writes through the
414        // sink half surface them.
415        if self.pending_flush {
416            match Pin::new(&mut self.inner).poll_flush(cx) {
417                Poll::Ready(_) => self.pending_flush = false,
418                Poll::Pending => {}
419            }
420        }
421
422        let result = match Pin::new(&mut self.inner).poll_next(cx) {
423            Poll::Ready(Some(Ok(msg))) => {
424                let message = Message::from(msg);
425                // A finished single-frame message bypasses Sockudo's fragment-only message cap
426                if exceeds_message_cap(self.max_message_size_bytes, &message) {
427                    Poll::Ready(Some(Err(TransportError::MessageTooLarge)))
428                } else {
429                    Poll::Ready(Some(Ok(message)))
430                }
431            }
432            Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(TransportError::from(e)))),
433            Poll::Ready(None) => Poll::Ready(None),
434            Poll::Pending => return Poll::Pending,
435        };
436
437        // Sockudo queues automatic Pong / close-response frames into the
438        // write buffer during poll_next. Nudge a flush so they reach the peer
439        // promptly even on a reader-only client; track a pending flush so the
440        // next poll retries when backpressure stalls the write socket.
441        match Pin::new(&mut self.inner).poll_flush(cx) {
442            Poll::Ready(_) => self.pending_flush = false,
443            Poll::Pending => self.pending_flush = true,
444        }
445
446        result
447    }
448}
449
450fn exceeds_message_cap(limit: Option<usize>, message: &Message) -> bool {
451    let Some(limit) = limit else {
452        return false;
453    };
454
455    let len = match message {
456        Message::Text(bytes) | Message::Binary(bytes) => bytes.len(),
457        Message::Ping(_) | Message::Pong(_) | Message::Close(_) => return false,
458    };
459
460    len > limit
461}
462
463impl<S> Sink<Message> for SockudoTransport<S>
464where
465    S: AsyncRead + AsyncWrite + Unpin,
466{
467    type Error = TransportError;
468
469    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
470        Pin::new(&mut self.inner)
471            .poll_ready(cx)
472            .map_err(TransportError::from)
473    }
474
475    fn start_send(mut self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> {
476        Pin::new(&mut self.inner)
477            .start_send(SockudoMessage::from(item))
478            .map_err(TransportError::from)
479    }
480
481    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
482        Pin::new(&mut self.inner)
483            .poll_flush(cx)
484            .map_err(TransportError::from)
485    }
486
487    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
488        Pin::new(&mut self.inner)
489            .poll_close(cx)
490            .map_err(TransportError::from)
491    }
492}
493
494const _: fn() = || {
495    fn assert_ws_transport<T: WsTransport>() {}
496    assert_ws_transport::<SockudoTransport<tokio::net::TcpStream>>();
497};
498
499#[cfg(test)]
500mod tests {
501    use bytes::Bytes;
502    use rstest::rstest;
503    #[cfg(not(feature = "turmoil"))]
504    use sockudo_ws::handshake::generate_accept_key;
505    #[cfg(not(feature = "turmoil"))]
506    use tokio::io::{AsyncRead, AsyncReadExt, AsyncWriteExt, duplex};
507
508    use super::*;
509
510    #[cfg(not(feature = "turmoil"))]
511    async fn read_http_request<S>(stream: &mut S) -> Vec<u8>
512    where
513        S: AsyncRead + Unpin,
514    {
515        let mut buf = Vec::new();
516        let mut chunk = [0u8; 256];
517
518        loop {
519            let n = stream.read(&mut chunk).await.unwrap();
520            assert!(n > 0, "HTTP request closed before headers completed");
521            buf.extend_from_slice(&chunk[..n]);
522            if buf.windows(4).any(|window| window == b"\r\n\r\n") {
523                return buf;
524            }
525        }
526    }
527
528    #[cfg(not(feature = "turmoil"))]
529    fn build_test_response(sec_websocket_key: &str, extra_bytes: &[u8]) -> Vec<u8> {
530        let accept = generate_accept_key(sec_websocket_key);
531        let mut response = format!(
532            concat!(
533                "HTTP/1.1 101 Switching Protocols\r\n",
534                "Upgrade: websocket\r\n",
535                "Connection: Upgrade\r\n",
536                "Sec-WebSocket-Accept: {}\r\n",
537                "\r\n",
538            ),
539            accept
540        )
541        .into_bytes();
542        response.extend_from_slice(extra_bytes);
543        response
544    }
545
546    #[cfg(not(feature = "turmoil"))]
547    fn extract_header<'a>(request: &'a str, name: &str) -> Option<&'a str> {
548        request.lines().find_map(|line| {
549            let (header_name, header_value) = line.split_once(':')?;
550            if header_name.eq_ignore_ascii_case(name) {
551                Some(header_value.trim())
552            } else {
553                None
554            }
555        })
556    }
557
558    #[rstest]
559    #[tokio::test]
560    #[cfg(not(feature = "turmoil"))]
561    async fn handshake_rejects_eof_before_response() {
562        let (mut client, mut server) = duplex(4096);
563
564        let peer = tokio::spawn(async move {
565            read_http_request(&mut server).await;
566        });
567
568        let error = client_handshake_with_headers(&mut client, "localhost", "/", &[])
569            .await
570            .unwrap_err();
571        peer.await.unwrap();
572
573        assert!(matches!(error, TransportError::ConnectionClosed));
574    }
575
576    #[rstest]
577    #[case::at_limit(MAX_HTTP_HEADER_SIZE, "connection closed")]
578    #[case::above_limit(MAX_HTTP_HEADER_SIZE + 1, "handshake failed: response too large")]
579    #[tokio::test]
580    #[cfg(not(feature = "turmoil"))]
581    async fn handshake_rejects_incomplete_headers_at_size_boundary(
582        #[case] response_size: usize,
583        #[case] expected: &str,
584    ) {
585        let (mut client, mut server) = duplex(MAX_HTTP_HEADER_SIZE * 2);
586
587        let peer = tokio::spawn(async move {
588            read_http_request(&mut server).await;
589            let mut response = b"HTTP/1.1 101 Switching Protocols\r\nX-Padding: ".to_vec();
590            response.resize(response_size, b'x');
591            server.write_all(&response).await.unwrap();
592        });
593
594        let error = client_handshake_with_headers(&mut client, "localhost", "/", &[])
595            .await
596            .unwrap_err();
597        peer.await.unwrap();
598
599        assert_eq!(error.to_string(), expected);
600    }
601
602    #[tokio::test]
603    #[cfg(not(feature = "turmoil"))]
604    async fn client_handshake_with_headers_sends_custom_headers() {
605        let (mut client, mut server) = duplex(4096);
606        let headers = vec![
607            ("ok-access-key".to_string(), "key-1".to_string()),
608            ("ok-access-passphrase".to_string(), "pass-1".to_string()),
609        ];
610
611        let server_task = tokio::spawn(async move {
612            let request = read_http_request(&mut server).await;
613            let request = String::from_utf8(request).unwrap();
614
615            assert!(request.starts_with("GET /ws/v5/public-sbe?instId=BTC-USDT HTTP/1.1\r\n"));
616            assert_eq!(extract_header(&request, "Host"), Some("ws.okx.com:8443"));
617            assert_eq!(extract_header(&request, "ok-access-key"), Some("key-1"));
618            assert_eq!(
619                extract_header(&request, "ok-access-passphrase"),
620                Some("pass-1")
621            );
622
623            let sec_websocket_key = extract_header(&request, "Sec-WebSocket-Key").unwrap();
624            let response = build_test_response(sec_websocket_key, &[]);
625            server.write_all(&response).await.unwrap();
626        });
627
628        let handshake = client_handshake_with_headers(
629            &mut client,
630            "ws.okx.com:8443",
631            "/ws/v5/public-sbe?instId=BTC-USDT",
632            &headers,
633        )
634        .await
635        .unwrap();
636
637        assert_eq!(handshake.path, "/ws/v5/public-sbe?instId=BTC-USDT");
638        assert!(handshake.leftover.is_none());
639        server_task.await.unwrap();
640    }
641
642    #[rstest]
643    #[cfg(not(feature = "turmoil"))]
644    #[case::host("Host")]
645    #[case::upgrade("Upgrade")]
646    #[case::connection("Connection")]
647    #[case::sec_websocket_key("Sec-WebSocket-Key")]
648    #[case::sec_websocket_version("Sec-WebSocket-Version")]
649    #[case::sec_websocket_protocol("Sec-WebSocket-Protocol")]
650    #[case::sec_websocket_extensions("Sec-WebSocket-Extensions")]
651    #[case::content_length("Content-Length")]
652    #[case::transfer_encoding("Transfer-Encoding")]
653    #[case::te("TE")]
654    #[case::trailer("Trailer")]
655    fn validate_extra_header_rejects_reserved_upgrade_headers(#[case] name: &str) {
656        let err = validate_extra_header(name, "value").unwrap_err();
657
658        assert!(matches!(
659            err,
660            SockudoError::InvalidHttp("reserved upgrade header not allowed in extra_headers")
661        ));
662    }
663
664    #[tokio::test]
665    #[cfg(not(feature = "turmoil"))]
666    async fn client_handshake_with_headers_rejects_missing_accept() {
667        let (mut client, mut server) = duplex(4096);
668
669        let server_task = tokio::spawn(async move {
670            let _request = read_http_request(&mut server).await;
671            server
672                .write_all(
673                    b"HTTP/1.1 101 Switching Protocols\r\n\
674                      Upgrade: websocket\r\n\
675                      Connection: Upgrade\r\n\
676                      \r\n",
677                )
678                .await
679                .unwrap();
680        });
681
682        let err = client_handshake_with_headers(&mut client, "example.com", "/ws", &[])
683            .await
684            .unwrap_err();
685
686        assert!(matches!(
687            err,
688            TransportError::Handshake(ref msg) if msg == "missing Sec-WebSocket-Accept"
689        ));
690        server_task.await.unwrap();
691    }
692
693    #[tokio::test]
694    #[cfg(not(feature = "turmoil"))]
695    async fn client_handshake_with_headers_preserves_rejected_status() {
696        let (mut client, mut server) = duplex(4096);
697
698        let server_task = tokio::spawn(async move {
699            let _request = read_http_request(&mut server).await;
700            server.write_all(b"HTTP/1.1 429\r\n\r\n").await.unwrap();
701        });
702
703        let err = client_handshake_with_headers(&mut client, "example.com", "/ws", &[])
704            .await
705            .unwrap_err();
706
707        assert!(matches!(err, TransportError::UpgradeRejected(429)));
708        server_task.await.unwrap();
709    }
710
711    #[tokio::test]
712    #[cfg(not(feature = "turmoil"))]
713    async fn client_handshake_with_headers_returns_leftover_bytes() {
714        let (mut client, mut server) = duplex(4096);
715        let extra = b"\x81\x05hello";
716
717        let server_task = tokio::spawn(async move {
718            let request = read_http_request(&mut server).await;
719            let request = String::from_utf8(request).unwrap();
720            let sec_websocket_key = extract_header(&request, "Sec-WebSocket-Key").unwrap();
721            let response = build_test_response(sec_websocket_key, extra);
722            server.write_all(&response).await.unwrap();
723        });
724
725        let handshake = client_handshake_with_headers(&mut client, "example.com", "/ws", &[])
726            .await
727            .unwrap();
728
729        assert_eq!(handshake.leftover.as_deref(), Some(extra.as_slice()));
730        server_task.await.unwrap();
731    }
732
733    #[tokio::test]
734    #[cfg(not(feature = "turmoil"))]
735    async fn prefixed_io_replays_leftover_before_socket() {
736        let (client, mut server) = duplex(4096);
737        let mut prefixed = PrefixedIo::new(client, Bytes::from_static(b"abc"));
738
739        let server_task = tokio::spawn(async move {
740            server.write_all(b"def").await.unwrap();
741        });
742
743        let mut buf = [0u8; 6];
744        prefixed.read_exact(&mut buf).await.unwrap();
745
746        assert_eq!(&buf, b"abcdef");
747        server_task.await.unwrap();
748    }
749
750    #[rstest]
751    fn round_trip_text() {
752        let original = SockudoMessage::Text(Bytes::from_static(b"hello"));
753        let neutral: Message = original.into();
754        assert!(neutral.is_text());
755        assert_eq!(neutral.as_bytes(), b"hello");
756
757        let back: SockudoMessage = neutral.into();
758        match back {
759            SockudoMessage::Text(b) => assert_eq!(&b[..], b"hello"),
760            other => panic!("expected text, was {other:?}"),
761        }
762    }
763
764    #[rstest]
765    fn round_trip_binary() {
766        let original = SockudoMessage::Binary(Bytes::from_static(&[1, 2, 3]));
767        let neutral: Message = original.into();
768        assert_eq!(neutral.as_bytes(), &[1, 2, 3]);
769
770        let back: SockudoMessage = neutral.into();
771        match back {
772            SockudoMessage::Binary(b) => assert_eq!(&b[..], &[1, 2, 3]),
773            other => panic!("expected binary, was {other:?}"),
774        }
775    }
776
777    #[rstest]
778    fn round_trip_ping_pong() {
779        let neutral: Message = SockudoMessage::Ping(Bytes::from_static(b"p")).into();
780        assert!(neutral.is_ping());
781
782        let neutral: Message = SockudoMessage::Pong(Bytes::from_static(b"q")).into();
783        assert!(neutral.is_pong());
784    }
785
786    #[rstest]
787    fn close_frame_round_trip() {
788        let original = SockudoMessage::Close(Some(SockudoCloseReason {
789            code: 1000,
790            reason: "bye".into(),
791        }));
792        let neutral: Message = original.into();
793        let Message::Close(Some(frame)) = &neutral else {
794            panic!("expected close frame");
795        };
796        assert_eq!(frame.code, 1000);
797        assert_eq!(frame.reason, "bye");
798
799        let back: SockudoMessage = neutral.into();
800        let SockudoMessage::Close(Some(reason)) = back else {
801            panic!("expected close frame");
802        };
803        assert_eq!(reason.code, 1000);
804        assert_eq!(reason.reason, "bye");
805    }
806
807    #[rstest]
808    fn error_translation_closed() {
809        let err: TransportError = SockudoError::ConnectionClosed.into();
810        assert!(matches!(err, TransportError::ConnectionClosed));
811    }
812
813    #[rstest]
814    fn error_translation_utf8() {
815        let err: TransportError = SockudoError::InvalidUtf8.into();
816        assert!(matches!(err, TransportError::InvalidUtf8));
817    }
818
819    #[rstest]
820    fn error_translation_handshake() {
821        let err: TransportError = SockudoError::HandshakeFailed("bad").into();
822        assert!(matches!(err, TransportError::Handshake(_)));
823    }
824
825    #[rstest]
826    #[case(SockudoError::HeartbeatTimeout, "WebSocket Pong deadline expired")]
827    #[case(SockudoError::IdleTimeout, "WebSocket inbound idle deadline expired")]
828    fn error_translation_timeouts_are_timed_out_io(
829        #[case] sockudo: SockudoError,
830        #[case] message: &str,
831    ) {
832        let err: TransportError = sockudo.into();
833        let TransportError::Io(io_err) = &err else {
834            panic!("expected I/O timeout, was: {err:?}");
835        };
836
837        assert_eq!(io_err.kind(), std::io::ErrorKind::TimedOut);
838        assert_eq!(io_err.to_string(), message);
839        assert_eq!(err.to_string(), format!("I/O error: {message}"));
840    }
841
842    // The log-capture harness is Linux-only for CI stability.
843    #[cfg(not(feature = "turmoil"))]
844    #[cfg(target_os = "linux")]
845    #[cfg(not(all(feature = "simulation", madsim)))]
846    mod handshake_logging {
847        use log::Level;
848        use rstest::rstest;
849        use tokio::io::{AsyncWriteExt, duplex};
850
851        use super::read_http_request;
852        use crate::{
853            logging::tests::capture_logs_for,
854            transport::{error::TransportError, sockudo::client_handshake_with_headers},
855        };
856
857        const LOG_TARGETS: &[&str] = &["nautilus_network::transport::sockudo"];
858        const HOST: &str = "ws.example.com:8443";
859        const PATH_SECRET: &str = "handshake-path-secret";
860
861        #[rstest]
862        #[case::retryable_bad_gateway("HTTP/1.1 502 Bad Gateway\r\n\r\n", Level::Warn, Some(502))]
863        #[case::retryable_rate_limited(
864            "HTTP/1.1 429 Too Many Requests\r\n\r\n",
865            Level::Warn,
866            Some(429)
867        )]
868        #[case::permanent_unauthorized(
869            "HTTP/1.1 401 Unauthorized\r\n\r\n",
870            Level::Error,
871            Some(401)
872        )]
873        #[case::permanent_not_found("HTTP/1.1 404 Not Found\r\n\r\n", Level::Error, Some(404))]
874        #[case::malformed_status_line("NOT-HTTP 502 Bad Gateway\r\n\r\n", Level::Error, None)]
875        #[case::missing_accept_key(
876            "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\r\n",
877            Level::Error,
878            None
879        )]
880        #[case::invalid_accept_key(
881            "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\
882             Sec-WebSocket-Accept: AAAAAAAAAAAAAAAAAAAAAAAAAAA=\r\n\r\n",
883            Level::Error,
884            None
885        )]
886        #[tokio::test]
887        async fn handshake_failure_log_level_follows_retry_policy(
888            #[case] response: &'static str,
889            #[case] expected_level: Level,
890            #[case] expected_status: Option<u16>,
891        ) {
892            let capture = capture_logs_for(LOG_TARGETS).await;
893            let (mut client, mut server) = duplex(4096);
894
895            let server_task = tokio::spawn(async move {
896                let _request = read_http_request(&mut server).await;
897                server.write_all(response.as_bytes()).await.unwrap();
898                server.flush().await.unwrap();
899            });
900
901            let err = client_handshake_with_headers(
902                &mut client,
903                HOST,
904                &format!("/ws?token={PATH_SECRET}"),
905                &[],
906            )
907            .await
908            .expect_err("handshake should fail");
909            server_task.await.unwrap();
910
911            match expected_status {
912                Some(status) => assert!(
913                    matches!(err, TransportError::UpgradeRejected(actual) if actual == status),
914                    "expected upgrade rejection {status}, was: {err:?}"
915                ),
916                None => assert!(
917                    matches!(err, TransportError::Handshake(_)),
918                    "expected a permanent handshake failure, was: {err:?}"
919                ),
920            }
921
922            let messages = capture.messages();
923            assert_eq!(
924                messages.len(),
925                1,
926                "expected exactly one handshake log, was: {messages:?}"
927            );
928
929            let (level, message) = &messages[0];
930            assert_eq!(*level, expected_level, "unexpected level for: {message}");
931            assert!(
932                message.contains(HOST),
933                "log should name the endpoint, was: {message}"
934            );
935
936            if let Some(status) = expected_status {
937                assert!(
938                    message.contains(&status.to_string()),
939                    "log should name the rejection status, was: {message}"
940                );
941            }
942
943            assert!(
944                !message.contains(PATH_SECRET),
945                "log must not carry the request path, was: {message}"
946            );
947        }
948    }
949}
nautilus_network/transport/sockudo.rs

nautilus_network/transport/
sockudo.rs
