diff --git a/bitreq/src/connection.rs b/bitreq/src/connection.rs index 5e3b2022..4972a177 100644 --- a/bitreq/src/connection.rs +++ b/bitreq/src/connection.rs @@ -322,12 +322,23 @@ impl AsyncConnection { tcp.write_all(proxy_request.as_bytes()).await?; tcp.flush().await?; + // Max proxy response size to prevent unbounded memory allocation + const MAX_PROXY_RESPONSE_SIZE: usize = 16 * 1024; let mut proxy_response = Vec::new(); - let mut buf = vec![0; 256]; + let mut buf = [0; 256]; + loop { let n = tcp.read(&mut buf).await?; + if n == 0 { + // EOF reached + break; + } proxy_response.extend_from_slice(&buf[..n]); - if n < 256 { + if proxy_response.len() > MAX_PROXY_RESPONSE_SIZE { + return Err(Error::ProxyConnect); + } + if n < buf.len() { + // Partial read indicates end of response break; } } @@ -510,6 +521,7 @@ impl AsyncConnection { request.config.method == Method::Head, request.config.max_headers_size, request.config.max_status_line_len, + request.config.max_body_size, ) .await?; @@ -655,13 +667,23 @@ impl Connection { write!(tcp, "{}", proxy.connect(params.host, params.port.port()))?; tcp.flush()?; + // Max proxy response size to prevent unbounded memory allocation + const MAX_PROXY_RESPONSE_SIZE: usize = 16 * 1024; let mut proxy_response = Vec::new(); + let mut buf = [0; 256]; loop { - let mut buf = vec![0; 256]; - let total = tcp.read(&mut buf)?; - proxy_response.append(&mut buf); - if total < 256 { + let n = tcp.read(&mut buf)?; + if n == 0 { + // EOF reached + break; + } + proxy_response.extend_from_slice(&buf[..n]); + if proxy_response.len() > MAX_PROXY_RESPONSE_SIZE { + return Err(Error::ProxyConnect); + } + if n < buf.len() { + // Partial read indicates end of response break; } } @@ -693,6 +715,7 @@ impl Connection { self.stream, request.config.max_headers_size, request.config.max_status_line_len, + request.config.max_body_size, )?; handle_redirects(request, response) }) diff --git a/bitreq/src/error.rs b/bitreq/src/error.rs index 64267bfc..3af0b181 100644 --- a/bitreq/src/error.rs +++ b/bitreq/src/error.rs @@ -58,13 +58,20 @@ pub enum Error { HttpsFeatureNotEnabled, /// The provided proxy information was not properly formatted. See /// [Proxy::new](crate::Proxy::new) for the valid format. + #[cfg(feature = "proxy")] BadProxy, /// The provided credentials were rejected by the proxy server. + #[cfg(feature = "proxy")] BadProxyCreds, /// The provided proxy credentials were malformed. + #[cfg(feature = "proxy")] ProxyConnect, /// The provided credentials were rejected by the proxy server. + #[cfg(feature = "proxy")] InvalidProxyCreds, + /// The response body size surpasses + /// [Request::with_max_body_size](crate::request::Request::with_max_body_size). + BodyOverflow, // TODO: Uncomment these two for 3.0 // /// The URL does not start with http:// or https://. // InvalidProtocol, @@ -102,10 +109,15 @@ impl fmt::Display for Error { TooManyRedirections => write!(f, "too many redirections (over the max)"), InvalidUtf8InResponse => write!(f, "response contained invalid utf-8 where valid utf-8 was expected"), HttpsFeatureNotEnabled => write!(f, "request url contains https:// but the https feature is not enabled"), + #[cfg(feature = "proxy")] BadProxy => write!(f, "the provided proxy information is malformed"), + #[cfg(feature = "proxy")] BadProxyCreds => write!(f, "the provided proxy credentials are malformed"), + #[cfg(feature = "proxy")] ProxyConnect => write!(f, "could not connect to the proxy server"), + #[cfg(feature = "proxy")] InvalidProxyCreds => write!(f, "the provided proxy credentials are invalid"), + BodyOverflow => write!(f, "the response body size surpassed max_body_size"), // TODO: Uncomment these two for 3.0 // InvalidProtocol => write!(f, "the url does not start with http:// or https://"), // InvalidProtocolInRedirect => write!(f, "got redirected to an absolute url which does not start with http:// or https://"), diff --git a/bitreq/src/lib.rs b/bitreq/src/lib.rs index 6d315a56..0af8fff7 100644 --- a/bitreq/src/lib.rs +++ b/bitreq/src/lib.rs @@ -59,7 +59,7 @@ //! //! ## `proxy` //! -//! This feature enables HTTP proxy support. See [Proxy]. +//! This feature enables HTTP proxy support. //! //! # Examples //! diff --git a/bitreq/src/request.rs b/bitreq/src/request.rs index b81539b3..9193ba80 100644 --- a/bitreq/src/request.rs +++ b/bitreq/src/request.rs @@ -94,6 +94,7 @@ pub struct Request { pub(crate) pipelining: bool, pub(crate) max_headers_size: Option, pub(crate) max_status_line_len: Option, + pub(crate) max_body_size: Option, max_redirects: usize, #[cfg(feature = "proxy")] pub(crate) proxy: Option, @@ -118,6 +119,7 @@ impl Request { pipelining: false, max_headers_size: None, max_status_line_len: None, + max_body_size: None, max_redirects: 100, #[cfg(feature = "proxy")] proxy: None, @@ -247,6 +249,23 @@ impl Request { self } + /// Sets the maximum size of the response body this request will + /// accept. + /// + /// If this limit is passed, the request will close the connection + /// and return an [Error::BodyOverflow] error. + /// + /// The maximum size is counted in bytes. + /// + /// `None` disables the cap, and may cause the program to use any + /// amount of memory if the server responds with a large (or + /// infinite) body. The default is None, so setting this + /// manually is recommended when talking to untrusted servers. + pub fn with_max_body_size>>(mut self, max_body_size: S) -> Request { + self.max_body_size = max_body_size.into(); + self + } + /// Sets the proxy to use. #[cfg(feature = "proxy")] pub fn with_proxy(mut self, proxy: Proxy) -> Request { @@ -282,10 +301,11 @@ impl Request { pub fn send(self) -> Result { let parsed_request = ParsedRequest::new(self)?; let is_head = parsed_request.config.method == Method::Head; + let max_body_size = parsed_request.config.max_body_size; let connection = Connection::new(parsed_request.connection_params(), parsed_request.timeout_at)?; let response = connection.send(parsed_request)?; - Response::create(response, is_head) + Response::create(response, is_head, max_body_size) } /// Sends this request to the host, loaded lazily. diff --git a/bitreq/src/response.rs b/bitreq/src/response.rs index 334610b0..eaf254f0 100644 --- a/bitreq/src/response.rs +++ b/bitreq/src/response.rs @@ -52,11 +52,18 @@ pub struct Response { impl Response { #[cfg(feature = "std")] - pub(crate) fn create(mut parent: ResponseLazy, is_head: bool) -> Result { + pub(crate) fn create( + mut parent: ResponseLazy, + is_head: bool, + max_body_size: Option, + ) -> Result { let mut body = Vec::new(); if !is_head && parent.status_code != 204 && parent.status_code != 304 { for byte in &mut parent { let (byte, length) = byte?; + if max_body_size.is_some_and(|max| body.len().saturating_add(length) > max) { + return Err(Error::BodyOverflow); + } body.reserve(length); body.push(byte); } @@ -79,6 +86,7 @@ impl Response { is_head: bool, max_headers_size: Option, max_status_line_len: Option, + max_body_size: Option, ) -> Result { use HttpStreamState::*; @@ -98,6 +106,10 @@ impl Response { EndOnClose => { while let Some(byte_result) = read_until_closed_async(&mut stream).await { let (byte, length) = byte_result?; + if max_body_size.is_some_and(|max| body.len().saturating_add(length) > max) + { + return Err(Error::BodyOverflow); + } body.reserve(length); body.push(byte); } @@ -107,6 +119,11 @@ impl Response { read_with_content_length_async(&mut stream, &mut length).await { let (byte, expected_length) = byte_result?; + if max_body_size + .is_some_and(|max| body.len().saturating_add(expected_length) > max) + { + return Err(Error::BodyOverflow); + } body.reserve(expected_length); body.push(byte); } @@ -123,6 +140,10 @@ impl Response { .await { let (byte, length) = byte_result?; + if max_body_size.is_some_and(|max| body.len().saturating_add(length) > max) + { + return Err(Error::BodyOverflow); + } body.reserve(length); body.push(byte); }, @@ -290,6 +311,8 @@ pub struct ResponseLazy { stream: HttpStreamBytes, state: HttpStreamState, max_trailing_headers_size: Option, + max_body_size: Option, + bytes_read: usize, } #[cfg(feature = "std")] @@ -301,6 +324,7 @@ impl ResponseLazy { stream: HttpStream, max_headers_size: Option, max_status_line_len: Option, + max_body_size: Option, ) -> Result { let mut stream = BufReader::with_capacity(BACKING_READ_BUFFER_LENGTH, stream).bytes(); let ResponseMetadata { @@ -319,6 +343,8 @@ impl ResponseLazy { stream, state, max_trailing_headers_size, + max_body_size, + bytes_read: 0, }) } @@ -333,6 +359,9 @@ impl ResponseLazy { stream: BufReader::with_capacity(1, http_stream).bytes(), state: HttpStreamState::EndOnClose, max_trailing_headers_size: None, + // Body was already fully loaded and size-checked by send_async + max_body_size: None, + bytes_read: 0, } } } @@ -343,7 +372,7 @@ impl Iterator for ResponseLazy { fn next(&mut self) -> Option { use HttpStreamState::*; - match self.state { + let result = match self.state { EndOnClose => read_until_closed(&mut self.stream), ContentLength(ref mut length) => read_with_content_length(&mut self.stream, length), Chunked(ref mut expecting_chunks, ref mut length, ref mut content_length) => @@ -355,7 +384,17 @@ impl Iterator for ResponseLazy { content_length, self.max_trailing_headers_size, ), + }; + + // Check body size limit before returning the byte + if let Some(Ok((_, expected_length))) = &result { + if self.max_body_size.is_some_and(|max| self.bytes_read + expected_length > max) { + return Some(Err(Error::BodyOverflow)); + } + self.bytes_read += 1; } + + result } } @@ -537,7 +576,7 @@ macro_rules! define_read_methods { return None; } *chunk_length = incoming_length; - *content_length += incoming_length; + *content_length = content_length.saturating_add(incoming_length); } if *chunk_length > 0 {