churust_core/
ws.rs

1//! WebSocket support (feature `ws`).
2//!
3//! A handler upgrades a request by taking the [`WebSocketUpgrade`] extractor and
4//! calling [`WebSocketUpgrade::on_upgrade`]:
5//!
6//! ```no_run
7//! use churust_core::Churust;
8//! use churust_core::ws::WebSocketUpgrade;
9//!
10//! # fn build() {
11//! Churust::server().routing(|r| {
12//!     r.get("/echo", |ws: WebSocketUpgrade| async move {
13//!         ws.on_upgrade(|mut sock| async move {
14//!             while let Some(Ok(msg)) = sock.recv().await {
15//!                 if sock.send(msg).await.is_err() { break; }
16//!             }
17//!         })
18//!     });
19//! });
20//! # }
21//! ```
22
23use crate::call::Call;
24use crate::error::{Error, Result};
25use crate::extract::FromCallParts;
26use crate::response::Response;
27use async_trait::async_trait;
28use futures_util::{SinkExt, StreamExt};
29use http::header::{
30    CONNECTION, SEC_WEBSOCKET_ACCEPT, SEC_WEBSOCKET_KEY, SEC_WEBSOCKET_PROTOCOL,
31    SEC_WEBSOCKET_VERSION, UPGRADE,
32};
33use http::{HeaderMap, HeaderValue, StatusCode};
34use hyper::upgrade::OnUpgrade;
35use hyper_util::rt::TokioIo;
36use std::future::Future;
37use std::sync::{Arc, Mutex};
38use tokio_tungstenite::tungstenite::protocol::Role;
39use tokio_tungstenite::WebSocketStream;
40
41/// A cloneable, takeable holder for hyper's pending connection upgrade. The
42/// engine inserts one into a [`Call`]'s extensions for WebSocket
43/// handshake requests; [`WebSocketUpgrade`] takes it back out.
44#[derive(Clone)]
45pub struct OnUpgradeHandle(Arc<Mutex<Option<OnUpgrade>>>);
46
47impl OnUpgradeHandle {
48    /// Wrap a pending upgrade.
49    pub fn new(on_upgrade: OnUpgrade) -> Self {
50        Self(Arc::new(Mutex::new(Some(on_upgrade))))
51    }
52
53    /// Take the upgrade future (can only succeed once).
54    pub(crate) fn take(&self) -> Option<OnUpgrade> {
55        self.0.lock().ok().and_then(|mut guard| guard.take())
56    }
57}
58
59impl std::fmt::Debug for OnUpgradeHandle {
60    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
61        f.write_str("OnUpgradeHandle")
62    }
63}
64
65/// True if the request headers request a WebSocket upgrade (`Connection:
66/// upgrade` + `Upgrade: websocket`, case-insensitive).
67pub(crate) fn is_upgrade_request(headers: &HeaderMap) -> bool {
68    let connection_upgrade = headers
69        .get(CONNECTION)
70        .and_then(|v| v.to_str().ok())
71        .map(|v| {
72            v.to_ascii_lowercase()
73                .split(',')
74                .any(|p| p.trim() == "upgrade")
75        })
76        .unwrap_or(false);
77    let upgrade_websocket = headers
78        .get(UPGRADE)
79        .and_then(|v| v.to_str().ok())
80        .map(|v| v.eq_ignore_ascii_case("websocket"))
81        .unwrap_or(false);
82    connection_upgrade && upgrade_websocket
83}
84
85use tokio_tungstenite::tungstenite::Message as TMessage;
86
87/// A WebSocket message. A deliberately small enum so user code never has to name
88/// `tungstenite` types.
89#[derive(Debug, Clone, PartialEq, Eq)]
90pub enum Message {
91    /// A UTF-8 text frame.
92    Text(String),
93    /// A binary frame.
94    Binary(Vec<u8>),
95    /// A ping control frame (payload echoed back by the peer as a pong).
96    Ping(Vec<u8>),
97    /// A pong control frame.
98    Pong(Vec<u8>),
99    /// A close frame (connection is closing).
100    Close,
101}
102
103impl From<Message> for TMessage {
104    fn from(m: Message) -> Self {
105        match m {
106            Message::Text(s) => TMessage::Text(s.into()),
107            Message::Binary(b) => TMessage::Binary(b.into()),
108            Message::Ping(b) => TMessage::Ping(b.into()),
109            Message::Pong(b) => TMessage::Pong(b.into()),
110            Message::Close => TMessage::Close(None),
111        }
112    }
113}
114
115impl From<TMessage> for Message {
116    fn from(m: TMessage) -> Self {
117        match m {
118            TMessage::Text(s) => Message::Text(s.to_string()),
119            TMessage::Binary(b) => Message::Binary(b.to_vec()),
120            TMessage::Ping(b) => Message::Ping(b.to_vec()),
121            TMessage::Pong(b) => Message::Pong(b.to_vec()),
122            TMessage::Close(_) => Message::Close,
123            // Raw frames are not surfaced to user code.
124            _ => Message::Close,
125        }
126    }
127}
128
129/// An established WebSocket connection. Obtained inside the
130/// [`WebSocketUpgrade::on_upgrade`] callback.
131pub struct WebSocket {
132    inner: WebSocketStream<TokioIo<hyper::upgrade::Upgraded>>,
133}
134
135impl WebSocket {
136    /// Receive the next message. `None` when the connection has closed.
137    pub async fn recv(&mut self) -> Option<Result<Message>> {
138        match self.inner.next().await {
139            Some(Ok(msg)) => Some(Ok(msg.into())),
140            Some(Err(e)) => Some(Err(Error::internal(format!("websocket recv: {e}")))),
141            None => None,
142        }
143    }
144
145    /// Send a message.
146    pub async fn send(&mut self, msg: Message) -> Result<()> {
147        self.inner
148            .send(msg.into())
149            .await
150            .map_err(|e| Error::internal(format!("websocket send: {e}")))
151    }
152
153    /// Convenience: send a text message.
154    pub async fn send_text(&mut self, text: impl Into<String>) -> Result<()> {
155        self.send(Message::Text(text.into())).await
156    }
157
158    /// Convenience: send a binary message.
159    pub async fn send_binary(&mut self, bytes: impl Into<Vec<u8>>) -> Result<()> {
160        self.send(Message::Binary(bytes.into())).await
161    }
162
163    /// Close the connection.
164    pub async fn close(&mut self) -> Result<()> {
165        self.inner
166            .close(None)
167            .await
168            .map_err(|e| Error::internal(format!("websocket close: {e}")))
169    }
170}
171
172/// Extractor that represents a pending WebSocket upgrade. A handler takes it as
173/// an argument, then calls [`on_upgrade`](WebSocketUpgrade::on_upgrade).
174///
175/// Extraction fails with **426 Upgrade Required** if the request is not a valid
176/// WebSocket handshake.
177pub struct WebSocketUpgrade {
178    on_upgrade: OnUpgrade,
179    accept_key: HeaderValue,
180    protocol: Option<HeaderValue>,
181}
182
183#[async_trait]
184impl FromCallParts for WebSocketUpgrade {
185    async fn from_call_parts(call: &mut Call) -> Result<Self> {
186        let version_ok = call
187            .header(SEC_WEBSOCKET_VERSION.as_str())
188            .map(|v| v == "13")
189            .unwrap_or(false);
190        if !is_upgrade_request(call.headers()) || !version_ok {
191            return Err(Error::new(
192                StatusCode::UPGRADE_REQUIRED,
193                "expected a WebSocket upgrade request",
194            )
195            .with_response_header(UPGRADE, HeaderValue::from_static("websocket")));
196        }
197
198        let key = call
199            .header(SEC_WEBSOCKET_KEY.as_str())
200            .ok_or_else(|| Error::bad_request("missing Sec-WebSocket-Key"))?;
201        let accept = tokio_tungstenite::tungstenite::handshake::derive_accept_key(key.as_bytes());
202        let accept_key =
203            HeaderValue::from_str(&accept).map_err(|_| Error::internal("invalid accept key"))?;
204
205        let protocol = call
206            .header(SEC_WEBSOCKET_PROTOCOL.as_str())
207            .and_then(|p| p.split(',').next())
208            .and_then(|p| HeaderValue::from_str(p.trim()).ok());
209
210        let handle = call.get::<OnUpgradeHandle>().ok_or_else(|| {
211            Error::new(
212                StatusCode::UPGRADE_REQUIRED,
213                "WebSocket upgrade unavailable (no pending connection upgrade)",
214            )
215        })?;
216        let on_upgrade = handle
217            .take()
218            .ok_or_else(|| Error::internal("WebSocket upgrade already consumed"))?;
219
220        Ok(WebSocketUpgrade {
221            on_upgrade,
222            accept_key,
223            protocol,
224        })
225    }
226}
227
228impl WebSocketUpgrade {
229    /// Finish the handshake: spawn a task that runs `callback` with the
230    /// established [`WebSocket`] once the upgrade completes, and return the
231    /// `101 Switching Protocols` response the engine will send.
232    pub fn on_upgrade<F, Fut>(self, callback: F) -> Response
233    where
234        F: FnOnce(WebSocket) -> Fut + Send + 'static,
235        Fut: Future<Output = ()> + Send + 'static,
236    {
237        let WebSocketUpgrade {
238            on_upgrade,
239            accept_key,
240            protocol,
241        } = self;
242
243        tokio::spawn(async move {
244            if let Ok(upgraded) = on_upgrade.await {
245                let stream =
246                    WebSocketStream::from_raw_socket(TokioIo::new(upgraded), Role::Server, None)
247                        .await;
248                callback(WebSocket { inner: stream }).await;
249            }
250        });
251
252        let mut res = Response::new(StatusCode::SWITCHING_PROTOCOLS);
253        res.headers
254            .insert(UPGRADE, HeaderValue::from_static("websocket"));
255        res.headers
256            .insert(CONNECTION, HeaderValue::from_static("upgrade"));
257        res.headers.insert(SEC_WEBSOCKET_ACCEPT, accept_key);
258        if let Some(p) = protocol {
259            res.headers.insert(SEC_WEBSOCKET_PROTOCOL, p);
260        }
261        res
262    }
263}
264
265#[cfg(test)]
266mod tests {
267    use super::*;
268    use crate::{Churust, TestClient};
269    use tokio_tungstenite::tungstenite::Message as TMessage;
270
271    #[tokio::test]
272    async fn plain_get_to_ws_route_is_426() {
273        let app = Churust::server()
274            .routing(|r| {
275                r.get("/ws", |ws: WebSocketUpgrade| async move {
276                    ws.on_upgrade(|_sock| async {})
277                });
278            })
279            .build();
280        // A normal GET (no upgrade headers, no captured handle) must be rejected.
281        let res = TestClient::new(app).get("/ws").send().await;
282        assert_eq!(res.status(), http::StatusCode::UPGRADE_REQUIRED);
283    }
284
285    #[test]
286    fn message_round_trips_through_tungstenite() {
287        let cases = [
288            Message::Text("hi".into()),
289            Message::Binary(vec![1, 2, 3]),
290            Message::Ping(vec![9]),
291            Message::Pong(vec![8]),
292            Message::Close,
293        ];
294        for m in cases {
295            let t: TMessage = m.clone().into();
296            let back: Message = t.into();
297            assert_eq!(m, back);
298        }
299    }
300
301    #[test]
302    fn accept_key_matches_rfc6455_example() {
303        // RFC 6455 §1.3 worked example.
304        let accept = tokio_tungstenite::tungstenite::handshake::derive_accept_key(
305            b"dGhlIHNhbXBsZSBub25jZQ==",
306        );
307        assert_eq!(accept, "s3pPLMBiTxaQ9kYGzzhZRbK+xOo=");
308    }
309}
churust_core/ws.rs

churust_core/
ws.rs
