Merge pull request #688 from hatoo/proxy-headers

Proxy headers

README.md

@@ -140,6 +140,8 @@ Options:
           HTTP method [default: GET]
   -H <HEADERS>
           Custom HTTP header. Examples: -H "foo: bar"
+      --proxy-header <PROXY_HEADERS>
+          Custom Proxy HTTP header. Examples: --proxy-header "foo: bar"
   -t <TIMEOUT>
           Timeout for each request. Default to infinite.
   -A <ACCEPT_HEADER>

src/client.rs

@@ -174,6 +174,7 @@ pub struct Client {
     pub url_generator: UrlGenerator,
     pub method: http::Method,
     pub headers: http::header::HeaderMap,
+    pub proxy_headers: http::header::HeaderMap,
     pub body: Option<&'static [u8]>,
     pub dns: Dns,
     pub timeout: Option<std::time::Duration>,
@@ -494,14 +495,21 @@ impl Client {
             let (dns_lookup, stream) = self.client(proxy_url, rng, self.is_proxy_http2()).await?;
             if url.scheme() == "https" {
                 // Do CONNECT request to proxy
-                let req = http::Request::builder()
-                    .method(Method::CONNECT)
-                    .uri(format!(
-                        "{}:{}",
-                        url.host_str().unwrap(),
-                        url.port_or_known_default().unwrap()
-                    ))
-                    .body(http_body_util::Full::default())?;
+                let req = {
+                    let mut builder =
+                        http::Request::builder()
+                            .method(Method::CONNECT)
+                            .uri(format!(
+                                "{}:{}",
+                                url.host_str().unwrap(),
+                                url.port_or_known_default().unwrap()
+                            ));
+                    *builder
+                        .headers_mut()
+                        .ok_or(ClientError::GetHeaderFromBuilderError)? =
+                        self.proxy_headers.clone();
+                    builder.body(http_body_util::Full::default())?
+                };
                 let res = if self.proxy_http_version == http::Version::HTTP_2 {
                     let mut send_request = stream.handshake_http2().await?;
                     send_request.send_request(req).await?
@@ -557,6 +565,12 @@ impl Client {
             aws_config.sign_request(self.method.as_str(), &mut headers, url, bytes)?
         }
 
+        if use_proxy {
+            for (key, value) in self.proxy_headers.iter() {
+                headers.insert(key, value.clone());
+            }
+        }
+
         *builder
             .headers_mut()
             .ok_or(ClientError::GetHeaderFromBuilderError)? = headers;
@@ -670,14 +684,21 @@ impl Client {
         if let Some(proxy_url) = &self.proxy_url {
             let (dns_lookup, stream) = self.client(proxy_url, rng, self.is_proxy_http2()).await?;
             if url.scheme() == "https" {
-                let req = http::Request::builder()
-                    .method(Method::CONNECT)
-                    .uri(format!(
-                        "{}:{}",
-                        url.host_str().unwrap(),
-                        url.port_or_known_default().unwrap()
-                    ))
-                    .body(http_body_util::Full::default())?;
+                let req = {
+                    let mut builder =
+                        http::Request::builder()
+                            .method(Method::CONNECT)
+                            .uri(format!(
+                                "{}:{}",
+                                url.host_str().unwrap(),
+                                url.port_or_known_default().unwrap()
+                            ));
+                    *builder
+                        .headers_mut()
+                        .ok_or(ClientError::GetHeaderFromBuilderError)? =
+                        self.proxy_headers.clone();
+                    builder.body(http_body_util::Full::default())?
+                };
                 let res = if self.proxy_http_version == http::Version::HTTP_2 {
                     let mut send_request = stream.handshake_http2().await?;
                     send_request.send_request(req).await?

src/db.rs

@@ -78,6 +78,7 @@ mod test_db {
             url_generator: UrlGenerator::new_static("http://example.com".parse().unwrap()),
             method: Method::GET,
             headers: HeaderMap::new(),
+            proxy_headers: HeaderMap::new(),
             body: None,
             dns: Dns {
                 resolver: hickory_resolver::AsyncResolver::tokio_from_system_conf().unwrap(),

src/main.rs

@@ -4,9 +4,12 @@ use clap::Parser;
 use crossterm::tty::IsTty;
 use hickory_resolver::config::{ResolverConfig, ResolverOpts};
 use humantime::Duration;
-use hyper::http::{
-    self,
-    header::{HeaderName, HeaderValue},
+use hyper::{
+    http::{
+        self,
+        header::{HeaderName, HeaderValue},
+    },
+    HeaderMap,
 };
 use printer::{PrintConfig, PrintMode};
 use rand_regex::Regex;
@@ -143,6 +146,11 @@ Note: If qps is specified, burst will be ignored",
     method: http::Method,
     #[arg(help = "Custom HTTP header. Examples: -H \"foo: bar\"", short = 'H')]
     headers: Vec<String>,
+    #[arg(
+        help = "Custom Proxy HTTP header. Examples: --proxy-header \"foo: bar\"",
+        long = "proxy-header"
+    )]
+    proxy_headers: Vec<String>,
     #[arg(help = "Timeout for each request. Default to infinite.", short = 't')]
     timeout: Option<humantime::Duration>,
     #[arg(help = "HTTP Accept Header.", short = 'A')]
@@ -501,13 +509,7 @@ async fn run() -> anyhow::Result<()> {
         for (k, v) in opts
             .headers
             .into_iter()
-            .map(|s| {
-                let header = s.splitn(2, ':').collect::<Vec<_>>();
-                anyhow::ensure!(header.len() == 2, anyhow::anyhow!("Parse header"));
-                let name = HeaderName::from_str(header[0])?;
-                let value = HeaderValue::from_str(header[1].trim_start_matches(' '))?;
-                Ok::<(HeaderName, HeaderValue), anyhow::Error>((name, value))
-            })
+            .map(|s| parse_header(s.as_str()))
             .collect::<anyhow::Result<Vec<_>>>()?
         {
             headers.insert(k, v);
@@ -516,6 +518,13 @@ async fn run() -> anyhow::Result<()> {
         headers
     };
 
+    let proxy_headers = {
+        opts.proxy_headers
+            .into_iter()
+            .map(|s| parse_header(s.as_str()))
+            .collect::<anyhow::Result<HeaderMap<_>>>()?
+    };
+
     let body: Option<&'static [u8]> = match (opts.body_string, opts.body_path) {
         (Some(body), _) => Some(Box::leak(body.into_boxed_str().into_boxed_bytes())),
         (_, Some(path)) => {
@@ -550,6 +559,7 @@ async fn run() -> anyhow::Result<()> {
         url_generator,
         method: opts.method,
         headers,
+        proxy_headers,
         body,
         dns: client::Dns {
             resolver,
@@ -946,3 +956,11 @@ impl Opts {
         }
     }
 }
+
+fn parse_header(s: &str) -> Result<(HeaderName, HeaderValue), anyhow::Error> {
+    let header = s.splitn(2, ':').collect::<Vec<_>>();
+    anyhow::ensure!(header.len() == 2, anyhow::anyhow!("Parse header"));
+    let name = HeaderName::from_str(header[0])?;
+    let value = HeaderValue::from_str(header[1].trim_start_matches(' '))?;
+    Ok::<(HeaderName, HeaderValue), anyhow::Error>((name, value))
+}

tests/tests.rs

@@ -755,22 +755,31 @@ where
 
             let proxy = proxy.clone();
             let service = service.clone();
+
+            let outer = service_fn(move |req| {
+                // Test --proxy-header option
+                assert_eq!(
+                    req.headers()
+                        .get("proxy-authorization")
+                        .unwrap()
+                        .to_str()
+                        .unwrap(),
+                    "test"
+                );
+
+                MitmProxy::wrap_service(proxy.clone(), service.clone()).call(req)
+            });
+
             tokio::spawn(async move {
                 if http2 {
                     let _ = hyper::server::conn::http2::Builder::new(TokioExecutor::new())
-                        .serve_connection(
-                            TokioIo::new(stream),
-                            MitmProxy::wrap_service(proxy, service),
-                        )
+                        .serve_connection(TokioIo::new(stream), outer)
                         .await;
                 } else {
                     let _ = hyper::server::conn::http1::Builder::new()
                         .preserve_header_case(true)
                         .title_case_headers(true)
-                        .serve_connection(
-                            TokioIo::new(stream),
-                            MitmProxy::wrap_service(proxy, service),
-                        )
+                        .serve_connection(TokioIo::new(stream), outer)
                         .with_upgrades()
                         .await;
                 }
@@ -800,6 +809,7 @@ async fn test_proxy_with_setting(https: bool, http2: bool, proxy_http2: bool) {
     let scheme = if https { "https" } else { "http" };
     proc.args(["--no-tui", "--debug", "--insecure", "-x"])
         .arg(format!("http://127.0.0.1:{proxy_port}/"))
+        .args(["--proxy-header", "proxy-authorization: test"])
         .arg(format!("{scheme}://example.com/"));
     if http2 {
         proc.arg("--http2");
@@ -808,18 +818,17 @@ async fn test_proxy_with_setting(https: bool, http2: bool, proxy_http2: bool) {
         proc.arg("--proxy-http2");
     }
 
+    // When std::process::Stdio::piped() is used, the wait_with_output() method will hang in Windows.
     proc.stdin(std::process::Stdio::null())
-        .stdout(std::process::Stdio::piped())
-        .stderr(std::process::Stdio::null());
-    let stdout = proc
-        .spawn()
-        .unwrap()
-        .wait_with_output()
-        .await
-        .unwrap()
-        .stdout;
-
-    assert!(String::from_utf8(stdout).unwrap().contains("Hello World"),);
+        .stdout(std::process::Stdio::inherit())
+        .stderr(std::process::Stdio::inherit());
+    // So, we test status code only for now.
+    assert!(proc.status().await.unwrap().success());
+    /*
+    let outputs = proc.spawn().unwrap().wait_with_output().await.unwrap();
+    let stdout = String::from_utf8(outputs.stdout).unwrap();
+    assert!(stdout.contains("Hello World"),);
+    */
 }
 
 #[tokio::test]