feat: pass context to should_intercept

omjadas · omjadas · commit 6987ee5e7534 · 2023-02-14T19:01:43.000+11:00
diff --git a/src/lib.rs b/src/lib.rs
@@ -115,7 +115,7 @@ pub trait HttpHandler: Clone + Send + Sync + 'static {
     }
 
     /// Whether a CONNECT request should be intercepted. Defaults to `true` for all requests.
-    async fn should_intercept(&mut self, _req: &Request<Body>) -> bool {
+    async fn should_intercept(&mut self, _ctx: &HttpContext, _req: &Request<Body>) -> bool {
         true
     }
 }
diff --git a/src/proxy/internal.rs b/src/proxy/internal.rs
@@ -69,6 +69,12 @@ where
     H: HttpHandler,
     W: WebSocketHandler,
 {
+    fn context(&self) -> HttpContext {
+        HttpContext {
+            client_addr: self.client_addr,
+        }
+    }
+
     #[instrument(
         skip_all,
         fields(
@@ -82,9 +88,7 @@ where
         mut self,
         req: Request<Body>,
     ) -> Result<Response<Body>, hyper::Error> {
-        let ctx = HttpContext {
-            client_addr: self.client_addr,
-        };
+        let ctx = self.context();
 
         let req = match self
             .http_handler
@@ -136,7 +140,11 @@ where
                                 bytes::Bytes::copy_from_slice(buffer[..bytes_read].as_ref()),
                             );
 
-                            if self.http_handler.should_intercept(&req).await {
+                            if self
+                                .http_handler
+                                .should_intercept(&self.context(), &req)
+                                .await
+                            {
                                 if buffer == *b"GET " {
                                     if let Err(e) =
                                         self.serve_stream(upgraded, Scheme::HTTP, authority).await
diff --git a/tests/common/mod.rs b/tests/common/mod.rs
@@ -305,7 +305,7 @@ impl HttpHandler for TestHandler {
         decode_response(res).unwrap()
     }
 
-    async fn should_intercept(&mut self, _req: &Request<Body>) -> bool {
+    async fn should_intercept(&mut self, _ctx: &HttpContext, _req: &Request<Body>) -> bool {
         self.should_intercept
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -115,7 +115,7 @@ pub trait HttpHandler: Clone + Send + Sync + 'static {`
`115`	`115`	`}`
`116`	`116`
`117`	`117`	/// Whether a CONNECT request should be intercepted. Defaults to `true` for all requests.
`118`		`- async fn should_intercept(&mut self, _req: &Request<Body>) -> bool {`
	`118`	`+ async fn should_intercept(&mut self, _ctx: &HttpContext, _req: &Request<Body>) -> bool {`
`119`	`119`	`true`
`120`	`120`	`}`
`121`	`121`	`}`
Original file line number	Diff line number	Diff line change
`@@ -305,7 +305,7 @@ impl HttpHandler for TestHandler {`
`305`	`305`	`decode_response(res).unwrap()`
`306`	`306`	`}`
`307`	`307`
`308`		`- async fn should_intercept(&mut self, _req: &Request<Body>) -> bool {`
	`308`	`+ async fn should_intercept(&mut self, _ctx: &HttpContext, _req: &Request<Body>) -> bool {`
`309`	`309`	`self.should_intercept`
`310`	`310`	`}`
`311`	`311`	`}`