From c6f095e11a8cf73288f9b5d9308258153ea876ef Mon Sep 17 00:00:00 2001
From: David Pedersen <david.pdrsn@gmail.com>
Date: Thu, 8 Jul 2021 11:45:29 +0200
Subject: [PATCH] fix(tonic): don't remove reserved headers in interceptor

`InterceptedService` would previously use `tonic::Request::into_http`
which removes reserved headers. That mean inner middleware in the stack
wouldn't be able to see those headers, which could result in errors.

Fixes https://github.com/hyperium/tonic/issues/700
---
 tonic/src/client/grpc.rs         |  3 +-
 tonic/src/request.rs             | 20 +++++++++++--
 tonic/src/service/interceptor.rs | 49 ++++++++++++++++++++++++++++----
 3 files changed, 62 insertions(+), 10 deletions(-)

diff --git a/tonic/src/client/grpc.rs b/tonic/src/client/grpc.rs
index c89c4c60f..2da5fbf14 100644
--- a/tonic/src/client/grpc.rs
+++ b/tonic/src/client/grpc.rs
@@ -4,6 +4,7 @@ use crate::{
     body::BoxBody,
     client::GrpcService,
     codec::{encode_client, Codec, Streaming},
+    request::SanitizeHeaders,
     Code, Request, Response, Status,
 };
 use futures_core::Stream;
@@ -247,7 +248,7 @@ impl<T> Grpc<T> {
             })
             .map(BoxBody::new);
 
-        let mut request = request.into_http(uri);
+        let mut request = request.into_http(uri, SanitizeHeaders::Yes);
 
         // Add the gRPC related HTTP headers
         request
diff --git a/tonic/src/request.rs b/tonic/src/request.rs
index 212a43b16..b5aae8c03 100644
--- a/tonic/src/request.rs
+++ b/tonic/src/request.rs
@@ -166,13 +166,20 @@ impl<T> Request<T> {
         Request::from_http_parts(parts, message)
     }
 
-    pub(crate) fn into_http(self, uri: http::Uri) -> http::Request<T> {
+    pub(crate) fn into_http(
+        self,
+        uri: http::Uri,
+        sanitize_headers: SanitizeHeaders,
+    ) -> http::Request<T> {
         let mut request = http::Request::new(self.message);
 
         *request.version_mut() = http::Version::HTTP_2;
         *request.method_mut() = http::Method::POST;
         *request.uri_mut() = uri;
-        *request.headers_mut() = self.metadata.into_sanitized_headers();
+        *request.headers_mut() = match sanitize_headers {
+            SanitizeHeaders::Yes => self.metadata.into_sanitized_headers(),
+            SanitizeHeaders::No => self.metadata.into_headers(),
+        };
         *request.extensions_mut() = self.extensions.into_http();
 
         request
@@ -412,6 +419,13 @@ fn duration_to_grpc_timeout(duration: Duration) -> String {
         .expect("duration is unrealistically large")
 }
 
+/// When converting a `tonic::Request` into a `http::Request` should reserved
+/// headers be removed?
+pub(crate) enum SanitizeHeaders {
+    Yes,
+    No,
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -427,7 +441,7 @@ mod tests {
                 .insert(*header, MetadataValue::from_static("invalid"));
         }
 
-        let http_request = r.into_http(Uri::default());
+        let http_request = r.into_http(Uri::default(), SanitizeHeaders::Yes);
         assert!(http_request.headers().is_empty());
     }
 
diff --git a/tonic/src/service/interceptor.rs b/tonic/src/service/interceptor.rs
index 831fe67e4..1158ee040 100644
--- a/tonic/src/service/interceptor.rs
+++ b/tonic/src/service/interceptor.rs
@@ -1,6 +1,6 @@
 //! gRPC interceptors which are a kind of middleware.
 
-use crate::Status;
+use crate::{request::SanitizeHeaders, Status};
 use pin_project::pin_project;
 use std::{
     fmt,
@@ -115,7 +115,7 @@ where
             Ok(req) => {
                 let (metadata, extensions, _) = req.into_parts();
                 let req = crate::Request::from_parts(metadata, extensions, msg);
-                let req = req.into_http(uri);
+                let req = req.into_http(uri, SanitizeHeaders::No);
                 ResponseFuture::future(self.inner.call(req))
             }
             Err(status) => ResponseFuture::error(status),
@@ -170,10 +170,7 @@ where
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         match self.project().kind.project() {
-            KindProj::Future(future) => {
-                let response = futures_core::ready!(future.poll(cx).map_err(Into::into)?);
-                Poll::Ready(Ok(response))
-            }
+            KindProj::Future(future) => future.poll(cx).map_err(Into::into),
             KindProj::Error(status) => {
                 let error = status.take().unwrap().into();
                 Poll::Ready(Err(error))
@@ -181,3 +178,43 @@ where
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    #[allow(unused_imports)]
+    use super::*;
+    use tower::ServiceExt;
+
+    #[tokio::test]
+    async fn doesnt_remove_headers() {
+        let svc = tower::service_fn(|request: http::Request<hyper::Body>| async move {
+            assert_eq!(
+                request
+                    .headers()
+                    .get("user-agent")
+                    .expect("missing in leaf service"),
+                "test-tonic"
+            );
+
+            Ok::<_, hyper::Error>(hyper::Response::new(hyper::Body::empty()))
+        });
+
+        let svc = InterceptedService::new(svc, |request: crate::Request<()>| {
+            assert_eq!(
+                request
+                    .metadata()
+                    .get("user-agent")
+                    .expect("missing in interceptor"),
+                "test-tonic"
+            );
+            Ok(request)
+        });
+
+        let request = http::Request::builder()
+            .header("user-agent", "test-tonic")
+            .body(hyper::Body::empty())
+            .unwrap();
+
+        svc.oneshot(request).await.unwrap();
+    }
+}