Skip to main content

h2ts_server/
handshake.rs

1//! Server-side WebSocket handshake (RFC 6455).
2//!
3//! It's small: parse the upgrade headers, compute `Sec-WebSocket-Accept` (SHA-1
4//! of the client key plus the RFC 6455 magic GUID, base64-encoded), reply `101`,
5//! and hand back hyper's upgraded connection as a raw byte stream. The framing on
6//! top of that stream is done by [`bridge`](crate::bridge) using wslay.
7
8use std::fmt;
9use std::future::Future;
10
11use base64::Engine as _;
12use bytes::Bytes;
13use http::header::{
14    HeaderName, HeaderValue, CONNECTION, SEC_WEBSOCKET_ACCEPT, SEC_WEBSOCKET_KEY,
15    SEC_WEBSOCKET_PROTOCOL, UPGRADE,
16};
17use http_body_util::Empty;
18use hyper::upgrade::Upgraded;
19use hyper::{Request, Response, StatusCode};
20use hyper_util::rt::TokioIo;
21use sha1::{Digest, Sha1};
22
23/// RFC 6455 §4.2.2: the magic GUID appended to the client key before hashing.
24const WS_GUID: &[u8] = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
25
26/// The subprotocol h2ts clients offer by default. [`accept`] echoes it when the
27/// client offered it and the handler didn't choose another (see [`accept_with`]).
28pub const DEFAULT_SUBPROTOCOL: &str = "h2ts";
29
30/// Options controlling how the handshake picks a subprotocol to echo when the
31/// handler's selector declines (see [`accept_with_options`]).
32#[derive(Debug, Clone, Copy, Default)]
33pub struct AcceptOptions {
34    /// When the selector returns `None` and the client did **not** offer
35    /// [`DEFAULT_SUBPROTOCOL`] (`h2ts`), accept the client's first offered
36    /// subprotocol (or, if it offered none, complete with no subprotocol)
37    /// instead of **rejecting** the handshake.
38    ///
39    /// Off by default: a client that speaks neither `h2ts` nor a
40    /// handler-selected protocol is rejected with
41    /// [`WebSocketError::UnsupportedSubprotocol`]. Turn this on for a
42    /// codec-agnostic tunnel that accepts whatever framing the client offered
43    /// (e.g. a raw/binary tunnel, or a websockify-style `binary` client).
44    pub allow_implicit_codec: bool,
45}
46
47/// The upgraded WebSocket connection, presented as a raw tokio byte stream
48/// (`AsyncRead + AsyncWrite`). Hand it to [`bridge`](crate::bridge),
49/// [`serve_h2`](crate::serve_h2), or [`WsByteStream`](crate::WsByteStream).
50pub type UpgradedIo = TokioIo<Upgraded>;
51
52/// Errors from the WebSocket handshake.
53#[derive(Debug)]
54pub enum WebSocketError {
55    /// The request was not a valid WebSocket upgrade (missing/invalid
56    /// `Upgrade` / `Connection` / `Sec-WebSocket-Key`).
57    NotUpgradeRequest,
58    /// The client offered no subprotocol the server accepts: it didn't offer
59    /// [`DEFAULT_SUBPROTOCOL`] (`h2ts`), the handler's selector declined, and
60    /// [`AcceptOptions::allow_implicit_codec`] was off. Reject it with a `400`
61    /// (see [`WebSocketError::rejection_response`]).
62    UnsupportedSubprotocol,
63    /// hyper failed to upgrade the connection after the `101` was sent.
64    Upgrade(hyper::Error),
65}
66
67impl WebSocketError {
68    /// The HTTP response to reject this handshake with, for the pre-upgrade
69    /// errors returned by [`accept`] / [`accept_with`] / [`accept_with_options`]:
70    /// `426 Upgrade Required` for a non-WebSocket request, `400 Bad Request` for
71    /// an unsupported subprotocol. (An [`Upgrade`](WebSocketError::Upgrade) error
72    /// happens *after* the `101` and has no meaningful rejection response; it maps
73    /// to `500` for completeness.)
74    ///
75    /// ```ignore
76    /// let (response, ws_fut) = match h2ts_server::accept(&mut req) {
77    ///     Ok(pair) => pair,
78    ///     Err(err) => return Ok(err.rejection_response()), // send the 4xx back
79    /// };
80    /// ```
81    pub fn rejection_response(&self) -> Response<Empty<Bytes>> {
82        let status = match self {
83            WebSocketError::NotUpgradeRequest => StatusCode::UPGRADE_REQUIRED,
84            WebSocketError::UnsupportedSubprotocol => StatusCode::BAD_REQUEST,
85            WebSocketError::Upgrade(_) => StatusCode::INTERNAL_SERVER_ERROR,
86        };
87        Response::builder()
88            .status(status)
89            .body(Empty::new())
90            .expect("static rejection response is well-formed")
91    }
92}
93
94impl fmt::Display for WebSocketError {
95    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
96        match self {
97            WebSocketError::NotUpgradeRequest => f.write_str("not a WebSocket upgrade request"),
98            WebSocketError::UnsupportedSubprotocol => {
99                f.write_str("client offered no supported subprotocol")
100            }
101            WebSocketError::Upgrade(e) => write!(f, "WebSocket upgrade failed: {e}"),
102        }
103    }
104}
105
106impl std::error::Error for WebSocketError {
107    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
108        match self {
109            WebSocketError::Upgrade(e) => Some(e),
110            WebSocketError::NotUpgradeRequest | WebSocketError::UnsupportedSubprotocol => None,
111        }
112    }
113}
114
115/// Whether `header` is present and lists `needle` as one of its comma-separated,
116/// case-insensitive tokens.
117fn header_lists<B>(request: &Request<B>, header: HeaderName, needle: &str) -> bool {
118    request
119        .headers()
120        .get(header)
121        .and_then(|v| v.to_str().ok())
122        .map(|list| list.split(',').any(|t| t.trim().eq_ignore_ascii_case(needle)))
123        .unwrap_or(false)
124}
125
126/// Whether `request` is a WebSocket upgrade: `Upgrade: websocket`,
127/// `Connection: upgrade`, and a `Sec-WebSocket-Key`.
128pub fn is_upgrade_request<B>(request: &Request<B>) -> bool {
129    header_lists(request, UPGRADE, "websocket")
130        && header_lists(request, CONNECTION, "upgrade")
131        && request.headers().contains_key(SEC_WEBSOCKET_KEY)
132}
133
134/// The WebSocket subprotocols the client offered, in order, whitespace-trimmed;
135/// empty if it offered none. Gives a handler full visibility into the offer so
136/// it can pass one to [`accept_with`].
137pub fn offered_protocols<B>(request: &Request<B>) -> Vec<&str> {
138    request
139        .headers()
140        .get(SEC_WEBSOCKET_PROTOCOL)
141        .and_then(|v| v.to_str().ok())
142        .map(|list| {
143            list.split(',')
144                .map(str::trim)
145                .filter(|s| !s.is_empty())
146                .collect()
147        })
148        .unwrap_or_default()
149}
150
151/// The outcome of the default subprotocol fallback, when the handler's selector
152/// declined.
153#[derive(Debug, PartialEq, Eq)]
154enum Fallback {
155    /// Complete the handshake, echoing this subprotocol (`Some`) or none (`None`).
156    Accept(Option<String>),
157    /// Reject the handshake: the client offered no acceptable subprotocol and
158    /// implicit codecs are disabled.
159    Reject,
160}
161
162/// The subprotocol decision when the handler's selector declined.
163///
164/// Prefers [`DEFAULT_SUBPROTOCOL`] (`h2ts`) whenever the client offered it (in
165/// its exact offered casing). Otherwise, with `allow_implicit_codec`, accept the
166/// client's first offered subprotocol (or none, if it offered nothing); without
167/// it, [`Fallback::Reject`] — we don't commit to a codec we don't understand.
168fn fallback_subprotocol(offered: &[&str], allow_implicit_codec: bool) -> Fallback {
169    if let Some(p) = offered
170        .iter()
171        .find(|p| p.eq_ignore_ascii_case(DEFAULT_SUBPROTOCOL))
172    {
173        return Fallback::Accept(Some(p.to_string()));
174    }
175    if allow_implicit_codec {
176        return Fallback::Accept(offered.first().map(|p| p.to_string()));
177    }
178    Fallback::Reject
179}
180
181/// Compute the `Sec-WebSocket-Accept` value for a client `Sec-WebSocket-Key`.
182fn accept_key(key: &[u8]) -> String {
183    let mut hasher = Sha1::new();
184    hasher.update(key);
185    hasher.update(WS_GUID);
186    base64::engine::general_purpose::STANDARD.encode(hasher.finalize())
187}
188
189/// Accept a WebSocket upgrade, choosing which offered subprotocol to echo, with
190/// explicit [`AcceptOptions`].
191///
192/// `select` is handed the full list of subprotocols the client offered (the same
193/// thing [`offered_protocols`] returns) and returns the one to echo in the `101`,
194/// or `None` to decline. When it declines, [`DEFAULT_SUBPROTOCOL`] (`h2ts`) is
195/// echoed if the client offered it; otherwise the handshake is **rejected** with
196/// [`WebSocketError::UnsupportedSubprotocol`] — unless
197/// [`AcceptOptions::allow_implicit_codec`] is set, which instead accepts the
198/// client's first offered codec (or completes with none if it offered nothing).
199/// Turn a rejection into a `4xx` to send back via
200/// [`WebSocketError::rejection_response`].
201///
202/// Per RFC 6455 a client that offered subprotocols will *fail* the connection if
203/// the server echoes one it did not offer, so `select` should return a member of
204/// the offered list (or `None`).
205///
206/// See [`accept_with`] (default options) and [`accept`] (the common case).
207// The returned `impl Future` can't be aliased away, so the tuple reads as
208// "complex"; it's just (response, upgrade-future).
209#[allow(clippy::type_complexity)]
210pub fn accept_with_options<B, F>(
211    request: &mut Request<B>,
212    select: F,
213    options: AcceptOptions,
214) -> Result<
215    (
216        Response<Empty<Bytes>>,
217        impl Future<Output = Result<UpgradedIo, WebSocketError>>,
218    ),
219    WebSocketError,
220>
221where
222    F: FnOnce(&[&str]) -> Option<String>,
223{
224    if !is_upgrade_request(request) {
225        return Err(WebSocketError::NotUpgradeRequest);
226    }
227    // Present because `is_upgrade_request` checked for it.
228    let key = request
229        .headers()
230        .get(SEC_WEBSOCKET_KEY)
231        .ok_or(WebSocketError::NotUpgradeRequest)?;
232    let accept_value = accept_key(key.as_bytes());
233
234    // Let the handler pick from the offered list; otherwise apply the default
235    // fallback (h2ts if offered, then optionally the first offered codec, else
236    // reject).
237    let offered = offered_protocols(request);
238    let decision = match select(&offered) {
239        Some(proto) => Fallback::Accept(Some(proto)),
240        None => fallback_subprotocol(&offered, options.allow_implicit_codec),
241    };
242    drop(offered); // release the immutable borrow before upgrading below
243    let chosen = match decision {
244        Fallback::Accept(proto) => proto,
245        Fallback::Reject => return Err(WebSocketError::UnsupportedSubprotocol),
246    };
247
248    // base64 output is always valid header-value ASCII, so this never fails.
249    let accept_header =
250        HeaderValue::from_str(&accept_value).expect("base64 is valid header ASCII");
251    let mut response = Response::builder()
252        .status(StatusCode::SWITCHING_PROTOCOLS)
253        .header(CONNECTION, HeaderValue::from_static("Upgrade"))
254        .header(UPGRADE, HeaderValue::from_static("websocket"))
255        .header(SEC_WEBSOCKET_ACCEPT, accept_header)
256        .body(Empty::<Bytes>::new())
257        .expect("static 101 response is well-formed");
258    if let Some(proto) = chosen {
259        // Skip a subprotocol that isn't a valid header value rather than fail.
260        if let Ok(value) = HeaderValue::from_str(&proto) {
261            response
262                .headers_mut()
263                .insert(SEC_WEBSOCKET_PROTOCOL, value);
264        }
265    }
266
267    // Capture the upgrade future now; it resolves once the 101 is flushed and
268    // hyper hands over the connection.
269    let on_upgrade = hyper::upgrade::on(&mut *request);
270    let fut = async move {
271        let upgraded = on_upgrade.await.map_err(WebSocketError::Upgrade)?;
272        Ok(TokioIo::new(upgraded))
273    };
274    Ok((response, fut))
275}
276
277/// Accept a WebSocket upgrade, choosing which offered subprotocol to echo.
278///
279/// Like [`accept_with_options`] with the default [`AcceptOptions`]: when `select`
280/// declines, [`DEFAULT_SUBPROTOCOL`] (`h2ts`) is echoed if the client offered it;
281/// otherwise the handshake is **rejected**
282/// ([`WebSocketError::UnsupportedSubprotocol`]). Use [`accept_with_options`] with
283/// `allow_implicit_codec` to instead accept the client's first offered codec.
284///
285/// ```ignore
286/// // Prefer "chat" if the client offered it, otherwise fall back to h2ts.
287/// let (response, ws_fut) = h2ts_server::accept_with(&mut req, |offered| {
288///     offered.iter().find(|p| p.eq_ignore_ascii_case("chat")).map(|p| p.to_string())
289/// })?;
290/// ```
291#[allow(clippy::type_complexity)]
292pub fn accept_with<B, F>(
293    request: &mut Request<B>,
294    select: F,
295) -> Result<
296    (
297        Response<Empty<Bytes>>,
298        impl Future<Output = Result<UpgradedIo, WebSocketError>>,
299    ),
300    WebSocketError,
301>
302where
303    F: FnOnce(&[&str]) -> Option<String>,
304{
305    accept_with_options(request, select, AcceptOptions::default())
306}
307
308/// Accept a WebSocket upgrade, requiring the [`DEFAULT_SUBPROTOCOL`] (`h2ts`).
309///
310/// Echoes `h2ts` when the client offered it; a client that doesn't offer `h2ts`
311/// is **rejected** with [`WebSocketError::UnsupportedSubprotocol`] (send
312/// [`rejection_response`](WebSocketError::rejection_response) back as a `400`).
313/// Use [`accept_with`] to select a different offered subprotocol, or
314/// [`accept_with_options`] with `allow_implicit_codec` to accept any codec.
315///
316/// On success, returns the `101 Switching Protocols` response to send back
317/// immediately, plus a future that resolves to the upgraded connection as a byte
318/// stream ([`UpgradedIo`]). Drive the response through your framework; spawn the
319/// future and hand the stream to [`bridge`](crate::bridge) or
320/// [`serve_h2`](crate::serve_h2). Control frames (ping/close) are handled
321/// downstream by wslay.
322///
323/// The outer HTTP/1 connection must be served with upgrades enabled
324/// (`http1::Builder::serve_connection(..).with_upgrades()`).
325#[allow(clippy::type_complexity)]
326pub fn accept<B>(
327    request: &mut Request<B>,
328) -> Result<
329    (
330        Response<Empty<Bytes>>,
331        impl Future<Output = Result<UpgradedIo, WebSocketError>>,
332    ),
333    WebSocketError,
334> {
335    accept_with(request, |_offered| None)
336}
337
338#[cfg(test)]
339mod tests {
340    use super::{
341        accept, accept_with, accept_with_options, fallback_subprotocol, is_upgrade_request,
342        offered_protocols, AcceptOptions, WebSocketError,
343    };
344    use http::header::{CONNECTION, SEC_WEBSOCKET_KEY, SEC_WEBSOCKET_PROTOCOL, UPGRADE};
345    use hyper::{Request, StatusCode};
346
347    fn with_protocol(protocol: Option<&str>) -> Request<()> {
348        let mut b = Request::builder();
349        if let Some(p) = protocol {
350            b = b.header(SEC_WEBSOCKET_PROTOCOL, p);
351        }
352        b.body(()).unwrap()
353    }
354
355    /// A well-formed WebSocket upgrade request offering `offer` (if any).
356    fn upgrade_request(offer: Option<&str>) -> Request<()> {
357        let mut b = Request::builder()
358            .header(UPGRADE, "websocket")
359            .header(CONNECTION, "Upgrade")
360            .header(SEC_WEBSOCKET_KEY, "dGhlIHNhbXBsZSBub25jZQ==");
361        if let Some(o) = offer {
362            b = b.header(SEC_WEBSOCKET_PROTOCOL, o);
363        }
364        b.body(()).unwrap()
365    }
366
367    #[test]
368    fn offered_protocols_parses_the_list_in_order() {
369        assert_eq!(offered_protocols(&with_protocol(Some("h2ts"))), ["h2ts"]);
370        assert_eq!(
371            offered_protocols(&with_protocol(Some("chat, h2ts, binary"))),
372            ["chat", "h2ts", "binary"]
373        );
374        assert_eq!(offered_protocols(&with_protocol(Some("  h2ts  "))), ["h2ts"]);
375        assert!(offered_protocols(&with_protocol(Some(""))).is_empty());
376        assert!(offered_protocols(&with_protocol(None)).is_empty());
377    }
378
379    #[test]
380    fn is_upgrade_request_requires_all_three_signals() {
381        let full = Request::builder()
382            .header(UPGRADE, "websocket")
383            .header(CONNECTION, "Upgrade")
384            .header(SEC_WEBSOCKET_KEY, "dGhlIHNhbXBsZSBub25jZQ==")
385            .body(())
386            .unwrap();
387        assert!(is_upgrade_request(&full));
388
389        // Connection may be a token list ("keep-alive, Upgrade").
390        let listed = Request::builder()
391            .header(UPGRADE, "websocket")
392            .header(CONNECTION, "keep-alive, Upgrade")
393            .header(SEC_WEBSOCKET_KEY, "dGhlIHNhbXBsZSBub25jZQ==")
394            .body(())
395            .unwrap();
396        assert!(is_upgrade_request(&listed));
397
398        // Missing the key.
399        let no_key = Request::builder()
400            .header(UPGRADE, "websocket")
401            .header(CONNECTION, "Upgrade")
402            .body(())
403            .unwrap();
404        assert!(!is_upgrade_request(&no_key));
405
406        // A plain request.
407        assert!(!is_upgrade_request(&Request::builder().body(()).unwrap()));
408    }
409
410    #[test]
411    fn fallback_prefers_h2ts_then_optionally_the_first_offered() {
412        use super::Fallback::{Accept, Reject};
413
414        // h2ts always wins when offered, in its exact casing, regardless of the flag.
415        assert_eq!(fallback_subprotocol(&["h2ts"], false), Accept(Some("h2ts".into())));
416        assert_eq!(
417            fallback_subprotocol(&["chat", "h2ts"], false),
418            Accept(Some("h2ts".into())),
419            "h2ts is preferred even when it isn't first"
420        );
421        assert_eq!(
422            fallback_subprotocol(&["H2TS"], false),
423            Accept(Some("H2TS".into())),
424            "matched case-insensitively but echoed in the offered casing"
425        );
426
427        // No h2ts, flag off: reject (an empty offer is also rejected).
428        assert_eq!(fallback_subprotocol(&["chat"], false), Reject);
429        assert_eq!(fallback_subprotocol(&["chat", "binary"], false), Reject);
430        assert_eq!(fallback_subprotocol(&[], false), Reject);
431
432        // No h2ts, flag on: implicitly accept the first offered codec...
433        assert_eq!(
434            fallback_subprotocol(&["chat", "binary"], true),
435            Accept(Some("chat".into()))
436        );
437        // ...and an empty offer completes with no subprotocol.
438        assert_eq!(fallback_subprotocol(&[], true), Accept(None));
439    }
440
441    #[test]
442    fn accept_echoes_h2ts_when_offered() {
443        let mut req = upgrade_request(Some("h2ts"));
444        let (resp, _fut) = accept(&mut req).unwrap();
445        assert_eq!(resp.status(), StatusCode::SWITCHING_PROTOCOLS);
446        assert_eq!(resp.headers().get(SEC_WEBSOCKET_PROTOCOL).unwrap(), "h2ts");
447
448        // Preferred even when offered alongside others.
449        let mut req = upgrade_request(Some("chat, h2ts"));
450        let (resp, _fut) = accept(&mut req).unwrap();
451        assert_eq!(resp.headers().get(SEC_WEBSOCKET_PROTOCOL).unwrap(), "h2ts");
452    }
453
454    #[test]
455    fn accept_rejects_when_h2ts_absent_by_default() {
456        // A non-h2ts offer and an empty offer both reject with a 400 response,
457        // rather than completing the handshake.
458        for offer in [Some("mystery"), Some("chat, binary"), None] {
459            let mut req = upgrade_request(offer);
460            let err = accept(&mut req).err().expect("should reject");
461            assert!(
462                matches!(&err, WebSocketError::UnsupportedSubprotocol),
463                "offer {offer:?}"
464            );
465            assert_eq!(err.rejection_response().status(), StatusCode::BAD_REQUEST);
466        }
467    }
468
469    #[test]
470    fn accept_with_honors_a_selection_even_without_h2ts() {
471        // The handler selecting an offered codec accepts it — no reject, no h2ts.
472        let mut req = upgrade_request(Some("chat, binary"));
473        let (resp, _fut) = accept_with(&mut req, |offered| {
474            offered.iter().find(|p| **p == "binary").map(|p| p.to_string())
475        })
476        .unwrap();
477        assert_eq!(resp.status(), StatusCode::SWITCHING_PROTOCOLS);
478        assert_eq!(resp.headers().get(SEC_WEBSOCKET_PROTOCOL).unwrap(), "binary");
479    }
480
481    #[test]
482    fn allow_implicit_codec_accepts_the_first_offered_codec() {
483        let opts = AcceptOptions {
484            allow_implicit_codec: true,
485        };
486        // Unknown codec: accepted, first offered echoed.
487        let mut req = upgrade_request(Some("mystery, other"));
488        let (resp, _fut) = accept_with_options(&mut req, |_| None, opts).unwrap();
489        assert_eq!(resp.status(), StatusCode::SWITCHING_PROTOCOLS);
490        assert_eq!(resp.headers().get(SEC_WEBSOCKET_PROTOCOL).unwrap(), "mystery");
491
492        // No offer: still accepted (permissive), nothing echoed.
493        let mut req = upgrade_request(None);
494        let (resp, _fut) = accept_with_options(&mut req, |_| None, opts).unwrap();
495        assert_eq!(resp.status(), StatusCode::SWITCHING_PROTOCOLS);
496        assert!(resp.headers().get(SEC_WEBSOCKET_PROTOCOL).is_none());
497    }
498
499    #[test]
500    fn accept_key_matches_rfc_6455_example() {
501        // RFC 6455 §1.3: key "dGhlIHNhbXBsZSBub25jZQ==" -> accept
502        // "s3pPLMBiTxaQ9kYGzzhZRbK+xOo=".
503        assert_eq!(
504            super::accept_key(b"dGhlIHNhbXBsZSBub25jZQ=="),
505            "s3pPLMBiTxaQ9kYGzzhZRbK+xOo="
506        );
507    }
508}