From 4835be515fca355ff818ab7a32a7d28c17194a18 Mon Sep 17 00:00:00 2001 From: Lucio Franco Date: Sun, 18 Aug 2019 01:47:40 -0400 Subject: [PATCH] Check response trailers --- tonic/src/client/grpc.rs | 23 +++++++++---- tonic/src/codec/decode.rs | 71 ++++++++++++++++++++++++++++++++++----- tonic/src/codec/mod.rs | 2 +- tonic/src/lib.rs | 2 +- tonic/src/server/grpc.rs | 6 ++-- 5 files changed, 85 insertions(+), 19 deletions(-) diff --git a/tonic/src/client/grpc.rs b/tonic/src/client/grpc.rs index f410d2f..1596dd1 100644 --- a/tonic/src/client/grpc.rs +++ b/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 Grpc { .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)) diff --git a/tonic/src/codec/decode.rs b/tonic/src/codec/decode.rs index 2cb2155..cfdfa0c 100644 --- a/tonic/src/codec/decode.rs +++ b/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( + decoder: T, + source: B, +) -> impl TryStream + 'static +where + T: Decoder + 'static, + T::Item: Unpin + 'static, + B: Body + 'static, + B::Error: Into, +{ + decode(decoder, source, Direction::Request) +} + +pub fn decode_response( + decoder: T, + source: B, + status: StatusCode, +) -> impl TryStream + 'static +where + T: Decoder + 'static, + T::Item: Unpin + 'static, + B: Body + 'static, + B::Error: Into, +{ + decode(decoder, source, Direction::Response(status)) +} + +pub fn decode_empty( + decoder: T, + source: B, +) -> impl TryStream + 'static +where + T: Decoder + 'static, + T::Item: Unpin + 'static, + B: Body + 'static, + B::Error: Into, +{ + decode(decoder, source, Direction::EmptyResponse) +} + pub struct Streaming { inner: Pin> + 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( +fn decode( mut decoder: T, mut source: B, + direction: Direction, ) -> impl TryStream + 'static where T: Decoder + '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, + }; } } } diff --git a/tonic/src/codec/mod.rs b/tonic/src/codec/mod.rs index 352bd95..0337786 100644 --- a/tonic/src/codec/mod.rs +++ b/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; diff --git a/tonic/src/lib.rs b/tonic/src/lib.rs index ab50e25..da4a27c 100644 --- a/tonic/src/lib.rs +++ b/tonic/src/lib.rs @@ -1,5 +1,5 @@ #![feature(async_await)] -#![recursion_limit = "256"] +#![recursion_limit = "512"] //! gRPC implementation diff --git a/tonic/src/server/grpc.rs b/tonic/src/server/grpc.rs index 710575d..c2fc273 100644 --- a/tonic/src/server/grpc.rs +++ b/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 + 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 + 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())), ) }