tide_disco/
socket.rs

1// Copyright (c) 2022 Espresso Systems (espressosys.com)
2// This file is part of the tide-disco library.
3
4// You should have received a copy of the MIT License
5// along with the tide-disco library. If not, see <https://mit-license.org/>.
6
7//! An interface for asynchronous communication with clients, using WebSockets.
8
9use crate::{
10    http::{content::Accept, mime},
11    request::{RequestParams, best_response_type},
12};
13use async_std::sync::Arc;
14use futures::{
15    FutureExt, Sink, SinkExt, Stream, StreamExt, TryFutureExt,
16    future::BoxFuture,
17    select, sink,
18    stream::BoxStream,
19    task::{Context, Poll},
20};
21use pin_project::pin_project;
22use serde::{Serialize, de::DeserializeOwned};
23use std::borrow::Cow;
24use std::fmt::Display;
25use std::marker::PhantomData;
26use std::pin::Pin;
27use tide_websockets::{
28    Message, WebSocketConnection,
29    tungstenite::protocol::frame::{CloseFrame, coding::CloseCode},
30};
31use vbs::{BinarySerializer, Serializer, version::StaticVersionType};
32
33pub use disco_types::error::SocketError;
34
35#[derive(Clone, Copy, Debug)]
36enum MessageType {
37    Binary,
38    Json,
39}
40
41/// A connection facilitating bi-directional, asynchronous communication with a client.
42///
43/// [Connection] implements [Stream], which can be used to receive `FromClient` messages from the
44/// client, and [Sink] which can be used to send `ToClient` messages to the client.
45#[pin_project]
46pub struct Connection<ToClient: ?Sized, FromClient, Error, VER: StaticVersionType> {
47    #[pin]
48    conn: WebSocketConnection,
49    // [Sink] wrapper around `conn`
50    sink: Pin<Box<dyn Send + Sink<Message, Error = SocketError<Error>>>>,
51    accept: MessageType,
52    #[allow(clippy::type_complexity)]
53    _phantom: PhantomData<fn(&ToClient, &FromClient, &Error, &VER) -> ()>,
54}
55
56impl<ToClient: ?Sized, FromClient: DeserializeOwned, E, VER: StaticVersionType> Stream
57    for Connection<ToClient, FromClient, E, VER>
58{
59    type Item = Result<FromClient, SocketError<E>>;
60
61    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
62        // Get a `Pin<&mut WebSocketConnection>` for the underlying connection, so we can use the
63        // `Stream` implementation of that field.
64        match self.project().conn.poll_next(cx) {
65            Poll::Ready(None) => Poll::Ready(None),
66            Poll::Ready(Some(Err(err))) => {
67                Poll::Ready(Some(Err(SocketError::WebSockets(err.to_string()))))
68            }
69            Poll::Ready(Some(Ok(msg))) => Poll::Ready(Some(match msg {
70                Message::Binary(bytes) => {
71                    Serializer::<VER>::deserialize(&bytes).map_err(SocketError::from)
72                }
73                Message::Text(s) => serde_json::from_str(&s).map_err(SocketError::from),
74                _ => Err(SocketError::UnsupportedMessageType),
75            })),
76            Poll::Pending => Poll::Pending,
77        }
78    }
79}
80
81impl<ToClient: Serialize + ?Sized, FromClient, E, VER: StaticVersionType> Sink<&ToClient>
82    for Connection<ToClient, FromClient, E, VER>
83{
84    type Error = SocketError<E>;
85
86    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
87        self.sink.as_mut().poll_ready(cx).map_err(SocketError::from)
88    }
89
90    fn start_send(mut self: Pin<&mut Self>, item: &ToClient) -> Result<(), Self::Error> {
91        let msg = match self.accept {
92            MessageType::Binary => Message::Binary(Serializer::<VER>::serialize(item)?),
93            MessageType::Json => Message::Text(serde_json::to_string(item)?),
94        };
95        self.sink.as_mut().start_send(msg)
96    }
97
98    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
99        self.sink.as_mut().poll_flush(cx).map_err(SocketError::from)
100    }
101
102    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
103        self.sink.as_mut().poll_close(cx).map_err(SocketError::from)
104    }
105}
106
107impl<ToClient: Serialize, FromClient, E, VER: StaticVersionType> Sink<ToClient>
108    for Connection<ToClient, FromClient, E, VER>
109{
110    type Error = SocketError<E>;
111
112    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
113        Sink::<&ToClient>::poll_ready(self, cx)
114    }
115
116    fn start_send(self: Pin<&mut Self>, item: ToClient) -> Result<(), Self::Error> {
117        self.start_send(&item)
118    }
119
120    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
121        Sink::<&ToClient>::poll_flush(self, cx)
122    }
123
124    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
125        Sink::<&ToClient>::poll_close(self, cx)
126    }
127}
128
129impl<ToClient: ?Sized, FromClient, E, VER: StaticVersionType>
130    Connection<ToClient, FromClient, E, VER>
131{
132    #[allow(clippy::result_large_err)]
133    fn new(accept: &Accept, conn: WebSocketConnection) -> Result<Self, SocketError<E>> {
134        let ty = best_response_type(accept, &[mime::JSON, mime::BYTE_STREAM])?;
135        let ty = if ty == mime::JSON {
136            MessageType::Json
137        } else if ty == mime::BYTE_STREAM {
138            MessageType::Binary
139        } else {
140            unreachable!()
141        };
142        Ok(Self {
143            sink: Self::sink(conn.clone()),
144            conn,
145            accept: ty,
146            _phantom: Default::default(),
147        })
148    }
149
150    /// Wrap a `WebSocketConnection` in a type that implements `Sink<Message>`.
151    fn sink(
152        conn: WebSocketConnection,
153    ) -> Pin<Box<dyn Send + Sink<Message, Error = SocketError<E>>>> {
154        Box::pin(sink::unfold(conn, |conn, msg| async move {
155            conn.send(msg)
156                .await
157                .map_err(|err| SocketError::WebSockets(err.to_string()))?;
158            Ok(conn)
159        }))
160    }
161}
162
163impl<ToClient: ?Sized, FromClient, E, VER: StaticVersionType> Clone
164    for Connection<ToClient, FromClient, E, VER>
165{
166    fn clone(&self) -> Self {
167        Self {
168            sink: Self::sink(self.conn.clone()),
169            conn: self.conn.clone(),
170            accept: self.accept,
171            _phantom: Default::default(),
172        }
173    }
174}
175
176pub(crate) type Handler<State, Error> = Box<
177    dyn 'static
178        + Send
179        + Sync
180        + Fn(RequestParams, WebSocketConnection, &State) -> BoxFuture<Result<(), SocketError<Error>>>,
181>;
182
183pub(crate) fn handler<State, Error, ToClient, FromClient, F, VER: StaticVersionType>(
184    f: F,
185) -> Handler<State, Error>
186where
187    F: 'static
188        + Send
189        + Sync
190        + Fn(
191            RequestParams,
192            Connection<ToClient, FromClient, Error, VER>,
193            &State,
194        ) -> BoxFuture<Result<(), Error>>,
195    State: 'static + Send + Sync,
196    ToClient: 'static + Serialize + ?Sized,
197    FromClient: 'static + DeserializeOwned,
198    Error: 'static + Send + Display,
199{
200    raw_handler(move |req, conn, state| {
201        f(req, conn, state)
202            .map_err(SocketError::AppSpecific)
203            .boxed()
204    })
205}
206
207struct StreamHandler<F, VER: StaticVersionType>(F, PhantomData<VER>);
208
209impl<F, VER: StaticVersionType> StreamHandler<F, VER> {
210    fn handle<'a, State, Error, Msg>(
211        &self,
212        req: RequestParams,
213        conn: Connection<Msg, (), Error, VER>,
214        state: &'a State,
215    ) -> BoxFuture<'a, Result<(), SocketError<Error>>>
216    where
217        F: 'static + Send + Sync + Fn(RequestParams, &State) -> BoxStream<Result<Msg, Error>>,
218        State: 'static + Send + Sync,
219        Msg: 'static + Serialize + Send + Sync,
220        Error: 'static + Send,
221        VER: 'static + Send + Sync,
222    {
223        let mut stream = (self.0)(req, state).fuse();
224        async move {
225            // Appease the borrow checker, this is a cheap clone
226            let (mut send, mut recv) = (conn.clone(), conn);
227
228            // Neither stream is documented to be cancel-safe, so we store the futures outside select
229            let mut item_fut = stream.next();
230            let mut client_fut = recv.next().fuse();
231
232            loop {
233                select! {
234                    item = item_fut => {
235                        match item {
236                            Some(msg) => {
237                                send.send(&msg.map_err(SocketError::AppSpecific)?).await?;
238                                item_fut = stream.next();
239                            }
240                            None => {
241                                break;
242                            }
243                        }
244                    }
245                    // We don't actually expect to receive anything from the client,
246                    // it is being polled only to handle connection closure by the client
247                    client_msg = client_fut => {
248                        client_fut = recv.next().fuse();
249                        match client_msg {
250                            None => return Ok(()),
251                            Some(Err(e)) => return Err(e),
252                            _ => {}
253                        }
254                    }
255                };
256            }
257            Ok(())
258        }
259        .boxed()
260    }
261}
262
263pub(crate) fn stream_handler<State, Error, Msg, F, VER>(f: F) -> Handler<State, Error>
264where
265    F: 'static + Send + Sync + Fn(RequestParams, &State) -> BoxStream<Result<Msg, Error>>,
266    State: 'static + Send + Sync,
267    Msg: 'static + Serialize + Send + Sync,
268    Error: 'static + Send + Display,
269    VER: 'static + Send + Sync + StaticVersionType,
270{
271    let handler: StreamHandler<F, VER> = StreamHandler(f, Default::default());
272    raw_handler(move |req, conn, state| handler.handle(req, conn, state))
273}
274
275fn raw_handler<State, Error, ToClient, FromClient, F, VER>(f: F) -> Handler<State, Error>
276where
277    F: 'static
278        + Send
279        + Sync
280        + Fn(
281            RequestParams,
282            Connection<ToClient, FromClient, Error, VER>,
283            &State,
284        ) -> BoxFuture<Result<(), SocketError<Error>>>,
285    State: 'static + Send + Sync,
286    ToClient: 'static + Serialize + ?Sized,
287    FromClient: 'static + DeserializeOwned,
288    Error: 'static + Send + Display,
289    VER: StaticVersionType,
290{
291    let close = |conn: WebSocketConnection, res: Result<(), SocketError<Error>>| async move {
292        // When the handler finishes, send a close message. If there was an error, include the error
293        // message.
294        let msg = res.as_ref().err().map(|err| CloseFrame {
295            code: CloseCode::Error,
296            reason: Cow::Owned(err.to_string()),
297        });
298        conn.send(Message::Close(msg))
299            .await
300            .map_err(|err| SocketError::WebSockets(err.to_string()))?;
301        res
302    };
303    Box::new(move |req, raw_conn, state| {
304        let accept = match req.accept() {
305            Ok(accept) => accept,
306            Err(err) => return close(raw_conn, Err(err.into())).boxed(),
307        };
308        let conn = match Connection::new(&accept, raw_conn.clone()) {
309            Ok(conn) => conn,
310            Err(err) => return close(raw_conn, Err(err)).boxed(),
311        };
312        f(req, conn, state)
313            .then(move |res| close(raw_conn, res))
314            .boxed()
315    })
316}
317
318struct MapErr<State, Error, F> {
319    handler: Handler<State, Error>,
320    map: Arc<F>,
321}
322
323impl<State, Error, F> MapErr<State, Error, F> {
324    fn handle<'a, Error2>(
325        &self,
326        req: RequestParams,
327        conn: WebSocketConnection,
328        state: &'a State,
329    ) -> BoxFuture<'a, Result<(), SocketError<Error2>>>
330    where
331        F: 'static + Send + Sync + Fn(Error) -> Error2,
332        State: 'static + Send + Sync,
333        Error: 'static,
334    {
335        let map = self.map.clone();
336        let fut = (self.handler)(req, conn, state);
337        async move { fut.await.map_err(|err| err.map_app_specific(&*map)) }.boxed()
338    }
339}
340
341pub(crate) fn map_err<State, Error, Error2>(
342    h: Handler<State, Error>,
343    f: impl 'static + Send + Sync + Fn(Error) -> Error2,
344) -> Handler<State, Error2>
345where
346    State: 'static + Send + Sync,
347    Error: 'static,
348{
349    let handler = MapErr {
350        handler: h,
351        map: Arc::new(f),
352    };
353    Box::new(move |req, conn, state| handler.handle(req, conn, state))
354}
355
356#[cfg(test)]
357mod tests {
358    use super::*;
359    use crate::{Api, App, Url, error::ServerError, testing::test_ws_client};
360    use async_std::task::{sleep, spawn};
361    use async_tungstenite::tungstenite::Message as TungsteniteMessage;
362    use futures::{StreamExt, stream};
363    use pin_project::pinned_drop;
364    use portpicker::pick_unused_port;
365    use std::{
366        sync::{
367            Arc,
368            atomic::{AtomicBool, Ordering},
369        },
370        time::Duration,
371    };
372    use vbs::version::StaticVersion;
373
374    type StaticVer01 = StaticVersion<0, 1>;
375
376    #[pin_project(PinnedDrop)]
377    struct DropStream<S: Stream> {
378        #[pin]
379        stream: S,
380        dropped: Arc<AtomicBool>,
381    }
382
383    impl<S: Stream> Stream for DropStream<S> {
384        type Item = S::Item;
385
386        fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
387            let stream = self.project().stream;
388            stream.poll_next(cx)
389        }
390    }
391
392    #[pinned_drop]
393    impl<S: Stream> PinnedDrop for DropStream<S> {
394        fn drop(self: Pin<&mut Self>) {
395            self.dropped.store(true, Ordering::SeqCst);
396        }
397    }
398
399    #[async_std::test]
400    async fn test_stream_handler_client_closure() {
401        // Setup: Create a simple API with a stream endpoint
402        let port = pick_unused_port().expect("No ports available");
403
404        let mut app = App::<(), ServerError>::with_state(());
405        let toml_content = r#"
406            [meta]
407            FORMAT_VERSION = "0.1.0"
408
409            [route.stream_test]
410            PATH = ["/stream"]
411            METHOD = "SOCKET"
412            "#;
413
414        let mut api =
415            Api::<(), ServerError, StaticVer01>::new(toml_content.parse::<toml::Value>().unwrap())
416                .unwrap();
417
418        // Register a stream handler that sends multiple messages and indicates
419        // whether it was dropped
420        let dropped = Arc::new(AtomicBool::new(false));
421        let _dropped = dropped.clone();
422        api.stream("stream_test", move |_req, _state| {
423            Box::pin(DropStream {
424                stream: stream::iter(0..).map(Result::Ok),
425                dropped: _dropped.clone(),
426            })
427        })
428        .unwrap();
429
430        app.register_module("test", api).unwrap();
431
432        // Start the server
433        spawn(async move {
434            app.serve(format!("127.0.0.1:{}", port), StaticVer01::instance())
435                .await
436                .unwrap();
437        });
438
439        // Give the server time to start
440        sleep(Duration::from_millis(500)).await;
441
442        // Connect as a client
443        let url = Url::parse(&format!("http://127.0.0.1:{}/test/stream", port)).unwrap();
444        let mut ws_stream = test_ws_client(url).await;
445
446        // Receive a few messages
447        let mut received_count = 0;
448        for _ in 0..5 {
449            if let Some(Ok(TungsteniteMessage::Text(msg))) = ws_stream.next().await {
450                let parsed: usize = serde_json::from_str(&msg).unwrap();
451                assert_eq!(parsed, received_count);
452                received_count += 1;
453            }
454        }
455
456        // Close the client connection
457        ws_stream
458            .close(None)
459            .await
460            .expect("Failed to close connection");
461
462        // Wait a bit to ensure the server processes the closure
463        sleep(Duration::from_millis(300)).await;
464
465        // The underlying stream should've been dropped
466        assert!(dropped.load(Ordering::SeqCst));
467    }
468}
tide_disco/socket.rs

tide_disco/
socket.rs
