diff --git a/tonic-interop/src/client.rs b/tonic-interop/src/client.rs index 064a194..a993aef 100644 --- a/tonic-interop/src/client.rs +++ b/tonic-interop/src/client.rs @@ -402,22 +402,21 @@ pub async fn custom_metadata(client: &mut Client, assertions: &mut Vec Grpc { M1: Send, M2: Send + Unpin + 'static, { - let (parts, body) = self.streaming(request, path, codec).await?.into_parts(); + let (mut parts, body) = self.streaming(request, path, codec).await?.into_parts(); futures_util::pin_mut!(body); @@ -86,6 +86,10 @@ impl Grpc { .await? .ok_or(Status::new(Code::Internal, "Missing response message."))?; + if let Some(trailers) = body.trailers().await? { + parts.merge(trailers); + } + Ok(Response::from_parts(parts, message)) } diff --git a/tonic/src/codec/decode.rs b/tonic/src/codec/decode.rs index 481b92f..b092dd5 100644 --- a/tonic/src/codec/decode.rs +++ b/tonic/src/codec/decode.rs @@ -91,10 +91,10 @@ impl Streaming { future::poll_fn(|cx| Pin::new(&mut *self).poll_next(cx)).await } - pub async fn trailers(mut self) -> Result, Status> { + pub async fn trailers(&mut self) -> Result, 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 { + if let Some(trailers) = self.trailers.take() { return Ok(Some(trailers)); } @@ -105,11 +105,10 @@ impl Streaming { // Since we call poll_trailers internally on poll_next we need to // check if it got cached again. - if let Some(trailers) = self.trailers { + if let Some(trailers) = self.trailers.take() { return Ok(Some(trailers)); } - // Trailers were not caught during poll_next and thus lets poll for // them manually. let map = diff --git a/tonic/src/metadata/map.rs b/tonic/src/metadata/map.rs index a1484dc..cab445d 100644 --- a/tonic/src/metadata/map.rs +++ b/tonic/src/metadata/map.rs @@ -1181,6 +1181,10 @@ impl MetadataMap { { key.remove(self) } + + pub(crate) fn merge(&mut self, other: MetadataMap) { + self.headers.extend(other.headers); + } } // ===== impl Iter ===== diff --git a/tonic/src/server/grpc.rs b/tonic/src/server/grpc.rs index 95eb7ad..aa43943 100644 --- a/tonic/src/server/grpc.rs +++ b/tonic/src/server/grpc.rs @@ -136,7 +136,14 @@ where .await? .ok_or(Status::new(Code::Internal, "Missing request message."))?; - Ok(Request::from_http_parts(parts, message)) + + let mut req = Request::from_http_parts(parts, message); + + if let Some(trailers) = stream.trailers().await? { + req.metadata_mut().merge(trailers); + } + + Ok(req) } fn map_request_streaming( @@ -170,8 +177,6 @@ where let body = encode_server(self.codec.encoder(), body.into_stream()); - // FIXME: try to return impl Trait? - // let body = Box::pin(body) as BoxStream; http::Response::from_parts(parts, BoxBody::new(body)) } Err(status) => {