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: do not panic while encoding oversized bodies #1142

Merged
merged 2 commits into from Nov 10, 2022
Merged
Show file tree
Hide file tree
Changes from 1 commit
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
15 changes: 11 additions & 4 deletions tonic/src/codec/encode.rs
Expand Up @@ -119,19 +119,26 @@ where
}

// now that we know length, we can write the header
Ok(finish_encoding(compression_encoding, buf))
finish_encoding(compression_encoding, buf)
}

fn finish_encoding(compression_encoding: Option<CompressionEncoding>, buf: &mut BytesMut) -> Bytes {
fn finish_encoding(
compression_encoding: Option<CompressionEncoding>,
buf: &mut BytesMut,
) -> Result<Bytes, Status> {
let len = buf.len() - HEADER_SIZE;
assert!(len <= std::u32::MAX as usize);
if len > std::u32::MAX as usize {
return Err(Status::resource_exhausted(format!(
"Cannot return body with more than 4GB of data but got {len} bytes"
)));
}
{
let mut buf = &mut buf[..HEADER_SIZE];
buf.put_u8(compression_encoding.is_some() as u8);
buf.put_u32(len as u32);
}

buf.split_to(len + HEADER_SIZE).freeze()
Ok(buf.split_to(len + HEADER_SIZE).freeze())
}

#[derive(Debug)]
Expand Down
31 changes: 31 additions & 0 deletions tonic/src/codec/prost.rs
Expand Up @@ -136,6 +136,37 @@ mod tests {
}
}

#[tokio::test]
async fn encode_too_big() {
let encoder = MockEncoder::default();

let msg = vec![0u8; u32::MAX as usize + 1];
LucioFranco marked this conversation as resolved.
Show resolved Hide resolved

let messages = std::iter::once(Ok::<_, Status>(msg));
let source = futures_util::stream::iter(messages);

let body = encode_server(
encoder,
source,
None,
SingleMessageCompressionOverride::default(),
);

futures_util::pin_mut!(body);

assert!(body.data().await.is_none());
assert_eq!(
body.trailers()
.await
.expect("no error polling trailers")
.expect("some trailers")
.get("grpc-status")
.expect("grpc-status header"),
"8"
);
assert!(body.is_end_stream());
}

#[derive(Debug, Clone, Default)]
struct MockEncoder;

Expand Down