| 1 | Upstream-status: <https://github.com/chayleaf/mitm-cache/pull/5>
|
| 2 |
|
| 3 | diff --git a/src/main.rs b/src/main.rs
|
| 4 | index d9f760f..215ed89 100644
|
| 5 | |
| 6 | |
| 7 | @@ -1,5 +1,6 @@
|
| 8 | use base64::Engine;
|
| 9 | use clap::{Parser, Subcommand};
|
| 10 | +use futures_util::StreamExt;
|
| 11 | use http_body_util::BodyExt;
|
| 12 | use hudsucker::{
|
| 13 | certificate_authority::RcgenAuthority,
|
| 14 | @@ -11,6 +12,7 @@ use hudsucker::{
|
| 15 | tokio_tungstenite::tungstenite::http::uri::Scheme,
|
| 16 | Body, HttpContext, HttpHandler, Proxy, RequestOrResponse,
|
| 17 | };
|
| 18 | +use hyper::Method;
|
| 19 | use hyper::{StatusCode, Uri};
|
| 20 | use hyper_rustls::{HttpsConnector, HttpsConnectorBuilder};
|
| 21 | use hyper_util::{
|
| 22 | @@ -137,11 +139,15 @@ impl HttpHandler for Handler {
|
| 23 | "GET" | "POST" | "HEAD" => {
|
| 24 | let original_url = req.uri().clone();
|
| 25 | println!("{req:?}");
|
| 26 | - let Ok(req) = decode_request(req) else {
|
| 27 | + let Ok(mut req) = decode_request(req) else {
|
| 28 | let mut res = Response::new("not found".into());
|
| 29 | *res.status_mut() = StatusCode::NOT_FOUND;
|
| 30 | return res.into();
|
| 31 | };
|
| 32 | + let is_head_request = req.method() == "HEAD";
|
| 33 | + if is_head_request {
|
| 34 | + *req.method_mut() = Method::GET;
|
| 35 | + }
|
| 36 | let (info, body) = req.into_parts();
|
| 37 | let Ok(req_body) = body.collect().await.map(|x| x.to_bytes())
|
| 38 | else {
|
| 39 | @@ -165,7 +171,7 @@ impl HttpHandler for Handler {
|
| 40 | req_body
|
| 41 | )])),
|
| 42 | );
|
| 43 | - let store_body_info = req.method() != "HEAD";
|
| 44 | + let store_body_info = true;
|
| 45 | let url = process_uri(original_url);
|
| 46 | if matches!(forget_regex, Some(x) if x.is_match(&url.to_string())) {
|
| 47 | forget = true;
|
| 48 | @@ -307,7 +313,20 @@ impl HttpHandler for Handler {
|
| 49 | }
|
| 50 | }
|
| 51 | });
|
| 52 | - Response::from_parts(info, ret_body)
|
| 53 | + if is_head_request {
|
| 54 | + // For HEAD requests, drain the channel
|
| 55 | + // output; otherwise rx will be dropped,
|
| 56 | + // which closes the channel and causes
|
| 57 | + // the tx.poll_ready call above to
|
| 58 | + // return an error
|
| 59 | + let mut stream = ret_body.into_data_stream();
|
| 60 | + while let Some(_) = stream.next().await {
|
| 61 | + // do nothing
|
| 62 | + }
|
| 63 | + Response::from_parts(info, Body::empty())
|
| 64 | + } else {
|
| 65 | + Response::from_parts(info, ret_body)
|
| 66 | + }
|
| 67 | } else {
|
| 68 | // remove hash headers to force the software to download this
|
| 69 | // so we get sha256
|