| 1 | //! Size-capped HTTP response body reads. |
| 2 | |
| 3 | use anyhow::Result; |
| 4 | |
| 5 | /// Read a whole response body, failing as soon as it would exceed `max_bytes`. |
| 6 | /// |
| 7 | /// A declared `Content-Length` over the cap is refused before anything is |
| 8 | /// read; chunked or length-less bodies are bounded while streaming, so a |
| 9 | /// server cannot make the caller buffer an unbounded body before a size check. |
| 10 | pub async fn read_response_body_capped( |
| 11 | response: reqwest::Response, |
| 12 | max_bytes: usize, |
| 13 | ) -> Result<Vec<u8>> { |
| 14 | use futures_util::StreamExt; |
| 15 | |
| 16 | if let Some(len) = response.content_length() |
| 17 | && len > max_bytes as u64 |
| 18 | { |
| 19 | anyhow::bail!("response body of {len} bytes exceeds {max_bytes} bytes — aborting"); |
| 20 | } |
| 21 | let mut stream = response.bytes_stream(); |
| 22 | let mut buf: Vec<u8> = Vec::new(); |
| 23 | while let Some(chunk) = stream.next().await { |
| 24 | let chunk = chunk.map_err(|e| anyhow::anyhow!("failed to read response body: {e}"))?; |
| 25 | if buf.len().saturating_add(chunk.len()) > max_bytes { |
| 26 | anyhow::bail!("response body exceeds {max_bytes} bytes — aborting"); |
| 27 | } |
| 28 | buf.extend_from_slice(&chunk); |
| 29 | } |
| 30 | Ok(buf) |
| 31 | } |
| 32 |