don't build a new response, modify the old one

This commit is contained in:
Brandon Thomas
2020-09-13 07:45:29 -04:00
parent fa05996e34
commit e4ee56e0dc
+13 -33
View File
@@ -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<Body>, proxy_hostname: &str, retry_count: u8)
-> Result<Response<Body>, ErrorResponse>
fn modify_response(mut response: Response<Body>, proxy_hostname: &str, retry_count: u8)
-> Response<Body>
{
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
}