Try to clean up code and to fix WebSocket problemsdocker-38

1 files changed, 34 insertions, 20 deletions

diff --git a/src/reverse_proxy.rs b/src/reverse_proxy.rs
index b5916be..99a7c98 100644
--- a/src/reverse_proxy.rs
+++ b/src/reverse_proxy.rs
@@ -48,9 +48,19 @@ fn remove_hop_headers(headers: &HeaderMap<HeaderValue>) -> HeaderMap<HeaderValue
 	result
 }
 
-fn create_proxied_response<B>(mut response: Response<B>) -> Response<B> {
-	*response.headers_mut() = remove_hop_headers(response.headers());
-	response
+fn copy_upgrade_headers(
+	old_headers: &HeaderMap<HeaderValue>,
+	new_headers: &mut HeaderMap<HeaderValue>,
+) -> Result<()> {
+	if let Some(conn) = old_headers.get(header::CONNECTION) {
+		if conn.to_str()?.to_lowercase() == "upgrade" {
+			if let Some(upgrade) = old_headers.get(header::UPGRADE) {
+				new_headers.insert(header::CONNECTION, "Upgrade".try_into()?);
+				new_headers.insert(header::UPGRADE, upgrade.clone());
+			}
+		}
+	}
+	Ok(())
 }
 
 fn forward_uri<B>(forward_url: &str, req: &Request<B>) -> Result<Uri> {
@@ -72,12 +82,13 @@ fn create_proxied_request<B>(
 		.uri(forward_uri(forward_url, &request)?)
 		.version(hyper::Version::HTTP_11);
 
-	let headers = builder.headers_mut().unwrap();
+	let old_headers = request.headers();
+	let new_headers = builder.headers_mut().unwrap();
 
-	*headers = remove_hop_headers(request.headers());
+	*new_headers = remove_hop_headers(old_headers);
 
 	// If request does not have host header, add it from original URI authority
-	if let header::Entry::Vacant(entry) = headers.entry(header::HOST) {
+	if let header::Entry::Vacant(entry) = new_headers.entry(header::HOST) {
 		if let Some(authority) = request.uri().authority() {
 			entry.insert(authority.as_str().parse()?);
 		}
@@ -86,7 +97,7 @@ fn create_proxied_request<B>(
 	// Concatenate cookie headers into single header
 	// (HTTP/2 allows several cookie headers, but we are proxying to HTTP/1.1 that does not)
 	let mut cookie_concat = vec![];
-	for cookie in headers.get_all(header::COOKIE) {
+	for cookie in new_headers.get_all(header::COOKIE) {
 		if !cookie_concat.is_empty() {
 			cookie_concat.extend(b"; ");
 		}
@@ -94,12 +105,12 @@ fn create_proxied_request<B>(
 	}
 	if !cookie_concat.is_empty() {
 		// insert clears the old value of COOKIE and inserts the concatenated version instead
-		headers.insert(header::COOKIE, cookie_concat.try_into()?);
+		new_headers.insert(header::COOKIE, cookie_concat.try_into()?);
 	}
 
 	// Add forwarding information in the headers
 	let x_forwarded_for_header_name = "x-forwarded-for";
-	match headers.entry(x_forwarded_for_header_name) {
+	match new_headers.entry(x_forwarded_for_header_name) {
 		header::Entry::Vacant(entry) => {
 			entry.insert(client_ip.to_string().parse()?);
 		}
@@ -110,24 +121,27 @@ fn create_proxied_request<B>(
 		}
 	}
 
-	headers.insert(
+	new_headers.insert(
 		HeaderName::from_bytes(b"x-forwarded-proto")?,
 		"https".try_into()?,
 	);
 
 	// Proxy upgrade requests properly
-	if let Some(conn) = request.headers().get(header::CONNECTION) {
-		if conn.to_str()?.to_lowercase() == "upgrade" {
-			if let Some(upgrade) = request.headers().get(header::UPGRADE) {
-				headers.insert(header::CONNECTION, "Upgrade".try_into()?);
-				headers.insert(header::UPGRADE, upgrade.clone());
-			}
-		}
-	}
+	copy_upgrade_headers(old_headers, new_headers)?;
 
 	Ok(builder.body(request.into_body())?)
 }
 
+fn create_proxied_response<B>(mut response: Response<B>) -> Result<Response<B>> {
+	let old_headers = response.headers();
+	let mut new_headers = remove_hop_headers(old_headers);
+
+	copy_upgrade_headers(old_headers, &mut new_headers)?;
+
+	*response.headers_mut() = new_headers;
+	Ok(response)
+}
+
 pub async fn call(
 	client_ip: IpAddr,
 	forward_uri: &str,
@@ -146,7 +160,7 @@ pub async fn call(
 
 	trace!("Inner response: {:?}", response);
 
-	let proxied_response = create_proxied_response(response);
+	let proxied_response = create_proxied_response(response)?;
 	Ok(proxied_response)
 }
 
@@ -173,7 +187,7 @@ pub async fn call_https(
 
 	trace!("Inner response (HTTPS): {:?}", response);
 
-	let proxied_response = create_proxied_response(response);
+	let proxied_response = create_proxied_response(response)?;
 	Ok(proxied_response)
 }