Check response trailers

tonic/src/client/grpc.rs

@@ -1,6 +1,6 @@
 use crate::{
     body::{Body, BoxBody},
-    codec::{decode, encode, Codec, Streaming},
+    codec::{decode_response, decode_empty, encode, Codec, Streaming},
     Code, GrpcService, Request, Response, Status,
 };
 use futures_core::Stream;
@@ -136,18 +136,29 @@ impl<T> Grpc<T> {
             .await
             .map_err(|err| Status::from_error(&*(err.into())))?;
 
-        // TODO: implement decode with status
-        let _status_code = response.status();
+        let status_code = response.status();
         let trailers_only_status = Status::from_header_map(response.headers());
 
-        if let Some(status) = trailers_only_status {
+        let expect_additional_trailers = if let Some(status) = trailers_only_status {
             if status.code() != Code::Ok {
                 return Err(status);
             }
-        }
+
+            true
+        } else {
+            false
+        };
 
         let response = response
-            .map(|b| decode(codec.decoder(), b).into_stream())
+            .map(|b| {
+                if expect_additional_trailers {
+                    future::Either::Left(
+                        decode_response(codec.decoder(), b, status_code).into_stream(),
+                    )
+                } else {
+                    future::Either::Right(decode_empty(codec.decoder(), b).into_stream())
+                }
+            })
             .map(Streaming::new);
 
         Ok(Response::from_http(response))

tonic/src/codec/decode.rs

@@ -2,12 +2,52 @@ use crate::{Code, Status};
 use bytes::{Buf, BufMut, BytesMut, IntoBuf};
 use futures_core::{Stream, TryStream};
 use futures_util::future;
-// use http::StatusCode;
+use http::StatusCode;
 use http_body::Body;
 use std::pin::Pin;
 use tokio_codec::Decoder;
 use tracing::{debug, trace};
 
+pub fn decode_request<T, B>(
+    decoder: T,
+    source: B,
+) -> impl TryStream<Ok = T::Item, Error = Status> + 'static
+where
+    T: Decoder<Error = Status> + 'static,
+    T::Item: Unpin + 'static,
+    B: Body + 'static,
+    B::Error: Into<crate::Error>,
+{
+    decode(decoder, source, Direction::Request)
+}
+
+pub fn decode_response<T, B>(
+    decoder: T,
+    source: B,
+    status: StatusCode,
+) -> impl TryStream<Ok = T::Item, Error = Status> + 'static
+where
+    T: Decoder<Error = Status> + 'static,
+    T::Item: Unpin + 'static,
+    B: Body + 'static,
+    B::Error: Into<crate::Error>,
+{
+    decode(decoder, source, Direction::Response(status))
+}
+
+pub fn decode_empty<T, B>(
+    decoder: T,
+    source: B,
+) -> impl TryStream<Ok = T::Item, Error = Status> + 'static
+where
+    T: Decoder<Error = Status> + 'static,
+    T::Item: Unpin + 'static,
+    B: Body + 'static,
+    B::Error: Into<crate::Error>,
+{
+    decode(decoder, source, Direction::EmptyResponse)
+}
+
 pub struct Streaming<T> {
     inner: Pin<Box<dyn Stream<Item = Result<T, Status>> + Send + 'static>>,
 }
@@ -34,15 +74,16 @@ enum State {
     ReadBody { compression: bool, len: usize },
 }
 
-// enum Direction {
-//     Request,
-//     Response(StatusCode),
-//     EmptyResponse,
-// }
+enum Direction {
+    Request,
+    Response(StatusCode),
+    EmptyResponse,
+}
 
-pub fn decode<T, B>(
+fn decode<T, B>(
     mut decoder: T,
     mut source: B,
+    direction: Direction,
 ) -> impl TryStream<Ok = T::Item, Error = Status> + 'static
 where
     T: Decoder<Error = Status> + 'static,
@@ -85,8 +126,22 @@ where
                     break;
                 }
             }
+        }
 
-            // TODO: poll_trailers for Response status code
+        if let Direction::Response(status) = direction {
+            let trailer = future::poll_fn(|cx| unsafe { std::pin::Pin::new_unchecked(&mut source) }.poll_trailers(cx));
+            let trailer = match trailer.await {
+                Ok(trailer) => {
+                    crate::status::infer_grpc_status(trailer, status)?;
+                },
+                Err(e) => {
+                    let err = e.into();
+                    debug!("decoder inner trailers error: {:?}", err);
+                    let status = Status::from_error(&*err);
+                    Err(status)?;
+                },
+                Ok(None) => return,
+            };
         }
     }
 }

tonic/src/codec/mod.rs

@@ -2,7 +2,7 @@ mod decode;
 mod encode;
 mod prost;
 
-pub use self::decode::{decode, Streaming};
+pub use self::decode::{decode_empty, decode_request, decode_response, Streaming};
 pub use self::encode::encode;
 pub use self::prost::ProstCodec;
 

tonic/src/lib.rs

@@ -1,5 +1,5 @@
 #![feature(async_await)]
-#![recursion_limit = "256"]
+#![recursion_limit = "512"]
 
 //! gRPC implementation
 

tonic/src/server/grpc.rs

@@ -1,6 +1,6 @@
 use crate::{
     body::{BoxBody, BytesBuf},
-    codec::{decode, encode, Codec, Streaming},
+    codec::{decode_request, encode, Codec, Streaming},
     server::{ClientStreamingService, ServerStreamingService, StreamingService, UnaryService},
     Code, Request, Response, Status,
 };
@@ -132,7 +132,7 @@ where
         B::Error: Into<crate::Error> + Send,
     {
         let (parts, body) = request.into_parts();
-        let stream = decode(self.codec.decoder(), body).into_stream();
+        let stream = decode_request(self.codec.decoder(), body).into_stream();
 
         futures_util::pin_mut!(stream);
 
@@ -154,7 +154,7 @@ where
         B::Error: Into<crate::Error> + Send,
     {
         Request::from_http(
-            request.map(|b| Streaming::new(decode(self.codec.decoder(), b).into_stream())),
+            request.map(|b| Streaming::new(decode_request(self.codec.decoder(), b).into_stream())),
         )
     }