nautilus_network/transport/
tungstenite.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//! `tokio-tungstenite` backend for the transport abstraction.
17//!
18//! Provides `From` conversions between the neutral [`Message`] and
19//! [`TransportError`] types and tungstenite's native types, plus the
20//! [`TungsteniteTransport<S>`] adapter that lifts a tungstenite
21//! `WebSocketStream<S>` into a backend-agnostic [`WsTransport`].
22//!
23//! The message conversions are structural (no payload copies): tungstenite
24//! stores payloads in `Bytes` and `Utf8Bytes`, which we re-wrap directly.
25
26use std::{
27    pin::Pin,
28    task::{Context, Poll},
29};
30
31use bytes::Bytes;
32use futures_util::{Sink, Stream};
33use tokio::io::{AsyncRead, AsyncWrite};
34use tokio_tungstenite::{
35    WebSocketStream,
36    tungstenite::{
37        self, Utf8Bytes,
38        protocol::{CloseFrame as TgCloseFrame, frame::coding::CloseCode},
39    },
40};
41
42use super::{
43    error::TransportError,
44    message::{CloseFrame, Message},
45    stream::WsTransport,
46};
47
48impl From<tungstenite::Message> for Message {
49    fn from(value: tungstenite::Message) -> Self {
50        match value {
51            tungstenite::Message::Text(text) => Self::Text(Bytes::from(text)),
52            tungstenite::Message::Binary(data) => Self::Binary(data),
53            tungstenite::Message::Ping(data) => Self::Ping(data),
54            tungstenite::Message::Pong(data) => Self::Pong(data),
55            tungstenite::Message::Close(frame) => Self::Close(frame.map(Into::into)),
56
57            // Tungstenite only emits Frame when explicitly constructed; treat as binary
58            tungstenite::Message::Frame(frame) => Self::Binary(frame.into_payload()),
59        }
60    }
61}
62
63impl TryFrom<Message> for tungstenite::Message {
64    type Error = TransportError;
65
66    /// Converts a neutral [`Message`] into a tungstenite [`tungstenite::Message`].
67    ///
68    /// Validates the `Text` payload as UTF-8 because tungstenite refuses to
69    /// transmit a Text frame whose body is not valid UTF-8. Other variants
70    /// are infallible.
71    ///
72    /// # Errors
73    ///
74    /// Returns an error if a `Text` payload is not valid UTF-8.
75    fn try_from(value: Message) -> Result<Self, Self::Error> {
76        Ok(match value {
77            Message::Text(bytes) => match Utf8Bytes::try_from(bytes) {
78                Ok(text) => Self::Text(text),
79                Err(_) => return Err(TransportError::InvalidUtf8),
80            },
81            Message::Binary(bytes) => Self::Binary(bytes),
82            Message::Ping(bytes) => Self::Ping(bytes),
83            Message::Pong(bytes) => Self::Pong(bytes),
84            Message::Close(frame) => Self::Close(frame.map(Into::into)),
85        })
86    }
87}
88
89impl From<TgCloseFrame> for CloseFrame {
90    fn from(value: TgCloseFrame) -> Self {
91        Self {
92            code: u16::from(value.code),
93            reason: value.reason.as_str().to_owned(),
94        }
95    }
96}
97
98impl From<CloseFrame> for TgCloseFrame {
99    fn from(value: CloseFrame) -> Self {
100        Self {
101            code: CloseCode::from(value.code),
102            reason: value.reason.into(),
103        }
104    }
105}
106
107impl From<tungstenite::Error> for TransportError {
108    fn from(value: tungstenite::Error) -> Self {
109        match value {
110            tungstenite::Error::ConnectionClosed | tungstenite::Error::AlreadyClosed => {
111                Self::ConnectionClosed
112            }
113            tungstenite::Error::Io(e) => Self::Io(e),
114            tungstenite::Error::Tls(e) => Self::Tls(e.to_string()),
115            tungstenite::Error::Capacity(e) => match e {
116                tungstenite::error::CapacityError::MessageTooLong { .. } => Self::MessageTooLarge,
117                e @ tungstenite::error::CapacityError::TooManyHeaders => Self::Other(e.to_string()),
118            },
119            tungstenite::Error::Protocol(
120                tungstenite::error::ProtocolError::ResetWithoutClosingHandshake,
121            ) => Self::ConnectionReset,
122            tungstenite::Error::Protocol(e) => Self::Protocol(e.to_string()),
123            tungstenite::Error::Utf8(_) => Self::InvalidUtf8,
124            tungstenite::Error::Url(e) => Self::InvalidUrl(e.to_string()),
125            tungstenite::Error::Http(resp) => Self::UpgradeRejected(resp.status().as_u16()),
126            tungstenite::Error::HttpFormat(e) => Self::Handshake(e.to_string()),
127            other => Self::Other(other.to_string()),
128        }
129    }
130}
131
132/// Adapter that lifts a `tokio-tungstenite` [`WebSocketStream<S>`] into a
133/// backend-agnostic [`WsTransport`].
134///
135/// Translates messages and errors to the neutral types on the way through
136/// `Stream::poll_next` and `Sink<Message>::start_send` / `poll_*`. The
137/// underlying stream is owned and forwarded to via pin projection.
138#[derive(Debug)]
139pub struct TungsteniteTransport<S> {
140    inner: WebSocketStream<S>,
141}
142
143impl<S> TungsteniteTransport<S> {
144    /// Wraps an established Tungstenite WebSocket stream.
145    #[inline]
146    #[must_use]
147    pub const fn new(inner: WebSocketStream<S>) -> Self {
148        Self { inner }
149    }
150
151    /// Consumes the adapter and returns the underlying stream.
152    #[inline]
153    pub fn into_inner(self) -> WebSocketStream<S> {
154        self.inner
155    }
156
157    /// Borrows the underlying stream.
158    #[inline]
159    pub const fn get_ref(&self) -> &WebSocketStream<S> {
160        &self.inner
161    }
162}
163
164impl<S> Stream for TungsteniteTransport<S>
165where
166    S: AsyncRead + AsyncWrite + Unpin,
167{
168    type Item = Result<Message, TransportError>;
169
170    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
171        match Pin::new(&mut self.inner).poll_next(cx) {
172            Poll::Ready(Some(Ok(msg))) => Poll::Ready(Some(Ok(Message::from(msg)))),
173            Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(TransportError::from(e)))),
174            Poll::Ready(None) => Poll::Ready(None),
175            Poll::Pending => Poll::Pending,
176        }
177    }
178}
179
180impl<S> Sink<Message> for TungsteniteTransport<S>
181where
182    S: AsyncRead + AsyncWrite + Unpin,
183{
184    type Error = TransportError;
185
186    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
187        Pin::new(&mut self.inner)
188            .poll_ready(cx)
189            .map_err(TransportError::from)
190    }
191
192    fn start_send(mut self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> {
193        let native = tungstenite::Message::try_from(item)?;
194        Pin::new(&mut self.inner)
195            .start_send(native)
196            .map_err(TransportError::from)
197    }
198
199    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
200        Pin::new(&mut self.inner)
201            .poll_flush(cx)
202            .map_err(TransportError::from)
203    }
204
205    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
206        Pin::new(&mut self.inner)
207            .poll_close(cx)
208            .map_err(TransportError::from)
209    }
210}
211
212const _: fn() = || {
213    fn assert_ws_transport<T: WsTransport>() {}
214    assert_ws_transport::<TungsteniteTransport<tokio::net::TcpStream>>();
215};
216
217#[cfg(test)]
218mod tests {
219    use bytes::Bytes;
220    use rstest::rstest;
221    use tokio_tungstenite::tungstenite::{self, Utf8Bytes};
222
223    use super::*;
224
225    #[rstest]
226    fn round_trip_text() {
227        let original = tungstenite::Message::Text(Utf8Bytes::from("hello"));
228        let neutral: Message = original.into();
229        assert!(neutral.is_text());
230        assert_eq!(neutral.as_bytes(), b"hello");
231
232        let back = tungstenite::Message::try_from(neutral).unwrap();
233        match back {
234            tungstenite::Message::Text(t) => assert_eq!(t.as_str(), "hello"),
235            other => panic!("expected text, was {other:?}"),
236        }
237    }
238
239    #[rstest]
240    fn try_from_text_rejects_invalid_utf8() {
241        let neutral = Message::Text(Bytes::from_static(&[0xFF, 0xFE]));
242        let err = tungstenite::Message::try_from(neutral).unwrap_err();
243        assert!(matches!(err, TransportError::InvalidUtf8));
244    }
245
246    #[rstest]
247    fn round_trip_binary() {
248        let original = tungstenite::Message::Binary(Bytes::from_static(&[1, 2, 3]));
249        let neutral: Message = original.into();
250        assert_eq!(neutral.as_bytes(), &[1, 2, 3]);
251
252        let back = tungstenite::Message::try_from(neutral).unwrap();
253        match back {
254            tungstenite::Message::Binary(b) => assert_eq!(&b[..], &[1, 2, 3]),
255            other => panic!("expected binary, was {other:?}"),
256        }
257    }
258
259    #[rstest]
260    fn round_trip_ping_pong() {
261        let ping = tungstenite::Message::Ping(Bytes::from_static(b"p"));
262        let neutral: Message = ping.into();
263        assert_eq!(neutral, Message::Ping(Bytes::from_static(b"p")));
264        assert_eq!(
265            tungstenite::Message::try_from(neutral).unwrap(),
266            tungstenite::Message::Ping(Bytes::from_static(b"p")),
267        );
268
269        let pong = tungstenite::Message::Pong(Bytes::from_static(b"q"));
270        let neutral: Message = pong.into();
271        assert_eq!(neutral, Message::Pong(Bytes::from_static(b"q")));
272        assert_eq!(
273            tungstenite::Message::try_from(neutral).unwrap(),
274            tungstenite::Message::Pong(Bytes::from_static(b"q")),
275        );
276    }
277
278    #[rstest]
279    fn close_frame_round_trip() {
280        let original = tungstenite::Message::Close(Some(TgCloseFrame {
281            code: CloseCode::Normal,
282            reason: "bye".into(),
283        }));
284
285        let neutral: Message = original.into();
286
287        let Message::Close(Some(frame)) = &neutral else {
288            panic!("expected close frame");
289        };
290
291        assert_eq!(frame.code, 1000);
292        assert_eq!(frame.reason, "bye");
293
294        let back = tungstenite::Message::try_from(neutral).unwrap();
295
296        let tungstenite::Message::Close(Some(frame)) = back else {
297            panic!("expected close frame");
298        };
299
300        assert_eq!(u16::from(frame.code), 1000);
301        assert_eq!(frame.reason.as_str(), "bye");
302    }
303
304    #[rstest]
305    fn error_translation_closed() {
306        let err: TransportError = tungstenite::Error::ConnectionClosed.into();
307        assert!(matches!(err, TransportError::ConnectionClosed));
308    }
309
310    #[rstest]
311    fn error_translation_reset_without_closing_handshake() {
312        let err: TransportError = tungstenite::Error::Protocol(
313            tungstenite::error::ProtocolError::ResetWithoutClosingHandshake,
314        )
315        .into();
316        assert!(matches!(err, TransportError::ConnectionReset));
317    }
318
319    #[rstest]
320    fn error_translation_utf8() {
321        let err: TransportError = tungstenite::Error::Utf8(String::from("bad")).into();
322        assert!(matches!(err, TransportError::InvalidUtf8));
323    }
324
325    #[rstest]
326    fn error_translation_message_too_long() {
327        let err: TransportError =
328            tungstenite::Error::Capacity(tungstenite::error::CapacityError::MessageTooLong {
329                size: 65,
330                max_size: 64,
331            })
332            .into();
333
334        assert!(matches!(err, TransportError::MessageTooLarge));
335    }
336
337    #[rstest]
338    fn error_translation_too_many_headers() {
339        let err: TransportError =
340            tungstenite::Error::Capacity(tungstenite::error::CapacityError::TooManyHeaders).into();
341
342        let TransportError::Other(message) = err else {
343            panic!("expected other error, was {err:?}");
344        };
345
346        assert_eq!(message, "Too many headers");
347    }
348
349    #[rstest]
350    fn error_translation_tls() {
351        let err: TransportError =
352            tungstenite::Error::Tls(tungstenite::error::TlsError::InvalidDnsName).into();
353
354        let TransportError::Tls(message) = err else {
355            panic!("expected TLS error, was {err:?}");
356        };
357
358        assert_eq!(message, "Invalid DNS name");
359    }
360
361    #[rstest]
362    fn error_translation_url() {
363        let err: TransportError =
364            tungstenite::Error::Url(tungstenite::error::UrlError::UnsupportedUrlScheme).into();
365
366        let TransportError::InvalidUrl(message) = err else {
367            panic!("expected invalid URL error, was {err:?}");
368        };
369
370        assert_eq!(message, "URL scheme not supported");
371    }
372}
nautilus_network/transport/tungstenite.rs

nautilus_network/transport/
tungstenite.rs
