Allow adding headers to switching protocols response

Signed-off-by: Jens Reidel <[email protected]>

Author: Gelbpunkt

src/server.rs

@@ -8,6 +8,7 @@
 use std::{future::poll_fn, io, pin::Pin};
 
 use futures_core::Stream;
+use http::{HeaderMap, HeaderName, HeaderValue, header};
 use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
 use tokio_util::codec::FramedRead;
 
@@ -20,12 +21,27 @@ use crate::{
 /// HTTP/1.1 400 Bad Request response payload.
 const BAD_REQUEST: &[u8] = b"HTTP/1.1 400 Bad Request\r\n\r\n";
 
+/// List of headers added by the server which will cause an error
+/// if added by the user:
+///
+/// - `host`
+/// - `upgrade`
+/// - `connection`
+/// - `sec-websocket-accept`
+pub const DISALLOWED_HEADERS: &[HeaderName] = &[
+    header::UPGRADE,
+    header::CONNECTION,
+    header::SEC_WEBSOCKET_ACCEPT,
+];
+
 /// Builder for WebSocket server connections.
 pub struct Builder {
     /// Configuration for the WebSocket stream.
     config: Config,
     /// Limits to impose on the WebSocket stream.
     limits: Limits,
+    /// Headers to be sent with the switching protocols response.
+    headers: HeaderMap,
 }
 
 impl Default for Builder {
@@ -42,6 +58,7 @@ impl Builder {
         Self {
             config: Config::default(),
             limits: Limits::default(),
+            headers: HeaderMap::new(),
         }
     }
 
@@ -61,6 +78,21 @@ impl Builder {
         self
     }
 
+    /// Adds an extra HTTP header to the switching protocols response.
+    ///
+    /// # Errors
+    ///
+    /// Returns [`Error::DisallowedHeader`] if the header is in
+    /// the [`DISALLOWED_HEADERS`] list.
+    pub fn add_header(mut self, name: HeaderName, value: HeaderValue) -> Result<Self, Error> {
+        if DISALLOWED_HEADERS.contains(&name) {
+            return Err(Error::DisallowedHeader);
+        }
+        self.headers.insert(name, value);
+
+        Ok(self)
+    }
+
     /// Perform a HTTP upgrade handshake on an already established stream and
     /// uses it to send and receive WebSocket messages.
     ///
@@ -71,12 +103,17 @@ impl Builder {
         &self,
         stream: S,
     ) -> Result<(http::Request<()>, WebSocketStream<S>), Error> {
-        let mut framed = FramedRead::new(stream, client_request::Codec {});
+        let mut framed = FramedRead::new(
+            stream,
+            client_request::Codec {
+                response_headers: &self.headers,
+            },
+        );
         let reply = poll_fn(|cx| Pin::new(&mut framed).poll_next(cx)).await;
 
         match reply {
             Some(Ok((request, response))) => {
-                framed.get_mut().write_all(response.as_bytes()).await?;
+                framed.get_mut().write_all(&response).await?;
                 Ok((
                     request,
                     WebSocketStream::from_framed(framed, Role::Server, self.config, self.limits),

src/upgrade/client_request.rs

@@ -3,14 +3,15 @@ use std::str::FromStr;
 
 use base64::{Engine, engine::general_purpose::STANDARD};
 use bytes::{Buf, BytesMut};
+use http::HeaderMap;
 use httparse::Request;
 use tokio_util::codec::Decoder;
 
 use crate::{sha::digest, upgrade::Error};
 
 /// A static HTTP/1.1 101 Switching Protocols response up until the
 /// `Sec-WebSocket-Accept` header value.
-const SWITCHING_PROTOCOLS_BODY: &str = "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: ";
+const SWITCHING_PROTOCOLS_BODY: &[u8] = b"HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: ";
 
 /// Returns whether an ASCII byte slice is contained in another one, ignoring
 /// captalization.
@@ -94,11 +95,14 @@ impl ClientRequest {
 /// It does not implement an [`Encoder`].
 ///
 /// [`Encoder`]: tokio_util::codec::Encoder
-pub struct Codec {}
+pub struct Codec<'a> {
+    /// List of headers to add to the Switching Protocols response.
+    pub response_headers: &'a HeaderMap,
+}
 
-impl Decoder for Codec {
+impl Decoder for Codec<'_> {
     type Error = crate::Error;
-    type Item = (http::Request<()>, String);
+    type Item = (http::Request<()>, Vec<u8>);
 
     fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
         let mut headers = [httparse::EMPTY_HEADER; 64];
@@ -151,11 +155,21 @@ impl Decoder for Codec {
 
         src.advance(request_len);
 
-        let mut resp = String::with_capacity(SWITCHING_PROTOCOLS_BODY.len() + ws_accept.len() + 4);
+        // Preallocate the size without extra headers
+        let mut resp = Vec::with_capacity(SWITCHING_PROTOCOLS_BODY.len() + ws_accept.len() + 4);
+
+        resp.extend_from_slice(SWITCHING_PROTOCOLS_BODY);
+        resp.extend_from_slice(ws_accept.as_bytes());
+        resp.extend_from_slice(b"\r\n");
+
+        for (name, value) in self.response_headers {
+            resp.extend_from_slice(name.as_str().as_bytes());
+            resp.extend_from_slice(b": ");
+            resp.extend_from_slice(value.as_bytes());
+            resp.extend_from_slice(b"\r\n");
+        }
 
-        resp.push_str(SWITCHING_PROTOCOLS_BODY);
-        resp.push_str(&ws_accept);
-        resp.push_str("\r\n\r\n");
+        resp.extend_from_slice(b"\r\n");
 
         Ok(Some((request, resp)))
     }