From d8e86d93d9e74ca2a87e5a1b33a6cbc65faca176 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20Ml=C3=A1dek?= <david.mladek.cz@gmail.com>
Date: Thu, 8 Feb 2024 22:54:17 +0100
Subject: [PATCH] util: add `Collected::to_bytes_mut`

---
 http-body-util/src/collected.rs | 19 ++++++++++++++++++-
 http-body-util/src/util.rs      | 14 ++++++++++----
 2 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/http-body-util/src/collected.rs b/http-body-util/src/collected.rs
index 31a02a4..5b2a061 100644
--- a/http-body-util/src/collected.rs
+++ b/http-body-util/src/collected.rs
@@ -4,7 +4,7 @@ use std::{
     task::{Context, Poll},
 };
 
-use bytes::{Buf, Bytes};
+use bytes::{Buf, Bytes, BytesMut};
 use http::HeaderMap;
 use http_body::{Body, Frame};
 
@@ -38,6 +38,11 @@ impl<B: Buf> Collected<B> {
         self.bufs.copy_to_bytes(self.bufs.remaining())
     }
 
+    /// Convert this body into a [`BytesMut`].
+    pub fn to_bytes_mut(mut self) -> BytesMut {
+        self.bufs.copy_to_bytes_mut(self.bufs.remaining())
+    }
+
     pub(crate) fn push_frame(&mut self, frame: Frame<B>) {
         let frame = match frame.into_data() {
             Ok(data) => {
@@ -125,6 +130,18 @@ mod tests {
         assert_eq!(&buf.copy_to_bytes(buf.remaining())[..], b"helloworld!");
     }
 
+    #[tokio::test]
+    async fn segmented_body_mut() {
+        let bufs = [&b"hello"[..], &b"world"[..], &b"!"[..]];
+        let body = StreamBody::new(stream::iter(bufs.map(Frame::data).map(Ok::<_, Infallible>)));
+
+        let buffered = body.collect().await.unwrap();
+
+        let mut buf = buffered.to_bytes_mut();
+
+        assert_eq!(&buf.copy_to_bytes(buf.remaining())[..], b"helloworld!");
+    }
+
     #[tokio::test]
     async fn delayed_segments() {
         let one = stream::once(async { Ok::<_, Infallible>(Frame::data(&b"hello "[..])) });
diff --git a/http-body-util/src/util.rs b/http-body-util/src/util.rs
index 89c4aaa..d19d25b 100644
--- a/http-body-util/src/util.rs
+++ b/http-body-util/src/util.rs
@@ -19,6 +19,14 @@ impl<T: Buf> BufList<T> {
     pub(crate) fn pop(&mut self) -> Option<T> {
         self.bufs.pop_front()
     }
+
+    #[inline]
+    pub(crate) fn copy_to_bytes_mut(&mut self, len: usize) -> BytesMut {
+        assert!(len <= self.remaining(), "`len` greater than remaining");
+        let mut bm = BytesMut::with_capacity(len);
+        bm.put(self.take(len));
+        bm
+    }
 }
 
 impl<T: Buf> Buf for BufList<T> {
@@ -77,10 +85,8 @@ impl<T: Buf> Buf for BufList<T> {
             }
             Some(front) if front.remaining() > len => front.copy_to_bytes(len),
             _ => {
-                assert!(len <= self.remaining(), "`len` greater than remaining");
-                let mut bm = BytesMut::with_capacity(len);
-                bm.put(self.take(len));
-                bm.freeze()
+                let bytes_mut = self.copy_to_bytes_mut(len);
+                bytes_mut.freeze()
             }
         }
     }