From 211daffcaa71ba7970c691274f4a08ef82a5c39c Mon Sep 17 00:00:00 2001
From: Zettroke <skorostnoy2000@yandex.ru>
Date: Sat, 11 Jun 2022 07:12:53 +0300
Subject: [PATCH] Fix chain remaining_mut(), allowing to chain growing buffer
 (#488)

---
 src/buf/buf_mut.rs  |  4 ++++
 src/buf/chain.rs    |  3 +--
 tests/test_chain.rs | 22 ++++++++++++++++++++++
 3 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/src/buf/buf_mut.rs b/src/buf/buf_mut.rs
index d43082ec0..a24b0c5d1 100644
--- a/src/buf/buf_mut.rs
+++ b/src/buf/buf_mut.rs
@@ -57,6 +57,10 @@ pub unsafe trait BufMut {
     /// Implementations of `remaining_mut` should ensure that the return value
     /// does not change unless a call is made to `advance_mut` or any other
     /// function that is documented to change the `BufMut`'s current position.
+    ///
+    /// # Note
+    ///
+    /// `remaining_mut` may return value smaller than actual available space.
     fn remaining_mut(&self) -> usize;
 
     /// Advance the internal cursor of the BufMut
diff --git a/src/buf/chain.rs b/src/buf/chain.rs
index 9ce5f23aa..78979a123 100644
--- a/src/buf/chain.rs
+++ b/src/buf/chain.rs
@@ -198,8 +198,7 @@ where
     fn remaining_mut(&self) -> usize {
         self.a
             .remaining_mut()
-            .checked_add(self.b.remaining_mut())
-            .unwrap()
+            .saturating_add(self.b.remaining_mut())
     }
 
     fn chunk_mut(&mut self) -> &mut UninitSlice {
diff --git a/tests/test_chain.rs b/tests/test_chain.rs
index 51ff83cdc..86114c70a 100644
--- a/tests/test_chain.rs
+++ b/tests/test_chain.rs
@@ -139,6 +139,28 @@ fn vectored_read() {
     }
 }
 
+#[test]
+fn chain_growing_buffer() {
+    let mut buff = [' ' as u8; 10];
+    let mut vec = b"wassup".to_vec();
+
+    let mut chained = (&mut buff[..]).chain_mut(&mut vec).chain_mut(Vec::new()); // Required for potential overflow because remaining_mut for Vec is isize::MAX - vec.len(), but for chain_mut is usize::MAX
+
+    chained.put_slice(b"hey there123123");
+
+    assert_eq!(&buff, b"hey there1");
+    assert_eq!(&vec, b"wassup23123");
+}
+
+#[test]
+fn chain_overflow_remaining_mut() {
+    let mut chained = Vec::<u8>::new().chain_mut(Vec::new()).chain_mut(Vec::new());
+
+    assert_eq!(chained.remaining_mut(), usize::MAX);
+    chained.put_slice(&[0; 256]);
+    assert_eq!(chained.remaining_mut(), usize::MAX);
+}
+
 #[test]
 fn chain_get_bytes() {
     let mut ab = Bytes::copy_from_slice(b"ab");