tachyon_web/ws/
mod.rs

1//! WebSocket support (RFC 6455), mirroring Axum's `extract::ws` API.
2//!
3//! Enabled via the `ws` feature. WebSocket upgrades happen over an already-accepted
4//! HTTP/1.1 connection — for TLS, that connection is already decrypted by the time it
5//! reaches the router, so `wss://` works automatically: bind with [`crate::Server::serve_https`]
6//! (or any of the other TLS entry points) exactly as you would for plain `ws://` with
7//! [`crate::Server::serve_http`]. There is nothing protocol-specific to configure.
8//!
9//! With the `http2` feature enabled, [`WebSocketUpgrade`] also accepts the RFC 8441 bootstrap:
10//! an HTTP/2 extended `CONNECT` request with a `:protocol` pseudo-header of `websocket`. Once
11//! the tunnel is established the wire format is identical RFC 6455 framing either way, so
12//! handlers don't need to care which transport a given [`WebSocket`] came from. HTTP/3 does not
13//! define a WebSocket bootstrap and is not supported.
14//!
15//! The `permessage-deflate` extension (RFC 7692) is negotiated automatically whenever the client
16//! offers it — disable it with [`WebSocketUpgrade::deflate`] or tune it with
17//! [`WebSocketUpgrade::deflate_config`].
18//!
19//! # Example
20//!
21//! ```rust,no_run
22//! use tachyon_web::ws::{WebSocket, WebSocketUpgrade};
23//! use tachyon_web::http::Response;
24//! use tachyon_web::http::response::Body;
25//! use tachyon_web::{Router, get};
26//!
27//! async fn handler(ws: WebSocketUpgrade) -> Response<Body> {
28//!     ws.on_upgrade(handle_socket)
29//! }
30//!
31//! async fn handle_socket(mut socket: WebSocket) {
32//!     while let Some(Ok(msg)) = socket.recv().await {
33//!         if socket.send(msg).await.is_err() {
34//!             break;
35//!         }
36//!     }
37//! }
38//!
39//! let _app: Router<()> = Router::new().route("/ws", get(handler));
40//! ```
41
42mod compat;
43mod deflate;
44mod socket;
45
46use crate::http::error::Error;
47use crate::http::response::Body;
48use hyper::header::{self, HeaderMap, HeaderName, HeaderValue};
49use hyper::http::request::Parts;
50use hyper::{Method, Response, StatusCode};
51use std::borrow::Cow;
52use std::future::Future;
53use tungstenite::handshake::derive_accept_key;
54
55pub use deflate::DeflateConfig;
56pub use socket::{Message, WebSocket};
57pub use tungstenite::protocol::{CloseFrame, WebSocketConfig, frame::coding::CloseCode};
58
59/// Extractor for establishing a WebSocket connection out of an HTTP/1.1 (or, with the `http2`
60/// feature, HTTP/2 extended-`CONNECT`) request.
61///
62/// See the [module docs](self) for an example.
63#[must_use]
64pub struct WebSocketUpgrade<F = DefaultOnFailedUpgrade> {
65    config: WebSocketConfig,
66    protocol: Option<HeaderValue>,
67    kind: UpgradeKind,
68    on_upgrade: hyper::upgrade::OnUpgrade,
69    on_failed_upgrade: F,
70    sec_websocket_protocol: Vec<HeaderValue>,
71    origin: Option<HeaderValue>,
72    deflate_offers: Vec<deflate::Offer>,
73    deflate_enabled: bool,
74    deflate_config: DeflateConfig,
75}
76
77/// Which bootstrap produced this upgrade, and the bits of state each one needs to finish the
78/// handshake: HTTP/1.1 needs the client's key to derive `Sec-WebSocket-Accept`; the RFC 8441
79/// HTTP/2 path needs nothing extra (there is no accept-key concept — see RFC 8441 §5).
80enum UpgradeKind {
81    Http1 { sec_websocket_key: HeaderValue },
82    #[cfg(feature = "http2")]
83    Http2,
84}
85
86impl<F> std::fmt::Debug for WebSocketUpgrade<F> {
87    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
88        f.debug_struct("WebSocketUpgrade")
89            .field("protocol", &self.protocol)
90            .field("sec_websocket_protocol", &self.sec_websocket_protocol)
91            .finish_non_exhaustive()
92    }
93}
94
95impl<F> WebSocketUpgrade<F> {
96    /// Read buffer capacity. The default value is 128 KiB.
97    pub const fn read_buffer_size(mut self, size: usize) -> Self {
98        self.config.read_buffer_size = size;
99        self
100    }
101
102    /// The target minimum size of the write buffer to reach before writing the data
103    /// to the underlying stream. The default value is 128 KiB.
104    ///
105    /// If set to `0`, each message is eagerly written to the underlying stream.
106    pub const fn write_buffer_size(mut self, size: usize) -> Self {
107        self.config.write_buffer_size = size;
108        self
109    }
110
111    /// The max size of the write buffer in bytes. The default value is unlimited.
112    pub const fn max_write_buffer_size(mut self, max: usize) -> Self {
113        self.config.max_write_buffer_size = max;
114        self
115    }
116
117    /// Set the maximum message size (defaults to 64 MiB).
118    pub const fn max_message_size(mut self, max: usize) -> Self {
119        self.config.max_message_size = Some(max);
120        self
121    }
122
123    /// Set the maximum frame size (defaults to 16 MiB).
124    pub const fn max_frame_size(mut self, max: usize) -> Self {
125        self.config.max_frame_size = Some(max);
126        self
127    }
128
129    /// Allow the server to accept unmasked frames (defaults to `false`).
130    pub const fn accept_unmasked_frames(mut self, accept: bool) -> Self {
131        self.config.accept_unmasked_frames = accept;
132        self
133    }
134
135    /// Enable or disable offering the `permessage-deflate` extension (RFC 7692) to the client.
136    /// Enabled by default; when the client doesn't offer it, this has no effect either way.
137    pub const fn deflate(mut self, enabled: bool) -> Self {
138        self.deflate_enabled = enabled;
139        self
140    }
141
142    /// Tune `permessage-deflate` parameters — context takeover and window size. See
143    /// [`DeflateConfig`]. No effect if [`deflate`](Self::deflate) is disabled or the client
144    /// didn't offer the extension.
145    pub const fn deflate_config(mut self, config: DeflateConfig) -> Self {
146        self.deflate_config = config;
147        self
148    }
149
150    /// Set the server's supported subprotocols, in decreasing order of preference.
151    ///
152    /// If any of them matches one the client requested (via `Sec-WebSocket-Protocol`),
153    /// the response includes a `Sec-WebSocket-Protocol` header naming it.
154    pub fn protocols<I>(mut self, protocols: I) -> Self
155    where
156        I: IntoIterator,
157        I::Item: Into<Cow<'static, str>>,
158    {
159        self.protocol = protocols.into_iter().map(Into::into).find_map(|proto| {
160            let value = match proto {
161                Cow::Owned(s) => HeaderValue::from_str(&s).ok()?,
162                Cow::Borrowed(s) => HeaderValue::from_static(s),
163            };
164            self.sec_websocket_protocol
165                .contains(&value)
166                .then_some(value)
167        });
168        self
169    }
170
171    /// The WebSocket subprotocols requested by the client, via `Sec-WebSocket-Protocol`.
172    pub fn requested_protocols(&self) -> impl Iterator<Item = &HeaderValue> {
173        self.sec_websocket_protocol.iter()
174    }
175
176    /// Return the selected WebSocket subprotocol, if [`protocols`](Self::protocols) matched one.
177    #[must_use]
178    pub const fn selected_protocol(&self) -> Option<&HeaderValue> {
179        self.protocol.as_ref()
180    }
181
182    /// The `Origin` header sent by the client, if any.
183    ///
184    /// Present for browser clients (and absent for most non-browser WebSocket clients).
185    /// Exposed for handlers that want to make their own origin decision inline.
186    #[must_use]
187    pub const fn origin(&self) -> Option<&HeaderValue> {
188        self.origin.as_ref()
189    }
190
191    /// Provide a callback invoked if completing the (background) connection upgrade fails.
192    ///
193    /// By default, failures are silently ignored.
194    pub fn on_failed_upgrade<C>(self, callback: C) -> WebSocketUpgrade<C>
195    where
196        C: OnFailedUpgrade,
197    {
198        WebSocketUpgrade {
199            config: self.config,
200            protocol: self.protocol,
201            kind: self.kind,
202            on_upgrade: self.on_upgrade,
203            on_failed_upgrade: callback,
204            sec_websocket_protocol: self.sec_websocket_protocol,
205            origin: self.origin,
206            deflate_offers: self.deflate_offers,
207            deflate_enabled: self.deflate_enabled,
208            deflate_config: self.deflate_config,
209        }
210    }
211
212    /// Finalize the upgrade, running `callback` with the [`WebSocket`] once the underlying
213    /// connection has actually switched protocols.
214    ///
215    /// The returned [`Response`] must be returned from the handler unmodified for the
216    /// upgrade to complete.
217    #[must_use = "the response from `on_upgrade` must be returned from the handler"]
218    pub fn on_upgrade<C, Fut>(self, callback: C) -> Response<Body>
219    where
220        C: FnOnce(WebSocket) -> Fut + Send + 'static,
221        Fut: Future<Output = ()> + Send + 'static,
222        F: OnFailedUpgrade,
223    {
224        let on_upgrade = self.on_upgrade;
225        let config = self.config;
226        let on_failed_upgrade = self.on_failed_upgrade;
227        let protocol = self.protocol.clone();
228        let agreement = self
229            .deflate_enabled
230            .then(|| deflate::negotiate(&self.deflate_offers, self.deflate_config))
231            .flatten();
232
233        tokio::spawn(async move {
234            let upgraded = match on_upgrade.await {
235                Ok(upgraded) => upgraded,
236                Err(err) => {
237                    on_failed_upgrade.call(Error::Internal(err.to_string()));
238                    return;
239                }
240            };
241            let io = hyper_util::rt::TokioIo::new(upgraded);
242            let deflate = agreement.map(deflate::PerMessageDeflate::new);
243            callback(WebSocket::new(io, protocol, config, deflate)).await;
244        });
245
246        let mut response = match self.kind {
247            UpgradeKind::Http1 { sec_websocket_key } => Response::builder()
248                .status(StatusCode::SWITCHING_PROTOCOLS)
249                .header(header::CONNECTION, HeaderValue::from_static("upgrade"))
250                .header(header::UPGRADE, HeaderValue::from_static("websocket"))
251                .header(
252                    header::SEC_WEBSOCKET_ACCEPT,
253                    derive_accept_key(sec_websocket_key.as_bytes()),
254                )
255                .body(Body::empty())
256                .unwrap_or_else(|_| Response::new(Body::empty())),
257            #[cfg(feature = "http2")]
258            UpgradeKind::Http2 => Response::builder()
259                .status(StatusCode::OK)
260                .body(Body::empty())
261                .unwrap_or_else(|_| Response::new(Body::empty())),
262        };
263
264        if let Some(protocol) = self.protocol {
265            response
266                .headers_mut()
267                .insert(header::SEC_WEBSOCKET_PROTOCOL, protocol);
268        }
269        if let Some(agreement) = agreement {
270            response.headers_mut().insert(
271                header::SEC_WEBSOCKET_EXTENSIONS,
272                deflate::agreement_header_value(agreement),
273            );
274        }
275        response
276    }
277}
278
279/// What to do when completing a WebSocket connection upgrade fails.
280///
281/// See [`WebSocketUpgrade::on_failed_upgrade`].
282pub trait OnFailedUpgrade: Send + 'static {
283    /// Handle the failure.
284    fn call(self, error: Error);
285}
286
287impl<F> OnFailedUpgrade for F
288where
289    F: FnOnce(Error) + Send + 'static,
290{
291    fn call(self, error: Error) {
292        self(error);
293    }
294}
295
296/// The default [`OnFailedUpgrade`]: silently ignores the error.
297#[non_exhaustive]
298#[derive(Debug)]
299pub struct DefaultOnFailedUpgrade;
300
301impl OnFailedUpgrade for DefaultOnFailedUpgrade {
302    fn call(self, _error: Error) {}
303}
304
305fn header_eq(headers: &HeaderMap, key: &HeaderName, value: &'static str) -> bool {
306    headers
307        .get(key)
308        .is_some_and(|h| h.as_bytes().eq_ignore_ascii_case(value.as_bytes()))
309}
310
311/// Case-insensitive, allocation-free substring search — `value` is always a lowercase ASCII
312/// literal at every call site, so a byte-window scan avoids `to_ascii_lowercase()`'s per-call
313/// heap allocation on the WebSocket upgrade path.
314fn header_contains(headers: &HeaderMap, key: &HeaderName, value: &'static str) -> bool {
315    let Some(header) = headers.get(key) else {
316        return false;
317    };
318    let haystack = header.as_bytes();
319    let needle = value.as_bytes();
320    !needle.is_empty()
321        && haystack
322            .windows(needle.len())
323            .any(|w| w.eq_ignore_ascii_case(needle))
324}
325
326impl<S> crate::routing::extract::FromRequest<S> for WebSocketUpgrade<DefaultOnFailedUpgrade>
327where
328    S: Send + Sync,
329{
330    type Rejection = Error;
331
332    async fn from_request(req: hyper::Request<Body>, state: &S) -> Result<Self, Self::Rejection> {
333        let (mut parts, _body) = req.into_parts();
334        Self::from_request_parts(&mut parts, state)
335    }
336}
337
338impl WebSocketUpgrade<DefaultOnFailedUpgrade> {
339    /// Build a [`WebSocketUpgrade`] from request parts, validating the RFC 6455 handshake
340    /// headers (or, for HTTP/2 with the `http2` feature enabled, the RFC 8441 extended-`CONNECT`
341    /// bootstrap) and pulling the pending [`hyper::upgrade::OnUpgrade`] out of the extensions.
342    ///
343    /// # Errors
344    ///
345    /// Returns a rejection if this isn't a well-formed WebSocket upgrade request.
346    pub fn from_request_parts<S>(parts: &mut Parts, _state: &S) -> Result<Self, Error> {
347        #[cfg(feature = "http2")]
348        if parts.version == hyper::Version::HTTP_2 {
349            return Self::from_h2_request_parts(parts);
350        }
351
352        if parts.version > hyper::Version::HTTP_11 {
353            return Err(Error::Rejection {
354                status: StatusCode::UPGRADE_REQUIRED,
355                message: "WebSocket upgrades require HTTP/1.1 (or HTTP/2 extended CONNECT, with the `http2` feature enabled)".to_string(),
356            });
357        }
358        if parts.method != Method::GET {
359            return Err(Error::Rejection {
360                status: StatusCode::METHOD_NOT_ALLOWED,
361                message: "Request method must be `GET`".to_string(),
362            });
363        }
364        if !header_contains(&parts.headers, &header::CONNECTION, "upgrade") {
365            return Err(Error::Rejection {
366                status: StatusCode::BAD_REQUEST,
367                message: "`Connection` header did not include 'upgrade'".to_string(),
368            });
369        }
370        if !header_eq(&parts.headers, &header::UPGRADE, "websocket") {
371            return Err(Error::Rejection {
372                status: StatusCode::BAD_REQUEST,
373                message: "`Upgrade` header did not include 'websocket'".to_string(),
374            });
375        }
376        if !header_eq(&parts.headers, &header::SEC_WEBSOCKET_VERSION, "13") {
377            return Err(Error::Rejection {
378                status: StatusCode::BAD_REQUEST,
379                message: "`Sec-WebSocket-Version` header did not include '13'".to_string(),
380            });
381        }
382        let sec_websocket_key = parts
383            .headers
384            .get(header::SEC_WEBSOCKET_KEY)
385            .cloned()
386            .ok_or_else(|| Error::Rejection {
387                status: StatusCode::BAD_REQUEST,
388                message: "`Sec-WebSocket-Key` header missing".to_string(),
389            })?;
390        let on_upgrade = parts
391            .extensions
392            .remove::<hyper::upgrade::OnUpgrade>()
393            .ok_or_else(|| Error::Rejection {
394                status: StatusCode::UPGRADE_REQUIRED,
395                message: "Request couldn't be upgraded: no upgrade state was present".to_string(),
396            })?;
397
398        let sec_websocket_protocol = parse_sec_websocket_protocol(&parts.headers);
399        let origin = parts.headers.get(header::ORIGIN).cloned();
400        let deflate_offers = deflate::parse_offers(&parts.headers);
401
402        Ok(Self {
403            config: WebSocketConfig::default(),
404            protocol: None,
405            kind: UpgradeKind::Http1 { sec_websocket_key },
406            on_upgrade,
407            on_failed_upgrade: DefaultOnFailedUpgrade,
408            sec_websocket_protocol,
409            origin,
410            deflate_offers,
411            deflate_enabled: true,
412            deflate_config: DeflateConfig::default(),
413        })
414    }
415
416    /// The RFC 8441 bootstrap: an HTTP/2 extended `CONNECT` request (`:method: CONNECT`,
417    /// `:protocol: websocket`). Unlike HTTP/1.1, there is no `Sec-WebSocket-Key`/`-Accept`
418    /// handshake — HTTP/2 already requires a validated transport, so RFC 8441 §5 drops it.
419    #[cfg(feature = "http2")]
420    fn from_h2_request_parts(parts: &mut Parts) -> Result<Self, Error> {
421        if parts.method != Method::CONNECT {
422            return Err(Error::Rejection {
423                status: StatusCode::UPGRADE_REQUIRED,
424                message: "WebSocket upgrades over HTTP/2 require the extended CONNECT method (RFC 8441)".to_string(),
425            });
426        }
427        let is_websocket_protocol = parts
428            .extensions
429            .get::<hyper::ext::Protocol>()
430            .is_some_and(|p| p.as_str().eq_ignore_ascii_case("websocket"));
431        if !is_websocket_protocol {
432            return Err(Error::Rejection {
433                status: StatusCode::BAD_REQUEST,
434                message: "`:protocol` pseudo-header must be `websocket`".to_string(),
435            });
436        }
437        if !header_eq(&parts.headers, &header::SEC_WEBSOCKET_VERSION, "13") {
438            return Err(Error::Rejection {
439                status: StatusCode::BAD_REQUEST,
440                message: "`Sec-WebSocket-Version` header did not include '13'".to_string(),
441            });
442        }
443        let on_upgrade = parts
444            .extensions
445            .remove::<hyper::upgrade::OnUpgrade>()
446            .ok_or_else(|| Error::Rejection {
447                status: StatusCode::UPGRADE_REQUIRED,
448                message: "Request couldn't be upgraded: no upgrade state was present".to_string(),
449            })?;
450
451        let sec_websocket_protocol = parse_sec_websocket_protocol(&parts.headers);
452        let origin = parts.headers.get(header::ORIGIN).cloned();
453        let deflate_offers = deflate::parse_offers(&parts.headers);
454
455        Ok(Self {
456            config: WebSocketConfig::default(),
457            protocol: None,
458            kind: UpgradeKind::Http2,
459            on_upgrade,
460            on_failed_upgrade: DefaultOnFailedUpgrade,
461            sec_websocket_protocol,
462            origin,
463            deflate_offers,
464            deflate_enabled: true,
465            deflate_config: DeflateConfig::default(),
466        })
467    }
468}
469
470fn parse_sec_websocket_protocol(headers: &HeaderMap) -> Vec<HeaderValue> {
471    headers
472        .get_all(header::SEC_WEBSOCKET_PROTOCOL)
473        .iter()
474        .flat_map(|val| val.as_bytes().split(|&b| b == b','))
475        .filter_map(|proto| HeaderValue::from_bytes(proto.trim_ascii()).ok())
476        .collect()
477}
478
479#[cfg(test)]
480mod tests {
481    #![allow(clippy::unwrap_used)]
482    use super::*;
483    use hyper::Request;
484
485    /// Builds a well-formed WebSocket-upgrade `Parts`, with a placeholder
486    /// `OnUpgrade` (this never drives a real upgrade in these tests, it only
487    /// needs to satisfy `WebSocketUpgrade::from_request_parts`'s extraction).
488    fn make_ws_parts() -> Parts {
489        let mut req = Request::builder()
490            .method(Method::GET)
491            .uri("/ws")
492            .header(header::CONNECTION, "upgrade")
493            .header(header::UPGRADE, "websocket")
494            .header(header::SEC_WEBSOCKET_VERSION, "13")
495            .header(header::SEC_WEBSOCKET_KEY, "dGhlIHNhbXBsZSBub25jZQ==")
496            .body(())
497            .unwrap();
498        let on_upgrade = hyper::upgrade::on(&mut req);
499        let (mut parts, ()) = req.into_parts();
500        let _ = parts.extensions.insert(on_upgrade);
501        parts
502    }
503
504    #[test]
505    fn rejects_http2_and_above() {
506        let mut parts = make_ws_parts();
507        parts.version = hyper::Version::HTTP_2;
508        let err = WebSocketUpgrade::from_request_parts(&mut parts, &()).unwrap_err();
509        assert!(matches!(
510            err,
511            Error::Rejection {
512                status: StatusCode::UPGRADE_REQUIRED,
513                ..
514            }
515        ));
516    }
517
518    #[test]
519    fn rejects_non_get_method() {
520        let mut parts = make_ws_parts();
521        parts.method = Method::POST;
522        let err = WebSocketUpgrade::from_request_parts(&mut parts, &()).unwrap_err();
523        assert!(matches!(
524            err,
525            Error::Rejection {
526                status: StatusCode::METHOD_NOT_ALLOWED,
527                ..
528            }
529        ));
530    }
531
532    #[test]
533    fn rejects_missing_connection_upgrade_token() {
534        let mut parts = make_ws_parts();
535        let _ = parts
536            .headers
537            .insert(header::CONNECTION, HeaderValue::from_static("keep-alive"));
538        let err = WebSocketUpgrade::from_request_parts(&mut parts, &()).unwrap_err();
539        assert!(matches!(
540            err,
541            Error::Rejection {
542                status: StatusCode::BAD_REQUEST,
543                ..
544            }
545        ));
546    }
547
548    #[test]
549    fn rejects_wrong_upgrade_header_value() {
550        let mut parts = make_ws_parts();
551        let _ = parts
552            .headers
553            .insert(header::UPGRADE, HeaderValue::from_static("h2c"));
554        let err = WebSocketUpgrade::from_request_parts(&mut parts, &()).unwrap_err();
555        assert!(matches!(
556            err,
557            Error::Rejection {
558                status: StatusCode::BAD_REQUEST,
559                ..
560            }
561        ));
562    }
563
564    #[test]
565    fn rejects_wrong_sec_websocket_version() {
566        let mut parts = make_ws_parts();
567        let _ = parts
568            .headers
569            .insert(header::SEC_WEBSOCKET_VERSION, HeaderValue::from_static("8"));
570        let err = WebSocketUpgrade::from_request_parts(&mut parts, &()).unwrap_err();
571        assert!(matches!(
572            err,
573            Error::Rejection {
574                status: StatusCode::BAD_REQUEST,
575                ..
576            }
577        ));
578    }
579
580    #[test]
581    fn rejects_missing_sec_websocket_key() {
582        let mut parts = make_ws_parts();
583        let _ = parts.headers.remove(header::SEC_WEBSOCKET_KEY);
584        let err = WebSocketUpgrade::from_request_parts(&mut parts, &()).unwrap_err();
585        assert!(matches!(
586            err,
587            Error::Rejection {
588                status: StatusCode::BAD_REQUEST,
589                ..
590            }
591        ));
592    }
593
594    #[test]
595    fn rejects_when_no_upgrade_state_is_present() {
596        // Built directly (not via `make_ws_parts`) so no `hyper::upgrade::on(&mut req)` was
597        // ever called — there's nothing in extensions for `from_request_parts` to remove.
598        let req = Request::builder()
599            .method(Method::GET)
600            .uri("/ws")
601            .header(header::CONNECTION, "upgrade")
602            .header(header::UPGRADE, "websocket")
603            .header(header::SEC_WEBSOCKET_VERSION, "13")
604            .header(header::SEC_WEBSOCKET_KEY, "dGhlIHNhbXBsZSBub25jZQ==")
605            .body(())
606            .unwrap();
607        let (mut parts, ()) = req.into_parts();
608        let err = WebSocketUpgrade::from_request_parts(&mut parts, &()).unwrap_err();
609        assert!(matches!(
610            err,
611            Error::Rejection {
612                status: StatusCode::UPGRADE_REQUIRED,
613                ..
614            }
615        ));
616    }
617
618    #[test]
619    fn builder_methods_configure_the_underlying_websocket_config() {
620        let mut parts = make_ws_parts();
621        let upgrade = WebSocketUpgrade::from_request_parts(&mut parts, &())
622            .unwrap()
623            .read_buffer_size(1024)
624            .write_buffer_size(2048)
625            .max_write_buffer_size(4096)
626            .max_message_size(8192)
627            .max_frame_size(16384)
628            .accept_unmasked_frames(true);
629
630        assert_eq!(upgrade.config.read_buffer_size, 1024);
631        assert_eq!(upgrade.config.write_buffer_size, 2048);
632        assert_eq!(upgrade.config.max_write_buffer_size, 4096);
633        assert_eq!(upgrade.config.max_message_size, Some(8192));
634        assert_eq!(upgrade.config.max_frame_size, Some(16384));
635        assert!(upgrade.config.accept_unmasked_frames);
636    }
637
638    #[test]
639    fn protocols_selects_a_requested_subprotocol_the_server_also_supports() {
640        let mut req = Request::builder()
641            .method(Method::GET)
642            .uri("/ws")
643            .header(header::CONNECTION, "upgrade")
644            .header(header::UPGRADE, "websocket")
645            .header(header::SEC_WEBSOCKET_VERSION, "13")
646            .header(header::SEC_WEBSOCKET_KEY, "dGhlIHNhbXBsZSBub25jZQ==")
647            .header(header::SEC_WEBSOCKET_PROTOCOL, "chat, superchat")
648            .body(())
649            .unwrap();
650        let on_upgrade = hyper::upgrade::on(&mut req);
651        let (mut parts, ()) = req.into_parts();
652        let _ = parts.extensions.insert(on_upgrade);
653
654        let upgrade = WebSocketUpgrade::from_request_parts(&mut parts, &()).unwrap();
655        let requested: Vec<_> = upgrade
656            .requested_protocols()
657            .map(|v| v.to_str().unwrap().to_string())
658            .collect();
659        assert_eq!(requested, vec!["chat", "superchat"]);
660
661        let upgrade = upgrade.protocols(["superchat"]);
662        assert_eq!(
663            upgrade.selected_protocol().unwrap().to_str().unwrap(),
664            "superchat"
665        );
666    }
667
668    #[test]
669    fn protocols_selects_none_when_nothing_matches() {
670        let mut parts = make_ws_parts();
671        let upgrade = WebSocketUpgrade::from_request_parts(&mut parts, &())
672            .unwrap()
673            .protocols(["some-protocol-the-client-never-asked-for"]);
674        assert!(upgrade.selected_protocol().is_none());
675    }
676
677    #[test]
678    fn origin_getter_reflects_the_request_header() {
679        let mut parts = make_ws_parts();
680        let _ = parts.headers.insert(
681            header::ORIGIN,
682            HeaderValue::from_static("https://example.com"),
683        );
684        let upgrade = WebSocketUpgrade::from_request_parts(&mut parts, &()).unwrap();
685        assert_eq!(
686            upgrade.origin().unwrap().to_str().unwrap(),
687            "https://example.com"
688        );
689    }
690
691    #[test]
692    fn origin_getter_is_none_when_absent() {
693        let mut parts = make_ws_parts();
694        let upgrade = WebSocketUpgrade::from_request_parts(&mut parts, &()).unwrap();
695        assert!(upgrade.origin().is_none());
696    }
697
698    #[test]
699    fn websocket_upgrade_debug_does_not_panic() {
700        let mut parts = make_ws_parts();
701        let upgrade = WebSocketUpgrade::from_request_parts(&mut parts, &()).unwrap();
702        assert!(format!("{upgrade:?}").contains("WebSocketUpgrade"));
703    }
704
705    #[test]
706    fn on_failed_upgrade_swaps_the_callback_type_and_preserves_config() {
707        let mut parts = make_ws_parts();
708        let upgrade = WebSocketUpgrade::from_request_parts(&mut parts, &())
709            .unwrap()
710            .max_message_size(1234)
711            .on_failed_upgrade(|_err: Error| {});
712        assert_eq!(upgrade.config.max_message_size, Some(1234));
713    }
714
715    #[test]
716    fn default_on_failed_upgrade_silently_ignores_the_error() {
717        // Just proves `call` doesn't panic — this is the "silently ignore" default.
718        DefaultOnFailedUpgrade.call(Error::Internal("boom".to_string()));
719    }
720
721    #[test]
722    fn deflate_is_offered_by_default_and_can_be_disabled() {
723        let mut parts = make_ws_parts();
724        let upgrade = WebSocketUpgrade::from_request_parts(&mut parts, &()).unwrap();
725        assert!(upgrade.deflate_enabled);
726        let upgrade = upgrade.deflate(false);
727        assert!(!upgrade.deflate_enabled);
728    }
729
730    #[test]
731    fn parses_permessage_deflate_offer_from_request() {
732        let mut req = Request::builder()
733            .method(Method::GET)
734            .uri("/ws")
735            .header(header::CONNECTION, "upgrade")
736            .header(header::UPGRADE, "websocket")
737            .header(header::SEC_WEBSOCKET_VERSION, "13")
738            .header(header::SEC_WEBSOCKET_KEY, "dGhlIHNhbXBsZSBub25jZQ==")
739            .header(
740                header::SEC_WEBSOCKET_EXTENSIONS,
741                "permessage-deflate; client_max_window_bits",
742            )
743            .body(())
744            .unwrap();
745        let on_upgrade = hyper::upgrade::on(&mut req);
746        let (mut parts, ()) = req.into_parts();
747        let _ = parts.extensions.insert(on_upgrade);
748
749        let upgrade = WebSocketUpgrade::from_request_parts(&mut parts, &()).unwrap();
750        assert_eq!(upgrade.deflate_offers.len(), 1);
751        let agreement = deflate::negotiate(&upgrade.deflate_offers, upgrade.deflate_config);
752        assert!(agreement.is_some());
753    }
754}
tachyon_web/ws/mod.rs

tachyon_web/ws/
mod.rs
