feat: aggregated MitM outbound connection

XOR-op · Nov 3, 2023 · c024191 · c024191
1 parent 411f66e
commit c024191
Showing 1 changed file with 25 additions and 10 deletions.
diff --git a/boltconn/src/intercept/https_intercept.rs b/boltconn/src/intercept/https_intercept.rs
@@ -6,12 +6,14 @@ use crate::intercept::modifier::Modifier;
 use crate::intercept::{sign_site_cert, ModifierContext};
 use crate::proxy::{ConnAbortHandle, ConnContext};
 use hyper::client::conn;
+use hyper::client::conn::SendRequest;
 use hyper::server::conn::Http;
 use hyper::service::service_fn;
 use hyper::{Body, Request, Response};
 use rcgen::Certificate as CaCertificate;
 use std::io;
 use std::sync::Arc;
+use tokio::sync::Mutex;
 use tokio_rustls::rustls::{Certificate, PrivateKey, ServerConfig, ServerName};
 use tokio_rustls::{TlsAcceptor, TlsConnector};
 
@@ -47,27 +49,38 @@ impl HttpsIntercept {
     }
 
     async fn proxy(
+        sender: Arc<Mutex<Option<SendRequest<Body>>>>,
         client_tls: TlsConnector,
         server_name: ServerName,
         creator: Arc<dyn Outbound>,
         modifier: Arc<dyn Modifier>,
         req: Request<Body>,
         ctx: ModifierContext,
     ) -> anyhow::Result<Response<Body>> {
-        let abort_handle = ConnAbortHandle::new();
-        abort_handle.fulfill(vec![]);
         let (req, fake_resp) = modifier.modify_request(req, &ctx).await?;
         if let Some(resp) = fake_resp {
             return Ok(resp);
         }
-        let (inbound, outbound) = Connector::new_pair(10);
-        let _handle = creator.spawn_tcp(inbound, abort_handle.clone());
-        let outbound = client_tls
-            .connect(server_name, DuplexChan::new(outbound))
-            .await?;
-        let (mut sender, connection) = conn::Builder::new().handshake(outbound).await?;
-        tokio::spawn(connection);
-        let resp = sender.send_request(req).await?;
+        let resp_future = {
+            let mut sender = sender.lock().await;
+            if sender.is_none() {
+                *sender = Some({
+                    let abort_handle = ConnAbortHandle::new();
+                    abort_handle.fulfill(vec![]);
+                    let (inbound, outbound) = Connector::new_pair(10);
+                    let _handle = creator.spawn_tcp(inbound, abort_handle.clone());
+                    let outbound = client_tls
+                        .connect(server_name, DuplexChan::new(outbound))
+                        .await?;
+                    let (sender, connection) = conn::Builder::new().handshake(outbound).await?;
+                    tokio::spawn(connection);
+                    sender
+                });
+            };
+            sender.as_mut().unwrap().send_request(req)
+        };
+
+        let resp = resp_future.await?;
         let resp = modifier.modify_response(resp, &ctx).await?;
         Ok(resp)
     }
@@ -87,9 +100,11 @@ impl HttpsIntercept {
         let server_name = ServerName::try_from(self.server_name.as_str())
             .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?;
         let id_gen = IdGenerator::default();
+        let sender = Arc::new(Mutex::new(None));
         let service = service_fn(|req| {
             // since sniffer is the middle part, async tasks should be cancelled properly
             Self::proxy(
+                sender.clone(),
                 client_tls.clone(),
                 server_name.clone(),
                 self.creator.clone(),