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

trillium_client/
websocket.rs
