From d6c095fe7bfa7992ef21adb420c1aacb0e0cf9fc Mon Sep 17 00:00:00 2001
From: Toshit Chawda <r58playz@gmail.com>
Date: Fri, 6 Sep 2024 22:24:56 -0700
Subject: [PATCH] make websocket errors more verbose

---
 client/src/lib.rs       | 37 +++++++++++++++++++------------------
 client/src/websocket.rs | 29 ++++++++++++++++++-----------
 2 files changed, 37 insertions(+), 29 deletions(-)

diff --git a/client/src/lib.rs b/client/src/lib.rs
index d705803..d8bc4b8 100644
--- a/client/src/lib.rs
+++ b/client/src/lib.rs
@@ -71,15 +71,15 @@ pub enum EpoxyError {
 	#[error("HTTP ToStr: {0:?} ({0})")]
 	ToStr(#[from] http::header::ToStrError),
 	#[cfg(feature = "full")]
-	#[error("Fastwebsockets: {0:?} ({0})")]
-	FastWebSockets(#[from] fastwebsockets::WebSocketError),
-	#[cfg(feature = "full")]
 	#[error("Pemfile: {0:?} ({0})")]
 	Pemfile(std::io::Error),
 	#[cfg(feature = "full")]
 	#[error("Webpki: {0:?} ({0})")]
 	Webpki(#[from] webpki::Error),
 
+	#[error("Wisp WebSocket failed to connect")]
+	WebSocketConnectFailed,
+
 	#[error("Custom wisp transport: {0}")]
 	WispTransport(String),
 	#[error("Invalid Wisp transport")]
@@ -89,6 +89,22 @@ pub enum EpoxyError {
 	#[error("Wisp transport already closed")]
 	WispTransportClosed,
 
+	#[cfg(feature = "full")]
+	#[error("Fastwebsockets: {0:?} ({0})")]
+	FastWebSockets(#[from] fastwebsockets::WebSocketError),
+	#[cfg(feature = "full")]
+	#[error("Invalid websocket response status code: {0} != {1}")]
+	WsInvalidStatusCode(u16, u16),
+	#[cfg(feature = "full")]
+	#[error("Invalid websocket upgrade header: {0:?} != \"websocket\"")]
+	WsInvalidUpgradeHeader(String),
+	#[cfg(feature = "full")]
+	#[error("Invalid websocket connection header: {0:?} != \"Upgrade\"")]
+	WsInvalidConnectionHeader(String),
+	#[cfg(feature = "full")]
+	#[error("Invalid websocket payload, only String/ArrayBuffer accepted")]
+	WsInvalidPayload,
+
 	#[error("Invalid URL scheme")]
 	InvalidUrlScheme,
 	#[error("No URL host found")]
@@ -99,22 +115,8 @@ pub enum EpoxyError {
 	InvalidRequestBody,
 	#[error("Invalid request")]
 	InvalidRequest,
-	#[error("Invalid websocket response status code")]
-	WsInvalidStatusCode,
-	#[error("Invalid websocket upgrade header")]
-	WsInvalidUpgradeHeader,
-	#[error("Invalid websocket connection header")]
-	WsInvalidConnectionHeader,
-	#[error("Invalid websocket payload")]
-	WsInvalidPayload,
 	#[error("Invalid payload")]
 	InvalidPayload,
-
-	#[error("Invalid certificate store")]
-	InvalidCertStore,
-	#[error("WebSocket failed to connect")]
-	WebSocketConnectFailed,
-
 	#[error("Failed to construct response headers object")]
 	ResponseHeadersFromEntriesFailed,
 	#[error("Failed to construct response object")]
@@ -185,7 +187,6 @@ enum EpoxyCompression {
 	Gzip,
 }
 
-
 // ugly hack. switch to serde-wasm-bindgen or a knockoff
 cfg_if! {
 	if #[cfg(feature = "full")] {
diff --git a/client/src/websocket.rs b/client/src/websocket.rs
index 4c399e4..a35fedc 100644
--- a/client/src/websocket.rs
+++ b/client/src/websocket.rs
@@ -205,27 +205,34 @@ impl EpoxyWebSocket {
 // https://github.com/snapview/tungstenite-rs/blob/314feea3055a93e585882fb769854a912a7e6dae/src/handshake/client.rs#L189
 fn verify(response: &Response<Incoming>) -> Result<(), EpoxyError> {
 	if response.status() != StatusCode::SWITCHING_PROTOCOLS {
-		return Err(EpoxyError::WsInvalidStatusCode);
+		return Err(EpoxyError::WsInvalidStatusCode(
+			response.status().as_u16(),
+			StatusCode::SWITCHING_PROTOCOLS.as_u16(),
+		));
 	}
 
 	let headers = response.headers();
 
-	if !headers
+	let upgrade_header = headers
 		.get(UPGRADE)
 		.and_then(|h| h.to_str().ok())
-		.map(|h| h.eq_ignore_ascii_case("websocket"))
-		.unwrap_or(false)
-	{
-		return Err(EpoxyError::WsInvalidUpgradeHeader);
+		.unwrap_or_default();
+
+	if !upgrade_header.eq_ignore_ascii_case("websocket") {
+		return Err(EpoxyError::WsInvalidUpgradeHeader(
+			upgrade_header.to_string(),
+		));
 	}
 
-	if !headers
+	let connection_header = headers
 		.get(CONNECTION)
 		.and_then(|h| h.to_str().ok())
-		.map(|h| h.eq_ignore_ascii_case("Upgrade"))
-		.unwrap_or(false)
-	{
-		return Err(EpoxyError::WsInvalidConnectionHeader);
+		.unwrap_or_default();
+
+	if !connection_header.eq_ignore_ascii_case("Upgrade") {
+		return Err(EpoxyError::WsInvalidConnectionHeader(
+			connection_header.to_string(),
+		));
 	}
 
 	Ok(())