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    /// Send a ping this often while [`WebSocket::recv`] is waiting (default
68    /// 30 s, `None` = never). Browsers answer automatically, which keeps
69    /// proxies from dropping idle connections.
70    pub ping_interval: Option<std::time::Duration>,
71    /// Close the connection when nothing at all (message, pong) arrived for this
72    /// long (default 90 s, `None` = never), so dead peers do not pile up
73    pub idle_timeout: Option<std::time::Duration>,
74}
75
76impl Default for WsOptions {
77    fn default() -> Self {
78        Self {
79            max_message_size: 1024 * 1024,
80            allowed_origins: Vec::new(),
81            ping_interval: Some(std::time::Duration::from_secs(30)),
82            idle_timeout: Some(std::time::Duration::from_secs(90)),
83        }
84    }
85}
86
87/// An upgraded connection
88pub struct WebSocket {
89    inner: WebSocketStream<TokioIo<Upgraded>>,
90    ping: Option<tokio::time::Interval>,
91    idle_timeout: Option<std::time::Duration>,
92    last_seen: std::time::Instant,
93}
94
95impl WebSocket {
96    fn new(inner: WebSocketStream<TokioIo<Upgraded>>, options: &WsOptions) -> Self {
97        let ping = options.ping_interval.map(|every| {
98            let mut i = tokio::time::interval_at(tokio::time::Instant::now() + every, every);
99            i.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
100            i
101        });
102        Self {
103            inner,
104            ping,
105            idle_timeout: options.idle_timeout,
106            last_seen: std::time::Instant::now(),
107        }
108    }
109
110    /// Next message, or `None` once the connection is closed (by either side, or
111    /// because the peer went silent for [`WsOptions::idle_timeout`])
112    ///
113    /// Keep-alive pings are sent from inside this call, so a handler should
114    /// spend most of its time waiting here (or in a `select!` that includes it).
115    pub async fn recv(&mut self) -> Option<Result<Message>> {
116        loop {
117            let Self {
118                inner,
119                ping,
120                idle_timeout,
121                last_seen,
122            } = self;
123            let tick = async {
124                match ping {
125                    Some(interval) => {
126                        interval.tick().await;
127                    }
128                    None => std::future::pending::<()>().await,
129                }
130            };
131            tokio::select! {
132                incoming = inner.next() => {
133                    *last_seen = std::time::Instant::now();
134                    return match incoming? {
135                        Ok(WsMessage::Text(t)) => Some(Ok(Message::Text(t.as_str().to_owned()))),
136                        Ok(WsMessage::Binary(b)) => Some(Ok(Message::Binary(b.to_vec()))),
137                        Ok(WsMessage::Ping(b)) => Some(Ok(Message::Ping(b.to_vec()))),
138                        Ok(WsMessage::Pong(b)) => Some(Ok(Message::Pong(b.to_vec()))),
139                        Ok(WsMessage::Close(_)) => Some(Ok(Message::Close)),
140                        Ok(WsMessage::Frame(_)) => continue,
141                        Err(tungstenite::Error::ConnectionClosed | tungstenite::Error::AlreadyClosed) => None,
142                        Err(e) => Some(Err(Error::bad_request(format!("WebSocket: {}", e)))),
143                    };
144                }
145                _ = tick => {
146                    if idle_timeout.is_some_and(|limit| last_seen.elapsed() >= limit) {
147                        let _ = inner.close(None).await;
148                        return None;
149                    }
150                    if inner.send(WsMessage::Ping(Vec::new().into())).await.is_err() {
151                        return None;
152                    }
153                }
154            }
155        }
156    }
157
158    /// Send `value` as a JSON text message
159    pub async fn send_json<T: serde::Serialize>(&mut self, value: &T) -> Result<()> {
160        let text = serde_json::to_string(value)
161            .map_err(|e| Error::internal(format!("WebSocket JSON: {}", e)))?;
162        self.send_text(text).await
163    }
164
165    /// Close with a status code (e.g. 1008 policy violation, 1009 message too
166    /// big, 1001 going away) and a short reason
167    pub async fn close_with(&mut self, code: u16, reason: &str) -> Result<()> {
168        use tungstenite::protocol::frame::{coding::CloseCode, CloseFrame};
169        self.inner
170            .close(Some(CloseFrame {
171                code: CloseCode::from(code),
172                reason: reason.to_string().into(),
173            }))
174            .await
175            .map_err(|e| Error::internal(format!("WebSocket close: {}", e)))
176    }
177
178    pub async fn send(&mut self, message: Message) -> Result<()> {
179        let msg = match message {
180            Message::Text(t) => WsMessage::text(t),
181            Message::Binary(b) => WsMessage::binary(b),
182            Message::Ping(b) => WsMessage::Ping(b.into()),
183            Message::Pong(b) => WsMessage::Pong(b.into()),
184            Message::Close => WsMessage::Close(None),
185        };
186        self.inner
187            .send(msg)
188            .await
189            .map_err(|e| Error::internal(format!("WebSocket send: {}", e)))
190    }
191
192    pub async fn send_text(&mut self, text: impl Into<String>) -> Result<()> {
193        self.send(Message::Text(text.into())).await
194    }
195
196    pub async fn send_binary(&mut self, data: impl Into<Vec<u8>>) -> Result<()> {
197        self.send(Message::Binary(data.into())).await
198    }
199
200    /// Start the closing handshake and flush it
201    pub async fn close(&mut self) -> Result<()> {
202        self.inner
203            .close(None)
204            .await
205            .map_err(|e| Error::internal(format!("WebSocket close: {}", e)))
206    }
207}
208
209/// Route handler performing the handshake; created by `App::ws`
210pub(crate) struct WsRoute<F> {
211    handler: Arc<F>,
212    options: WsOptions,
213}
214
215impl<F> WsRoute<F> {
216    pub(crate) fn new(handler: F, options: WsOptions) -> Self {
217        Self {
218            handler: Arc::new(handler),
219            options,
220        }
221    }
222}
223
224fn header<'a>(req: &'a Request, name: &str) -> Option<&'a str> {
225    req.headers.get(name).and_then(|v| v.to_str().ok())
226}
227
228/// Token in a comma separated header, across repeated header lines
229fn has_token(req: &Request, name: &str, token: &str) -> bool {
230    req.headers
231        .get_all(name)
232        .iter()
233        .filter_map(|v| v.to_str().ok())
234        .any(|v| v.split(',').any(|t| t.trim().eq_ignore_ascii_case(token)))
235}
236
237/// Same host as the request, or explicitly allowed
238fn origin_allowed(req: &Request, options: &WsOptions) -> bool {
239    let Some(origin) = header(req, "origin") else {
240        return true;
241    };
242    if options.allowed_origins.iter().any(|o| o == origin) {
243        return true;
244    }
245    let authority = origin
246        .split_once("://")
247        .map(|(_, rest)| rest)
248        .unwrap_or(origin);
249    header(req, "host")
250        .map(|h| h.eq_ignore_ascii_case(authority))
251        .unwrap_or(false)
252}
253
254#[async_trait]
255impl<F, Fut> Handler for WsRoute<F>
256where
257    F: Fn(Request, WebSocket) -> Fut + Send + Sync + 'static,
258    Fut: Future<Output = ()> + Send + 'static,
259{
260    async fn call(&self, mut req: Request) -> Result<Response> {
261        if !has_token(&req, "connection", "upgrade")
262            || !header(&req, "upgrade").is_some_and(|v| v.eq_ignore_ascii_case("websocket"))
263        {
264            return Err(Error::bad_request("Expected a WebSocket upgrade request"));
265        }
266        if header(&req, "sec-websocket-version") != Some("13") {
267            return Ok(
268                Response::from_error(Error::custom(426, "Unsupported WebSocket version"))
269                    .with_header("sec-websocket-version", "13"),
270            );
271        }
272        let Some(key) = header(&req, "sec-websocket-key").filter(|k| !k.is_empty()) else {
273            return Err(Error::bad_request("Missing Sec-WebSocket-Key"));
274        };
275        let accept = derive_accept_key(key.as_bytes());
276        if !origin_allowed(&req, &self.options) {
277            return Err(Error::forbidden("Origin not allowed"));
278        }
279        let Some(on_upgrade) = req.extensions.remove::<OnUpgrade>() else {
280            return Err(Error::internal("Connection cannot be upgraded"));
281        };
282
283        let handler = self.handler.clone();
284        let options = self.options.clone();
285        let config = WebSocketConfig::default()
286            .max_message_size(Some(self.options.max_message_size))
287            .max_frame_size(Some(self.options.max_message_size));
288        tokio::spawn(async move {
289            let Ok(upgraded) = on_upgrade.await else {
290                return;
291            };
292            let inner = WebSocketStream::from_raw_socket(
293                TokioIo::new(upgraded),
294                Role::Server,
295                Some(config),
296            )
297            .await;
298            handler(req, WebSocket::new(inner, &options)).await;
299        });
300
301        let mut res = Response::new(http::StatusCode::SWITCHING_PROTOCOLS)
302            .with_header("upgrade", "websocket")
303            .with_header("connection", "Upgrade")
304            .with_header("sec-websocket-accept", accept);
305        res.body = bytes::Bytes::new();
306        Ok(res)
307    }
308}
ullun/websocket.rs

ullun/
websocket.rs
