diff --git a/tonic-web/src/call.rs b/tonic-web/src/call.rs index c53c99720..353e644e5 100644 --- a/tonic-web/src/call.rs +++ b/tonic-web/src/call.rs @@ -242,33 +242,41 @@ where cx: &mut Context<'_>, ) -> Poll>> { if self.client && self.direction == Direction::Decode { - let buf = ready!(self.as_mut().poll_decode(cx)); - - return if let Some(Ok(mut buf)) = buf { - // We found some trailers so extract them since we - // want to return them via `poll_trailers`. - if let Some(len) = find_trailers(&buf[..]) { - // Extract up to len of where the trailers are at - let msg_buf = buf.copy_to_bytes(len); - match decode_trailers_frame(buf) { - Ok(Some(trailers)) => { - self.project().trailers.replace(trailers); + let mut me = self.as_mut(); + + loop { + let buf = ready!(me.as_mut().poll_decode(cx)); + + return if let Some(Ok(incoming_buf)) = buf { + let buf = &mut me.as_mut().project().buf; + + buf.put(incoming_buf); + + match find_trailers(&buf[..]) { + FindTrailers::Trailer(len) => { + // Extract up to len of where the trailers are at + let msg_buf = buf.copy_to_bytes(len); + match decode_trailers_frame(buf.split().freeze()) { + Ok(Some(trailers)) => { + self.project().trailers.replace(trailers); + } + Err(e) => return Poll::Ready(Some(Err(e))), + _ => {} + } + + if msg_buf.has_remaining() { + return Poll::Ready(Some(Ok(msg_buf))); + } else { + return Poll::Ready(None); + } } - Err(e) => return Poll::Ready(Some(Err(e))), - _ => {} + FindTrailers::IncompleteBuf => continue, + FindTrailers::Done => Poll::Ready(Some(Ok(buf.split().freeze()))), } - - if msg_buf.has_remaining() { - return Poll::Ready(Some(Ok(msg_buf))); - } else { - return Poll::Ready(None); - } - } - - Poll::Ready(Some(Ok(buf))) - } else { - Poll::Ready(buf) - }; + } else { + Poll::Ready(buf) + }; + } } match self.direction { @@ -413,7 +421,7 @@ fn make_trailers_frame(trailers: HeaderMap) -> Vec { /// its location in the original buf. If `None` is returned we did /// not find a trailers in this buffer either because its incomplete /// or the buffer jsut contained grpc message frames. -fn find_trailers(buf: &[u8]) -> Option { +fn find_trailers(buf: &[u8]) -> FindTrailers { let mut len = 0; let mut temp_buf = &buf[..]; @@ -421,13 +429,13 @@ fn find_trailers(buf: &[u8]) -> Option { // To check each frame, there must be at least GRPC_HEADER_SIZE // amount of bytes available otherwise the buffer is incomplete. if temp_buf.is_empty() || temp_buf.len() < GRPC_HEADER_SIZE { - return None; + return FindTrailers::Done; } let header = temp_buf.get_u8(); if header == GRPC_WEB_TRAILERS_BIT { - return Some(len); + return FindTrailers::Trailer(len); } let msg_len = temp_buf.get_u32(); @@ -437,13 +445,20 @@ fn find_trailers(buf: &[u8]) -> Option { // If the msg len of a non-grpc-web trailer frame is larger than // the overall buffer we know within that buffer there are no trailers. if len > buf.len() { - return None; + return FindTrailers::IncompleteBuf; } temp_buf = &buf[len as usize..]; } } +#[derive(Debug, PartialEq, Eq)] +enum FindTrailers { + Trailer(usize), + IncompleteBuf, + Done, +} + #[cfg(test)] mod tests { use super::*; @@ -494,7 +509,7 @@ mod tests { let out = find_trailers(&buf[..]); - assert_eq!(out, Some(0)); + assert_eq!(out, FindTrailers::Trailer(0)); } #[test] @@ -511,7 +526,7 @@ mod tests { let out = find_trailers(&buf[..]); - assert_eq!(out, Some(81)); + assert_eq!(out, FindTrailers::Trailer(81)); let trailers = decode_trailers_frame(Bytes::copy_from_slice(&buf[81..])) .unwrap() @@ -519,4 +534,40 @@ mod tests { let status = trailers.get("grpc-status").unwrap(); assert_eq!(status.to_str().unwrap(), "0") } + + #[test] + fn find_trailers_buffered_incomplete_message() { + let buf = vec![ + 0, 0, 0, 9, 238, 10, 233, 19, 18, 230, 19, 10, 9, 10, 1, 120, 26, 4, 84, 69, 88, 84, + 18, 60, 10, 58, 10, 56, 3, 0, 0, 0, 44, 0, 0, 0, 0, 0, 0, 0, 116, 104, 105, 115, 32, + 118, 97, 108, 117, 101, 32, 119, 97, 115, 32, 119, 114, 105, 116, 116, 101, 110, 32, + 118, 105, 97, 32, 119, 114, 105, 116, 101, 32, 100, 101, 108, 101, 103, 97, 116, 105, + 111, 110, 33, 18, 62, 10, 60, 10, 58, 3, 0, 0, 0, 46, 0, 0, 0, 0, 0, 0, 0, 116, 104, + 105, 115, 32, 118, 97, 108, 117, 101, 32, 119, 97, 115, 32, 119, 114, 105, 116, 116, + 101, 110, 32, 98, 121, 32, 97, 110, 32, 101, 109, 98, 101, 100, 100, 101, 100, 32, 114, + 101, 112, 108, 105, 99, 97, 33, 18, 62, 10, 60, 10, 58, 3, 0, 0, 0, 46, 0, 0, 0, 0, 0, + 0, 0, 116, 104, 105, 115, 32, 118, 97, 108, 117, 101, 32, 119, 97, 115, 32, 119, 114, + 105, 116, 116, 101, 110, 32, 98, 121, 32, 97, 110, 32, 101, 109, 98, 101, 100, 100, + 101, 100, 32, 114, 101, 112, 108, 105, 99, 97, 33, 18, 62, 10, 60, 10, 58, 3, 0, 0, 0, + 46, 0, 0, 0, 0, 0, 0, 0, 116, 104, 105, 115, 32, 118, 97, 108, 117, 101, 32, 119, 97, + 115, 32, 119, 114, 105, 116, 116, 101, 110, 32, 98, 121, 32, 97, 110, 32, 101, 109, 98, + 101, 100, 100, 101, 100, 32, 114, 101, 112, 108, 105, 99, 97, 33, 18, 62, 10, 60, 10, + 58, 3, 0, 0, 0, 46, 0, 0, 0, 0, 0, 0, 0, 116, 104, 105, 115, 32, 118, 97, 108, 117, + 101, 32, 119, 97, 115, 32, 119, 114, 105, 116, 116, 101, 110, 32, 98, 121, 32, 97, 110, + 32, 101, 109, 98, 101, 100, 100, 101, 100, 32, 114, 101, 112, 108, 105, 99, 97, 33, 18, + 62, 10, 60, 10, 58, 3, 0, 0, 0, 46, 0, 0, 0, 0, 0, 0, 0, 116, 104, 105, 115, 32, 118, + 97, 108, 117, 101, 32, 119, 97, 115, 32, 119, 114, 105, 116, 116, 101, 110, 32, 98, + 121, 32, 97, 110, 32, 101, 109, 98, 101, 100, 100, 101, 100, 32, 114, 101, 112, 108, + 105, 99, 97, 33, 18, 62, 10, 60, 10, 58, 3, 0, 0, 0, 46, 0, 0, 0, 0, 0, 0, 0, 116, 104, + 105, 115, 32, 118, 97, 108, 117, 101, 32, 119, 97, 115, 32, 119, 114, 105, 116, 116, + 101, 110, 32, 98, 121, 32, 97, 110, 32, 101, 109, 98, 101, 100, 100, 101, 100, 32, 114, + 101, 112, 108, 105, 99, 97, 33, 18, 62, 10, 60, 10, 58, 3, 0, 0, 0, 46, 0, 0, 0, 0, 0, + 0, 0, 116, 104, 105, 115, 32, 118, 97, 108, 117, 101, 32, 119, 97, 115, 32, 119, 114, + 105, 116, 116, 101, 110, 32, 98, 121, 32, + ]; + + let out = find_trailers(&buf[..]); + + assert_eq!(out, FindTrailers::IncompleteBuf); + } }