[bitreq] Fix unbounded proxy / body reads

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)
         })

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://"),

bitreq/src/lib.rs

@@ -59,7 +59,7 @@
 //!
 //! ## `proxy`
 //!
-//! This feature enables HTTP proxy support. See [Proxy].
+//! This feature enables HTTP proxy support.
 //!
 //! # Examples
 //!

bitreq/src/request.rs

@@ -94,6 +94,7 @@ pub struct Request {
     pub(crate) pipelining: bool,
     pub(crate) max_headers_size: Option<usize>,
     pub(crate) max_status_line_len: Option<usize>,
+    pub(crate) max_body_size: Option<usize>,
     max_redirects: usize,
     #[cfg(feature = "proxy")]
     pub(crate) proxy: Option<Proxy>,
@@ -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<S: Into<Option<usize>>>(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<Response, Error> {
         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.

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<Response, Error> {
+    pub(crate) fn create(
+        mut parent: ResponseLazy,
+        is_head: bool,
+        max_body_size: Option<usize>,
+    ) -> Result<Response, Error> {
         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<usize>,
         max_status_line_len: Option<usize>,
+        max_body_size: Option<usize>,
     ) -> Result<Response, Error> {
         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<usize>,
+    max_body_size: Option<usize>,
+    bytes_read: usize,
 }
 
 #[cfg(feature = "std")]
@@ -301,6 +324,7 @@ impl ResponseLazy {
         stream: HttpStream,
         max_headers_size: Option<usize>,
         max_status_line_len: Option<usize>,
+        max_body_size: Option<usize>,
     ) -> Result<ResponseLazy, Error> {
         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<Self::Item> {
         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 {