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::{
12        Connection, SecWebsocketAccept, SecWebsocketKey, SecWebsocketVersion,
13        Upgrade as UpgradeHeader,
14    },
15    Method, Status, Upgrade, Version,
16};
17pub use trillium_websockets::Message;
18use trillium_websockets::{Role, websocket_accept_hash, websocket_key};
19
20impl Conn {
21    fn set_websocket_upgrade_headers_h1(&mut self) {
22        let headers = self.request_headers_mut();
23        headers.try_insert(UpgradeHeader, "websocket");
24        headers.try_insert(Connection, "upgrade");
25        headers.try_insert(SecWebsocketVersion, "13");
26        headers.try_insert(SecWebsocketKey, websocket_key());
27    }
28
29    /// Attempt to transform this `Conn` into a [`WebSocketConn`].
30    ///
31    /// This is an *execution* method: calling it on a conn that has already been awaited
32    /// returns [`ErrorKind::AlreadyExecuted`]. Build the conn, then call this — don't await
33    /// it yourself first.
34    ///
35    /// Protocol selection follows the conn's [`http_version`][Conn::http_version] hint:
36    /// `Http2` uses the extended-CONNECT bootstrap (RFC 8441); the default uses an h1
37    /// `Upgrade` handshake (RFC 6455). If the peer is h2 but doesn't advertise
38    /// `SETTINGS_ENABLE_CONNECT_PROTOCOL`, the upgrade hard-errors — there is no silent
39    /// fallback to h1 from a non-capable h2 peer.
40    ///
41    /// HTTP/3 (RFC 9220) extended CONNECT is not yet supported on the client. The h3 transport
42    /// would need to wrap post-upgrade bytes in h3 DATA frames on both client and server before
43    /// the byte channel will round-trip, and that framing layer doesn't exist yet. A `Http3`
44    /// hint here surfaces as `ErrorKind::ExtendedConnectUnsupported`.
45    pub async fn into_websocket(self) -> Result<WebSocketConn, WebSocketUpgradeError> {
46        self.into_websocket_with_config(WebSocketConfig::default())
47            .await
48    }
49
50    /// Like [`Conn::into_websocket`] but with a caller-supplied [`WebSocketConfig`].
51    pub async fn into_websocket_with_config(
52        self,
53        config: WebSocketConfig,
54    ) -> Result<WebSocketConn, WebSocketUpgradeError> {
55        if self.status().is_some() {
56            return Err(WebSocketUpgradeError::new(self, ErrorKind::AlreadyExecuted));
57        }
58
59        match self.http_version() {
60            Version::Http2 => self.into_websocket_extended_connect(config).await,
61            Version::Http3 => Err(WebSocketUpgradeError::new(
62                self,
63                ErrorKind::ExtendedConnectUnsupported,
64            )),
65            _ => self.into_websocket_h1(config).await,
66        }
67    }
68
69    async fn into_websocket_h1(
70        mut self,
71        config: WebSocketConfig,
72    ) -> Result<WebSocketConn, WebSocketUpgradeError> {
73        self.set_websocket_upgrade_headers_h1();
74        if let Err(e) = (&mut self).await {
75            return Err(WebSocketUpgradeError::new(self, e.into()));
76        }
77        let status = self.status().expect("Response did not include status");
78        if status != Status::SwitchingProtocols {
79            return Err(WebSocketUpgradeError::new(self, ErrorKind::Status(status)));
80        }
81        let key = self
82            .request_headers()
83            .get_str(SecWebsocketKey)
84            .expect("Request did not include Sec-WebSocket-Key");
85        let accept_key = websocket_accept_hash(key);
86        if self.response_headers().get_str(SecWebsocketAccept) != Some(&accept_key) {
87            return Err(WebSocketUpgradeError::new(self, ErrorKind::InvalidAccept));
88        }
89        let peer_ip = self.peer_addr().map(|addr| addr.ip());
90        let mut conn = WebSocketConn::new(Upgrade::from(self), Some(config), Role::Client).await;
91        conn.set_peer_ip(peer_ip);
92        Ok(conn)
93    }
94
95    async fn into_websocket_extended_connect(
96        mut self,
97        config: WebSocketConfig,
98    ) -> Result<WebSocketConn, WebSocketUpgradeError> {
99        // RFC 8441 §4 / RFC 9220 §3: the upgrade carries `Sec-WebSocket-Version: 13` and the
100        // optional `Sec-WebSocket-Protocol`, but skips the `Sec-WebSocket-Key` /
101        // `Sec-WebSocket-Accept` SHA1 dance — those are h1-only artifacts. The
102        // `Connection: upgrade` / `Upgrade: websocket` headers are likewise h1-only and would
103        // be stripped by `finalize_headers_h2` / `_h3` even if we set them.
104        self.request_headers_mut()
105            .try_insert(SecWebsocketVersion, "13");
106        self.set_method(Method::Connect);
107        self.protocol = Some(Cow::Borrowed("websocket"));
108
109        // The peer-capability gate (RFC 8441 §3 — server must have advertised
110        // `SETTINGS_ENABLE_CONNECT_PROTOCOL` before the client may send a `:protocol`
111        // HEADERS) lives inside the h2 client send path, where it can park on the peer's
112        // first SETTINGS *before* putting any HEADERS on the wire. A "not supported"
113        // outcome surfaces here as `Error::ExtendedConnectUnsupported`.
114        if let Err(e) = (&mut self).await {
115            let kind = match e {
116                trillium_http::Error::ExtendedConnectUnsupported => {
117                    ErrorKind::ExtendedConnectUnsupported
118                }
119                other => other.into(),
120            };
121            return Err(WebSocketUpgradeError::new(self, kind));
122        }
123
124        let status = self.status().expect("Response did not include status");
125        if status != Status::Ok {
126            return Err(WebSocketUpgradeError::new(self, ErrorKind::Status(status)));
127        }
128
129        let peer_ip = self.peer_addr().map(|addr| addr.ip());
130        let mut conn = WebSocketConn::new(Upgrade::from(self), Some(config), Role::Client).await;
131        conn.set_peer_ip(peer_ip);
132        Ok(conn)
133    }
134}
135
136/// The kind of error that occurred when attempting a websocket upgrade
137#[derive(thiserror::Error, Debug)]
138#[non_exhaustive]
139pub enum ErrorKind {
140    /// an HTTP error attempting to make the request
141    #[error(transparent)]
142    Http(#[from] trillium_http::Error),
143
144    /// Response didn't have the expected status (101 Switching Protocols for h1, 200 OK for
145    /// h2/h3 extended CONNECT).
146    #[error("Unexpected response status {0} for websocket upgrade")]
147    Status(Status),
148
149    /// Response Sec-WebSocket-Accept was missing or invalid; generally a server bug
150    #[error("Response Sec-WebSocket-Accept was missing or invalid")]
151    InvalidAccept,
152
153    /// `into_websocket` was called on a `Conn` that had already been executed (its status is
154    /// already set). The websocket upgrade *is* the execution; build the conn and call
155    /// `into_websocket` directly without awaiting first.
156    #[error(
157        "Conn::into_websocket called after execution — build the conn and await into_websocket \
158         instead of awaiting the conn separately"
159    )]
160    AlreadyExecuted,
161
162    /// h2 peer did not advertise `SETTINGS_ENABLE_CONNECT_PROTOCOL = 1`, so the extended-CONNECT
163    /// bootstrap (RFC 8441) is not available on this connection.
164    ///
165    /// Also surfaced when the conn was hinted as `Version::Http3`: client-side WebSocket-over-h3
166    /// (RFC 9220) requires h3 DATA-frame wrapping for the post-upgrade byte channel and that
167    /// framing layer doesn't exist yet.
168    #[error("peer does not support extended CONNECT, or h3 client websocket framing is missing")]
169    ExtendedConnectUnsupported,
170}
171
172/// An attempted upgrade to a WebSocket failed.
173///
174/// You can transform this back into the Conn with [`From::from`]/[`Into::into`], if you need to
175/// look at the server response.
176#[derive(Debug)]
177pub struct WebSocketUpgradeError {
178    /// The kind of error that occurred
179    pub kind: ErrorKind,
180    conn: Box<Conn>,
181}
182
183impl WebSocketUpgradeError {
184    fn new(conn: Conn, kind: ErrorKind) -> Self {
185        let conn = Box::new(conn);
186        Self { conn, kind }
187    }
188}
189
190impl From<WebSocketUpgradeError> for Conn {
191    fn from(value: WebSocketUpgradeError) -> Self {
192        *value.conn
193    }
194}
195
196impl Deref for WebSocketUpgradeError {
197    type Target = Conn;
198
199    fn deref(&self) -> &Self::Target {
200        &self.conn
201    }
202}
203impl DerefMut for WebSocketUpgradeError {
204    fn deref_mut(&mut self) -> &mut Self::Target {
205        &mut self.conn
206    }
207}
208
209impl Error for WebSocketUpgradeError {}
210
211impl Display for WebSocketUpgradeError {
212    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
213        self.kind.fmt(f)
214    }
215}
trillium_client/websocket.rs

trillium_client/
websocket.rs
