diff --git a/tonic/src/codec/encode.rs b/tonic/src/codec/encode.rs index 60a8548ce..d94a1f0ba 100644 --- a/tonic/src/codec/encode.rs +++ b/tonic/src/codec/encode.rs @@ -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, buf: &mut BytesMut) -> Bytes { +fn finish_encoding( + compression_encoding: Option, + buf: &mut BytesMut, +) -> Result { 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)] diff --git a/tonic/src/codec/prost.rs b/tonic/src/codec/prost.rs index 678ab96d8..2facddcef 100644 --- a/tonic/src/codec/prost.rs +++ b/tonic/src/codec/prost.rs @@ -136,6 +136,39 @@ mod tests { } } + // skip on windows because CI stumbles over our 4GB allocation + #[cfg(not(target_family = "windows"))] + #[tokio::test] + async fn encode_too_big() { + let encoder = MockEncoder::default(); + + let msg = vec![0u8; u32::MAX as usize + 1]; + + 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;