From 91c26048e0f949dc4417e9514469fbf28aca218e Mon Sep 17 00:00:00 2001
From: Darius <darius.clark@proton.me>
Date: Fri, 1 Dec 2023 23:34:40 -0500
Subject: [PATCH] feat: Allow configurable request/response sizes for json and
 cbor codec.

---
 protocols/request-response/CHANGELOG.md |  2 ++
 protocols/request-response/src/cbor.rs  | 35 +++++++++++++++++++------
 protocols/request-response/src/json.rs  | 35 +++++++++++++++++++------
 3 files changed, 56 insertions(+), 16 deletions(-)

diff --git a/protocols/request-response/CHANGELOG.md b/protocols/request-response/CHANGELOG.md
index d53ff479ee2..c146a40c68f 100644
--- a/protocols/request-response/CHANGELOG.md
+++ b/protocols/request-response/CHANGELOG.md
@@ -1,5 +1,7 @@
 ## 0.26.1
 
+- Allow configurable request and response sizes for `json` and `cbor` codec.
+  See [PR XXXX](https://github.com/libp2p/rust-libp2p/pull/XXXX).
 - Derive `PartialOrd` and `Ord` for `{Out,In}boundRequestId`.
   See [PR 4956](https://github.com/libp2p/rust-libp2p/pull/4956).
 
diff --git a/protocols/request-response/src/cbor.rs b/protocols/request-response/src/cbor.rs
index f371f6149dc..ddf11feb29f 100644
--- a/protocols/request-response/src/cbor.rs
+++ b/protocols/request-response/src/cbor.rs
@@ -52,18 +52,19 @@ mod codec {
     use serde::{de::DeserializeOwned, Serialize};
     use std::{collections::TryReserveError, convert::Infallible, io, marker::PhantomData};
 
-    /// Max request size in bytes
-    const REQUEST_SIZE_MAXIMUM: u64 = 1024 * 1024;
-    /// Max response size in bytes
-    const RESPONSE_SIZE_MAXIMUM: u64 = 10 * 1024 * 1024;
-
     pub struct Codec<Req, Resp> {
+        /// Max request size in bytes
+        request_size_maximum: u64,
+        /// Max response size in bytes
+        response_size_maximum: u64,
         phantom: PhantomData<(Req, Resp)>,
     }
 
     impl<Req, Resp> Default for Codec<Req, Resp> {
         fn default() -> Self {
             Codec {
+                request_size_maximum: 1024 * 1024,
+                response_size_maximum: 10 * 1024 * 1024,
                 phantom: PhantomData,
             }
         }
@@ -71,7 +72,21 @@ mod codec {
 
     impl<Req, Resp> Clone for Codec<Req, Resp> {
         fn clone(&self) -> Self {
-            Self::default()
+            Self {
+                request_size_maximum: self.request_size_maximum,
+                response_size_maximum: self.response_size_maximum,
+                phantom: PhantomData,
+            }
+        }
+    }
+
+    impl<Req, Resp> Codec<Req, Resp> {
+        pub fn new(request_size_maximum: u64, response_size_maximum: u64) -> Self {
+            Self {
+                request_size_maximum,
+                response_size_maximum,
+                ..Default::default()
+            }
         }
     }
 
@@ -91,7 +106,9 @@ mod codec {
         {
             let mut vec = Vec::new();
 
-            io.take(REQUEST_SIZE_MAXIMUM).read_to_end(&mut vec).await?;
+            io.take(self.request_size_maximum)
+                .read_to_end(&mut vec)
+                .await?;
 
             cbor4ii::serde::from_slice(vec.as_slice()).map_err(decode_into_io_error)
         }
@@ -102,7 +119,9 @@ mod codec {
         {
             let mut vec = Vec::new();
 
-            io.take(RESPONSE_SIZE_MAXIMUM).read_to_end(&mut vec).await?;
+            io.take(self.response_size_maximum)
+                .read_to_end(&mut vec)
+                .await?;
 
             cbor4ii::serde::from_slice(vec.as_slice()).map_err(decode_into_io_error)
         }
diff --git a/protocols/request-response/src/json.rs b/protocols/request-response/src/json.rs
index 0b3d634573b..6a9533633b7 100644
--- a/protocols/request-response/src/json.rs
+++ b/protocols/request-response/src/json.rs
@@ -50,18 +50,19 @@ mod codec {
     use serde::{de::DeserializeOwned, Serialize};
     use std::{io, marker::PhantomData};
 
-    /// Max request size in bytes
-    const REQUEST_SIZE_MAXIMUM: u64 = 1024 * 1024;
-    /// Max response size in bytes
-    const RESPONSE_SIZE_MAXIMUM: u64 = 10 * 1024 * 1024;
-
     pub struct Codec<Req, Resp> {
+        /// Max request size in bytes
+        request_size_maximum: u64,
+        /// Max response size in bytes
+        response_size_maximum: u64,
         phantom: PhantomData<(Req, Resp)>,
     }
 
     impl<Req, Resp> Default for Codec<Req, Resp> {
         fn default() -> Self {
             Codec {
+                request_size_maximum: 1024 * 1024,
+                response_size_maximum: 10 * 1024 * 1024,
                 phantom: PhantomData,
             }
         }
@@ -69,7 +70,21 @@ mod codec {
 
     impl<Req, Resp> Clone for Codec<Req, Resp> {
         fn clone(&self) -> Self {
-            Self::default()
+            Self {
+                request_size_maximum: self.request_size_maximum,
+                response_size_maximum: self.response_size_maximum,
+                phantom: self.phantom,
+            }
+        }
+    }
+
+    impl<Req, Resp> Codec<Req, Resp> {
+        pub fn new(request_size_maximum: u64, response_size_maximum: u64) -> Self {
+            Self {
+                request_size_maximum,
+                response_size_maximum,
+                ..Default::default()
+            }
         }
     }
 
@@ -89,7 +104,9 @@ mod codec {
         {
             let mut vec = Vec::new();
 
-            io.take(REQUEST_SIZE_MAXIMUM).read_to_end(&mut vec).await?;
+            io.take(self.request_size_maximum)
+                .read_to_end(&mut vec)
+                .await?;
 
             Ok(serde_json::from_slice(vec.as_slice())?)
         }
@@ -100,7 +117,9 @@ mod codec {
         {
             let mut vec = Vec::new();
 
-            io.take(RESPONSE_SIZE_MAXIMUM).read_to_end(&mut vec).await?;
+            io.take(self.response_size_maximum)
+                .read_to_end(&mut vec)
+                .await?;
 
             Ok(serde_json::from_slice(vec.as_slice())?)
         }