feat(ext/websocket): add header support to WebSocketStream (#11887)

1 files changed, 29 insertions, 0 deletions

diff --git a/ext/websocket/lib.rs b/ext/websocket/lib.rs
index 4796eddc6..544423066 100644
--- a/ext/websocket/lib.rs
+++ b/ext/websocket/lib.rs
@@ -1,6 +1,7 @@
 // Copyright 2018-2021 the Deno authors. All rights reserved. MIT license.
 
 use deno_core::error::invalid_hostname;
+use deno_core::error::type_error;
 use deno_core::error::AnyError;
 use deno_core::futures::stream::SplitSink;
 use deno_core::futures::stream::SplitStream;
@@ -11,6 +12,7 @@ use deno_core::op_async;
 use deno_core::op_sync;
 use deno_core::url;
 use deno_core::AsyncRefCell;
+use deno_core::ByteString;
 use deno_core::CancelFuture;
 use deno_core::CancelHandle;
 use deno_core::Extension;
@@ -20,6 +22,8 @@ use deno_core::Resource;
 use deno_core::ResourceId;
 use deno_core::ZeroCopyBuf;
 use deno_tls::create_client_config;
+use http::header::HeaderName;
+use http::HeaderValue;
 use http::Method;
 use http::Request;
 use http::Uri;
@@ -215,6 +219,7 @@ pub struct CreateArgs {
   url: String,
   protocols: String,
   cancel_handle: Option<ResourceId>,
+  headers: Option<Vec<(ByteString, ByteString)>>,
 }
 
 #[derive(Serialize)]
@@ -267,6 +272,30 @@ where
     request = request.header("Sec-WebSocket-Protocol", args.protocols);
   }
 
+  if let Some(headers) = args.headers {
+    for (key, value) in headers {
+      let name = HeaderName::from_bytes(&key)
+        .map_err(|err| type_error(err.to_string()))?;
+      let v = HeaderValue::from_bytes(&value)
+        .map_err(|err| type_error(err.to_string()))?;
+
+      let is_disallowed_header = matches!(
+        name,
+        http::header::HOST
+          | http::header::SEC_WEBSOCKET_ACCEPT
+          | http::header::SEC_WEBSOCKET_EXTENSIONS
+          | http::header::SEC_WEBSOCKET_KEY
+          | http::header::SEC_WEBSOCKET_PROTOCOL
+          | http::header::SEC_WEBSOCKET_VERSION
+          | http::header::UPGRADE
+          | http::header::CONNECTION
+      );
+      if !is_disallowed_header {
+        request = request.header(name, v);
+      }
+    }
+  }
+
   let request = request.body(())?;
   let domain = &uri.host().unwrap().to_string();
   let port = &uri.port_u16().unwrap_or(match uri.scheme_str() {