fix panic when upgrading connections incorrectlyv1.1.0

2 files changed, 8 insertions, 5 deletions

diff --git a/src/error.rs b/src/error.rs
index b78999f..1675b3c 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -31,6 +31,8 @@ pub enum ServiceError {
     BadAuth,
     #[error("bad base64: {0}")]
     BadBase64(#[from] base64::DecodeError),
+    #[error("connection upgrade failed")]
+    UpgradeFailed,
     #[error("impossible error")]
     Other,
 }
diff --git a/src/filters/proxy.rs b/src/filters/proxy.rs
index 40ebf17..0e1a556 100644
--- a/src/filters/proxy.rs
+++ b/src/filters/proxy.rs
@@ -75,13 +75,14 @@ pub async fn proxy_request(
     };
 
     if do_upgrade {
-        let on_upgrade_upstream = resp.extensions_mut().remove::<OnUpgrade>();
+        let on_upgrade_upstream = resp
+            .extensions_mut()
+            .remove::<OnUpgrade>()
+            .ok_or(ServiceError::UpgradeFailed)?;
+        let on_upgrade_downstream = on_upgrade_downstream.ok_or(ServiceError::UpgradeFailed)?;
         tokio::task::spawn(async move {
             debug!("about to upgrade connection, sending empty response");
-            match (
-                on_upgrade_upstream.unwrap().await,
-                on_upgrade_downstream.unwrap().await,
-            ) {
+            match (on_upgrade_upstream.await, on_upgrade_downstream.await) {
                 (Ok(upgraded_upstream), Ok(upgraded_downstream)) => {
                     debug!("upgrade successful");
                     match tokio::io::copy_bidirectional(