flowcat_core/transport/
ws_media.rs

1// SPDX-License-Identifier: Apache-2.0
2//
3//! `WsMediaTransport` — a [`MediaSocket`] over a `tokio-tungstenite` WebSocket.
4//!
5//! Flowcat does not own the socket: it is **generic over any tungstenite-message
6//! stream/sink** so a host can plug in whatever it has. The standalone client
7//! path ([`WsMediaTransport::connect`]) uses `tokio_tungstenite::connect_async`;
8//! an embedder can adapt its own inbound (e.g. axum) WS into the same
9//! `Stream<Item = Result<Message, _>> + Sink<Message>` shape via a thin wrapper
10//! and construct a `WsMediaTransport` over that. See DESIGN.md "Crate layout".
11
12use std::fmt::Display;
13
14use async_trait::async_trait;
15use futures_util::{SinkExt, StreamExt};
16use tokio_tungstenite::tungstenite::Message;
17
18use crate::error::FlowcatError;
19use crate::transport::socket::MediaSocket;
20use crate::types::WsIn;
21
22/// A [`MediaSocket`] backed by any tungstenite-`Message` WebSocket stream/sink.
23///
24/// `S` is the underlying duplex socket. It is intentionally generic so Flowcat
25/// stays host-agnostic: both a `tokio_tungstenite::WebSocketStream` (client) and
26/// an axum-WS adapter (server) satisfy
27/// `Stream<Item = Result<Message, E>> + Sink<Message, Error = E>`.
28///
29/// `WebSocketStream` uses one error type (`tungstenite::Error`) for *both* the
30/// `Stream` and the `Sink`, so a single `E` bound is sufficient and keeps the
31/// `From`/error handling simple.
32pub struct WsMediaTransport<S> {
33    socket: S,
34}
35
36impl<S> WsMediaTransport<S> {
37    /// Wrap an already-established duplex WS stream/sink (server path, mocks,
38    /// or any pre-built socket) as a `MediaSocket`.
39    pub fn new(socket: S) -> Self {
40        Self { socket }
41    }
42
43    /// Consume the transport and return the underlying socket.
44    pub fn into_inner(self) -> S {
45        self.socket
46    }
47}
48
49/// The concrete socket type produced by [`tokio_tungstenite::connect_async`].
50type ClientSocket =
51    tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;
52
53impl WsMediaTransport<ClientSocket> {
54    /// Connect to `url` (`ws://` or `wss://`) and wrap the socket (client path).
55    ///
56    /// This is the standalone/CLI path; servers should construct via
57    /// [`WsMediaTransport::new`] over their own adapted socket instead.
58    pub async fn connect(url: &str) -> Result<Self, FlowcatError> {
59        let (socket, _resp) = tokio_tungstenite::connect_async(url)
60            .await
61            .map_err(|e| FlowcatError::Transport(format!("connect_async({url}): {e}")))?;
62        Ok(Self::new(socket))
63    }
64}
65
66#[async_trait]
67impl<S, E> MediaSocket for WsMediaTransport<S>
68where
69    S: futures_util::Stream<Item = Result<Message, E>>
70        + futures_util::Sink<Message, Error = E>
71        + Unpin
72        + Send,
73    E: Display + Send,
74{
75    async fn recv(&mut self) -> Option<WsIn> {
76        // Loop so control frames (Ping/Pong) and the rare protocol `Frame` don't
77        // surface as actionable inbound events — keep reading until we hit a
78        // Text/Binary/Close, a stream error, or end-of-stream.
79        loop {
80            match self.socket.next().await {
81                Some(Ok(Message::Text(t))) => return Some(WsIn::Text(t.as_str().to_owned())),
82                Some(Ok(Message::Binary(b))) => return Some(WsIn::Binary(b.to_vec())),
83                Some(Ok(Message::Close(_))) => return Some(WsIn::Close),
84                // Keepalives + raw frames carry no media for Flowcat; ignore.
85                Some(Ok(Message::Ping(_)))
86                | Some(Ok(Message::Pong(_)))
87                | Some(Ok(Message::Frame(_))) => continue,
88                // A transport error is terminal for the call: treat like a peer
89                // close so the pipeline finalizes cleanly.
90                Some(Err(_)) => return Some(WsIn::Close),
91                None => return None,
92            }
93        }
94    }
95
96    async fn send_text(&mut self, s: String) -> Result<(), FlowcatError> {
97        self.socket
98            .send(Message::text(s))
99            .await
100            .map_err(|e| FlowcatError::Transport(format!("send_text: {e}")))
101    }
102
103    async fn send_binary(&mut self, b: Vec<u8>) -> Result<(), FlowcatError> {
104        self.socket
105            .send(Message::binary(b))
106            .await
107            .map_err(|e| FlowcatError::Transport(format!("send_binary: {e}")))
108    }
109}
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114    use futures_util::stream;
115    use std::collections::VecDeque;
116    use std::pin::Pin;
117    use std::task::{Context, Poll};
118
119    /// A tiny duplex mock: a canned inbound stream + an outbound capture buffer.
120    /// Implements `Stream<Item = Result<Message, String>>` and
121    /// `Sink<Message, Error = String>` so it matches the `MediaSocket` bounds.
122    struct MockSocket {
123        inbound: stream::Iter<std::vec::IntoIter<Result<Message, String>>>,
124        sent: std::sync::Arc<std::sync::Mutex<VecDeque<Message>>>,
125    }
126
127    impl MockSocket {
128        fn new(
129            inbound: Vec<Result<Message, String>>,
130        ) -> (Self, std::sync::Arc<std::sync::Mutex<VecDeque<Message>>>) {
131            let sent = std::sync::Arc::new(std::sync::Mutex::new(VecDeque::new()));
132            (
133                Self {
134                    inbound: stream::iter(inbound),
135                    sent: sent.clone(),
136                },
137                sent,
138            )
139        }
140    }
141
142    impl futures_util::Stream for MockSocket {
143        type Item = Result<Message, String>;
144        fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
145            Pin::new(&mut self.inbound).poll_next(cx)
146        }
147    }
148
149    impl futures_util::Sink<Message> for MockSocket {
150        type Error = String;
151        fn poll_ready(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), String>> {
152            Poll::Ready(Ok(()))
153        }
154        fn start_send(self: Pin<&mut Self>, item: Message) -> Result<(), String> {
155            self.sent.lock().unwrap().push_back(item);
156            Ok(())
157        }
158        fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), String>> {
159            Poll::Ready(Ok(()))
160        }
161        fn poll_close(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), String>> {
162            Poll::Ready(Ok(()))
163        }
164    }
165
166    #[tokio::test]
167    async fn recv_maps_text_binary_close_and_skips_keepalives() {
168        let inbound = vec![
169            Ok(Message::Ping(bytes::Bytes::from_static(b"hb"))), // skipped
170            Ok(Message::text("hello")),
171            Ok(Message::Pong(bytes::Bytes::new())), // skipped
172            Ok(Message::binary(vec![1u8, 2, 3])),
173            Ok(Message::Close(None)),
174        ];
175        let (mock, _sent) = MockSocket::new(inbound);
176        let mut t = WsMediaTransport::new(mock);
177
178        assert_eq!(t.recv().await, Some(WsIn::Text("hello".to_string())));
179        assert_eq!(t.recv().await, Some(WsIn::Binary(vec![1, 2, 3])));
180        assert_eq!(t.recv().await, Some(WsIn::Close));
181        // After the canned stream is exhausted, recv yields None.
182        assert_eq!(t.recv().await, None);
183    }
184
185    #[tokio::test]
186    async fn recv_treats_stream_error_as_close() {
187        let inbound = vec![Err("boom".to_string()), Ok(Message::text("unreached"))];
188        let (mock, _sent) = MockSocket::new(inbound);
189        let mut t = WsMediaTransport::new(mock);
190        assert_eq!(t.recv().await, Some(WsIn::Close));
191    }
192
193    #[tokio::test]
194    async fn send_text_and_binary_emit_matching_messages() {
195        let (mock, sent) = MockSocket::new(vec![]);
196        let mut t = WsMediaTransport::new(mock);
197
198        t.send_text("ping".to_string()).await.unwrap();
199        t.send_binary(vec![9, 8, 7]).await.unwrap();
200
201        let q = sent.lock().unwrap();
202        match &q[0] {
203            Message::Text(s) => assert_eq!(s.as_str(), "ping"),
204            other => panic!("expected Text, got {other:?}"),
205        }
206        match &q[1] {
207            Message::Binary(b) => assert_eq!(b.as_ref(), &[9u8, 8, 7]),
208            other => panic!("expected Binary, got {other:?}"),
209        }
210    }
211}
flowcat_core/transport/ws_media.rs

flowcat_core/transport/
ws_media.rs
