From fa05996e3466de60316318350ee876aff9fd5f21 Mon Sep 17 00:00:00 2001
From: Brandon Thomas <echelon@gmail.com>
Date: Sun, 13 Sep 2020 07:03:01 -0400
Subject: [PATCH] clean up response modification

---
 src/speak_proxy.rs | 60 ++++++++++++++++++++++++++++++----------------
 1 file changed, 40 insertions(+), 20 deletions(-)

diff --git a/src/speak_proxy.rs b/src/speak_proxy.rs
index a91c4ed..ab716a7 100644
--- a/src/speak_proxy.rs
+++ b/src/speak_proxy.rs
@@ -75,26 +75,7 @@ pub async fn speak_proxy_with_retry(
     match result {
       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());
+        return modify_response(good_response, proxy_hostname, attempt);
       },
       Err(_) => {},
     }
@@ -147,3 +128,42 @@ async fn speak_proxy(
     Err(_result) => Err(()), // TODO: Surface this?
   }
 }
+
+fn modify_response(response: Response<Body>, proxy_hostname: &str, retry_count: u8)
+  -> Result<Response<Body>, ErrorResponse>
+{
+  let status = response.status();
+
+  let response_headers = {
+    let mut response_headers = response.headers().clone();
+
+    let proxy_hostname = HeaderValue::from_str(proxy_hostname)
+      .ok()
+      .unwrap_or(HeaderValue::from_static("proxy-unknown"));
+
+    let retry_count = HeaderValue::from_str(&retry_count.to_string())
+      .ok()
+      .unwrap_or(HeaderValue::from_static("unknown"));
+
+    response_headers.append(HeaderName::from_static("x-proxy-hostname"), proxy_hostname);
+    response_headers.append(HeaderName::from_static("x-retry-count"), retry_count);
+    response_headers
+  };
+
+  let (_parts, body) = response.into_parts();
+
+  let mut response_builder = Response::builder()
+    .status(status);
+
+  for header in response_headers {
+    if let Some(header_name) = header.0.as_ref() {
+      response_builder = response_builder.header(header_name, header.1);
+    }
+  }
+
+  response_builder.body(body)
+    .map_err(|_| ErrorResponse {
+      error_type: ErrorType::ProxyError,
+      error_description: "could not modify response".to_string()
+    })
+}
\ No newline at end of file