trillium_client/
websocket.rs

1//! Support for client-side WebSockets
2
3use crate::{Conn, WebSocketConfig, WebSocketConn};
4use std::{
5    borrow::Cow,
6    error::Error,
7    fmt::{self, Display},
8    ops::{Deref, DerefMut},
9};
10use trillium_http::{
11    KnownHeaderName::{SecWebsocketAccept, SecWebsocketKey, SecWebsocketVersion},
12    Status, Upgrade, Version,
13};
14pub use trillium_websockets::Message;
15use trillium_websockets::{Role, websocket_accept_hash};
16
17impl Conn {
18    /// Attempt to transform this `Conn` into a [`WebSocketConn`].
19    ///
20    /// This is an *execution* method: calling it on a conn that has already been awaited
21    /// returns [`ErrorKind::AlreadyExecuted`]. Build the conn, then call this — don't await
22    /// it yourself first.
23    ///
24    /// The handshake is an `Upgrade` over HTTP/1.1 and an extended CONNECT over HTTP/2 and
25    /// HTTP/3. Which one is sent follows the same protocol selection as any other request,
26    /// with one addition: a peer that speaks h2 or h3 but does not support extended CONNECT
27    /// is retried as an HTTP/1.1 upgrade, unless
28    /// [`strict_http_version`](Conn::strict_http_version) is on, in which case it yields
29    /// [`ErrorKind::ExtendedConnectUnsupported`]. See the crate-level [Protocol
30    /// selection][crate#protocol-selection] documentation.
31    pub async fn into_websocket(self) -> Result<WebSocketConn, WebSocketUpgradeError> {
32        self.upgrade_to_websocket(None).await
33    }
34
35    /// Like [`Conn::into_websocket`] but with a caller-supplied [`WebSocketConfig`].
36    pub async fn into_websocket_with_config(
37        self,
38        config: WebSocketConfig,
39    ) -> Result<WebSocketConn, WebSocketUpgradeError> {
40        self.upgrade_to_websocket(Some(config)).await
41    }
42
43    async fn upgrade_to_websocket(
44        mut self,
45        config: Option<WebSocketConfig>,
46    ) -> Result<WebSocketConn, WebSocketUpgradeError> {
47        if self.status().is_some() {
48            return Err(WebSocketUpgradeError::new(self, ErrorKind::AlreadyExecuted));
49        }
50
51        // Only the protocol-neutral parts of the handshake are set here. The h1-only
52        // `Upgrade`/`Connection`/`Sec-WebSocket-Key` headers are added when the request is
53        // rendered for h1, so a request that goes out as an extended CONNECT never carries them.
54        self.protocol = Some(Cow::Borrowed("websocket"));
55        self.request_headers_mut()
56            .try_insert(SecWebsocketVersion, "13");
57
58        if let Err(e) = (&mut self).await {
59            let kind = match e {
60                trillium_http::Error::ExtendedConnectUnsupported => {
61                    ErrorKind::ExtendedConnectUnsupported
62                }
63                other => other.into(),
64            };
65            return Err(WebSocketUpgradeError::new(self, kind));
66        }
67
68        let status = self.status().expect("Response did not include status");
69        match self.http_version() {
70            Version::Http2 | Version::Http3 => {
71                if status != Status::Ok {
72                    return Err(WebSocketUpgradeError::new(self, ErrorKind::Status(status)));
73                }
74            }
75            _ => {
76                if status != Status::SwitchingProtocols {
77                    return Err(WebSocketUpgradeError::new(self, ErrorKind::Status(status)));
78                }
79                let key = self
80                    .request_headers()
81                    .get_str(SecWebsocketKey)
82                    .expect("h1 websocket request did not include Sec-WebSocket-Key");
83                let accept_key = websocket_accept_hash(key);
84                if self.response_headers().get_str(SecWebsocketAccept) != Some(&accept_key) {
85                    return Err(WebSocketUpgradeError::new(self, ErrorKind::InvalidAccept));
86                }
87            }
88        }
89
90        let peer_ip = self.peer_addr().map(|addr| addr.ip());
91        let mut conn = WebSocketConn::new(Upgrade::from(self), config, Role::Client).await;
92        conn.set_peer_ip(peer_ip);
93        Ok(conn)
94    }
95}
96
97/// The kind of error that occurred when attempting a websocket upgrade
98#[derive(thiserror::Error, Debug)]
99#[non_exhaustive]
100pub enum ErrorKind {
101    /// an HTTP error attempting to make the request
102    #[error(transparent)]
103    Http(#[from] trillium_http::Error),
104
105    /// Response didn't have the expected status (101 Switching Protocols for h1, 200 OK for
106    /// h2/h3 extended CONNECT).
107    #[error("Unexpected response status {0} for websocket upgrade")]
108    Status(Status),
109
110    /// Response Sec-WebSocket-Accept was missing or invalid; generally a server bug
111    #[error("Response Sec-WebSocket-Accept was missing or invalid")]
112    InvalidAccept,
113
114    /// `into_websocket` was called on a `Conn` that had already been executed (its status is
115    /// already set). The websocket upgrade *is* the execution; build the conn and call
116    /// `into_websocket` directly without awaiting first.
117    #[error(
118        "Conn::into_websocket called after execution — build the conn and await into_websocket \
119         instead of awaiting the conn separately"
120    )]
121    AlreadyExecuted,
122
123    /// The h2 or h3 peer did not advertise `SETTINGS_ENABLE_CONNECT_PROTOCOL = 1`, so the
124    /// extended-CONNECT bootstrap (RFC 8441 over h2, RFC 9220 over h3) is not available on this
125    /// connection. Only surfaced with [`strict_http_version`](Conn::strict_http_version) on;
126    /// otherwise the client retries as an HTTP/1.1 upgrade instead.
127    #[error("peer does not support extended CONNECT")]
128    ExtendedConnectUnsupported,
129}
130
131/// An attempted upgrade to a WebSocket failed.
132///
133/// You can transform this back into the Conn with [`From::from`]/[`Into::into`], if you need to
134/// look at the server response.
135#[derive(Debug)]
136pub struct WebSocketUpgradeError {
137    /// The kind of error that occurred
138    pub kind: ErrorKind,
139    conn: Box<Conn>,
140}
141
142impl WebSocketUpgradeError {
143    fn new(conn: Conn, kind: ErrorKind) -> Self {
144        let conn = Box::new(conn);
145        Self { conn, kind }
146    }
147}
148
149impl From<WebSocketUpgradeError> for Conn {
150    fn from(value: WebSocketUpgradeError) -> Self {
151        *value.conn
152    }
153}
154
155impl Deref for WebSocketUpgradeError {
156    type Target = Conn;
157
158    fn deref(&self) -> &Self::Target {
159        &self.conn
160    }
161}
162impl DerefMut for WebSocketUpgradeError {
163    fn deref_mut(&mut self) -> &mut Self::Target {
164        &mut self.conn
165    }
166}
167
168impl Error for WebSocketUpgradeError {}
169
170impl Display for WebSocketUpgradeError {
171    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
172        self.kind.fmt(f)
173    }
174}
trillium_client/websocket.rs

trillium_client/
websocket.rs
