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
, proxy_hostname: &str, retry_count: u8) + -> Result