Add builder for additional header values

src/client.rs

@@ -1,12 +1,13 @@
 //! Methods to connect to a WebSocket as a client.
 
 use std::{
+    convert::TryFrom,
     io::{Read, Write},
     net::{SocketAddr, TcpStream, ToSocketAddrs},
     result::Result as StdResult,
 };
 
-use http::{request::Parts, Uri};
+use http::{request::Parts, HeaderName, Uri};
 use log::*;
 
 use url::Url;
@@ -265,3 +266,71 @@ impl<'h, 'b> IntoClientRequest for httparse::Request<'h, 'b> {
         Request::from_httparse(self)
     }
 }
+
+/// Builder for a custom [`IntoClientRequest`] with options to add
+/// custom additional headers and sub protocols.
+///
+/// # Example
+///
+/// ```rust no_run
+/// # use crate::*;
+///
+/// let uri: Uri = "ws://localhost:3012/socket".parse().unwrap();
+/// let token = "my_jwt_token";
+/// let builder = ClientRequestBuilder::new(uri)
+///     .with_header("Authorization", format!("Bearer {token}"))
+///     .with_sub_protocol("my_sub_protocol"");
+/// let socket = connect(builder).unwrap();
+/// ```
+#[derive(Debug, Clone)]
+pub struct ClientRequestBuilder {
+    uri: Uri,
+    /// Additional [`Request`] handshake headers
+    additional_headers: Vec<(String, String)>,
+    /// Handsake subprotocols
+    subprotocols: Vec<String>,
+}
+
+impl ClientRequestBuilder {
+    /// Initializes an empty request builder
+    #[must_use]
+    pub const fn new(uri: Uri) -> Self {
+        Self { uri, additional_headers: Vec::new(), subprotocols: Vec::new() }
+    }
+
+    /// Adds (`key`, `value`) as an additional header to the handshake request
+    pub fn with_header<K, V>(mut self, key: K, value: V) -> Self
+    where
+        K: Into<String>,
+        V: Into<String>,
+    {
+        self.additional_headers.push((key.into(), value.into()));
+        self
+    }
+
+    /// Adds `protocol` to the handshake request subprotocols (`Sec-WebSocket-Protocol`)
+    pub fn with_sub_protocol<P>(mut self, protocol: P) -> Self
+    where
+        P: Into<String>,
+    {
+        self.subprotocols.push(protocol.into());
+        self
+    }
+}
+
+impl IntoClientRequest for ClientRequestBuilder {
+    fn into_client_request(self) -> Result<Request> {
+        let mut request = self.uri.into_client_request()?;
+        let headers = request.headers_mut();
+        for (k, v) in self.additional_headers {
+            let key = HeaderName::try_from(k)?;
+            let value = v.parse()?;
+            headers.append(key, value);
+        }
+        if !self.subprotocols.is_empty() {
+            let protocols = self.subprotocols.join(", ").parse()?;
+            headers.append("Sec-WebSocket-Protocol", protocols);
+        }
+        Ok(request)
+    }
+}

src/lib.rs

@@ -39,7 +39,7 @@ pub use crate::{
 
 #[cfg(feature = "handshake")]
 pub use crate::{
-    client::{client, connect},
+    client::{client, connect, ClientRequestBuilder},
     handshake::{client::ClientHandshake, server::ServerHandshake, HandshakeError},
     server::{accept, accept_hdr, accept_hdr_with_config, accept_with_config},
 };