From 591f870631c7918d75881b4441d5094983622d79 Mon Sep 17 00:00:00 2001 From: Jesse White Date: Sat, 9 Dec 2023 20:26:45 -0500 Subject: [PATCH] add the ability to set the protocol(s) when creating a websocket client --- .gitignore | 5 ++++- worker/src/websocket.rs | 20 +++++++++++++++++++- 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/.gitignore b/.gitignore index 4f91bedb..dcb11df6 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 8b5a386b..d0a8e654 100644 --- a/worker/src/websocket.rs +++ b/worker/src/websocket.rs @@ -63,7 +63,18 @@ impl WebSocket { /// /// Response::error("never got a message echoed back :(", 500) /// ``` - pub async fn connect(mut url: Url) -> Result { + pub async fn connect(url: Url) -> Result { + return 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>) -> Result { let scheme: String = match url.scheme() { "ws" => "http".into(), "wss" => "https".into(), @@ -77,6 +88,13 @@ 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() {