From dcaa13546787c982f947e6ec3b4fb3418f488f18 Mon Sep 17 00:00:00 2001
From: Or Ricon <rikonor@gmail.com>
Date: Thu, 23 Jun 2022 20:25:37 +0300
Subject: [PATCH] feat: set max response size for http transport (#41)

* set max response size for http transport

* handle responze size exceeded limit

* change max size
---
 Cargo.lock  |  7 +++++--
 Cargo.toml  |  2 +-
 src/main.rs | 18 +++++++++++++++---
 3 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 28f2643..ba018eb 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -879,17 +879,19 @@ dependencies = [
 
 [[package]]
 name = "ic-agent"
-version = "0.17.0"
+version = "0.17.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6a0cabf758d04a2389ffba0700bd7099de9b5cd47a04255063de1b0f9aac1f6e"
+checksum = "f084a526233000cc9c08ab1cb95e59a03d300fae638c185785d1daa1256cc9c8"
 dependencies = [
  "async-trait",
  "base32",
  "base64",
  "byteorder",
+ "futures-util",
  "garcon",
  "hex",
  "http",
+ "http-body",
  "hyper-rustls",
  "ic-types",
  "k256",
@@ -1782,6 +1784,7 @@ dependencies = [
  "tokio",
  "tokio-native-tls",
  "tokio-rustls",
+ "tokio-util 0.6.9",
  "url",
  "wasm-bindgen",
  "wasm-bindgen-futures",
diff --git a/Cargo.toml b/Cargo.toml
index 8155cfb..a4c49b8 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -30,7 +30,7 @@ http-body = "0.4.5"
 hyper = { version = "0.14", features = ["full"] }
 hyper-rustls = { version = "0.23", features = [ "webpki-roots" ] }
 hyper-tls = "0.5"
-ic-agent = { version = "0.17" }
+ic-agent = "0.17.1"
 ic-utils = { version = "0.17", features = ["raw"] }
 lazy-regex = "2"
 opentelemetry = "0.17.0"
diff --git a/src/main.rs b/src/main.rs
index 66fb828..9871389 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -71,6 +71,7 @@ const KB: usize = 1024;
 const MB: usize = 1024 * KB;
 
 const REQUEST_BODY_SIZE_LIMIT: usize = 10 * MB;
+const RESPONSE_BODY_SIZE_LIMIT: usize = 10 * MB;
 
 /// Resolve overrides for [`reqwest::ClientBuilder::resolve()`]
 /// `ic0.app=[::1]:9090`
@@ -359,6 +360,7 @@ async fn forward_request(
         // leak here because a user could use `dfx` to get the same reply.
         match result {
             Ok((http_response,)) => Ok(http_response),
+
             Err(AgentError::ReplicaError {
                 reject_code,
                 reject_message,
@@ -366,6 +368,7 @@ async fn forward_request(
                 .status(StatusCode::INTERNAL_SERVER_ERROR)
                 .body(format!(r#"Replica Error ({}): "{}""#, reject_code, reject_message).into())
                 .unwrap())),
+
             Err(AgentError::HttpError(HttpErrorPayload {
                 status: 451,
                 content_type,
@@ -376,6 +379,12 @@ async fn forward_request(
                 .status(451)
                 .body(content.into())
                 .unwrap())),
+
+            Err(AgentError::ResponseSizeExceededLimit()) => Err(Ok(Response::builder()
+                .status(StatusCode::INSUFFICIENT_STORAGE)
+                .body("Response size exceeds limit".into())
+                .unwrap())),
+
             Err(e) => Err(Err(e.into())),
         }
     }
@@ -800,14 +809,17 @@ async fn handle_request(
             not_found()
         }
     } else {
+        let transport = ReqwestHttpReplicaV2Transport::create_with_client(replica_url, client)
+            .expect("failed to create transport")
+            .with_max_response_body_size(RESPONSE_BODY_SIZE_LIMIT);
+
         let agent = Arc::new(
             ic_agent::Agent::builder()
-                .with_transport(
-                    ReqwestHttpReplicaV2Transport::create_with_client(replica_url, client).unwrap(),
-                )
+                .with_transport(transport)
                 .build()
                 .expect("Could not create agent..."),
         );
+
         if fetch_root_key && agent.fetch_root_key().await.is_err() {
             unable_to_fetch_root_key()
         } else {