diff options
| -rw-r--r-- | src/error.rs | 2 | ||||
| -rw-r--r-- | src/filters/proxy.rs | 11 |
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( |