From c3536eb138d61a7a5e720ae5846c98fb7e573400 Mon Sep 17 00:00:00 2001
From: glendc <contact@glendc.com>
Date: Fri, 20 Sep 2024 20:01:11 +0200
Subject: [PATCH] allways remove illegal h2 headers in h2

---
 rama-http-backend/src/client/svc.rs | 83 +++++++++++++++++------------
 1 file changed, 48 insertions(+), 35 deletions(-)

diff --git a/rama-http-backend/src/client/svc.rs b/rama-http-backend/src/client/svc.rs
index 2b37e69e..29e8344d 100644
--- a/rama-http-backend/src/client/svc.rs
+++ b/rama-http-backend/src/client/svc.rs
@@ -60,25 +60,16 @@ fn sanitize_client_req_header<S, B>(
     ctx: &mut Context<S>,
     req: Request<B>,
 ) -> Result<Request<B>, BoxError> {
-    Ok(match req.method() {
-        &Method::CONNECT => {
-            // CONNECT
-            if req.uri().host().is_none() {
-                return Err(OpaqueError::from_display("missing host in CONNECT request").into());
-            }
-            req
-        }
-        _ => {
-            // [HTTP/1.1] GET | HEAD | POST | PUT | DELETE | OPTIONS | TRACE | PATCH
-            if !ctx.contains::<ProxyAddress>()
-                && req.uri().host().is_some()
-                && req.version() <= Version::HTTP_11
-            {
-                // ensure request context is defined prior to doing this, as otherwise we can get issues
-                let _ = ctx.get_or_try_insert_with_ctx::<RequestContext, _>(|ctx| {
-                    (ctx, &req).try_into()
-                })?;
+    // logic specific to this method
+    if req.method() == Method::CONNECT && req.uri().host().is_none() {
+        return Err(OpaqueError::from_display("missing host in CONNECT request").into());
+    }
 
+    // logic specific to http versions
+    Ok(match req.version() {
+        Version::HTTP_09 | Version::HTTP_10 | Version::HTTP_11 => {
+            // remove authority and scheme for non-connect requests
+            if !ctx.contains::<ProxyAddress>() && req.uri().host().is_some() {
                 tracing::trace!(
                     "remove authority and scheme from non-connect direct http(~1) request"
                 );
@@ -93,6 +84,7 @@ fn sanitize_client_req_header<S, B>(
                     uri_parts.path_and_query = Some(PathAndQuery::from_static("/"));
                 }
 
+                // add required host header if not defined
                 if !parts.headers.contains_key(HOST) {
                     parts
                         .headers
@@ -101,8 +93,26 @@ fn sanitize_client_req_header<S, B>(
 
                 parts.uri = rama_http_types::Uri::from_parts(uri_parts)?;
                 Request::from_parts(parts, body)
-            } else if req.uri().host().is_none() && req.version() >= Version::HTTP_2 {
-                // [h2/h3] GET | HEAD | POST | PUT | DELETE | OPTIONS | TRACE | PATCH
+            } else if !req.headers().contains_key(HOST) {
+                let authority = req
+                    .uri()
+                    .authority()
+                    .ok_or_else(|| {
+                        OpaqueError::from_display("missing authority in uri and missing host")
+                    })?
+                    .clone();
+                let mut req = req;
+                req.headers_mut()
+                    .typed_insert(rama_http_types::headers::Host::from(authority));
+                req
+            } else {
+                req
+            }
+        }
+        Version::HTTP_2 => {
+            // set scheme/host if not defined as otherwise pseudo
+            // headers won't be possible to be set in the h2 crate
+            let mut req = if req.uri().host().is_none() {
                 let request_ctx = ctx.get::<RequestContext>().ok_or_else(|| {
                     OpaqueError::from_display("[h2+] add scheme/host: missing RequestCtx")
                         .into_boxed()
@@ -112,8 +122,6 @@ fn sanitize_client_req_header<S, B>(
                     http_version = ?req.version(),
                     "defining authority and scheme to non-connect direct http request"
                 );
-                // setting this information is required for libs such as `h2`
-                // in order to define the pseudo headers correctly
 
                 let (mut parts, body) = req.into_parts();
                 let mut uri_parts = parts.uri.into_parts();
@@ -135,26 +143,31 @@ fn sanitize_client_req_header<S, B>(
                         .context("use RequestContext.authority as http authority")?,
                 );
 
-                for illegal_h2_header in [
-                    &CONNECTION,
-                    &TRANSFER_ENCODING,
-                    &PROXY_CONNECTION,
-                    &UPGRADE,
-                    &KEEP_ALIVE,
-                    &HOST,
-                ] {
-                    if let Some(header) = parts.headers.remove(illegal_h2_header) {
-                        tracing::trace!(?header, "removed illegal (~http1) header from h2 request");
-                    }
-                }
-
                 parts.uri = rama_http_types::Uri::from_parts(uri_parts)
                     .context("create http uri from parts")?;
 
                 Request::from_parts(parts, body)
             } else {
                 req
+            };
+
+            // remove illegal headers
+            for illegal_h2_header in [
+                &CONNECTION,
+                &TRANSFER_ENCODING,
+                &PROXY_CONNECTION,
+                &UPGRADE,
+                &KEEP_ALIVE,
+                &HOST,
+            ] {
+                if let Some(header) = req.headers_mut().remove(illegal_h2_header) {
+                    tracing::trace!(?header, "removed illegal (~http1) header from h2 request");
+                }
             }
+
+            req
         }
+        Version::HTTP_3 => req,
+        _ => req,
     })
 }