diff --git a/changelog.d/19498.misc b/changelog.d/19498.misc new file mode 100644 index 00000000000..7d048c283d6 --- /dev/null +++ b/changelog.d/19498.misc @@ -0,0 +1 @@ +Pre-allocate the buffer based on the expected `Content-Length` with the Rust HTTP client. diff --git a/rust/src/http_client.rs b/rust/src/http_client.rs index 9bbdff8b454..dd37a10426b 100644 --- a/rust/src/http_client.rs +++ b/rust/src/http_client.rs @@ -16,6 +16,7 @@ use std::{collections::HashMap, future::Future, sync::OnceLock}; use anyhow::Context; use futures::TryStreamExt; +use headers::HeaderMapExt; use once_cell::sync::OnceCell; use pyo3::{create_exception, exceptions::PyException, prelude::*}; use reqwest::RequestBuilder; @@ -235,8 +236,47 @@ impl HttpClient { let status = response.status(); + // Find the expected `Content-Length` so we can pre-allocate the buffer + // necessary to read the response. It's expected that not every request will + // have a `Content-Length` header. + // + // `response.content_length()` does exist but the "value does not directly + // represents the value of the `Content-Length` header, but rather the size + // of the response’s body" + // (https://docs.rs/reqwest/latest/reqwest/struct.Response.html#method.content_length) + // and we want to avoid reading the entire body at this point because we + // purposely stream it below until the `response_limit`. + let content_length = { + let content_length = response + .headers() + .typed_get::() + // We need a `usize` for the `Vec::with_capacity(...)` usage below + .and_then(|content_length| content_length.0.try_into().ok()); + + // Sanity check that the request isn't too large from the information + // they told us (may be inaccurate so we also check below as we actually + // read the bytes) + if let Some(content_length_bytes) = content_length { + if content_length_bytes > response_limit { + Err(anyhow::anyhow!( + "Response size (defined by `Content-Length`) too large" + ))?; + } + } + + content_length + }; + + // Stream the response to avoid allocating a giant object on the server + // above our expected `response_limit`. let mut stream = response.bytes_stream(); - let mut buffer = Vec::new(); + // Pre-allocate the buffer based on the expected `Content-Length` + let mut buffer = Vec::with_capacity( + content_length + // Default to pre-allocating nothing when the request doesn't have a + // `Content-Length` header + .unwrap_or(0), + ); while let Some(chunk) = stream.try_next().await.context("reading body")? { if buffer.len() + chunk.len() > response_limit { Err(anyhow::anyhow!("Response size too large"))?;