From bfab56d42d776f1cef8a2e82f5eeb13c143130bc Mon Sep 17 00:00:00 2001
From: Lucio Franco <luciofranco14@gmail.com>
Date: Wed, 28 Aug 2019 20:09:43 -0400
Subject: [PATCH] Add trailing metadata and fix test

---
 tonic-interop/src/client.rs | 12 +++++--
 tonic-interop/test.sh       |  4 +--
 tonic/src/codec/decode.rs   | 64 ++++++++++++++++++++++++-------------
 tonic/src/status.rs         |  2 +-
 4 files changed, 54 insertions(+), 28 deletions(-)

diff --git a/tonic-interop/src/client.rs b/tonic-interop/src/client.rs
index 7d46f57..064a194 100644
--- a/tonic-interop/src/client.rs
+++ b/tonic-interop/src/client.rs
@@ -418,6 +418,7 @@ pub async fn custom_metadata(client: &mut Client, assertions: &mut Vec<TestAsser
     //     format!("result={:?}", response.metadata().get_bin(key1))
     // ));
 
+
     let response = client
         .full_duplex_call(req_stream)
         .await
@@ -428,10 +429,17 @@ pub async fn custom_metadata(client: &mut Client, assertions: &mut Vec<TestAsser
         response.metadata().get(key1) == Some(&value1),
         format!("result={:?}", response.metadata().get(key1))
     ));
+
+    let mut stream = response.into_inner();
+
+    // while let Some(_) = stream.next().await {}
+
+    let trailers = stream.trailers().await.unwrap().unwrap();
+
     assertions.push(test_assert!(
         "metadata bin must match in unary",
-        response.metadata().get_bin(key2) == Some(&value2),
-        format!("result={:?}", response.metadata().get_bin(key1))
+        trailers.get_bin(key2) == Some(&value2),
+        format!("result={:?}", trailers.get_bin(key1))
     ));
 }
 
diff --git a/tonic-interop/test.sh b/tonic-interop/test.sh
index 43ef7af..bacf8d2 100755
--- a/tonic-interop/test.sh
+++ b/tonic-interop/test.sh
@@ -1,6 +1,4 @@
  cargo run -p tonic-interop --bin client -- \
  --test_case=empty_unary,large_unary,client_streaming,server_streaming,ping_pong,\
 empty_stream,status_code_and_message,special_status_message,unimplemented_method,\
-unimplemented_service
-
-# DISABLED: ,custom_metadata
+unimplemented_service,custom_metadata
diff --git a/tonic/src/codec/decode.rs b/tonic/src/codec/decode.rs
index 998bd5e..481b92f 100644
--- a/tonic/src/codec/decode.rs
+++ b/tonic/src/codec/decode.rs
@@ -19,6 +19,7 @@ pub struct Streaming<T> {
     state: State,
     direction: Direction,
     buf: BytesMut,
+    trailers: Option<MetadataMap>,
 }
 
 impl<T> Unpin for Streaming<T> {}
