From d85d7e6c390e5e6c2febca02c3c50ab7b9a8a937 Mon Sep 17 00:00:00 2001
From: Jun Kurihara <junkurihara@users.noreply.github.com>
Date: Fri, 15 Dec 2023 15:36:00 +0900
Subject: [PATCH] use hyper::upgrade::on

---
 rpxy-lib/src/message_handler/handler_main.rs | 21 +++++++++++---------
 rpxy-lib/src/message_handler/http_result.rs  | 13 ++++++------
 2 files changed, 18 insertions(+), 16 deletions(-)

diff --git a/rpxy-lib/src/message_handler/handler_main.rs b/rpxy-lib/src/message_handler/handler_main.rs
index ceb5db4..5fb978d 100644
--- a/rpxy-lib/src/message_handler/handler_main.rs
+++ b/rpxy-lib/src/message_handler/handler_main.rs
@@ -143,7 +143,8 @@ where
 
     // Upgrade in request header
     let upgrade_in_request = extract_upgrade(req.headers());
-    let request_upgraded = req.extensions_mut().remove::<hyper::upgrade::OnUpgrade>();
+    // let request_upgraded = req.extensions_mut().remove::<hyper::upgrade::OnUpgrade>();
+    let req_on_upgrade = hyper::upgrade::on(&mut req);
 
     // Build request from destination information
     let _context = match self.generate_request_forwarded(
@@ -209,19 +210,21 @@ where
         upgrade_in_response, upgrade_in_request
       )));
     }
-    let Some(request_upgraded) = request_upgraded else {
-      return Err(HttpError::NoUpgradeExtensionInRequest);
-    };
-    let Some(onupgrade) = res_backend.extensions_mut().remove::<hyper::upgrade::OnUpgrade>() else {
-      return Err(HttpError::NoUpgradeExtensionInResponse);
-    };
+    // let Some(request_upgraded) = request_upgraded else {
+    //   return Err(HttpError::NoUpgradeExtensionInRequest);
+    // };
+
+    // let Some(onupgrade) = res_backend.extensions_mut().remove::<hyper::upgrade::OnUpgrade>() else {
+    //   return Err(HttpError::NoUpgradeExtensionInResponse);
+    // };
+    let res_on_upgrade = hyper::upgrade::on(&mut res_backend);
 
     self.globals.runtime_handle.spawn(async move {
-      let mut response_upgraded = TokioIo::new(onupgrade.await.map_err(|e| {
+      let mut response_upgraded = TokioIo::new(res_on_upgrade.await.map_err(|e| {
         error!("Failed to upgrade response: {}", e);
         RpxyError::FailedToUpgradeResponse(e.to_string())
       })?);
-      let mut request_upgraded = TokioIo::new(request_upgraded.await.map_err(|e| {
+      let mut request_upgraded = TokioIo::new(req_on_upgrade.await.map_err(|e| {
         error!("Failed to upgrade request: {}", e);
         RpxyError::FailedToUpgradeRequest(e.to_string())
       })?);
diff --git a/rpxy-lib/src/message_handler/http_result.rs b/rpxy-lib/src/message_handler/http_result.rs
index ec48200..691c087 100644
--- a/rpxy-lib/src/message_handler/http_result.rs
+++ b/rpxy-lib/src/message_handler/http_result.rs
@@ -32,11 +32,10 @@ pub enum HttpError {
 
   #[error("Failed to upgrade connection: {0}")]
   FailedToUpgrade(String),
-  #[error("Request does not have an upgrade extension")]
-  NoUpgradeExtensionInRequest,
-  #[error("Response does not have an upgrade extension")]
-  NoUpgradeExtensionInResponse,
-
+  // #[error("Request does not have an upgrade extension")]
+  // NoUpgradeExtensionInRequest,
+  // #[error("Response does not have an upgrade extension")]
+  // NoUpgradeExtensionInResponse,
   #[error(transparent)]
   Other(#[from] anyhow::Error),
 }
@@ -54,8 +53,8 @@ impl From<HttpError> for StatusCode {
       HttpError::FailedToAddSetCookeInResponse(_) => StatusCode::INTERNAL_SERVER_ERROR,
       HttpError::FailedToGenerateDownstreamResponse(_) => StatusCode::INTERNAL_SERVER_ERROR,
       HttpError::FailedToUpgrade(_) => StatusCode::INTERNAL_SERVER_ERROR,
-      HttpError::NoUpgradeExtensionInRequest => StatusCode::BAD_REQUEST,
-      HttpError::NoUpgradeExtensionInResponse => StatusCode::BAD_GATEWAY,
+      // HttpError::NoUpgradeExtensionInRequest => StatusCode::BAD_REQUEST,
+      // HttpError::NoUpgradeExtensionInResponse => StatusCode::BAD_GATEWAY,
       _ => StatusCode::INTERNAL_SERVER_ERROR,
     }
   }