From 0aec6e563c93ff71e1e53b000e3e8635d7a4e46d Mon Sep 17 00:00:00 2001
From: Toshit Chawda <r58playz@gmail.com>
Date: Tue, 2 Jul 2024 20:54:21 -0700
Subject: [PATCH] ws close code and reason, ws headers

---
 client/Cargo.toml       |  2 +-
 client/build.sh         |  2 +-
 client/demo.js          |  3 ++-
 client/package.json     |  2 +-
 client/src/lib.rs       |  3 ++-
 client/src/websocket.rs | 21 ++++++++++++++++-----
 6 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/client/Cargo.toml b/client/Cargo.toml
index 7d1408d..6d122c9 100644
--- a/client/Cargo.toml
+++ b/client/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "epoxy-client"
-version = "2.0.0"
+version = "2.0.1"
 edition = "2021"
 
 [lib]
diff --git a/client/build.sh b/client/build.sh
index 210c84c..3bbdeb2 100755
--- a/client/build.sh
+++ b/client/build.sh
@@ -2,7 +2,7 @@
 set -euo pipefail
 shopt -s inherit_errexit
 
-mkdir out/
+mkdir out/ || true
 rm -r pkg/ || true
 mkdir pkg/
 
diff --git a/client/demo.js b/client/demo.js
index bde2e80..a485b5d 100644
--- a/client/demo.js
+++ b/client/demo.js
@@ -182,8 +182,9 @@ onmessage = async (msg) => {
 		);
 		let ws = await epoxy_client.connect_websocket(
 			handlers,
-			"wss://echo.websocket.events",
+			"ws://localhost:5000",
 			[],
+			{ "x-header": "abc" },
 		);
 		while (true) {
 			log("sending `data`");
diff --git a/client/package.json b/client/package.json
index 1020c31..3de9e3a 100644
--- a/client/package.json
+++ b/client/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@mercuryworkshop/epoxy-tls",
-  "version": "2.0.0-3",
+  "version": "2.0.1-1",
   "description": "A wasm library for using raw encrypted tls/ssl/https/websocket streams on the browser",
   "scripts": {
     "build": "./build.sh"
diff --git a/client/src/lib.rs b/client/src/lib.rs
index 750f3c7..1fe3f75 100644
--- a/client/src/lib.rs
+++ b/client/src/lib.rs
@@ -245,8 +245,9 @@ impl EpoxyClient {
         handlers: EpoxyHandlers,
         url: String,
         protocols: Vec<String>,
+		headers: JsValue,
     ) -> Result<EpoxyWebSocket, EpoxyError> {
-        EpoxyWebSocket::connect(self, handlers, url, protocols).await
+        EpoxyWebSocket::connect(self, handlers, url, protocols, headers).await
     }
 
     pub async fn connect_tcp(
diff --git a/client/src/websocket.rs b/client/src/websocket.rs
index a08a15e..67d1a73 100644
--- a/client/src/websocket.rs
+++ b/client/src/websocket.rs
@@ -3,7 +3,7 @@ use std::{str::from_utf8, sync::Arc};
 use base64::{prelude::BASE64_STANDARD, Engine};
 use bytes::Bytes;
 use fastwebsockets::{
-    CloseCode, FragmentCollectorRead, Frame, OpCode, Payload, Role, WebSocket, WebSocketWrite,
+    FragmentCollectorRead, Frame, OpCode, Payload, Role, WebSocket, WebSocketWrite,
 };
 use futures_util::lock::Mutex;
 use getrandom::getrandom;
@@ -17,12 +17,12 @@ use hyper::{
     body::Incoming,
     upgrade::{self, Upgraded},
 };
-use js_sys::{ArrayBuffer, Function, Uint8Array};
+use js_sys::{ArrayBuffer, Function, Object, Uint8Array};
 use tokio::io::WriteHalf;
 use wasm_bindgen::{prelude::*, JsError, JsValue};
 use wasm_bindgen_futures::spawn_local;
 
-use crate::{tokioio::TokioIo, EpoxyClient, EpoxyError, EpoxyHandlers, HttpBody};
+use crate::{tokioio::TokioIo, utils::entries_of_object, EpoxyClient, EpoxyError, EpoxyHandlers, HttpBody};
 
 #[wasm_bindgen]
 pub struct EpoxyWebSocket {
@@ -37,6 +37,7 @@ impl EpoxyWebSocket {
         handlers: EpoxyHandlers,
         url: String,
         protocols: Vec<String>,
+		headers: JsValue,
     ) -> Result<Self, EpoxyError> {
         let EpoxyHandlers {
             onopen,
@@ -66,6 +67,16 @@ impl EpoxyWebSocket {
                 request = request.header(SEC_WEBSOCKET_PROTOCOL, protocols.join(","));
             }
 
+			if web_sys::Headers::instanceof(&headers) && let Ok(entries) = Object::from_entries(&headers) {
+				for header in entries_of_object(&entries) {
+					request = request.header(&header[0], &header[1]);
+				}
+			} else if headers.is_truthy() {
+				for header in entries_of_object(&headers.into()) {
+					request = request.header(&header[0], &header[1]);
+				}
+			}
+
             let request = request.body(HttpBody::new(Bytes::new()))?;
 
             let mut response = client.client.request(request).await?;
@@ -170,12 +181,12 @@ impl EpoxyWebSocket {
         }
     }
 
-    pub async fn close(&self) -> Result<(), EpoxyError> {
+    pub async fn close(&self, code: u16, reason: String) -> Result<(), EpoxyError> {
         let ret = self
             .tx
             .lock()
             .await
-            .write_frame(Frame::close(CloseCode::Normal.into(), b""))
+            .write_frame(Frame::close(code, reason.as_bytes()))
             .await;
         match ret {
             Ok(ok) => Ok(ok),