nautilus_network/websocket/
types.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//! Shared transport halves, message handlers, and writer commands.
17//!
18//! [`MessageReader`] and the internal writer use [`BoxedWsTransport`], keeping lifecycle code
19//! independent of the selected transport backend. [`MessageHandler`] receives messages without
20//! connection identity, while [`EpochMessageHandler`] attaches the owning reader's connection epoch
21//! to application messages and reconnect notifications.
22//!
23//! The channel handlers convert neutral [`Message`] values to Tungstenite messages at the adapter
24//! compatibility boundary. `WriterCommand` remains internal so all sink access, connection swaps,
25//! and reconnect buffering stay serialized in the writer task.
26
27use std::{fmt::Debug, sync::Arc};
28
29use futures_util::stream::{SplitSink, SplitStream};
30
31use crate::{
32    error::SendError,
33    transport::{BoxedWsTransport, CloseFrame, Message},
34};
35
36/// Sink half of the active WebSocket transport.
37///
38/// Backed by [`BoxedWsTransport`], so the writer is decoupled from the concrete
39/// backend stream type. Sends are keyed off the neutral [`Message`] enum.
40pub(crate) type MessageWriter = SplitSink<BoxedWsTransport, Message>;
41
42/// Stream half of the active WebSocket transport.
43///
44/// Backed by [`BoxedWsTransport`], yielding neutral [`Message`] values regardless
45/// of the underlying backend.
46pub type MessageReader = SplitStream<BoxedWsTransport>;
47
48/// Function type for handling WebSocket messages.
49///
50/// Handler-mode clients spawn an internal task to read messages and pass them to this handler.
51/// This enables automatic reconnection because the client can replace the reader internally.
52/// [`WebSocketClient::stream_builder`](super::WebSocketClient::stream_builder) returns a
53/// [`MessageReader`] instead.
54pub type MessageHandler = Arc<dyn Fn(Message) + Send + Sync>;
55
56/// Function type for handling WebSocket messages with connection ownership.
57///
58/// The initial connection has epoch `0`. Each replacement connection increments the epoch, and
59/// every application or reconnect-control message carries the epoch of its originating reader.
60pub type EpochMessageHandler = Arc<dyn Fn(u64, Message) + Send + Sync>;
61
62/// Function type for handling WebSocket ping messages.
63pub type PingHandler = Arc<dyn Fn(Vec<u8>) + Send + Sync>;
64
65/// Function type for handling WebSocket ping messages with connection ownership.
66pub type EpochPingHandler = Arc<dyn Fn(u64, Vec<u8>) + Send + Sync>;
67
68/// Handles a peer close frame and decides whether automatic reconnection continues.
69///
70/// Returns `true` to use the client's reconnect policy, or `false` to disconnect without retrying.
71/// The callback runs synchronously on the reader task before it ends. Close frames are delivered
72/// only to this callback, not to the message handler.
73pub type CloseHandler = Arc<dyn Fn(Option<CloseFrame>) -> bool + Send + Sync>;
74
75/// Creates a channel-based message handler.
76///
77/// The handler accepts neutral [`Message`] values. The receiver yields
78/// `tokio_tungstenite::tungstenite::Message` values for adapter compatibility.
79#[must_use]
80pub fn channel_message_handler() -> (
81    MessageHandler,
82    tokio::sync::mpsc::UnboundedReceiver<tokio_tungstenite::tungstenite::Message>,
83) {
84    let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
85
86    let handler: MessageHandler = Arc::new(move |msg: Message| {
87        match tokio_tungstenite::tungstenite::Message::try_from(msg) {
88            Ok(legacy) => {
89                if let Err(e) = tx.send(legacy) {
90                    log::debug!("Failed to send message to channel: {e}");
91                }
92            }
93            Err(e) => log::debug!("Dropping message that failed legacy conversion: {e}"),
94        }
95    });
96
97    (handler, rx)
98}
99
100/// Creates a channel-based message handler carrying each reader's connection epoch.
101///
102/// The receiver preserves the handler's `(connection_epoch, message)` order.
103#[must_use]
104pub fn channel_epoch_message_handler() -> (
105    EpochMessageHandler,
106    tokio::sync::mpsc::UnboundedReceiver<(u64, tokio_tungstenite::tungstenite::Message)>,
107) {
108    let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
109
110    let handler: EpochMessageHandler =
111        Arc::new(
112            move |epoch, msg: Message| match tokio_tungstenite::tungstenite::Message::try_from(msg)
113            {
114                Ok(legacy) => {
115                    if let Err(e) = tx.send((epoch, legacy)) {
116                        log::debug!("Failed to send connection message to channel: {e}");
117                    }
118                }
119                Err(e) => log::debug!("Dropping message that failed legacy conversion: {e}"),
120            },
121        );
122
123    (handler, rx)
124}
125
126/// A command processed by the WebSocket writer task.
127pub(crate) enum WriterCommand {
128    /// Replaces the writer after reconnection.
129    Update(MessageWriter, tokio::sync::oneshot::Sender<u64>),
130    /// Sends a message to the server.
131    Send(Message),
132    /// Sends a connection keepalive. Never replayed on a replacement connection.
133    Heartbeat(Message),
134    /// Sends once if the active connection still owns the message.
135    SendOnConnection {
136        message: Message,
137        connection_epoch: u64,
138        response_tx: tokio::sync::oneshot::Sender<Result<(), SendError>>,
139    },
140    /// Sends a pong if the active connection still owns the ping that caused it.
141    SendPongOnConnection {
142        data: Vec<u8>,
143        connection_epoch: u64,
144    },
145}
146
147impl Debug for WriterCommand {
148    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
149        match self {
150            Self::Update(_, _) => f.debug_tuple("Update").field(&"<writer>").finish(),
151            Self::Send(msg) => f.debug_tuple("Send").field(msg).finish(),
152            Self::Heartbeat(msg) => f.debug_tuple("Heartbeat").field(msg).finish(),
153            Self::SendOnConnection { message, .. } => {
154                f.debug_tuple("SendOnConnection").field(message).finish()
155            }
156            Self::SendPongOnConnection { data, .. } => {
157                f.debug_tuple("SendPongOnConnection").field(data).finish()
158            }
159        }
160    }
161}
162
163#[cfg(test)]
164mod tests {
165    use bytes::Bytes;
166    use rstest::rstest;
167    use tokio_tungstenite::tungstenite::Message as TgMessage;
168
169    use super::*;
170    use crate::RECONNECTED;
171
172    #[rstest]
173    fn channel_handler_drops_invalid_utf8_text_without_panic() {
174        let (handler, mut rx) = channel_message_handler();
175
176        // Invalid UTF-8 in a Text frame must not propagate or panic,
177        // the shim logs and drops it when Utf8Bytes::try_from fails.
178        handler(Message::Text(Bytes::from_static(&[0xFF, 0xFE])));
179        handler(Message::Binary(Bytes::from_static(b"ok")));
180
181        let received = rx.try_recv().expect("binary should arrive");
182        assert!(matches!(received, TgMessage::Binary(ref b) if b.as_ref() == b"ok"));
183        assert!(rx.try_recv().is_err(), "no further messages expected");
184    }
185
186    #[rstest]
187    fn channel_handler_forwards_valid_text() {
188        let (handler, mut rx) = channel_message_handler();
189        handler(Message::text("hello"));
190
191        let received = rx.try_recv().expect("text should arrive");
192        match received {
193            TgMessage::Text(t) => assert_eq!(t.as_str(), "hello"),
194            other => panic!("expected text, was {other:?}"),
195        }
196    }
197
198    #[rstest]
199    fn epoch_message_handler_forwards_epoch_and_text() {
200        let (handler, mut rx) = channel_epoch_message_handler();
201        handler(7, Message::text("hello"));
202
203        let (connection_epoch, received) = rx.try_recv().expect("text should arrive");
204        assert_eq!(connection_epoch, 7);
205
206        match received {
207            TgMessage::Text(text) => assert_eq!(text.as_str(), "hello"),
208            other => panic!("expected text, was {other:?}"),
209        }
210    }
211
212    #[rstest]
213    fn channel_handler_forwards_reconnected_text_to_rust_consumers() {
214        let (handler, mut rx) = channel_message_handler();
215        handler(Message::text(RECONNECTED));
216
217        let received = rx.try_recv().expect("reconnect text should arrive");
218        match received {
219            TgMessage::Text(t) => assert_eq!(t.as_str(), RECONNECTED),
220            other => panic!("expected text, was {other:?}"),
221        }
222    }
223
224    #[rstest]
225    fn writer_command_send_debug_includes_message() {
226        let cmd = WriterCommand::Send(Message::text("hi"));
227        let formatted = format!("{cmd:?}");
228        assert!(
229            formatted.starts_with("Send("),
230            "unexpected debug output: {formatted}"
231        );
232        assert!(
233            formatted.contains("Text"),
234            "debug output should retain the message variant: {formatted}"
235        );
236        assert!(
237            formatted.contains("hi"),
238            "debug output should retain the payload bytes: {formatted}"
239        );
240    }
241}
nautilus_network/websocket/types.rs

nautilus_network/websocket/
types.rs
