From f096394fe5af787565c35113d51b3aad20502590 Mon Sep 17 00:00:00 2001
From: Darius Clark <dariusc93@users.noreply.github.com>
Date: Tue, 7 Jan 2025 10:12:35 -0500
Subject: [PATCH] feat(request-response): allow custom sizes for `json` and
 `cbor` codec

resolves #5791.

Pull-Request: #5792.
---
 protocols/request-response/CHANGELOG.md |  3 ++
 protocols/request-response/src/cbor.rs  | 39 ++++++++++++++++++++-----
 protocols/request-response/src/json.rs  | 39 ++++++++++++++++++++-----
 3 files changed, 65 insertions(+), 16 deletions(-)

diff --git a/protocols/request-response/CHANGELOG.md b/protocols/request-response/CHANGELOG.md
index 15cb0c91797..34fc27b7432 100644
--- a/protocols/request-response/CHANGELOG.md
+++ b/protocols/request-response/CHANGELOG.md
@@ -3,6 +3,9 @@
 - Add connection id to the events emitted by a request-response `Behaviour`.
   See [PR 5719](https://github.com/libp2p/rust-libp2p/pull/5719).
 
+- Allow configurable request and response sizes for `json` and `cbor` codec.
+  See [PR 5792](https://github.com/libp2p/rust-libp2p/pull/5792).
+
 ## 0.27.1
 
 - Deprecate `void` crate.
diff --git a/protocols/request-response/src/cbor.rs b/protocols/request-response/src/cbor.rs
index 744d94cb961..eac1944bb09 100644
--- a/protocols/request-response/src/cbor.rs
+++ b/protocols/request-response/src/cbor.rs
@@ -55,18 +55,19 @@ mod codec {
     use libp2p_swarm::StreamProtocol;
     use serde::{de::DeserializeOwned, Serialize};
 
-    /// 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,
             }
         }
@@ -74,7 +75,25 @@ 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> {
+        /// Sets the limit for request size in bytes.
+        pub fn set_request_size_maximum(mut self, request_size_maximum: u64) -> Self {
+            self.request_size_maximum = request_size_maximum;
+            self
+        }
+
+        /// Sets the limit for response size in bytes.
+        pub fn set_response_size_maximum(mut self, response_size_maximum: u64) -> Self {
+            self.response_size_maximum = response_size_maximum;
+            self
         }
     }
 
@@ -94,7 +113,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)
         }
@@ -105,7 +126,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 9bd5b8c6df9..f151b16bf5f 100644
--- a/protocols/request-response/src/json.rs
+++ b/protocols/request-response/src/json.rs
@@ -54,18 +54,19 @@ mod codec {
     use libp2p_swarm::StreamProtocol;
     use serde::{de::DeserializeOwned, Serialize};
 
-    /// 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,
             }
         }
@@ -73,7 +74,25 @@ 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> {
+        /// Sets the limit for request size in bytes.
+        pub fn set_request_size_maximum(mut self, request_size_maximum: u64) -> Self {
+            self.request_size_maximum = request_size_maximum;
+            self
+        }
+
+        /// Sets the limit for response size in bytes.
+        pub fn set_response_size_maximum(mut self, response_size_maximum: u64) -> Self {
+            self.response_size_maximum = response_size_maximum;
+            self
         }
     }
 
@@ -93,7 +112,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())?)
         }
@@ -104,7 +125,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())?)
         }