fix(tonic): respect max_message_size when decompressing a message (#2484)

This patch fixes a bug where the configured max_message_size is not
respected when a payload is compressed, leading the possibility of over
allocating and exhausting memory over the configured maximum message
size when decompressing the message.

## Motivation

When compression is enabled, and a compressed message is sent or
received, the configured max_message_size (for both clients and servers)
is only checked against the compressed message size and that limit is
not respected for the resulting uncompressed message, which can lead to
resource exhaustion.

## Solution

Respect the configured, or default, max_message_size limit while
decompressing a message, returning an error if the resultant
decompressed message would exceed the limit.

Author: bmwill

tests/compression/src/compressing_request.rs

@@ -234,3 +234,71 @@ async fn client_mark_compressed_without_header_server_enabled(encoding: Compress
         "protocol error: received message with compressed-flag but no grpc-encoding was specified"
     );
 }
+
+util::parametrized_tests! {
+    limit_decoded_message_size,
+    zstd: CompressionEncoding::Zstd,
+    gzip: CompressionEncoding::Gzip,
+    deflate: CompressionEncoding::Deflate,
+}
+
+#[cfg(test)]
+async fn limit_decoded_message_size(encoding: CompressionEncoding) {
+    use prost::Message;
+
+    let under_limit_request = SomeData {
+        data: [0_u8; UNCOMPRESSED_MIN_BODY_SIZE].to_vec(),
+    };
+    let limit = under_limit_request.encoded_len();
+    let over_limit_request = SomeData {
+        data: [0_u8; 1 + UNCOMPRESSED_MIN_BODY_SIZE].to_vec(),
+    };
+
+    let (client, server) = tokio::io::duplex(UNCOMPRESSED_MIN_BODY_SIZE * 10);
+
+    let svc = test_server::TestServer::new(Svc::default())
+        .accept_compressed(encoding)
+        .max_decoding_message_size(limit);
+
+    let request_bytes_counter = Arc::new(AtomicUsize::new(0));
+
+    tokio::spawn({
+        let request_bytes_counter = request_bytes_counter.clone();
+        async move {
+            Server::builder()
+                .layer(
+                    ServiceBuilder::new()
+                        .layer(
+                            ServiceBuilder::new()
+                                .layer(measure_request_body_size_layer(request_bytes_counter))
+                                .into_inner(),
+                        )
+                        .into_inner(),
+                )
+                .add_service(svc)
+                .serve_with_incoming(tokio_stream::once(Ok::<_, std::io::Error>(server)))
+                .await
+                .unwrap();
+        }
+    });
+
+    let mut client =
+        test_client::TestClient::new(mock_io_channel(client).await).send_compressed(encoding);
+
+    for _ in 0..3 {
+        // compressed messages that are under or exactly at the limit are successful.
+        client
+            .compress_input_unary(under_limit_request.clone())
+            .await
+            .unwrap();
+        let bytes_sent = request_bytes_counter.load(SeqCst);
+        assert!(bytes_sent < UNCOMPRESSED_MIN_BODY_SIZE);
+
+        // compressed messages that are over the limit are fail with resource exhausted
+        let status = client
+            .compress_input_unary(over_limit_request.clone())
+            .await
+            .unwrap_err();
+        assert_eq!(status.code(), tonic::Code::ResourceExhausted);
+    }
+}

tests/compression/src/compressing_response.rs

@@ -488,3 +488,78 @@ async fn disabling_compression_on_response_from_client_stream(encoding: Compress
     let bytes_sent = response_bytes_counter.load(SeqCst);
     assert!(bytes_sent > UNCOMPRESSED_MIN_BODY_SIZE);
 }
