From 37e665d768d98c31208791f1fcc28d032cc9c96c Mon Sep 17 00:00:00 2001 From: hatoo Date: Thu, 7 Nov 2024 21:49:49 +0900 Subject: [PATCH] tweak --- examples/https.rs | 36 ++++----- src/lib.rs | 190 +++++++++++++++++++++++----------------------- 2 files changed, 113 insertions(+), 113 deletions(-) diff --git a/examples/https.rs b/examples/https.rs index c972dc2..1c61ebb 100644 --- a/examples/https.rs +++ b/examples/https.rs @@ -106,31 +106,27 @@ async fn main() { let tls_acceptor = tls_acceptor.clone(); tokio::spawn(async move { - let service = service_fn(move |req| { - let client = client.clone(); + let client = client.clone(); + let service = MitmProxy::wrap_service( + proxy.clone(), + service_fn(move |req| { + let client = client.clone(); + async move { + let uri = req.uri().clone(); - MitmProxy::wrap_service( - proxy.clone(), - req, - service_fn(move |req| { - let client = client.clone(); - async move { - let uri = req.uri().clone(); + // You can modify request here + // or You can just return response anywhere - // You can modify request here - // or You can just return response anywhere + let (res, _upgrade) = client.send_request(req).await?; - let (res, _upgrade) = client.send_request(req).await?; + println!("{} -> {}", uri, res.status()); - println!("{} -> {}", uri, res.status()); + // You can modify response here - // You can modify response here - - Ok::<_, http_mitm_proxy::default_client::Error>(res) - } - }), - ) - }); + Ok::<_, http_mitm_proxy::default_client::Error>(res) + } + }), + ); let stream = tls_acceptor.accept(stream).await.unwrap(); hyper::server::conn::http1::Builder::new() diff --git a/src/lib.rs b/src/lib.rs index feaa905..5e533d9 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -87,9 +87,7 @@ impl + Send + Sync + 'static> MitmProxy { .title_case_headers(true) .serve_connection( TokioIo::new(stream), - service_fn(|req| { - Self::wrap_service(proxy.clone(), req, service.clone()) - }), + Self::wrap_service(proxy.clone(), service.clone()), ) .with_upgrades() .await @@ -106,112 +104,118 @@ impl + Send + Sync + 'static> MitmProxy { /// See `examples/https.rs` for usage. /// If you want to serve simple HTTP proxy server, you can use `bind` method instead. /// `bind` will call this method internally. - pub async fn wrap_service( + pub fn wrap_service( proxy: Arc, - req: Request, - mut service: S, - ) -> Result>, E2> + service: S, + ) -> impl HttpService, Error = E2, Future: Send> where - S: HttpService - + Send - + Sync - + Clone - + 'static, + S: HttpService + Send + Clone + 'static, B: Body + Send + Sync + 'static, E: std::error::Error + Send + Sync + 'static, E2: std::error::Error + Send + Sync + 'static, { - if req.method() == Method::CONNECT { - // https - let Some(connect_authority) = req.uri().authority().cloned() else { - tracing::error!( - "Bad CONNECT request: {}, Reason: Invalid Authority", - req.uri() - ); - return Ok(no_body(StatusCode::BAD_REQUEST)); - }; + service_fn(move |req| { + let proxy = proxy.clone(); + let mut service = service.clone(); - tokio::spawn(async move { - let Ok(client) = hyper::upgrade::on(req).await else { - tracing::error!( - "Bad CONNECT request: {}, Reason: Invalid Upgrade", - connect_authority - ); - return; - }; - if let Some(server_config) = - proxy.server_config(connect_authority.host().to_string(), true) - { - let server_config = match server_config { - Ok(server_config) => server_config, - Err(err) => { - tracing::error!( - "Failed to create server config for {}, {}", - connect_authority.host(), - err - ); - return; - } + async move { + if req.method() == Method::CONNECT { + // https + let Some(connect_authority) = req.uri().authority().cloned() else { + tracing::error!( + "Bad CONNECT request: {}, Reason: Invalid Authority", + req.uri() + ); + return Ok(no_body(StatusCode::BAD_REQUEST)); }; - let server_config = Arc::new(server_config); - let tls_acceptor = tokio_rustls::TlsAcceptor::from(server_config); - let client = match tls_acceptor.accept(TokioIo::new(client)).await { - Ok(client) => client, - Err(err) => { + + tokio::spawn(async move { + let Ok(client) = hyper::upgrade::on(req).await else { tracing::error!( - "Failed to accept TLS connection for {}, {}", - connect_authority.host(), - err + "Bad CONNECT request: {}, Reason: Invalid Upgrade", + connect_authority ); return; - } - }; - let f = move |mut req: Request<_>| { - let connect_authority = connect_authority.clone(); - let mut service = service.clone(); + }; + if let Some(server_config) = + proxy.server_config(connect_authority.host().to_string(), true) + { + let server_config = match server_config { + Ok(server_config) => server_config, + Err(err) => { + tracing::error!( + "Failed to create server config for {}, {}", + connect_authority.host(), + err + ); + return; + } + }; + let server_config = Arc::new(server_config); + let tls_acceptor = tokio_rustls::TlsAcceptor::from(server_config); + let client = match tls_acceptor.accept(TokioIo::new(client)).await { + Ok(client) => client, + Err(err) => { + tracing::error!( + "Failed to accept TLS connection for {}, {}", + connect_authority.host(), + err + ); + return; + } + }; + let f = move |mut req: Request<_>| { + let connect_authority = connect_authority.clone(); + let mut service = service.clone(); + + async move { + inject_authority(&mut req, connect_authority.clone()); + service.call(req).await + } + }; + let res = if client.get_ref().1.alpn_protocol() == Some(b"h2") { + server::conn::http2::Builder::new(TokioExecutor::new()) + .serve_connection(TokioIo::new(client), service_fn(f)) + .await + } else { + server::conn::http1::Builder::new() + .preserve_header_case(true) + .title_case_headers(true) + .serve_connection(TokioIo::new(client), service_fn(f)) + .with_upgrades() + .await + }; - async move { - inject_authority(&mut req, connect_authority.clone()); - service.call(req).await + if let Err(_err) = res { + // Suppress error because if we serving HTTPS proxy server and forward to HTTPS server, it will always error when closing connection. + // tracing::error!("Error in proxy: {}", err); + } + } else { + let Ok(mut server) = + TcpStream::connect(connect_authority.as_str()).await + else { + tracing::error!("Failed to connect to {}", connect_authority); + return; + }; + let _ = tokio::io::copy_bidirectional( + &mut TokioIo::new(client), + &mut server, + ) + .await; } - }; - let res = if client.get_ref().1.alpn_protocol() == Some(b"h2") { - server::conn::http2::Builder::new(TokioExecutor::new()) - .serve_connection(TokioIo::new(client), service_fn(f)) - .await - } else { - server::conn::http1::Builder::new() - .preserve_header_case(true) - .title_case_headers(true) - .serve_connection(TokioIo::new(client), service_fn(f)) - .with_upgrades() - .await - }; + }); - if let Err(_err) = res { - // Suppress error because if we serving HTTPS proxy server and forward to HTTPS server, it will always error when closing connection. - // tracing::error!("Error in proxy: {}", err); - } + Ok(Response::new( + http_body_util::Empty::new() + .map_err(|never: std::convert::Infallible| match never {}) + .boxed(), + )) } else { - let Ok(mut server) = TcpStream::connect(connect_authority.as_str()).await - else { - tracing::error!("Failed to connect to {}", connect_authority); - return; - }; - let _ = - tokio::io::copy_bidirectional(&mut TokioIo::new(client), &mut server).await; + // http + service.call(req).await.map(|res| res.map(|b| b.boxed())) } - }); - - Ok(Response::new( - http_body_util::Empty::new() - .map_err(|never: std::convert::Infallible| match never {}) - .boxed(), - )) - } else { - // http - service.call(req).await.map(|res| res.map(|b| b.boxed())) - } + } + }) } fn get_certified_key(&self, host: String) -> Option {