tokio_websockets/
client.rs

1//! Implementation of a WebSocket client.
2//!
3//! This can be used in three ways:
4//!   - By letting the library connect to a remote URI and performing a HTTP/1.1
5//!     Upgrade handshake, via [`Builder::connect`]
6//!   - By letting the library perform a HTTP/1.1 Upgrade handshake on an
7//!     established stream, via [`Builder::connect_on`]
8//!   - By performing the handshake yourself and then using
9//!     [`Builder::take_over`] to let it take over a WebSocket stream
10use std::{future::poll_fn, io, pin::Pin, str::FromStr};
11
12use base64::{Engine, engine::general_purpose};
13use futures_core::Stream;
14use http::{
15    HeaderMap, HeaderValue, Uri,
16    header::{self, HeaderName},
17};
18use tokio::{
19    io::{AsyncRead, AsyncWrite, AsyncWriteExt},
20    net::TcpStream,
21};
22use tokio_util::codec::FramedRead;
23
24use crate::{
25    Connector, Error, MaybeTlsStream, WebSocketStream,
26    proto::{Config, Limits, Role},
27    resolver::{self, Resolver},
28    upgrade::{self, server_response},
29};
30
31/// Generates a new, random 16-byte WebSocket key and encodes it as base64.
32pub(crate) fn make_key() -> [u8; 24] {
33    let mut key_base64 = [0; 24];
34    let key_bytes = crate::rand::get_key();
35
36    // SAFETY: We know that 16 bytes will be 24 bytes base64-encoded
37    unsafe {
38        general_purpose::STANDARD
39            .encode_slice(key_bytes, &mut key_base64)
40            .unwrap_unchecked()
41    };
42
43    key_base64
44}
45
46/// Guesses the port to connect on for a URI. If none is specified, port 443
47/// will be used for TLS, 80 for plain HTTP.
48fn default_port(uri: &Uri) -> Option<u16> {
49    if let Some(port) = uri.port_u16() {
50        return Some(port);
51    }
52
53    let scheme = uri.scheme_str();
54
55    match scheme {
56        Some("https" | "wss") => Some(443),
57        Some("http" | "ws") => Some(80),
58        _ => None,
59    }
60}
61
62/// List of headers added by the client which will cause an error
63/// if added by the user:
64///
65/// - `host`
66/// - `upgrade`
67/// - `connection`
68/// - `sec-websocket-key`
69/// - `sec-websocket-version`
70pub const DISALLOWED_HEADERS: &[HeaderName] = &[
71    header::HOST,
72    header::UPGRADE,
73    header::CONNECTION,
74    header::SEC_WEBSOCKET_KEY,
75    header::SEC_WEBSOCKET_VERSION,
76];
77
78/// Builds a HTTP/1.1 Upgrade request for a URI with extra headers and a
79/// WebSocket key.
80fn build_request(uri: &Uri, key: &[u8], headers: &HeaderMap) -> Vec<u8> {
81    let mut buf = Vec::new();
82
83    buf.extend_from_slice(b"GET ");
84    buf.extend_from_slice(uri.path().as_bytes());
85
86    if let Some(query) = uri.query() {
87        buf.extend_from_slice(b"?");
88        buf.extend_from_slice(query.as_bytes());
89    }
90
91    buf.extend_from_slice(b" HTTP/1.1\r\n");
92
93    if let Some(host) = uri.host() {
94        buf.extend_from_slice(b"Host: ");
95        buf.extend_from_slice(host.as_bytes());
96
97        if let Some(port) = uri.port_u16() {
98            buf.extend_from_slice(b":");
99            buf.extend_from_slice(port.to_string().as_bytes());
100        }
101
102        buf.extend_from_slice(b"\r\n");
103    }
104
105    buf.extend_from_slice(b"Upgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Key: ");
106    buf.extend_from_slice(key);
107    buf.extend_from_slice(b"\r\nSec-WebSocket-Version: 13\r\n");
108
109    for (name, value) in headers {
110        buf.extend_from_slice(name.as_str().as_bytes());
111        buf.extend_from_slice(b": ");
112        buf.extend_from_slice(value.as_bytes());
113        buf.extend_from_slice(b"\r\n");
114    }
115
116    buf.extend_from_slice(b"\r\n");
117
118    buf
119}
120
121/// Builder for WebSocket client connections.
122pub struct Builder<'a, R: Resolver = resolver::Gai> {
123    /// URI to connect to, required unless connecting to an established
124    /// WebSocket stream.
125    uri: Option<Uri>,
126    /// A TLS connector to use for the connection. If not set and required, a
127    /// new one will be created.
128    connector: Option<&'a Connector>,
129    /// A DNS resolver to use for looking up the hostname.
130    resolver: R,
131    /// Configuration for the WebSocket stream.
132    config: Config,
133    /// Limits to impose on the WebSocket stream.
134    limits: Limits,
135    /// Headers to be sent with the upgrade request.
136    headers: HeaderMap,
137}
138
139impl Builder<'_> {
140    /// Creates a [`Builder`] with all defaults that is not configured to
141    /// connect to any server.
142    #[must_use]
143    pub fn new() -> Self {
144        Self {
145            uri: None,
146            connector: None,
147            resolver: resolver::Gai,
148            config: Config::default(),
149            limits: Limits::default(),
150            headers: HeaderMap::new(),
151        }
152    }
153
154    /// Creates a [`Builder`] that connects to a given URI. This URI must use
155    /// the `ws` or `wss` schemes.
156    ///
157    /// This method never fails as the URI has already been parsed.
158    #[must_use]
159    pub fn from_uri(uri: Uri) -> Self {
160        Self {
161            uri: Some(uri),
162            connector: None,
163            resolver: resolver::Gai,
164            config: Config::default(),
165            limits: Limits::default(),
166            headers: HeaderMap::new(),
167        }
168    }
169}
170
171impl<'a, R: Resolver> Builder<'a, R> {
172    /// Sets the [`Uri`] to connect to. This URI must use the `ws` or `wss`
173    /// schemes.
174    ///
175    /// # Errors
176    ///
177    /// This method returns a [`http::uri::InvalidUri`] error if URI parsing
178    /// fails.
179    pub fn uri(mut self, uri: &str) -> Result<Self, http::uri::InvalidUri> {
180        self.uri = Some(Uri::from_str(uri)?);
181
182        Ok(self)
183    }
184
185    /// Sets the TLS connector for the client.
186    ///
187    /// By default, the client will create a new one for each connection instead
188    /// of reusing one.
189    #[must_use]
190    pub fn connector(mut self, connector: &'a Connector) -> Self {
191        self.connector = Some(connector);
192
193        self
194    }
195
196    /// Sets the DNS resolver for the client.
197    ///
198    /// By default, the client will use the [`Gai`] resolver, a wrapper around
199    /// the blocking `getaddrinfo` syscall.
200    ///
201    /// [`Gai`]: resolver::Gai
202    #[must_use]
203    pub fn resolver<NewR: Resolver>(self, resolver: NewR) -> Builder<'a, NewR> {
204        let Builder {
205            uri,
206            connector,
207            resolver: _,
208            config,
209            limits,
210            headers,
211        } = self;
212
213        Builder {
214            uri,
215            connector,
216            resolver,
217            config,
218            limits,
219            headers,
220        }
221    }
222
223    /// Sets the configuration for the WebSocket stream.
224    #[must_use]
225    pub fn config(mut self, config: Config) -> Self {
226        self.config = config;
227
228        self
229    }
230
231    /// Sets the limits for the WebSocket stream.
232    #[must_use]
233    pub fn limits(mut self, limits: Limits) -> Self {
234        self.limits = limits;
235
236        self
237    }
238
239    /// Adds an extra HTTP header to the handshake request.
240    ///
241    /// # Errors
242    ///
243    /// Returns [`Error::DisallowedHeader`] if the header is in
244    /// the [`DISALLOWED_HEADERS`] list.
245    pub fn add_header(mut self, name: HeaderName, value: HeaderValue) -> Result<Self, Error> {
246        if DISALLOWED_HEADERS.contains(&name) {
247            return Err(Error::DisallowedHeader);
248        }
249        self.headers.insert(name, value);
250
251        Ok(self)
252    }
253
254    /// Establishes a connection to the WebSocket server. This requires a URI to
255    /// be configured via [`Builder::uri`].
256    ///
257    /// # Errors
258    ///
259    /// This method returns an [`Error`] if connecting to the server fails or no
260    /// URI has been configured.
261    pub async fn connect(
262        &self,
263    ) -> Result<
264        (
265            WebSocketStream<MaybeTlsStream<TcpStream>>,
266            upgrade::Response,
267        ),
268        Error,
269    > {
270        let uri = self.uri.as_ref().ok_or(Error::NoUriConfigured)?;
271        // Uri::host contains square brackets around IPv6 addresses, which is required
272        // by the RFC: https://datatracker.ietf.org/doc/html/rfc3986#section-3.2.2
273        // These, however, do not resolve.
274        let host = uri
275            .host()
276            .ok_or(Error::CannotResolveHost)?
277            .trim_start_matches('[')
278            .trim_end_matches(']');
279        let port = default_port(uri).unwrap_or(80);
280        let addr = self.resolver.resolve(host, port).await?;
281
282        let stream = TcpStream::connect(&addr).await?;
283
284        let stream = if uri.scheme_str() == Some("wss") {
285            if let Some(connector) = self.connector {
286                connector.wrap(host, stream).await?
287            } else {
288                let connector = Connector::new()?;
289
290                connector.wrap(host, stream).await?
291            }
292        } else if uri.scheme_str() == Some("ws") {
293            Connector::Plain.wrap(host, stream).await?
294        } else {
295            return Err(Error::UnsupportedScheme);
296        };
297
298        self.connect_on(stream).await
299    }
300
301    /// Takes over an already established stream and uses it to send and receive
302    /// WebSocket messages. This requires a URI to be configured via
303    /// [`Builder::uri`].
304    ///
305    /// This method assumes that the TLS connection has already been
306    /// established, if needed. It sends an HTTP upgrade request and waits
307    /// for an HTTP Switching Protocols response before proceeding.
308    ///
309    /// # Errors
310    ///
311    /// This method returns an [`Error`] if writing or reading from the stream
312    /// fails or no URI has been configured.
313    pub async fn connect_on<S: AsyncRead + AsyncWrite + Unpin>(
314        &self,
315        mut stream: S,
316    ) -> Result<(WebSocketStream<S>, upgrade::Response), Error> {
317        let uri = self.uri.as_ref().ok_or(Error::NoUriConfigured)?;
318
319        let key_base64 = make_key();
320
321        let upgrade_codec = server_response::Codec::new(&key_base64);
322        let request = build_request(uri, &key_base64, &self.headers);
323        stream.write_all(&request).await?;
324        stream.flush().await?;
325
326        let mut framed = FramedRead::new(stream, upgrade_codec);
327        let res = poll_fn(|cx| Pin::new(&mut framed).poll_next(cx))
328            .await
329            .ok_or(Error::Io(io::ErrorKind::UnexpectedEof.into()))??;
330
331        Ok((
332            WebSocketStream::from_framed(framed, Role::Client, self.config, self.limits),
333            res,
334        ))
335    }
336
337    /// Takes over an already established stream that has already performed a
338    /// HTTP upgrade handshake and uses it to send and receive WebSocket
339    /// messages.
340    ///
341    /// This method will not perform a TLS handshake or a HTTP upgrade
342    /// handshake, it assumes the stream is ready to use for writing and
343    /// reading the WebSocket protocol.
344    pub fn take_over<S: AsyncRead + AsyncWrite + Unpin>(&self, stream: S) -> WebSocketStream<S> {
345        WebSocketStream::from_raw_stream(stream, Role::Client, self.config, self.limits)
346    }
347}
348
349impl Default for Builder<'_> {
350    fn default() -> Self {
351        Self::new()
352    }
353}
354
355#[cfg(test)]
356mod tests {
357    use futures_util::StreamExt;
358    use static_assertions::assert_impl_all;
359
360    use super::Builder;
361    use crate::{Error, proto::ProtocolError};
362
363    assert_impl_all!(Builder: Send, Sync);
364
365    #[tokio::test]
366    async fn control_payload_limit_receive() {
367        #[rustfmt::skip]
368        let overlong: [&[u8]; 3]= [
369            // Ping with 126 byte payload.
370            &[137, 126, 0, 126, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
371            // Pong with 126 byte payload.
372            &[138, 126, 0, 126, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
373            // Close with 124 byte reason.
374            &[136, 126, 0, 126, 3, 232, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97, 97]
375        ];
376
377        for mut message in overlong {
378            let mut stream =
379                Builder::new().take_over(tokio::io::join(&mut message, tokio::io::empty()));
380
381            let result = stream.next().await.unwrap();
382            assert!(matches!(
383                result,
384                Err(Error::Protocol(ProtocolError::InvalidPayloadLength))
385            ));
386        }
387    }
388}
tokio_websockets/client.rs

tokio_websockets/
client.rs
