From d12494cc9c3f3347193d684f2a39bfb6c3cafcaa Mon Sep 17 00:00:00 2001
From: Olivier Pinon <op@impero.com>
Date: Thu, 17 Feb 2022 13:30:42 +0100
Subject: [PATCH] Add IntoResponse impl for BytesMut and Chain (#767)

* Add IntoResponse impl for BytesMut and Chain

Co-authored-by: David Pedersen <david.pdrsn@gmail.com>

* Add CHANGELOG entry

Co-authored-by: David Pedersen <david.pdrsn@gmail.com>
---
 axum-core/CHANGELOG.md        |  2 +
 axum-core/src/response/mod.rs | 91 +++++++++++++++++++++++++++++++++--
 2 files changed, 90 insertions(+), 3 deletions(-)

diff --git a/axum-core/CHANGELOG.md b/axum-core/CHANGELOG.md
index fb3f36b5..0d56e641 100644
--- a/axum-core/CHANGELOG.md
+++ b/axum-core/CHANGELOG.md
@@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 # Unreleased
 
 - **added:** Add `IntoResponseHeaders` trait ([#644])
+- **added:** Implement `IntoResponse` for `bytes::BytesMut` and `bytes::Chain<T, U>` ([#767])
 - **breaking:** Using `HeaderMap` as an extractor will no longer remove the headers and thus
   they'll still be accessible to other extractors, such as `axum::extract::Json`. Instead
   `HeaderMap` will clone the headers. You should prefer to use `TypedHeader` to extract only the
@@ -39,6 +40,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 [#644]: https://github.com/tokio-rs/axum/pull/644
 [#698]: https://github.com/tokio-rs/axum/pull/698
 [#699]: https://github.com/tokio-rs/axum/pull/699
+[#767]: https://github.com/tokio-rs/axum/pull/767
 
 # 0.1.1 (06. December, 2021)
 
diff --git a/axum-core/src/response/mod.rs b/axum-core/src/response/mod.rs
index 166b301b..f72b4708 100644
--- a/axum-core/src/response/mod.rs
+++ b/axum-core/src/response/mod.rs
@@ -8,16 +8,22 @@ use crate::{
     body::{boxed, BoxBody},
     BoxError,
 };
-use bytes::Bytes;
+use bytes::{buf::Chain, Buf, Bytes, BytesMut};
 use http::{
     header::{self, HeaderMap, HeaderName, HeaderValue},
     StatusCode,
 };
 use http_body::{
     combinators::{MapData, MapErr},
-    Empty, Full,
+    Empty, Full, SizeHint,
+};
+use std::{
+    borrow::Cow,
+    convert::Infallible,
+    iter,
+    pin::Pin,
+    task::{Context, Poll},
 };
-use std::{borrow::Cow, convert::Infallible, iter};
 
 /// Type alias for [`http::Response`] whose body type defaults to [`BoxBody`], the most common body
 /// type used with axum.
@@ -292,6 +298,85 @@ impl IntoResponse for Bytes {
     }
 }
 
+impl IntoResponse for BytesMut {
+    fn into_response(self) -> Response {
+        self.freeze().into_response()
+    }
+}
+
+impl<T, U> IntoResponse for Chain<T, U>
+where
+    T: Buf + Unpin + Send + 'static,
+    U: Buf + Unpin + Send + 'static,
+{
+    fn into_response(self) -> Response {
+        let (first, second) = self.into_inner();
+        let mut res = Response::new(boxed(BytesChainBody {
+            first: Some(first),
+            second: Some(second),
+        }));
+        res.headers_mut().insert(
+            header::CONTENT_TYPE,
+            HeaderValue::from_static(mime::APPLICATION_OCTET_STREAM.as_ref()),
+        );
+        res
+    }
+}
+
+struct BytesChainBody<T, U> {
+    first: Option<T>,
+    second: Option<U>,
+}
+
+impl<T, U> http_body::Body for BytesChainBody<T, U>
+where
+    T: Buf + Unpin,
+    U: Buf + Unpin,
+{
+    type Data = Bytes;
+    type Error = Infallible;
+
+    fn poll_data(
+        mut self: Pin<&mut Self>,
+        _cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<Self::Data, Self::Error>>> {
+        if let Some(mut buf) = self.first.take() {
+            let bytes = buf.copy_to_bytes(buf.remaining());
+            return Poll::Ready(Some(Ok(bytes)));
+        }
+
+        if let Some(mut buf) = self.second.take() {
+            let bytes = buf.copy_to_bytes(buf.remaining());
+            return Poll::Ready(Some(Ok(bytes)));
+        }
+
+        Poll::Ready(None)
+    }
+
+    fn poll_trailers(
+        self: Pin<&mut Self>,
+        _cx: &mut Context<'_>,
+    ) -> Poll<Result<Option<HeaderMap>, Self::Error>> {
+        Poll::Ready(Ok(None))
+    }
+
+    fn is_end_stream(&self) -> bool {
+        self.first.is_none() && self.second.is_none()
+    }
+
+    fn size_hint(&self) -> SizeHint {
+        match (self.first.as_ref(), self.second.as_ref()) {
+            (Some(first), Some(second)) => {
+                let total_size = first.remaining() + second.remaining();
+                SizeHint::with_exact(total_size as u64)
+            }
+            (Some(buf), None) => SizeHint::with_exact(buf.remaining() as u64),
+            (None, Some(buf)) => SizeHint::with_exact(buf.remaining() as u64),
+            (None, None) => SizeHint::with_exact(0),
+        }
+    }
+}
+
 impl IntoResponse for &'static [u8] {
     fn into_response(self) -> Response {
         let mut res = Response::new(boxed(Full::from(self)));