ullun/
websocket.rs

1//! WebSocket endpoints (feature `websocket`)
2//!
3//! ```no_run
4//! use ullun::prelude::*;
5//! use ullun::websocket::{Message, WebSocket};
6//!
7//! # async fn run() {
8//! App::new()
9//!     .ws("/echo/{room}", |req: Request, mut socket: WebSocket| async move {
10//!         let room = req.params.get_opt("room").unwrap_or("lobby").to_string();
11//!         let _ = socket.send_text(format!("welcome to {room}")).await;
12//!         while let Some(Ok(msg)) = socket.recv().await {
13//!             match msg {
14//!                 Message::Text(t) => { let _ = socket.send_text(t).await; }
15//!                 Message::Close => break,
16//!                 _ => {}
17//!             }
18//!         }
19//!     })
20//!     .run("127.0.0.1:3000")
21//!     .await
22//!     .unwrap();
23//! # }
24//! ```
25//!
26//! The handshake is validated (version 13, key, `Connection`/`Upgrade`) and,
27//! because browsers do not apply the same-origin policy to WebSockets, a
28//! cross-site `Origin` is rejected with `403` unless it is listed in
29//! [`WsOptions::allowed_origins`]. Clients without an `Origin` header (native
30//! apps, CLI tools) are accepted; authenticate them in middleware.
31
32use crate::error::{Error, Result};
33use crate::handler::Handler;
34use crate::request::Request;
35use crate::response::Response;
36use async_trait::async_trait;
37use futures_util::{SinkExt, StreamExt};
38use hyper::upgrade::{OnUpgrade, Upgraded};
39use hyper_util::rt::TokioIo;
40use std::future::Future;
41use std::sync::Arc;
42use tokio_tungstenite::tungstenite::handshake::derive_accept_key;
43use tokio_tungstenite::tungstenite::protocol::{Role, WebSocketConfig};
44use tokio_tungstenite::tungstenite::{self, Message as WsMessage};
45use tokio_tungstenite::WebSocketStream;
46
47/// A WebSocket message
48#[derive(Debug, Clone, PartialEq, Eq)]
49pub enum Message {
50    Text(String),
51    Binary(Vec<u8>),
52    /// Pings are answered automatically; they are surfaced for information
53    Ping(Vec<u8>),
54    Pong(Vec<u8>),
55    /// The peer started the closing handshake
56    Close,
57}
58
59/// Per-endpoint settings
60#[derive(Debug, Clone)]
61pub struct WsOptions {
62    /// Largest accepted message and frame in bytes (default 1 MiB)
63    pub max_message_size: usize,
64    /// Extra `Origin` values (e.g. `https://app.example.com`) that may connect
65    /// besides the host serving the endpoint
66    pub allowed_origins: Vec<String>,
67}
68
69impl Default for WsOptions {
70    fn default() -> Self {
71        Self {
72            max_message_size: 1024 * 1024,
73            allowed_origins: Vec::new(),
74        }
75    }
76}
77
78/// An upgraded connection
79pub struct WebSocket {
80    inner: WebSocketStream<TokioIo<Upgraded>>,
81}
82
83impl WebSocket {
84    /// Next message, or `None` once the connection is closed
85    pub async fn recv(&mut self) -> Option<Result<Message>> {
86        loop {
87            return match self.inner.next().await? {
88                Ok(WsMessage::Text(t)) => Some(Ok(Message::Text(t.as_str().to_owned()))),
89                Ok(WsMessage::Binary(b)) => Some(Ok(Message::Binary(b.to_vec()))),
90                Ok(WsMessage::Ping(b)) => Some(Ok(Message::Ping(b.to_vec()))),
91                Ok(WsMessage::Pong(b)) => Some(Ok(Message::Pong(b.to_vec()))),
92                Ok(WsMessage::Close(_)) => Some(Ok(Message::Close)),
93                Ok(WsMessage::Frame(_)) => continue,
94                Err(tungstenite::Error::ConnectionClosed | tungstenite::Error::AlreadyClosed) => {
95                    None
96                }
97                Err(e) => Some(Err(Error::bad_request(format!("WebSocket: {}", e)))),
98            };
99        }
100    }
101
102    pub async fn send(&mut self, message: Message) -> Result<()> {
103        let msg = match message {
104            Message::Text(t) => WsMessage::text(t),
105            Message::Binary(b) => WsMessage::binary(b),
106            Message::Ping(b) => WsMessage::Ping(b.into()),
107            Message::Pong(b) => WsMessage::Pong(b.into()),
108            Message::Close => WsMessage::Close(None),
109        };
110        self.inner
111            .send(msg)
112            .await
113            .map_err(|e| Error::internal(format!("WebSocket send: {}", e)))
114    }
115
116    pub async fn send_text(&mut self, text: impl Into<String>) -> Result<()> {
117        self.send(Message::Text(text.into())).await
118    }
119
120    pub async fn send_binary(&mut self, data: impl Into<Vec<u8>>) -> Result<()> {
121        self.send(Message::Binary(data.into())).await
122    }
123
124    /// Start the closing handshake and flush it
125    pub async fn close(&mut self) -> Result<()> {
126        self.inner
127            .close(None)
128            .await
129            .map_err(|e| Error::internal(format!("WebSocket close: {}", e)))
130    }
131}
132
133/// Route handler performing the handshake; created by `App::ws`
134pub(crate) struct WsRoute<F> {
135    handler: Arc<F>,
136    options: WsOptions,
137}
138
139impl<F> WsRoute<F> {
140    pub(crate) fn new(handler: F, options: WsOptions) -> Self {
141        Self {
142            handler: Arc::new(handler),
143            options,
144        }
145    }
146}
147
148fn header<'a>(req: &'a Request, name: &str) -> Option<&'a str> {
149    req.headers.get(name).and_then(|v| v.to_str().ok())
150}
151
152/// Token in a comma separated header, across repeated header lines
153fn has_token(req: &Request, name: &str, token: &str) -> bool {
154    req.headers
155        .get_all(name)
156        .iter()
157        .filter_map(|v| v.to_str().ok())
158        .any(|v| v.split(',').any(|t| t.trim().eq_ignore_ascii_case(token)))
159}
160
161/// Same host as the request, or explicitly allowed
162fn origin_allowed(req: &Request, options: &WsOptions) -> bool {
163    let Some(origin) = header(req, "origin") else {
164        return true;
165    };
166    if options.allowed_origins.iter().any(|o| o == origin) {
167        return true;
168    }
169    let authority = origin
170        .split_once("://")
171        .map(|(_, rest)| rest)
172        .unwrap_or(origin);
173    header(req, "host")
174        .map(|h| h.eq_ignore_ascii_case(authority))
175        .unwrap_or(false)
176}
177
178#[async_trait]
179impl<F, Fut> Handler for WsRoute<F>
180where
181    F: Fn(Request, WebSocket) -> Fut + Send + Sync + 'static,
182    Fut: Future<Output = ()> + Send + 'static,
183{
184    async fn call(&self, mut req: Request) -> Result<Response> {
185        if !has_token(&req, "connection", "upgrade")
186            || !header(&req, "upgrade").is_some_and(|v| v.eq_ignore_ascii_case("websocket"))
187        {
188            return Err(Error::bad_request("Expected a WebSocket upgrade request"));
189        }
190        if header(&req, "sec-websocket-version") != Some("13") {
191            return Ok(
192                Response::from_error(Error::custom(426, "Unsupported WebSocket version"))
193                    .with_header("sec-websocket-version", "13"),
194            );
195        }
196        let Some(key) = header(&req, "sec-websocket-key").filter(|k| !k.is_empty()) else {
197            return Err(Error::bad_request("Missing Sec-WebSocket-Key"));
198        };
199        let accept = derive_accept_key(key.as_bytes());
200        if !origin_allowed(&req, &self.options) {
201            return Err(Error::forbidden("Origin not allowed"));
202        }
203        let Some(on_upgrade) = req.extensions.remove::<OnUpgrade>() else {
204            return Err(Error::internal("Connection cannot be upgraded"));
205        };
206
207        let handler = self.handler.clone();
208        let config = WebSocketConfig::default()
209            .max_message_size(Some(self.options.max_message_size))
210            .max_frame_size(Some(self.options.max_message_size));
211        tokio::spawn(async move {
212            let Ok(upgraded) = on_upgrade.await else {
213                return;
214            };
215            let inner = WebSocketStream::from_raw_socket(
216                TokioIo::new(upgraded),
217                Role::Server,
218                Some(config),
219            )
220            .await;
221            handler(req, WebSocket { inner }).await;
222        });
223
224        let mut res = Response::new(http::StatusCode::SWITCHING_PROTOCOLS)
225            .with_header("upgrade", "websocket")
226            .with_header("connection", "Upgrade")
227            .with_header("sec-websocket-accept", accept);
228        res.body = bytes::Bytes::new();
229        Ok(res)
230    }
231}
ullun/websocket.rs

ullun/
websocket.rs
