Fix URI in request headers to match target service

Author: ameba23

src/lib.rs

@@ -163,9 +163,20 @@ impl ProxyServer {
 
         // Setup a request handler
         let service = service_fn(move |mut req| {
+            // Change the request uri to match the target service
+            let new_uri = rewrite_uri_authority(req.uri(), &target);
+            *req.uri_mut() = new_uri;
+            let headers = req.headers_mut();
+
+            // Add or update the HOST header
+            if let Ok(host_header_value) = target.parse() {
+                headers.insert(http::header::HOST, host_header_value);
+            } else {
+                error!("Failed to encode host as header value: {target}");
+            }
+
             // If we have measurements, from the remote peer, add them to the request header
             let measurements = measurements.clone();
-            let headers = req.headers_mut();
             if let Some(measurements) = measurements {
                 match measurements.to_header_format() {
                     Ok(header_value) => {
@@ -508,6 +519,29 @@ impl ProxyClient {
     }
 }
 
+/// Given an http request, update the host in the URI when proxying
+fn rewrite_uri_authority(old_uri: &http::Uri, new_authority: &str) -> http::Uri {
+    let path_and_query = old_uri
+        .path_and_query()
+        .map(|pq| pq.as_str())
+        .unwrap_or("/");
+
+    let scheme = old_uri.scheme_str().unwrap_or("http");
+
+    match http::Uri::builder()
+        .scheme(scheme)
+        .authority(new_authority)
+        .path_and_query(path_and_query)
+        .build()
+    {
+        Ok(new_uri) => new_uri,
+        Err(err) => {
+            tracing::error!("Failed to build URI - falling back to old one: {err}");
+            old_uri.clone()
+        }
+    }
+}
+
 /// An error when running a proxy client or server
 #[derive(Error, Debug)]
 pub enum ProxyError {