From 6025dd50cad526b56eaec6e10c72a7e0feb32279 Mon Sep 17 00:00:00 2001
From: Martin Aeschlimann <martinae@microsoft.com>
Date: Sun, 19 Jun 2022 01:20:59 +0200
Subject: [PATCH] add x-forwarded-host header

---
 src/lib.rs         | 32 +++++++++++++++++++-------------
 tests/test_http.rs | 21 +++++++++++++++++++++
 2 files changed, 40 insertions(+), 13 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 76d9146..39c1c6a 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -140,6 +140,7 @@ lazy_static! {
     ];
 
     static ref X_FORWARDED_FOR: HeaderName = HeaderName::from_static("x-forwarded-for");
+    static ref X_FORWARDED_HOST: HeaderName = HeaderName::from_static("x-forwarded-host");
 }
 
 #[derive(Debug)]
@@ -334,7 +335,7 @@ fn create_proxied_request<B>(
     debug!("Setting headers of proxied request");
 
     // remove the original HOST header. It will be set by the client that sends the request
-    request.headers_mut().remove(HOST);
+    let original_host = request.headers_mut().remove(HOST);
 
     *request.uri_mut() = uri;
 
@@ -361,22 +362,27 @@ fn create_proxied_request<B>(
     }
 
     // Add forwarding information in the headers
+    let forwarded_for_value = client_ip.to_string().parse()?;
     match request.headers_mut().entry(&*X_FORWARDED_FOR) {
         hyper::header::Entry::Vacant(entry) => {
-            debug!("X-Fowraded-for header was vacant");
-            entry.insert(client_ip.to_string().parse()?);
+            debug!("x-forwarded-for header was vacant");
+            entry.insert(forwarded_for_value);
         }
+        hyper::header::Entry::Occupied(mut entry) => {
+            debug!("x-forwarded-for header was occupied");
+            entry.append(forwarded_for_value);
+        }
+    }
 
-        hyper::header::Entry::Occupied(entry) => {
-            debug!("X-Fowraded-for header was occupied");
-            let client_ip_str = client_ip.to_string();
-            let mut addr =
-                String::with_capacity(entry.get().as_bytes().len() + 2 + client_ip_str.len());
-
-            addr.push_str(std::str::from_utf8(entry.get().as_bytes()).unwrap());
-            addr.push(',');
-            addr.push(' ');
-            addr.push_str(&client_ip_str);
+    if let Some(host) = original_host {
+        match request.headers_mut().entry(&*X_FORWARDED_HOST) {
+            hyper::header::Entry::Vacant(entry) => {
+                debug!("x-forwarded-host header was vacant");
+                entry.insert(host.to_owned());
+            }
+            hyper::header::Entry::Occupied(mut _entry) => {
+                debug!("x-forwarded-host header was occupied");
+            }
         }
     }
 
diff --git a/tests/test_http.rs b/tests/test_http.rs
index d28108e..9749d18 100644
--- a/tests/test_http.rs
+++ b/tests/test_http.rs
@@ -90,6 +90,27 @@ async fn test_get(ctx: &mut ProxyTestContext) {
         HOST,
         format!("127.0.0.1:{}", ctx.http_back.port).parse().unwrap(),
     );
+
+    ctx.http_back.add(
+        HandlerBuilder::new("/foo")
+            .status_code(StatusCode::OK)
+            .headers(headers)
+            .build(),
+    );
+    let resp = Client::new().get(ctx.uri("/foo")).await.unwrap();
+    assert_eq!(200, resp.status());
+}
+
+#[test_context(ProxyTestContext)]
+#[tokio::test]
+async fn test_headers(ctx: &mut ProxyTestContext) {
+    let mut headers = HeaderMap::new();
+    headers.insert("x-forwarded-for", "127.0.0.1".parse().unwrap());
+    headers.insert(
+        "x-forwarded-host",
+        format!("localhost:{}", ctx.port).parse().unwrap(),
+    );
+
     ctx.http_back.add(
         HandlerBuilder::new("/foo")
             .status_code(StatusCode::OK)