From bd38ea52537cd2851813525636a897d826e71ae0 Mon Sep 17 00:00:00 2001
From: nitn3lav <77448526+nitn3lav@users.noreply.github.com>
Date: Sat, 30 Mar 2024 15:44:29 +0100
Subject: [PATCH] Fix Axum panic caused by missing upgrade header & HTTP 400 on
 error (#75)

---
 src/upgrade.rs | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/src/upgrade.rs b/src/upgrade.rs
index 5f1d1fc..4a1e6cd 100644
--- a/src/upgrade.rs
+++ b/src/upgrade.rs
@@ -76,26 +76,29 @@ impl<S> axum_core::extract::FromRequestParts<S> for IncomingUpgrade
 where
   S: Sync,
 {
-  type Rejection = ();
+  type Rejection = hyper::StatusCode;
 
   async fn from_request_parts(
     parts: &mut http::request::Parts,
     _state: &S,
   ) -> Result<Self, Self::Rejection> {
-    let key = parts.headers.get("Sec-WebSocket-Key").ok_or(())?;
+    let key = parts
+      .headers
+      .get("Sec-WebSocket-Key")
+      .ok_or(hyper::StatusCode::BAD_REQUEST)?;
     if parts
       .headers
       .get("Sec-WebSocket-Version")
       .map(|v| v.as_bytes())
       != Some(b"13")
     {
-      return Err(());
+      return Err(hyper::StatusCode::BAD_REQUEST);
     }
 
     let on_upgrade = parts
       .extensions
       .remove::<hyper::upgrade::OnUpgrade>()
-      .unwrap();
+      .ok_or(hyper::StatusCode::BAD_REQUEST)?;
     Ok(Self {
       on_upgrade,
       key: sec_websocket_protocol(key.as_bytes()),