From e6bb066fa7f97232330a5766d09823317942c80c Mon Sep 17 00:00:00 2001
From: Jesse White <jesse.white.oss@gmail.com>
Date: Wed, 28 Feb 2024 11:10:17 -0400
Subject: [PATCH] add the ability to set the protocol(s) when creating a
 websocket client (#424)

* add the ability to set the protocol(s) when creating a websocket client

* fmt and clippy

---------

Co-authored-by: Kevin Flansburg <kflansburg@cloudflare.com>
---
 .gitignore              |  5 ++++-
 worker/src/websocket.rs | 24 +++++++++++++++++++++++-
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/.gitignore b/.gitignore
index 4f91bedb9..dcb11df6e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -133,4 +133,7 @@ dist
 .yarn/unplugged
 .yarn/build-state.yml
 .yarn/install-state.gz
-.pnp.*
\ No newline at end of file
+.pnp.*
+
+# RustRover
+.idea
diff --git a/worker/src/websocket.rs b/worker/src/websocket.rs
index 65056e2cd..a79869ce6 100644
--- a/worker/src/websocket.rs
+++ b/worker/src/websocket.rs
@@ -63,7 +63,21 @@ impl WebSocket {
     ///
     /// Response::error("never got a message echoed back :(", 500)
     /// ```
-    pub async fn connect(mut url: Url) -> Result<WebSocket> {
+    pub async fn connect(url: Url) -> Result<WebSocket> {
+        WebSocket::connect_with_protocols(url, None).await
+    }
+
+    /// Attempts to establish a [`WebSocket`] connection to the provided [`Url`] and protocol.
+    ///
+    /// # Example:
+    /// ```rust,ignore
+    /// let ws = WebSocket::connect_with_protocols("wss://echo.zeb.workers.dev/".parse()?, Some(vec!["GiggleBytes"])).await?;
+    ///
+    /// ```
+    pub async fn connect_with_protocols(
+        mut url: Url,
+        protocols: Option<Vec<&str>>,
+    ) -> Result<WebSocket> {
         let scheme: String = match url.scheme() {
             "ws" => "http".into(),
             "wss" => "https".into(),
@@ -77,6 +91,14 @@ impl WebSocket {
         let mut req = Request::new(url.as_str(), Method::Get)?;
         req.headers_mut()?.set("upgrade", "websocket")?;
 
+        match protocols {
+            None => {}
+            Some(v) => {
+                req.headers_mut()?
+                    .set("Sec-WebSocket-Protocol", v.join(",").as_str())?;
+            }
+        }
+
         let res = Fetch::Request(req).send().await?;
 
         match res.websocket() {