@@ -44,14 +45,7 @@ impl<T> Streaming<T> {
         B::Error: Into<crate::Error>,
         D: Decoder<Item = T, Error = Status> + Send + 'static,
     {
-        Self {
-            decoder: Box::new(decoder),
-            body: BoxBody::map_from(body),
-            state: State::ReadHeader,
-            direction: Direction::Response(status_code),
-            // FIXME: update this with a reasonable size
-            buf: BytesMut::with_capacity(1024 * 1024),
-        }
+        Self::new(decoder, body, Direction::Response(status_code))
     }
 
     pub fn new_empty<B, D>(decoder: D, body: B) -> Self
@@ -61,17 +55,19 @@ impl<T> Streaming<T> {
         B::Error: Into<crate::Error>,
         D: Decoder<Item = T, Error = Status> + Send + 'static,
     {
-        Self {
-            decoder: Box::new(decoder),
-            body: BoxBody::map_from(body),
-            state: State::ReadHeader,
-            direction: Direction::EmptyResponse,
-            // FIXME: update this with a reasonable size
-            buf: BytesMut::with_capacity(1024 * 1024),
-        }
+        Self::new(decoder, body, Direction::EmptyResponse)
     }
 
     pub fn new_request<B, D>(decoder: D, body: B) -> Self
+    where
+        B: Body + Send + 'static,
+        B::Data: Into<Bytes>,
+        B::Error: Into<crate::Error>,
+        D: Decoder<Item = T, Error = Status> + Send + 'static,
+    {
+        Self::new(decoder, body, Direction::Request)
+    }
+    fn new<B, D>(decoder: D, body: B, direction: Direction) -> Self
     where
         B: Body + Send + 'static,
         B::Data: Into<Bytes>,
@@ -82,23 +78,45 @@ impl<T> Streaming<T> {
             decoder: Box::new(decoder),
             body: BoxBody::map_from(body),
             state: State::ReadHeader,
-            direction: Direction::Request,
+            direction,
             // FIXME: update this with a reasonable size
             buf: BytesMut::with_capacity(1024 * 1024),
+            trailers: None,
         }
     }
 }
 
 impl<T> Streaming<T> {
-    // pub async fn message(&mut self) -> Option<Result<T::Item, Status>> {
-    //     future::poll_fn(|cx| Pin::new(&mut *self).poll_next(cx)).await
-    // }
+    pub async fn message(&mut self) -> Option<Result<T, Status>> {
+        future::poll_fn(|cx| Pin::new(&mut *self).poll_next(cx)).await
+    }
 
-    pub async fn trailers(&mut self) -> Result<Option<MetadataMap>, Status> {
+    pub async fn trailers(mut self) -> Result<Option<MetadataMap>, Status> {
+        // Shortcut to see if we already pulled the trailers in the stream step
+        // we need to do that so that the stream can error on trailing grpc-status
+        if let Some(trailers) = self.trailers {
+            return Ok(Some(trailers));
+        }
+
+        // To fetch the trailers we must clear the body and drop it.
+        while let Some(res) = self.message().await {
+            res?;
+        }
+
+        // Since we call poll_trailers internally on poll_next we need to
+        // check if it got cached again.
+        if let Some(trailers) = self.trailers {
+            return Ok(Some(trailers));
+        }
+
+
+        // Trailers were not caught during poll_next and thus lets poll for
+        // them manually.
         let map =
             future::poll_fn(|cx| unsafe { Pin::new_unchecked(&mut self.body) }.poll_trailers(cx))
                 .await
                 .map_err(|e| Status::from_error(&e))?;
+
         Ok(map.map(MetadataMap::from_headers))
     }
 
@@ -205,8 +223,10 @@ impl<T> Stream for Streaming<T> {
         if let Direction::Response(status) = self.direction {
             match ready!(unsafe { Pin::new_unchecked(&mut self.body) }.poll_trailers(cx)) {
                 Ok(trailer) => {
-                    if let Err(e) = crate::status::infer_grpc_status(trailer, status) {
+                    if let Err(e) = crate::status::infer_grpc_status(trailer.as_ref(), status) {
                         return Some(Err(e)).into();
+                    } else {
+                        self.trailers = trailer.map(MetadataMap::from_headers);
                     }
                 }
                 Err(e) => {
diff --git a/tonic/src/status.rs b/tonic/src/status.rs
index 49ec386..519b75f 100644
--- a/tonic/src/status.rs
+++ b/tonic/src/status.rs
@@ -312,7 +312,7 @@ impl Error for Status {}
 /// Take the `Status` value from `trailers` if it is available, else from `status_code`.
 ///
 pub(crate) fn infer_grpc_status(
-    trailers: Option<HeaderMap>,
+    trailers: Option<&HeaderMap>,
     status_code: http::StatusCode,
 ) -> Result<(), Status> {
     if let Some(trailers) = trailers {