From f2b0ec1c466ec83731f9bd590021968bf3aaf2bd Mon Sep 17 00:00:00 2001
From: Brandon Thomas <echelon@gmail.com>
Date: Sun, 13 Sep 2020 06:37:24 -0400
Subject: [PATCH] add proxy hostname to response (but this isn't elegant/safe)

---
 Cargo.lock         | 20 +++++++++++++++++++-
 Cargo.toml         |  3 ++-
 src/main.rs        |  9 +++++++++
 src/speak_proxy.rs | 25 ++++++++++++++++++++++++-
 4 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index dcd869e..812794e 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -225,6 +225,17 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "hostname"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3c731c3e10504cc8ed35cfe2f1db4c9274c3d35fa486e3b31df46f068ef3e867"
+dependencies = [
+ "libc",
+ "match_cfg",
+ "winapi 0.3.9",
+]
+
 [[package]]
 name = "http"
 version = "0.2.1"
@@ -350,6 +361,12 @@ dependencies = [
  "cfg-if",
 ]
 
+[[package]]
+name = "match_cfg"
+version = "0.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ffbee8634e0d45d258acb448e7eaab3fce7a0a467395d4d9f228e3c1f01fb2e4"
+
 [[package]]
 name = "memchr"
 version = "2.3.3"
@@ -795,11 +812,12 @@ checksum = "e604eb7b43c06650e854be16a2a03155743d3752dd1c943f6829e26b7a36e382"
 
 [[package]]
 name = "tts-service-proxy"
-version = "0.0.1"
+version = "0.0.2"
 dependencies = [
  "anyhow",
  "env_logger",
  "futures",
+ "hostname",
  "hyper",
  "hyper-reverse-proxy",
  "log",
diff --git a/Cargo.toml b/Cargo.toml
index 8d9aa86..4b1a26f 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -16,6 +16,7 @@ path = "src/main.rs"
 anyhow = "1.0.31"
 env_logger = "0.7.1"
 futures = "0.3"
+hostname = "0.3.1"
 hyper = "0.13"
 hyper-reverse-proxy = { git = "https://github.com/felipenoris/hyper-reverse-proxy.git", commit = "d109fcbd629c62eeb5c34bec1074d737e7b58529" }
 log = "0.4.8"
@@ -25,4 +26,4 @@ serde = "1.0.104"
 serde_derive = "1.0.104"
 serde_json = "1.0.55"
 tokio = { version = "0.2", features = ["full"] }
-toml = "0.5.6"
+toml = "0.5.6"
\ No newline at end of file
diff --git a/src/main.rs b/src/main.rs
index d3f0fd1..4423c43 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -135,6 +135,7 @@ pub async fn handle(
         nr_transaction,
         server_params.max_retries,
         server_params.retry_wait_ms,
+        &server_params.server_hostname,
         "/speak").await;
       match result {
         Ok(response) => Ok(response),
@@ -150,6 +151,7 @@ pub async fn handle(
         nr_transaction,
         server_params.max_retries,
         server_params.retry_wait_ms,
+        &server_params.server_hostname,
         "/speak_spectrogram").await;
       match result {
         Ok(response) => Ok(response),
@@ -186,6 +188,7 @@ pub async fn handle(
 pub struct ServerParams {
   pub max_retries: u8,
   pub retry_wait_ms: u64,
+  pub server_hostname: String,
 }
 
 #[tokio::main]
@@ -227,9 +230,15 @@ async fn main() {
 
   info!("Proxy configs: {:?}", proxy_configs);
 
+  let server_hostname = hostname::get()
+    .ok()
+    .and_then(|h| h.into_string().ok())
+    .unwrap_or("proxy-unknown".to_string());
+
   let server_params = ServerParams {
     max_retries,
     retry_wait_ms,
+    server_hostname,
   };
 
   let mut route_map = HashMap::new();
diff --git a/src/speak_proxy.rs b/src/speak_proxy.rs
index 98a689d..a91c4ed 100644
--- a/src/speak_proxy.rs
+++ b/src/speak_proxy.rs
@@ -18,6 +18,7 @@ pub async fn speak_proxy_with_retry(
   _transaction: MaybeNewRelicTransaction,
   max_retries: u8,
   retry_wait_ms: u64,
+  proxy_hostname: &str,
   endpoint: &'static str) -> Result<Response<Body>, ErrorResponse>
 {
   let mut attempt : u8 = 0;
@@ -72,7 +73,29 @@ pub async fn speak_proxy_with_retry(
       endpoint).await;
 
     match result {
-      Ok(good_response) => return Ok(good_response),
+      Ok(good_response) => {
+        // All we want to do is add our special header
+        let proxy_hostname_header = HeaderValue::from_str(proxy_hostname)
+          .ok()
+          .unwrap_or(HeaderValue::from_static("proxy-unknown"));
+
+        let status = good_response.status();
+
+        let mut response_headers = good_response.headers().clone();
+        response_headers.append(HeaderName::from_static("x-proxy-hostname"), proxy_hostname_header);
+
+        let (_parts, body) = good_response.into_parts();
+
+        let mut response_builder = Response::builder()
+          .status(status);
+
+        for header in response_headers {
+          response_builder = response_builder.header(header.0.unwrap(), header.1);
+        }
+
+        return Ok(response_builder.body(body)
+          .unwrap());
+      },
       Err(_) => {},
     }