diff --git a/src/speak_proxy.rs b/src/speak_proxy.rs
index ab716a7..4cb0ddd 100644
--- a/src/speak_proxy.rs
+++ b/src/speak_proxy.rs
@@ -75,7 +75,7 @@ pub async fn speak_proxy_with_retry(
match result {
Ok(good_response) => {
// All we want to do is add our special header
- return modify_response(good_response, proxy_hostname, attempt);
+ return Ok(modify_response(good_response, proxy_hostname, attempt));
},
Err(_) => {},
}
@@ -129,41 +129,21 @@ async fn speak_proxy(
}
}
-fn modify_response(response: Response
, proxy_hostname: &str, retry_count: u8)
- -> Result, ErrorResponse>
+fn modify_response(mut response: Response, proxy_hostname: &str, retry_count: u8)
+ -> Response
{
- let status = response.status();
+ let headers = response.headers_mut();
- 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 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"));
- let retry_count = HeaderValue::from_str(&retry_count.to_string())
- .ok()
- .unwrap_or(HeaderValue::from_static("unknown"));
+ headers.append(HeaderName::from_static("x-proxy-hostname"), proxy_hostname);
+ headers.append(HeaderName::from_static("x-retry-count"), retry_count);
- 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()
- })
+ response
}
\ No newline at end of file