+
+util::parametrized_tests! {
+    limit_decoded_message_size,
+    zstd: CompressionEncoding::Zstd,
+    gzip: CompressionEncoding::Gzip,
+    deflate: CompressionEncoding::Deflate,
+}
+
+#[cfg(test)]
+async fn limit_decoded_message_size(encoding: CompressionEncoding) {
+    use prost::Message;
+
+    let under_limit_request = SomeData {
+        data: [0_u8; UNCOMPRESSED_MIN_BODY_SIZE].to_vec(),
+    };
+    let limit = under_limit_request.encoded_len();
+
+    let (client, server) = tokio::io::duplex(UNCOMPRESSED_MIN_BODY_SIZE * 10);
+
+    let svc = test_server::TestServer::new(Svc::default()).send_compressed(encoding);
+
+    let response_bytes_counter = Arc::new(AtomicUsize::new(0));
+
+    tokio::spawn({
+        let response_bytes_counter = response_bytes_counter.clone();
+        async move {
+            Server::builder()
+                .layer(
+                    ServiceBuilder::new()
+                        .layer(MapResponseBodyLayer::new(move |body| {
+                            util::CountBytesBody {
+                                inner: body,
+                                counter: response_bytes_counter.clone(),
+                            }
+                        }))
+                        .into_inner(),
+                )
+                .add_service(svc)
+                .serve_with_incoming(tokio_stream::once(Ok::<_, std::io::Error>(server)))
+                .await
+                .unwrap();
+        }
+    });
+
+    let expected = match encoding {
+        CompressionEncoding::Gzip => "gzip",
+        CompressionEncoding::Zstd => "zstd",
+        CompressionEncoding::Deflate => "deflate",
+        _ => panic!("unexpected encoding {encoding:?}"),
+    };
+
+    // compressed messages that are under or exactly at the limit are successful.
+    let mut under_limit_client = test_client::TestClient::new(mock_io_channel(client).await)
+        .accept_compressed(encoding)
+        .max_decoding_message_size(limit);
+
+    for _ in 0..3 {
+        let res = under_limit_client.compress_output_unary(()).await.unwrap();
+        assert_eq!(res.metadata().get("grpc-encoding").unwrap(), expected);
+        let bytes_sent = response_bytes_counter.load(SeqCst);
+        assert!(bytes_sent < UNCOMPRESSED_MIN_BODY_SIZE);
+    }
+
+    // compressed messages that are over the limit are fail with resource exhausted
+    let mut over_limit_client = under_limit_client.max_decoding_message_size(limit - 1);
+
+    for _ in 0..3 {
+        let status = over_limit_client
+            .compress_output_unary(())
+            .await
+            .unwrap_err();
+
+        assert_eq!(status.code(), tonic::Code::ResourceExhausted);
+    }
+}

tonic/src/codec/compression.rs

@@ -256,14 +256,17 @@ pub(crate) fn compress(
 pub(crate) fn decompress(
     settings: CompressionSettings,
     compressed_buf: &mut BytesMut,
-    out_buf: &mut BytesMut,
+    mut out_buf: bytes::buf::Limit<&mut BytesMut>,
     len: usize,
 ) -> Result<(), std::io::Error> {
     let buffer_growth_interval = settings.buffer_growth_interval;
     let estimate_decompressed_len = len * 2;
-    let capacity =
-        ((estimate_decompressed_len / buffer_growth_interval) + 1) * buffer_growth_interval;
-    out_buf.reserve(capacity);
+    let capacity = std::cmp::min(
+        bytes::buf::Limit::limit(&out_buf),
+        ((estimate_decompressed_len / buffer_growth_interval) + 1) * buffer_growth_interval,
+    );
+
+    out_buf.get_mut().reserve(capacity);
 
     #[cfg(any(feature = "gzip", feature = "deflate", feature = "zstd"))]
     let mut out_writer = out_buf.writer();

tonic/src/codec/decode.rs

@@ -211,16 +211,25 @@ impl StreamingInner {
 
             let decode_buf = if let Some(encoding) = compression {
                 self.decompress_buf.clear();
+                let limit = self
+                    .max_message_size
+                    .unwrap_or(DEFAULT_MAX_RECV_MESSAGE_SIZE);
+                let limited_out_buf = (&mut self.decompress_buf).limit(limit);
 
                 if let Err(err) = decompress(
                     CompressionSettings {
                         encoding,
                         buffer_growth_interval: buffer_settings.buffer_size,
                     },
                     &mut self.buf,
-                    &mut self.decompress_buf,
+                    limited_out_buf,
                     len,
                 ) {
+                    if matches!(err.kind(), std::io::ErrorKind::WriteZero) {
+                        return Err(Status::resource_exhausted(format!(
+                            "Error decompressing: size limit, of {limit} bytes, exceeded while decompressing message"
+                        )));
+                    }
                     let message = if let Direction::Response(status) = self.direction {
                         format!(
                             "Error decompressing: {err}, while receiving response with status: {status}"