feat(request-response): allow custom sizes for json and cbor codec

resolves #5791. Pull-Request: #5792.
libp2p · Jan 7, 2025 · f096394 · f096394
1 parent 34ac476
commit f096394
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 16 deletions.
diff --git 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
@@ -55,26 +55,45 @@ 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,
             }
         }
     }
 
     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
@@ -54,26 +54,45 @@ 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,
             }
         }
     }
 
     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())?)
         }