fix: do not panic while encoding oversized bodies (#1142)

Fixes #1141.

Author: crepererum

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<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)]

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;