Skip to content

Commit

Permalink
fix(web): Buffer incomplete messages
Browse files Browse the repository at this point in the history 
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`.
  • Loading branch information
@LucioFranco
LucioFranco committed Sep 21, 2023
1 parent 6af8d5c commit 38d6988
Showing 1 changed file with 82 additions and 31 deletions.
113 changes: 82 additions & 31 deletions tonic-web/src/call.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -242,33 +242,41 @@ where
cx: &mut Context<'_>,
) -> Poll<Option<Result<Self::Data, Self::Error>>> {
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 {
Expand Down Expand Up @@ -413,21 +421,21 @@ fn make_trailers_frame(trailers: HeaderMap) -> Vec<u8> {
/// 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<usize> {
fn find_trailers(buf: &[u8]) -> FindTrailers {
let mut len = 0;
let mut temp_buf = &buf[..];

loop {
// 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();
Expand All @@ -437,13 +445,20 @@ fn find_trailers(buf: &[u8]) -> Option<usize> {
// 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::*;
Expand Down Expand Up @@ -494,7 +509,7 @@ mod tests {

let out = find_trailers(&buf[..]);

assert_eq!(out, Some(0));
assert_eq!(out, FindTrailers::Trailer(0));
}

#[test]
Expand All @@ -511,12 +526,48 @@ 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()
.unwrap();
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);
}
}

0 comments on commit 38d6988

Please sign in to comment.