Skip to content

Commit

Permalink
tweak
Browse files Browse the repository at this point in the history
  • Loading branch information
@hatoo
hatoo committed Nov 7, 2024
1 parent 41623a2 commit 37e665d
Show file tree
Hide file tree
Showing 2 changed files with 113 additions and 113 deletions.
36 changes: 16 additions & 20 deletions examples/https.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
190 changes: 97 additions & 93 deletions src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -87,9 +87,7 @@ impl<C: Borrow<rcgen::CertifiedKey> + Send + Sync + 'static> MitmProxy<C> {
.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
Expand All @@ -106,112 +104,118 @@ impl<C: Borrow<rcgen::CertifiedKey> + Send + Sync + 'static> MitmProxy<C> {
/// 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<S, B, E, E2>(
pub fn wrap_service<S, B, E, E2>(
proxy: Arc<Self>,
req: Request<Incoming>,
mut service: S,
) -> Result<Response<BoxBody<Bytes, E>>, E2>
service: S,
) -> impl HttpService<Incoming, ResBody = BoxBody<Bytes, E>, Error = E2, Future: Send>
where
S: HttpService<Incoming, ResBody = B, Error = E2, Future: Send>
+ Send
+ Sync
+ Clone
+ 'static,
S: HttpService<Incoming, ResBody = B, Error = E2, Future: Send> + Send + Clone + 'static,
B: Body<Data = Bytes, Error = E> + 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<CertifiedKeyDer> {
Expand Down

0 comments on commit 37e665d

Please sign in to comment.