51 lines
1.5 KiB
Rust
51 lines
1.5 KiB
Rust
use futures::StreamExt;
|
|
use num_traits::ToPrimitive;
|
|
|
|
use crate::Err;
|
|
|
|
/// Reads the response body while enforcing a maximum size limit to prevent
|
|
/// memory exhaustion.
|
|
pub async fn limit_read(response: reqwest::Response, max_size: u64) -> crate::Result<Vec<u8>> {
|
|
if response.content_length().is_some_and(|len| len > max_size) {
|
|
return Err!(BadServerResponse("Response too large"));
|
|
}
|
|
let mut data = Vec::new();
|
|
let mut reader = response.bytes_stream();
|
|
|
|
while let Some(chunk) = reader.next().await {
|
|
let chunk = chunk?;
|
|
data.extend_from_slice(&chunk);
|
|
|
|
if data.len() > max_size.to_usize().expect("max_size must fit in usize") {
|
|
return Err!(BadServerResponse("Response too large"));
|
|
}
|
|
}
|
|
|
|
Ok(data)
|
|
}
|
|
|
|
/// Reads the response body as text while enforcing a maximum size limit to
|
|
/// prevent memory exhaustion.
|
|
pub async fn limit_read_text(
|
|
response: reqwest::Response,
|
|
max_size: u64,
|
|
) -> crate::Result<String> {
|
|
let text = String::from_utf8(limit_read(response, max_size).await?)?;
|
|
Ok(text)
|
|
}
|
|
|
|
#[allow(async_fn_in_trait)]
|
|
pub trait LimitReadExt {
|
|
async fn limit_read(self, max_size: u64) -> crate::Result<Vec<u8>>;
|
|
async fn limit_read_text(self, max_size: u64) -> crate::Result<String>;
|
|
}
|
|
|
|
impl LimitReadExt for reqwest::Response {
|
|
async fn limit_read(self, max_size: u64) -> crate::Result<Vec<u8>> {
|
|
limit_read(self, max_size).await
|
|
}
|
|
|
|
async fn limit_read_text(self, max_size: u64) -> crate::Result<String> {
|
|
limit_read_text(self, max_size).await
|
|
}
|
|
}
|