From e2218460e3176672eb2be07ba785ca7fdf4cbc13 Mon Sep 17 00:00:00 2001
From: Abdullah Alyan <AbdullahAlyan@yahoo.com>
Date: Sun, 21 Apr 2024 19:07:00 +0300
Subject: [PATCH] feat: Add option in client builder to ignore content length

---
 src/client.rs             | 15 +++++++++++++++
 src/proto/connection.rs   |  2 ++
 src/proto/streams/mod.rs  |  3 +++
 src/proto/streams/recv.rs | 13 +++++++++----
 src/server.rs             |  1 +
 5 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index 25b151f53..b74d1deee 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -344,6 +344,9 @@ pub struct Builder {
     ///
     /// When this gets exceeded, we issue GOAWAYs.
     local_max_error_reset_streams: Option<usize>,
+
+    /// Ignore content length and only stop receiving when peer close stream
+    ignore_content_length: bool,
 }
 
 #[derive(Debug)]
@@ -654,6 +657,7 @@ impl Builder {
             settings: Default::default(),
             stream_id: 1.into(),
             local_max_error_reset_streams: Some(proto::DEFAULT_LOCAL_RESET_COUNT_MAX),
+            ignore_content_length: false,
         }
     }
 
@@ -1136,6 +1140,16 @@ impl Builder {
         self
     }
 
+    /// Setup whether to ignore the content length when receiving
+    ///
+    /// If content length is ignored, receiving only stops when server closes stream
+    pub fn ignore_content_length(&mut self, ignore: bool) -> &mut Self {
+        self.ignore_content_length = ignore;
+        self
+    }
+
+
+
     /// Sets the first stream ID to something other than 1.
     #[cfg(feature = "unstable")]
     pub fn initial_stream_id(&mut self, stream_id: u32) -> &mut Self {
@@ -1326,6 +1340,7 @@ where
                 remote_reset_stream_max: builder.pending_accept_reset_stream_max,
                 local_error_reset_streams_max: builder.local_max_error_reset_streams,
                 settings: builder.settings.clone(),
+                ignore_content_length: builder.ignore_content_length,
             },
         );
         let send_request = SendRequest {
diff --git a/src/proto/connection.rs b/src/proto/connection.rs
index 5969bb841..bebe8ad7a 100644
--- a/src/proto/connection.rs
+++ b/src/proto/connection.rs
@@ -83,6 +83,7 @@ pub(crate) struct Config {
     pub remote_reset_stream_max: usize,
     pub local_error_reset_streams_max: Option<usize>,
     pub settings: frame::Settings,
+    pub ignore_content_length: bool,
 }
 
 #[derive(Debug)]
@@ -123,6 +124,7 @@ where
                     .max_concurrent_streams()
                     .map(|max| max as usize),
                 local_max_error_reset_streams: config.local_error_reset_streams_max,
+                ignore_content_length: config.ignore_content_length,
             }
         }
         let streams = Streams::new(streams_config(&config));
diff --git a/src/proto/streams/mod.rs b/src/proto/streams/mod.rs
index c4a832342..7903f9c9b 100644
--- a/src/proto/streams/mod.rs
+++ b/src/proto/streams/mod.rs
@@ -72,4 +72,7 @@ pub struct Config {
     ///
     /// When this gets exceeded, we issue GOAWAYs.
     pub local_max_error_reset_streams: Option<usize>,
+
+    /// Ignore the content length header and only stop receiving when peer close stream
+    pub ignore_content_length: bool,
 }
diff --git a/src/proto/streams/recv.rs b/src/proto/streams/recv.rs
index 46cb87cd0..6cb120ddc 100644
--- a/src/proto/streams/recv.rs
+++ b/src/proto/streams/recv.rs
@@ -59,6 +59,9 @@ pub(super) struct Recv {
 
     /// If extended connect protocol is enabled.
     is_extended_connect_protocol_enabled: bool,
+
+    /// Ignore the content length header and only stop when peer close stream
+    ignore_content_length: bool,
 }
 
 #[derive(Debug)]
@@ -107,6 +110,7 @@ impl Recv {
             refused: None,
             is_push_enabled: config.local_push_enabled,
             is_extended_connect_protocol_enabled: config.extended_connect_protocol_enabled,
+            ignore_content_length: config.ignore_content_length,
         }
     }
 
@@ -171,7 +175,7 @@ impl Recv {
             counts.inc_num_recv_streams(stream);
         }
 
-        if !stream.content_length.is_head() {
+        if !self.ignore_content_length && !stream.content_length.is_head() {
             use super::stream::ContentLength;
             use http::header;
 
@@ -341,7 +345,7 @@ impl Recv {
         // Transition the state
         stream.state.recv_close()?;
 
-        if stream.ensure_content_length_zero().is_err() {
+        if !self.ignore_content_length && stream.ensure_content_length_zero().is_err() {
             proto_err!(stream: "recv_trailers: content-length is not zero; stream={:?};",  stream.id);
             return Err(Error::library_reset(stream.id, Reason::PROTOCOL_ERROR));
         }
@@ -616,7 +620,8 @@ impl Recv {
             return Err(Error::library_reset(stream.id, Reason::FLOW_CONTROL_ERROR));
         }
 
-        if stream.dec_content_length(frame.payload().len()).is_err() {
+        if !self.ignore_content_length && stream.dec_content_length(frame.payload().len()).is_err()
+        {
             proto_err!(stream:
                 "recv_data: content-length overflow; stream={:?}; len={:?}",
                 stream.id,
@@ -626,7 +631,7 @@ impl Recv {
         }
 
         if frame.is_end_stream() {
-            if stream.ensure_content_length_zero().is_err() {
+            if !self.ignore_content_length && stream.ensure_content_length_zero().is_err() {
                 proto_err!(stream:
                     "recv_data: content-length underflow; stream={:?}; len={:?}",
                     stream.id,
diff --git a/src/server.rs b/src/server.rs
index 4f8722269..7c86e66ac 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -1385,6 +1385,7 @@ where
                                 .builder
                                 .local_max_error_reset_streams,
                             settings: self.builder.settings.clone(),
+                            ignore_content_length: false,
                         },
                     );