feat: pass CONNECT requests through handler

omjadas · Aug 9, 2022 · 35aca05 · 35aca05
1 parent e6432db
commit 35aca05
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 30 deletions.
diff --git a/src/proxy/internal.rs b/src/proxy/internal.rs
@@ -70,16 +70,10 @@ where
             client_addr = %self.client_addr,
         )
     )]
-    pub(crate) async fn proxy(self, req: Request<Body>) -> Result<Response<Body>, hyper::Error> {
-        if req.method() == Method::CONNECT {
-            self.process_connect(req).await
-        } else {
-            self.process_request(req).await
-        }
-    }
-
-    #[instrument(skip_all)]
-    async fn process_request(mut self, req: Request<Body>) -> Result<Response<Body>, hyper::Error> {
+    pub(crate) async fn proxy(
+        mut self,
+        req: Request<Body>,
+    ) -> Result<Response<Body>, hyper::Error> {
         let ctx = HttpContext {
             client_addr: self.client_addr,
         };
@@ -94,24 +88,26 @@ where
             RequestOrResponse::Response(res) => return Ok(res),
         };
 
-        if hyper_tungstenite::is_upgrade_request(&req) {
-            return Ok(self.upgrade_websocket(req));
-        }
-
-        let res = self
-            .client
-            .request(normalize_request(req))
-            .instrument(info_span!("proxy_request"))
-            .await?;
+        if req.method() == Method::CONNECT {
+            self.process_connect(req)
+        } else if hyper_tungstenite::is_upgrade_request(&req) {
+            Ok(self.upgrade_websocket(req))
+        } else {
+            let res = self
+                .client
+                .request(normalize_request(req))
+                .instrument(info_span!("proxy_request"))
+                .await?;
 
-        Ok(self
-            .http_handler
-            .handle_response(&ctx, res)
-            .instrument(info_span!("handle_response"))
-            .await)
+            Ok(self
+                .http_handler
+                .handle_response(&ctx, res)
+                .instrument(info_span!("handle_response"))
+                .await)
+        }
     }
 
-    async fn process_connect(self, mut req: Request<Body>) -> Result<Response<Body>, hyper::Error> {
+    fn process_connect(self, mut req: Request<Body>) -> Result<Response<Body>, hyper::Error> {
         let span = info_span!("process_connect");
         let fut = async move {
             match hyper::upgrade::on(&mut req).await {
@@ -288,7 +284,7 @@ where
                 req = Request::from_parts(parts, body);
             };
 
-            self.clone().process_request(req)
+            self.clone().proxy(req)
         });
 
         Http::new()

diff --git a/tests/openssl_ca.rs b/tests/openssl_ca.rs
@@ -35,7 +35,7 @@ async fn https_rustls() {
         .unwrap();
 
     assert_eq!(res.status(), 200);
-    assert_eq!(http_handler.request_counter.load(Ordering::Relaxed), 1);
+    assert_eq!(http_handler.request_counter.load(Ordering::Relaxed), 2);
     assert_eq!(http_handler.response_counter.load(Ordering::Relaxed), 1);
 
     stop_server.send(()).unwrap();
@@ -61,7 +61,7 @@ async fn https_native_tls() {
         .unwrap();
 
     assert_eq!(res.status(), 200);
-    assert_eq!(http_handler.request_counter.load(Ordering::Relaxed), 1);
+    assert_eq!(http_handler.request_counter.load(Ordering::Relaxed), 2);
     assert_eq!(http_handler.response_counter.load(Ordering::Relaxed), 1);
 
     stop_server.send(()).unwrap();

diff --git a/tests/rcgen_ca.rs b/tests/rcgen_ca.rs
@@ -41,7 +41,7 @@ async fn https_rustls() {
         .unwrap();
 
     assert_eq!(res.status(), 200);
-    assert_eq!(http_handler.request_counter.load(Ordering::Relaxed), 1);
+    assert_eq!(http_handler.request_counter.load(Ordering::Relaxed), 2);
     assert_eq!(http_handler.response_counter.load(Ordering::Relaxed), 1);
 
     stop_server.send(()).unwrap();
@@ -67,7 +67,7 @@ async fn https_native_tls() {
         .unwrap();
 
     assert_eq!(res.status(), 200);
-    assert_eq!(http_handler.request_counter.load(Ordering::Relaxed), 1);
+    assert_eq!(http_handler.request_counter.load(Ordering::Relaxed), 2);
     assert_eq!(http_handler.response_counter.load(Ordering::Relaxed), 1);
 
     stop_server.send(()).unwrap();