From 6cc706ac87cfa63d6c1039f090f753e1fd23aa4e Mon Sep 17 00:00:00 2001 From: Lucio Franco Date: Thu, 21 Sep 2023 15:19:04 -0400 Subject: [PATCH] fix(web): Buffer incomplete messages Previously, if we did not find trailers in the incoming message and the message len was longer than the actual buffer we returned the buffer. This would mean we would fail to parse incomplete message buffers. This fixes `find_trailers` to support three return types, `Trailers(len)`, `IncompleteBuf`, `Done`. Previously, there was only `Some(len)` and `None`. --- tonic-web/src/call.rs | 148 ++++++++++++++++++++---------------------- 1 file changed, 72 insertions(+), 76 deletions(-) diff --git a/tonic-web/src/call.rs b/tonic-web/src/call.rs index 4f0f4a335..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() @@ -521,57 +536,38 @@ mod tests { } #[test] - fn find_trailers_buffered_case2() { + fn find_trailers_buffered_incomplete_message() { let buf = vec![ - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 48, 4, 2, 105, 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, 48, 3, 2, 105, 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, 48, 2, 2, 105, 116, 104, + 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, 48, 1, 2, 105, 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, 128, 0, 0, 0, 15, 103, 114, 112, 99, 45, 115, 116, 97, 116, 117, 115, 58, 48, 13, - 10, + 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, Some(81)); - - let trailers = decode_trailers_frame(Bytes::copy_from_slice(&buf[81..])) - .unwrap() - .unwrap(); - let status = trailers.get("grpc-status").unwrap(); - assert_eq!(status.to_str().unwrap(), "0") + assert_eq!(out, FindTrailers::IncompleteBuf); } }