Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix(web): Trailers decode bug #1528

Merged
merged 1 commit into from
Sep 21, 2023
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
113 changes: 82 additions & 31 deletions tonic-web/src/call.rs
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);
}
}