Skip to content

Commit

Permalink
refactor
Browse files Browse the repository at this point in the history
  • Loading branch information
hatoo committed Jun 17, 2024
1 parent 2c2d9d8 commit a633d31
Show file tree
Hide file tree
Showing 3 changed files with 12 additions and 40 deletions.
3 changes: 2 additions & 1 deletion examples/dev_proxy.rs
Original file line number Diff line number Diff line change
Expand Up @@ -106,7 +106,8 @@ async fn main() {
if req.uri().host() == Some("dev.example") {
req.headers_mut().insert(
hyper::header::HOST,
hyper::header::HeaderValue::from_static("127.0.0.1"),
hyper::header::HeaderValue::from_maybe_shared(format!("127.0.0.1:{}", port))
.unwrap(),
);

let mut parts = req.uri().clone().into_parts();
Expand Down
35 changes: 6 additions & 29 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -154,7 +154,7 @@ impl<C: Borrow<rcgen::CertifiedKey> + Send + Sync + 'static> MitmProxy<C> {

async fn proxy<B>(
proxy: Arc<MitmProxyImpl<C>>,
mut req: Request<hyper::body::Incoming>,
req: Request<hyper::body::Incoming>,
tx: UnboundedSender<Communication<B>>,
client_addr: std::net::SocketAddr,
) -> Result<Response<BoxBody<Bytes, Arc<hyper::Error>>>, hyper::Error>
Expand All @@ -163,15 +163,6 @@ impl<C: Borrow<rcgen::CertifiedKey> + Send + Sync + 'static> MitmProxy<C> {
B::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
{
let original_uri = req.uri().clone();
if req.method() == Method::CONNECT {
// https
let mut parts = req.uri().clone().into_parts();
parts.scheme = Some(hyper::http::uri::Scheme::HTTPS);
// Dummy path, to avoid error
parts.path_and_query = Some(hyper::http::uri::PathAndQuery::from_static("/"));

*req.uri_mut() = Uri::from_parts(parts).unwrap();
}
let (req_back_tx, req_back_rx) = futures::channel::oneshot::channel();
let (res_tx, res_rx) = futures::channel::oneshot::channel();
let (upgrade_tx, upgrade_rx) = futures::channel::oneshot::channel();
Expand All @@ -190,30 +181,17 @@ impl<C: Borrow<rcgen::CertifiedKey> + Send + Sync + 'static> MitmProxy<C> {
};

if req.method() == Method::CONNECT {
// Modified CONNECT request is ignored
// HTTPS connection
let uri = req.uri().clone();
let uri = original_uri;
let Some(authority) = uri.authority().cloned() else {
tracing::error!("Bad CONNECT request: {}, Reason: Invalid Authority", uri);
return Ok(no_body(StatusCode::BAD_REQUEST));
};
let Some(original_authority) = original_uri.authority().cloned() else {
tracing::error!(
"Bad CONNECT request: {}, Reason: Invalid Authority",
original_uri
);
return Ok(no_body(StatusCode::BAD_REQUEST));
};
let Some(host) = uri.host().map(str::to_string) else {
tracing::error!("Bad CONNECT request: {}, Reason: Invalid Host", uri);
return Ok(no_body(StatusCode::BAD_REQUEST));
};
let Some(original_host) = original_uri.host().map(str::to_string) else {
tracing::error!(
"Bad CONNECT request: {}, Reason: Invalid Host",
original_uri
);
return Ok(no_body(StatusCode::BAD_REQUEST));
};
tokio::spawn(async move {
let Ok(client) = hyper::upgrade::on(req).await else {
tracing::error!("Bad CONNECT request: {}, Reason: Invalid Upgrade", uri);
Expand All @@ -223,9 +201,9 @@ impl<C: Borrow<rcgen::CertifiedKey> + Send + Sync + 'static> MitmProxy<C> {
if let Some(root_cert) = proxy.root_cert.as_ref() {
let Ok(server_config) =
// Even if URL is modified by middleman, we should sign with original host name to communicate client.
server_config(original_host.to_string(), root_cert.borrow())
server_config(host.to_string(), root_cert.borrow())
else {
tracing::error!("Failed to create server config for {}", original_host);
tracing::error!("Failed to create server config for {}", host);
return;
};
// TODO: Cache server_config
Expand All @@ -242,7 +220,6 @@ impl<C: Borrow<rcgen::CertifiedKey> + Send + Sync + 'static> MitmProxy<C> {
.serve_connection(
TokioIo::new(client),
service_fn(move |mut req| {
let original_authority = original_authority.clone();
let tx = tx.clone();
let authority = authority.clone();
let host = host.clone();
Expand All @@ -255,7 +232,7 @@ impl<C: Borrow<rcgen::CertifiedKey> + Send + Sync + 'static> MitmProxy<C> {
let (upgrade_tx, upgrade_rx) =
futures::channel::oneshot::channel();

inject_authority(&mut req, original_authority);
inject_authority(&mut req, authority.clone());
let _ = tx.unbounded_send(Communication {
client_addr,
request: req,
Expand Down
14 changes: 4 additions & 10 deletions tests/test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -539,12 +539,9 @@ async fn test_tls_modify_url() {

let response = tokio::spawn(setup.client.get("https://example.com/").send());

let mut comm = setup.proxy.next().await.unwrap();
let comm = setup.proxy.next().await.unwrap();
assert_eq!(comm.request.method(), hyper::Method::CONNECT);
assert_eq!(comm.request.uri().to_string(), "https://example.com:443/");
*comm.request.uri_mut() = format!("https://127.0.0.1:{}/", setup.server_port)
.parse()
.unwrap();
assert_eq!(comm.request.uri().to_string(), "example.com:443");
comm.request_back.send(comm.request).unwrap();

let mut comm = setup.proxy.next().await.unwrap();
Expand Down Expand Up @@ -576,12 +573,9 @@ async fn test_tls_modify_url_http() {

let response = tokio::spawn(setup.client.get("https://example.com/").send());

let mut comm = setup.proxy.next().await.unwrap();
let comm = setup.proxy.next().await.unwrap();
assert_eq!(comm.request.method(), hyper::Method::CONNECT);
assert_eq!(comm.request.uri().to_string(), "https://example.com:443/");
*comm.request.uri_mut() = format!("http://127.0.0.1:{}/", setup.server_port)
.parse()
.unwrap();
assert_eq!(comm.request.uri().to_string(), "example.com:443");
comm.request_back.send(comm.request).unwrap();

let mut comm = setup.proxy.next().await.unwrap();
Expand Down

0 comments on commit a633d31

Please sign in to comment.