px_wsdom_axum/
lib.rs

1//! Integration code for conveniently using WSDOM with the Axum web framework.
2//!
3//! This library provides only one function: [socket_to_browser].
4
5use std::{
6    pin::Pin,
7    task::{Context, Poll},
8};
9
10use axum::extract::ws::{Message, WebSocket};
11use futures_util::{Future, Sink, Stream, StreamExt};
12use pin_project_lite::pin_project;
13use wsdom_core::Browser;
14
15pin_project! {
16    /// Future type returned from [socket_to_browser].
17    pub struct ToBrowserFuture<Fut: Future> {
18        #[pin] ws: WebSocket,
19        #[pin] fut: Fut,
20        browser: Browser,
21        output: Option<Fut::Output>
22    }
23}
24
25impl<Fut> Future for ToBrowserFuture<Fut>
26where
27    Fut: Future,
28{
29    type Output = Output<Fut::Output>;
30
31    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
32        let mut this = self.project();
33        match this.ws.as_mut().poll_next(cx) {
34            Poll::Ready(Some(Ok(Message::Text(message)))) => {
35                this.browser.receive_incoming_message(message);
36            }
37            Poll::Ready(None | Some(Ok(Message::Close(_)))) => {
38                return Poll::Ready(Output::ConnectionClosed);
39            }
40            Poll::Ready(Some(Ok(_))) => {}
41            Poll::Ready(Some(Err(e))) => {
42                return Poll::Ready(Output::AxumError(e));
43            }
44            Poll::Pending => {}
45        }
46        match this.ws.as_mut().poll_ready(cx) {
47            Poll::Ready(Ok(_)) => {
48                match this.browser.poll_next_unpin(cx) {
49                    Poll::Ready(Some(message)) => {
50                        match this.ws.as_mut().start_send(Message::Text(message)) {
51                            Ok(_) => {}
52                            Err(e) => return Poll::Ready(Output::AxumError(e)),
53                        }
54                    }
55                    Poll::Ready(None) => {
56                        if let Some(err) = this.browser.take_error() {
57                            return Poll::Ready(Output::WsdomError(err));
58                        } else {
59                            return Poll::Pending;
60                        }
61                    }
62                    Poll::Pending => {}
63                }
64                match this.ws.as_mut().poll_flush(cx) {
65                    Poll::Ready(Err(e)) => {
66                        return Poll::Ready(Output::AxumError(e));
67                    }
68                    Poll::Ready(Ok(_)) => {
69                        if let Some(output) = this.output.take() {
70                            return Poll::Ready(Output::Done(output));
71                        }
72                    }
73                    _ => {}
74                }
75            }
76            Poll::Ready(Err(e)) => {
77                return Poll::Ready(Output::AxumError(e));
78            }
79            Poll::Pending => {}
80        }
81        if this.output.is_none() {
82            if let Poll::Ready(t) = this.fut.poll(cx) {
83                *this.output = Some(t);
84            }
85        }
86        Poll::Pending
87    }
88}
89
90/// Output type of [ToBrowserFuture].
91pub enum Output<T> {
92    /// The inner function (the second argument passed to `socket_to_browser`) completed with this result.
93    Done(T),
94    /// The WebSocket connection was closed, for whatever reason.
95    ConnectionClosed,
96    /// Axum raised an error.
97    AxumError(axum::Error),
98    /// WSDOM raised an error.
99    WsdomError(wsdom_core::Error),
100}
101
102/// Get a [Browser] from an [axum::WebSocket] object.
103///
104/// You should first follow [axum's tutorial on how to obtain the WebSocket](https://docs.rs/axum/latest/axum/extract/ws/index.html).
105/// Once you have the WebSocket object,
106/// pass it to `socket_to_browser` along with an async function/closure that takes a Browser as argument,
107/// then await the returned Future.
108///
109/// ```rust
110/// # use wsdom_core::Browser;
111/// use wsdom_axum::socket_to_browser;
112/// use axum::extract::{WebSocketUpgrade, ws::WebSocket};
113/// use axum::response::Response;
114/// async fn axum_handler(wsu: WebSocketUpgrade) -> Response {
115///     wsu.on_upgrade(|ws: WebSocket| async move {
116///         socket_to_browser(ws, app).await;
117///     })
118/// }
119/// async fn app(browser: Browser) {
120///     // do things...
121/// }
122/// ````
123#[must_use = "the return type is a Future and should be .awaited"]
124pub fn socket_to_browser<Func, Fut>(ws: WebSocket, f: Func) -> ToBrowserFuture<Fut>
125where
126    Func: FnOnce(Browser) -> Fut,
127    Fut: Future,
128{
129    let browser = Browser::new();
130    ToBrowserFuture {
131        fut: f(browser.clone()),
132        ws,
133        browser,
134        output: None,
135    }
136}
px_wsdom_axum/lib.rs

px_wsdom_axum/
lib.rs
