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::{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    /// Convert 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 [`TransportError::InvalidUtf8`] if a `Text` payload is not
75    /// valid UTF-8.
76    fn try_from(value: Message) -> Result<Self, Self::Error> {
77        Ok(match value {
78            Message::Text(bytes) => match Utf8Bytes::try_from(bytes) {
79                Ok(text) => Self::Text(text),
80                Err(_) => return Err(TransportError::InvalidUtf8),
81            },
82            Message::Binary(bytes) => Self::Binary(bytes),
83            Message::Ping(bytes) => Self::Ping(bytes),
84            Message::Pong(bytes) => Self::Pong(bytes),
85            Message::Close(frame) => Self::Close(frame.map(Into::into)),
86        })
87    }
88}
89
90impl From<TgCloseFrame> for CloseFrame {
91    fn from(value: TgCloseFrame) -> Self {
92        Self {
93            code: u16::from(value.code),
94            reason: value.reason.as_str().to_owned(),
95        }
96    }
97}
98
99impl From<CloseFrame> for TgCloseFrame {
100    fn from(value: CloseFrame) -> Self {
101        Self {
102            code: CloseCode::from(value.code),
103            reason: value.reason.into(),
104        }
105    }
106}
107
108impl From<tungstenite::Error> for TransportError {
109    fn from(value: tungstenite::Error) -> Self {
110        match value {
111            tungstenite::Error::ConnectionClosed | tungstenite::Error::AlreadyClosed => {
112                Self::ConnectionClosed
113            }
114            tungstenite::Error::Io(e) => Self::Io(e),
115            tungstenite::Error::Tls(e) => Self::Tls(e.to_string()),
116            tungstenite::Error::Capacity(e) => match e {
117                tungstenite::error::CapacityError::MessageTooLong { .. } => Self::MessageTooLarge,
118                e @ tungstenite::error::CapacityError::TooManyHeaders => Self::Other(e.to_string()),
119            },
120            tungstenite::Error::Protocol(
121                tungstenite::error::ProtocolError::ResetWithoutClosingHandshake,
122            ) => Self::ConnectionReset,
123            tungstenite::Error::Protocol(e) => Self::Protocol(e.to_string()),
124            tungstenite::Error::Utf8(_) => Self::InvalidUtf8,
125            tungstenite::Error::Url(e) => Self::InvalidUrl(e.to_string()),
126            tungstenite::Error::Http(resp) => {
127                Self::Handshake(format!("HTTP status {}", resp.status()))
128            }
129            tungstenite::Error::HttpFormat(e) => Self::Handshake(e.to_string()),
130            other => Self::Other(other.to_string()),
131        }
132    }
133}
134
135/// Adapter that lifts a `tokio-tungstenite` [`WebSocketStream<S>`] into a
136/// backend-agnostic [`WsTransport`].
137///
138/// Translates messages and errors to the neutral types on the way through
139/// `Stream::poll_next` and `Sink<Message>::start_send` / `poll_*`. The
140/// underlying stream is owned and forwarded to via pin projection.
141#[derive(Debug)]
142pub struct TungsteniteTransport<S> {
143    inner: WebSocketStream<S>,
144}
145
146impl<S> TungsteniteTransport<S> {
147    /// Wrap an established tungstenite WebSocket stream.
148    #[inline]
149    #[must_use]
150    pub const fn new(inner: WebSocketStream<S>) -> Self {
151        Self { inner }
152    }
153
154    /// Consume the adapter and return the underlying stream.
155    #[inline]
156    pub fn into_inner(self) -> WebSocketStream<S> {
157        self.inner
158    }
159
160    /// Borrow the underlying stream.
161    #[inline]
162    pub const fn get_ref(&self) -> &WebSocketStream<S> {
163        &self.inner
164    }
165}
166
167impl<S> Stream for TungsteniteTransport<S>
168where
169    S: AsyncRead + AsyncWrite + Unpin,
170{
171    type Item = Result<Message, TransportError>;
172
173    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
174        match Pin::new(&mut self.inner).poll_next(cx) {
175            Poll::Ready(Some(Ok(msg))) => Poll::Ready(Some(Ok(Message::from(msg)))),
176            Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(TransportError::from(e)))),
177            Poll::Ready(None) => Poll::Ready(None),
178            Poll::Pending => Poll::Pending,
179        }
180    }
181}
182
183impl<S> Sink<Message> for TungsteniteTransport<S>
184where
185    S: AsyncRead + AsyncWrite + Unpin,
186{
187    type Error = TransportError;
188
189    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
190        Pin::new(&mut self.inner)
191            .poll_ready(cx)
192            .map_err(TransportError::from)
193    }
194
195    fn start_send(mut self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> {
196        let native = tungstenite::Message::try_from(item)?;
197        Pin::new(&mut self.inner)
198            .start_send(native)
199            .map_err(TransportError::from)
200    }
201
202    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
203        Pin::new(&mut self.inner)
204            .poll_flush(cx)
205            .map_err(TransportError::from)
206    }
207
208    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
209        Pin::new(&mut self.inner)
210            .poll_close(cx)
211            .map_err(TransportError::from)
212    }
213}
214
215const _: fn() = || {
216    fn assert_ws_transport<T: WsTransport>() {}
217    assert_ws_transport::<TungsteniteTransport<tokio::net::TcpStream>>();
218};
219
220#[cfg(test)]
221mod tests {
222    use bytes::Bytes;
223    use rstest::rstest;
224    use tokio_tungstenite::tungstenite::{self, Utf8Bytes};
225
226    use super::*;
227
228    #[rstest]
229    fn round_trip_text() {
230        let original = tungstenite::Message::Text(Utf8Bytes::from("hello"));
231        let neutral: Message = original.into();
232        assert!(neutral.is_text());
233        assert_eq!(neutral.as_bytes(), b"hello");
234
235        let back = tungstenite::Message::try_from(neutral).unwrap();
236        match back {
237            tungstenite::Message::Text(t) => assert_eq!(t.as_str(), "hello"),
238            other => panic!("expected text, was {other:?}"),
239        }
240    }
241
242    #[rstest]
243    fn try_from_text_rejects_invalid_utf8() {
244        let neutral = Message::Text(Bytes::from_static(&[0xFF, 0xFE]));
245        let err = tungstenite::Message::try_from(neutral).unwrap_err();
246        assert!(matches!(err, TransportError::InvalidUtf8));
247    }
248
249    #[rstest]
250    fn round_trip_binary() {
251        let original = tungstenite::Message::Binary(Bytes::from_static(&[1, 2, 3]));
252        let neutral: Message = original.into();
253        assert_eq!(neutral.as_bytes(), &[1, 2, 3]);
254
255        let back = tungstenite::Message::try_from(neutral).unwrap();
256        match back {
257            tungstenite::Message::Binary(b) => assert_eq!(&b[..], &[1, 2, 3]),
258            other => panic!("expected binary, was {other:?}"),
259        }
260    }
261
262    #[rstest]
263    fn round_trip_ping_pong() {
264        let ping = tungstenite::Message::Ping(Bytes::from_static(b"p"));
265        let neutral: Message = ping.into();
266        assert!(neutral.is_ping());
267
268        let pong = tungstenite::Message::Pong(Bytes::from_static(b"q"));
269        let neutral: Message = pong.into();
270        assert!(neutral.is_pong());
271    }
272
273    #[rstest]
274    fn close_frame_round_trip() {
275        let original = tungstenite::Message::Close(Some(TgCloseFrame {
276            code: CloseCode::Normal,
277            reason: "bye".into(),
278        }));
279        let neutral: Message = original.into();
280        let Message::Close(Some(frame)) = &neutral else {
281            panic!("expected close frame");
282        };
283        assert_eq!(frame.code, 1000);
284        assert_eq!(frame.reason, "bye");
285
286        let back = tungstenite::Message::try_from(neutral).unwrap();
287        let tungstenite::Message::Close(Some(frame)) = back else {
288            panic!("expected close frame");
289        };
290        assert_eq!(u16::from(frame.code), 1000);
291        assert_eq!(frame.reason.as_str(), "bye");
292    }
293
294    #[rstest]
295    fn error_translation_closed() {
296        let err: TransportError = tungstenite::Error::ConnectionClosed.into();
297        assert!(matches!(err, TransportError::ConnectionClosed));
298    }
299
300    #[rstest]
301    fn error_translation_reset_without_closing_handshake() {
302        let err: TransportError = tungstenite::Error::Protocol(
303            tungstenite::error::ProtocolError::ResetWithoutClosingHandshake,
304        )
305        .into();
306        assert!(matches!(err, TransportError::ConnectionReset));
307    }
308
309    #[rstest]
310    fn error_translation_utf8() {
311        let err: TransportError = tungstenite::Error::Utf8(String::from("bad")).into();
312        assert!(matches!(err, TransportError::InvalidUtf8));
313    }
314}
nautilus_network/transport/tungstenite.rs

nautilus_network/transport/
tungstenite.rs
