siderite_core/
ws.rs

1//! WebSocket upgrade and connection handling.
2
3use crate::body::Body;
4use crate::error::ApiError;
5use crate::extract::FromRequestParts;
6use crate::response::{IntoResponse, Response};
7use bytes::Bytes;
8use http::request::Parts;
9use siderite_openapi::{Operation, SchemaRegistry};
10use std::future::Future;
11use thiserror::Error;
12
13/// Extractor that upgrades an HTTP request to a WebSocket.
14///
15/// Failed upgrades (missing headers, wrong method, non-upgradable
16/// connection) become `400` or `426` problem details responses.
17pub struct WebSocketUpgrade {
18    inner: axum::extract::ws::WebSocketUpgrade,
19    owner: crate::server::tasks::TaskHandle,
20    permit: tokio::sync::OwnedSemaphorePermit,
21}
22
23impl std::fmt::Debug for WebSocketUpgrade {
24    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
25        f.debug_struct("WebSocketUpgrade").finish_non_exhaustive()
26    }
27}
28
29impl FromRequestParts for WebSocketUpgrade {
30    async fn from_request_parts(parts: &mut Parts) -> Result<Self, ApiError> {
31        use axum::extract::FromRequestParts as _;
32        match axum::extract::ws::WebSocketUpgrade::from_request_parts(parts, &()).await {
33            Ok(inner) => {
34                let owner = parts
35                    .extensions
36                    .get::<crate::server::tasks::TaskHandle>()
37                    .cloned()
38                    .ok_or_else(upgrade_unavailable)?;
39                let permit = owner.reserve().ok_or_else(upgrade_unavailable)?;
40                Ok(Self {
41                    inner,
42                    owner,
43                    permit,
44                })
45            }
46            Err(rejection) => Err(ApiError::new(rejection.status(), rejection.body_text())),
47        }
48    }
49}
50
51impl WebSocketUpgrade {
52    /// Limit incoming message size in bytes (axum default is 64 MiB).
53    #[must_use]
54    pub fn max_message_size(self, max: usize) -> Self {
55        Self {
56            inner: self.inner.max_message_size(max),
57            owner: self.owner,
58            permit: self.permit,
59        }
60    }
61
62    /// Complete the upgrade and run `callback` on the established socket.
63    ///
64    /// The returned [`WebSocketResponse`] must be returned from the handler.
65    #[must_use]
66    pub fn on_upgrade<F, Fut>(self, callback: F) -> WebSocketResponse
67    where
68        F: FnOnce(WebSocket) -> Fut + Send + 'static,
69        Fut: Future<Output = ()> + Send + 'static,
70    {
71        let Self {
72            inner,
73            owner,
74            permit,
75        } = self;
76        let (established, socket) = tokio::sync::oneshot::channel();
77        if !owner.spawn(permit, async move {
78            if let Ok(socket) = socket.await {
79                callback(WebSocket { inner: socket }).await;
80            }
81        }) {
82            let response = upgrade_unavailable().into_response();
83            return WebSocketResponse(response.map(Body::into_inner));
84        }
85        let response = inner.on_upgrade(move |socket| async move {
86            // Axum's handoff owns only the channel. Pending callbacks and
87            // established sockets remain under the lifespan supervisor.
88            let _ = established.send(socket);
89        });
90        WebSocketResponse(response)
91    }
92}
93
94fn upgrade_unavailable() -> ApiError {
95    ApiError::new(
96        http::StatusCode::SERVICE_UNAVAILABLE,
97        "WebSocket admission is unavailable.",
98    )
99    .with_header(
100        http::header::RETRY_AFTER,
101        http::HeaderValue::from_static("1"),
102    )
103}
104
105/// Response produced by [`WebSocketUpgrade::on_upgrade`].
106pub struct WebSocketResponse(axum::response::Response);
107
108impl IntoResponse for WebSocketResponse {
109    fn into_response(self) -> Response {
110        self.0.map(Body::from_inner)
111    }
112
113    fn describe(op: &mut Operation, _registry: &mut SchemaRegistry) {
114        op.add_response("101", "Switching Protocols", None);
115    }
116}
117
118/// An established WebSocket connection.
119pub struct WebSocket {
120    inner: axum::extract::ws::WebSocket,
121}
122
123impl std::fmt::Debug for WebSocket {
124    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
125        f.debug_struct("WebSocket").finish_non_exhaustive()
126    }
127}
128
129impl WebSocket {
130    /// Receive the next message, or `None` if the connection has closed.
131    pub async fn recv(&mut self) -> Option<Result<Message, WsError>> {
132        self.inner.recv().await.map(|result| {
133            result
134                .map(Message::from_engine)
135                .map_err(WsError::from_engine)
136        })
137    }
138
139    /// Send a message.
140    ///
141    /// # Errors
142    /// Returns [`WsError`] if the connection is closed or the send fails.
143    pub async fn send(&mut self, message: Message) -> Result<(), WsError> {
144        self.inner
145            .send(message.into_engine())
146            .await
147            .map_err(WsError::from_engine)
148    }
149
150    /// Send a close frame with `code` and `reason`.
151    ///
152    /// # Errors
153    /// Returns [`WsError`] if the connection is closed or the send fails.
154    pub async fn close(&mut self, code: u16, reason: impl Into<String>) -> Result<(), WsError> {
155        self.send(Message::Close(Some(CloseFrame {
156            code,
157            reason: reason.into(),
158        })))
159        .await
160    }
161}
162
163/// A WebSocket message.
164#[derive(Debug, Clone, PartialEq, Eq)]
165pub enum Message {
166    /// UTF-8 text.
167    Text(String),
168    /// Binary payload.
169    Binary(Bytes),
170    /// Ping (payload at most 125 bytes).
171    Ping(Bytes),
172    /// Pong (payload at most 125 bytes).
173    Pong(Bytes),
174    /// Close frame.
175    Close(Option<CloseFrame>),
176}
177
178/// Payload of a WebSocket close frame.
179#[derive(Debug, Clone, PartialEq, Eq)]
180pub struct CloseFrame {
181    /// Close status code.
182    pub code: u16,
183    /// Close reason.
184    pub reason: String,
185}
186
187impl Message {
188    fn from_engine(message: axum::extract::ws::Message) -> Self {
189        match message {
190            axum::extract::ws::Message::Text(text) => Self::Text(text.to_string()),
191            axum::extract::ws::Message::Binary(data) => Self::Binary(data),
192            axum::extract::ws::Message::Ping(data) => Self::Ping(data),
193            axum::extract::ws::Message::Pong(data) => Self::Pong(data),
194            axum::extract::ws::Message::Close(frame) => {
195                Self::Close(frame.map(|frame| CloseFrame {
196                    code: frame.code,
197                    reason: frame.reason.to_string(),
198                }))
199            }
200        }
201    }
202
203    fn into_engine(self) -> axum::extract::ws::Message {
204        match self {
205            Self::Text(text) => axum::extract::ws::Message::Text(text.into()),
206            Self::Binary(data) => axum::extract::ws::Message::Binary(data),
207            Self::Ping(data) => axum::extract::ws::Message::Ping(data),
208            Self::Pong(data) => axum::extract::ws::Message::Pong(data),
209            Self::Close(frame) => axum::extract::ws::Message::Close(frame.map(|frame| {
210                axum::extract::ws::CloseFrame {
211                    code: frame.code,
212                    reason: frame.reason.into(),
213                }
214            })),
215        }
216    }
217}
218
219/// Error sending or receiving a WebSocket message.
220#[derive(Debug, Error)]
221#[error("{0}")]
222pub struct WsError(String);
223
224impl WsError {
225    fn from_engine(err: axum::Error) -> Self {
226        Self(err.to_string())
227    }
228}
229
230#[cfg(test)]
231#[allow(clippy::unwrap_used)]
232mod tests {
233    use super::*;
234    use crate::app::App;
235    use crate::routing::get;
236    use std::time::Duration;
237    use tokio::io::{AsyncReadExt, AsyncWriteExt};
238    use tokio::net::TcpStream;
239
240    #[test]
241    fn message_conversions_round_trip() {
242        let messages = [
243            Message::Text("hello".into()),
244            Message::Binary(Bytes::from_static(b"bin")),
245            Message::Ping(Bytes::from_static(b"ping")),
246            Message::Pong(Bytes::from_static(b"pong")),
247            Message::Close(None),
248            Message::Close(Some(CloseFrame {
249                code: 1000,
250                reason: "bye".into(),
251            })),
252        ];
253        for message in messages {
254            let back = Message::from_engine(message.clone().into_engine());
255            assert_eq!(message, back);
256        }
257    }
258
259    async fn echo(ws: WebSocketUpgrade) -> WebSocketResponse {
260        ws.on_upgrade(|mut socket| async move {
261            while let Some(Ok(message)) = socket.recv().await {
262                if socket.send(message).await.is_err() {
263                    break;
264                }
265            }
266        })
267    }
268
269    #[tokio::test]
270    async fn websocket_http11_upgrade_round_trip() {
271        let app = App::new().route("/ws", get(echo));
272        let (router, _hooks) = app.build().unwrap();
273        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
274        let addr = listener.local_addr().unwrap();
275        tokio::spawn(async move {
276            let _ = axum::serve(listener, router).await;
277        });
278
279        let mut stream = connect_with_retry(addr).await;
280        stream
281            .write_all(
282                b"GET /ws HTTP/1.1\r\n\
283Host: localhost\r\n\
284Upgrade: websocket\r\n\
285Connection: Upgrade\r\n\
286Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n\
287Sec-WebSocket-Version: 13\r\n\
288\r\n",
289            )
290            .await
291            .unwrap();
292
293        let headers = read_http_headers(&mut stream).await;
294        assert!(
295            headers.starts_with("HTTP/1.1 101"),
296            "expected 101 Switching Protocols, got: {headers}"
297        );
298        assert!(headers.to_ascii_lowercase().contains("upgrade: websocket"));
299
300        // Masked text frame "hi".
301        stream
302            .write_all(&[0x81, 0x82, 0x01, 0x02, 0x03, 0x04, b'h' ^ 0x01, b'i' ^ 0x02])
303            .await
304            .unwrap();
305
306        let mut buf = [0_u8; 4];
307        stream.read_exact(&mut buf).await.unwrap();
308        assert_eq!(buf, [0x81, 0x02, b'h', b'i']);
309    }
310
311    async fn connect_with_retry(addr: std::net::SocketAddr) -> TcpStream {
312        let deadline = tokio::time::Instant::now() + Duration::from_secs(2);
313        loop {
314            match TcpStream::connect(addr).await {
315                Ok(stream) => return stream,
316                Err(err) if tokio::time::Instant::now() < deadline => {
317                    tokio::time::sleep(Duration::from_millis(10)).await;
318                    let _ = err;
319                }
320                Err(err) => panic!("connect: {err}"),
321            }
322        }
323    }
324
325    async fn read_http_headers(stream: &mut TcpStream) -> String {
326        let mut buf = Vec::new();
327        let mut byte = [0_u8; 1];
328        loop {
329            stream.read_exact(&mut byte).await.unwrap();
330            buf.push(byte[0]);
331            if buf.ends_with(b"\r\n\r\n") {
332                break;
333            }
334            if buf.len() > 8192 {
335                panic!("HTTP response headers too large");
336            }
337        }
338        String::from_utf8_lossy(&buf).into_owned()
339    }
340}
siderite_core/ws.rs

siderite_core/
ws.rs
