feat: add with_max_response_body_size to ReqwestHttpReplicaV2Transport (

#349)

CHANGELOG.md

@@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+[agent-rs/349](https://github.com/dfinity/agent-rs/pull/349) feat: add with_max_response_body_size to ReqwestHttpReplicaV2Tran
+
 ## [0.17.0] - 2022-05-19
 
 Updated dependencies.  Some had breaking changes: k256 0.11, pkcs 0.9, and sec1 0.3.

ic-agent/Cargo.toml

@@ -21,6 +21,7 @@ byteorder = "1.3.2"
 garcon = { version = "0.2", features = ["async"] }
 hex = "0.4.0"
 http = "0.2.6"
+http-body = "0.4.5"
 hyper-rustls = { version = "0.23.0", features = [ "webpki-roots" ] }
 ic-types = "0.3.0"
 k256 = { version = "0.11", features = ["pem"] }
@@ -38,11 +39,12 @@ thiserror = "1.0.30"
 url = "2.1.0"
 pkcs8 = { version = "0.9", features = ["std"] }
 sec1 = { version = "0.3", features = ["pem"]}
+futures-util = "0.3.21"
 
 [dependencies.reqwest]
 version = "0.11.7"
 default-features = false
-features = [ "blocking", "json", "rustls-tls" ]
+features = [ "blocking", "json", "rustls-tls", "stream" ]
 optional = true
 
 [dependencies.pem]

ic-agent/src/agent/agent_error.rs

@@ -168,6 +168,10 @@ pub enum AgentError {
     #[error("Missing replica transport in the Agent Builder.")]
     MissingReplicaTransport(),
 
+    /// The response size exceeded the provided limit.
+    #[error("Response size exceeded limit.")]
+    ResponseSizeExceededLimit(),
+
     /// An unknown error occurred during communication with the replica.
     #[error("An error happened during communication with the replica: {0}")]
     TransportError(Box<dyn std::error::Error + Send + Sync>),

ic-agent/src/agent/http_transport.rs

@@ -4,6 +4,7 @@
 pub use reqwest;
 
 use crate::{agent::agent_error::HttpErrorPayload, ic_types::Principal, AgentError, RequestId};
+use futures_util::StreamExt;
 use hyper_rustls::ConfigBuilderExt;
 use reqwest::Method;
 use std::{future::Future, pin::Pin, sync::Arc};
@@ -32,6 +33,7 @@ pub struct ReqwestHttpReplicaV2Transport {
     url: reqwest::Url,
     client: reqwest::Client,
     password_manager: Option<Arc<dyn PasswordManager>>,
+    max_response_body_size: Option<usize>,
 }
 
 const IC0_DOMAIN: &str = "ic0.app";
@@ -77,24 +79,31 @@ impl ReqwestHttpReplicaV2Transport {
                 .map_err(|_| AgentError::InvalidReplicaUrl(url.clone()))?,
             client,
             password_manager: None,
+            max_response_body_size: None,
         })
     }
 
     /// Sets a password manager to use with HTTP authentication.
     pub fn with_password_manager<P: 'static + PasswordManager>(self, password_manager: P) -> Self {
         ReqwestHttpReplicaV2Transport {
             password_manager: Some(Arc::new(password_manager)),
-            url: self.url,
-            client: self.client,
+            ..self
         }
     }
 
     /// Same as [`with_password_manager`], but providing the Arc so one does not have to be created.
     pub fn with_arc_password_manager(self, password_manager: Arc<dyn PasswordManager>) -> Self {
         ReqwestHttpReplicaV2Transport {
             password_manager: Some(password_manager),
-            url: self.url,
-            client: self.client,
+            ..self
+        }
+    }
+
+    /// Sets a max response body size limit
+    pub fn with_max_response_body_size(self, max_response_body_size: usize) -> Self {
+        ReqwestHttpReplicaV2Transport {
+            max_response_body_size: Some(max_response_body_size),
+            ..self
         }
     }
 
@@ -142,13 +151,35 @@ impl ReqwestHttpReplicaV2Transport {
 
         let http_status = response.status();
         let response_headers = response.headers().clone();
-        let bytes = response
-            .bytes()
-            .await
-            .map_err(|x| AgentError::TransportError(Box::new(x)))?
-            .to_vec();
 
-        Ok((http_status, response_headers, bytes))
+        // Size Check (Content-Length)
+        if matches!(self
+            .max_response_body_size
+            .zip(response.content_length()), Some((size_limit, content_length)) if content_length as usize > size_limit)
+        {
+            return Err(AgentError::ResponseSizeExceededLimit());
+        }
+
+        let mut body: Vec<u8> = response
+            .content_length()
+            .map_or_else(Vec::new, |n| Vec::with_capacity(n as usize));
+
+        let mut stream = response.bytes_stream();
+
+        while let Some(chunk) = stream.next().await {
+            let chunk = chunk.map_err(|x| AgentError::TransportError(Box::new(x)))?;
+
+            // Size Check (Body Size)
+            if matches!(self
+                .max_response_body_size, Some(size_limit) if body.len() + chunk.len() > size_limit)
+            {
+                return Err(AgentError::ResponseSizeExceededLimit());
+            }
+
+            body.extend_from_slice(chunk.as_ref());
+        }
+
+        Ok((http_status, response_headers, body))
     }
 
     async fn execute(