tokio_websockets/
server.rs

1//! Implementation of a WebSocket server.
2//!
3//! This can be used in two ways:
4//!   - By letting the library perform a HTTP/1.1 Upgrade handshake on an
5//!     established stream, via [`Builder::accept`]
6//!   - By performing the handshake yourself and then using [`Builder::serve`]
7//!     to let it take over a WebSocket stream
8use std::{future::poll_fn, io, pin::Pin};
9
10use futures_core::Stream;
11use http::{HeaderMap, HeaderName, HeaderValue, header};
12use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
13use tokio_util::codec::FramedRead;
14
15use crate::{
16    Error, WebSocketStream,
17    proto::{Config, Limits, Role},
18    upgrade::client_request,
19};
20
21/// HTTP/1.1 400 Bad Request response payload.
22const BAD_REQUEST: &[u8] = b"HTTP/1.1 400 Bad Request\r\n\r\n";
23
24/// List of headers added by the server which will cause an error
25/// if added by the user:
26///
27/// - `host`
28/// - `upgrade`
29/// - `connection`
30/// - `sec-websocket-accept`
31pub const DISALLOWED_HEADERS: &[HeaderName] = &[
32    header::UPGRADE,
33    header::CONNECTION,
34    header::SEC_WEBSOCKET_ACCEPT,
35];
36
37/// Builder for WebSocket server connections.
38pub struct Builder {
39    /// Configuration for the WebSocket stream.
40    config: Config,
41    /// Limits to impose on the WebSocket stream.
42    limits: Limits,
43    /// Headers to be sent with the switching protocols response.
44    headers: HeaderMap,
45}
46
47impl Default for Builder {
48    fn default() -> Self {
49        Self::new()
50    }
51}
52
53impl Builder {
54    /// Creates a [`Builder`] that can be used to create a [`WebSocketStream`]
55    /// to receive messages at the server end.
56    #[must_use]
57    pub fn new() -> Self {
58        Self {
59            config: Config::default(),
60            limits: Limits::default(),
61            headers: HeaderMap::new(),
62        }
63    }
64
65    /// Sets the configuration for the WebSocket stream.
66    #[must_use]
67    pub fn config(mut self, config: Config) -> Self {
68        self.config = config;
69
70        self
71    }
72
73    /// Sets the limits for the WebSocket stream.
74    #[must_use]
75    pub fn limits(mut self, limits: Limits) -> Self {
76        self.limits = limits;
77
78        self
79    }
80
81    /// Adds an extra HTTP header to the switching protocols response.
82    ///
83    /// # Errors
84    ///
85    /// Returns [`Error::DisallowedHeader`] if the header is in
86    /// the [`DISALLOWED_HEADERS`] list.
87    pub fn add_header(mut self, name: HeaderName, value: HeaderValue) -> Result<Self, Error> {
88        if DISALLOWED_HEADERS.contains(&name) {
89            return Err(Error::DisallowedHeader);
90        }
91        self.headers.insert(name, value);
92
93        Ok(self)
94    }
95
96    /// Perform a HTTP upgrade handshake on an already established stream and
97    /// uses it to send and receive WebSocket messages.
98    ///
99    /// # Errors
100    ///
101    /// This method returns an [`Error`] if the handshake fails.
102    pub async fn accept<S: AsyncRead + AsyncWrite + Unpin>(
103        &self,
104        stream: S,
105    ) -> Result<(http::Request<()>, WebSocketStream<S>), Error> {
106        let mut framed = FramedRead::new(
107            stream,
108            client_request::Codec {
109                response_headers: &self.headers,
110            },
111        );
112        let reply = poll_fn(|cx| Pin::new(&mut framed).poll_next(cx)).await;
113
114        match reply {
115            Some(Ok((request, response))) => {
116                framed.get_mut().write_all(&response).await?;
117                Ok((
118                    request,
119                    WebSocketStream::from_framed(framed, Role::Server, self.config, self.limits),
120                ))
121            }
122            Some(Err(e)) => {
123                framed.get_mut().write_all(BAD_REQUEST).await?;
124
125                Err(e)
126            }
127            None => Err(Error::Io(io::ErrorKind::UnexpectedEof.into())),
128        }
129    }
130
131    /// Takes over an already established stream and uses it to send and receive
132    /// WebSocket messages.
133    ///
134    /// This does not perform a HTTP upgrade handshake.
135    pub fn serve<S: AsyncRead + AsyncWrite + Unpin>(&self, stream: S) -> WebSocketStream<S> {
136        WebSocketStream::from_raw_stream(stream, Role::Server, self.config, self.limits)
137    }
138}
tokio_websockets/server.rs

tokio_websockets/
server.rs
