diff --git a/src/sources/util/http/encoding.rs b/src/sources/util/http/encoding.rs index 5d90ca0c20189..beb418c5a7b38 100644 --- a/src/sources/util/http/encoding.rs +++ b/src/sources/util/http/encoding.rs @@ -60,13 +60,12 @@ pub(crate) fn limited_body(max_body_size: usize) -> BoxedFilter<(Bytes,)> { max_body_size, ))) } else { - Ok(()) + Ok(declared) } }) - .untuple_one() .and(warp::body::stream()) - .and_then(move |body| async move { - collect_body_with_limit(body, max_body_size) + .and_then(move |declared: Option, body| async move { + collect_body_with_limit(body, max_body_size, declared) .await .map_err(warp::reject::custom) }) @@ -178,19 +177,35 @@ fn decompress_snappy( Ok(decoded.into()) } +/// Cap on the initial buffer allocation when pre-sizing from a `Content-Length` header. +/// Prevents stalling clients from reserving large amounts of memory before sending any bytes. #[cfg(any( feature = "sources-utils-http-prelude", feature = "sources-opentelemetry", test ))] -async fn collect_body_with_limit(body: S, max_body_size: usize) -> Result +const MAX_INITIAL_BODY_CAPACITY: usize = 8 * 1024 * 1024; // 8 MiB + +#[cfg(any( + feature = "sources-utils-http-prelude", + feature = "sources-opentelemetry", + test +))] +async fn collect_body_with_limit( + body: S, + max_body_size: usize, + declared_len: Option, +) -> Result where S: futures_util::Stream>, B: Buf, { futures_util::pin_mut!(body); - let mut bytes = BytesMut::new(); + let capacity = declared_len + .and_then(|len| usize::try_from(len).ok()) + .map_or(0, |len| len.min(max_body_size).min(MAX_INITIAL_BODY_CAPACITY)); + let mut bytes = BytesMut::with_capacity(capacity); while let Some(chunk) = body.next().await { let chunk = chunk.map_err(|error| { ErrorMessage::new( @@ -368,7 +383,7 @@ mod tests { Ok::<_, warp::Error>(Bytes::from_static(b" world")), ]); - let collected = collect_body_with_limit(body, 11).await.unwrap(); + let collected = collect_body_with_limit(body, 11, Some(11)).await.unwrap(); assert_eq!(collected, Bytes::from_static(b"hello world")); } @@ -379,7 +394,7 @@ mod tests { Ok::<_, warp::Error>(Bytes::from_static(b" world")), ]); - let err = collect_body_with_limit(body, 5) + let err = collect_body_with_limit(body, 5, None) .await .expect_err("should reject"); assert_eq!(err.status_code(), StatusCode::PAYLOAD_TOO_LARGE);