WIP

axum-core/src/extract/request_parts.rs

@@ -1,10 +1,14 @@
 use super::{rejection::*, FromRequest, FromRequestParts, Request};
-use crate::{body::Body, RequestExt};
+use crate::body::Body;
 use async_trait::async_trait;
-use bytes::Bytes;
+use bytes::{Buf as _, BufMut, Bytes, BytesMut};
 use http::{request::Parts, Extensions, HeaderMap, Method, Uri, Version};
+use http_body::Body as _;
 use http_body_util::BodyExt;
-use std::convert::Infallible;
+use std::{
+    convert::Infallible,
+    pin::{pin, Pin},
+};
 
 #[async_trait]
 impl<S> FromRequest<S> for Request
@@ -71,6 +75,40 @@ where
     }
 }
 
+#[async_trait]
+impl<S> FromRequest<S> for BytesMut
+where
+    S: Send + Sync,
+{
+    type Rejection = BytesRejection;
+
+    async fn from_request(req: Request, _: &S) -> Result<Self, Self::Rejection> {
+        let body = pin!(req.into_body());
+        let mut bytes = BytesMut::new();
+        body_to_bytes_mut(body, &mut bytes).await?;
+        Ok(bytes)
+    }
+}
+
+async fn body_to_bytes_mut(
+    mut body: Pin<&mut Body>,
+    bytes: &mut BytesMut,
+) -> Result<(), BytesRejection> {
+    while let Some(frame) = body
+        .frame()
+        .await
+        .transpose()
+        .map_err(FailedToBufferBody::from_err)?
+    {
+        let Ok(data) = frame.into_data() else {
+            return Ok(());
+        };
+        bytes.put(data);
+    }
+
+    Ok(())
+}
+
 #[async_trait]
 impl<S> FromRequest<S> for Bytes
 where
@@ -79,14 +117,47 @@ where
     type Rejection = BytesRejection;
 
     async fn from_request(req: Request, _: &S) -> Result<Self, Self::Rejection> {
-        let bytes = req
-            .into_limited_body()
-            .collect()
+        let mut body = pin!(req.into_body());
+
+        // If there's only 1 chunk, we can just return Buf::to_bytes()
+        let first_chunk = if let Some(frame) = body
+            .frame()
             .await
+            .transpose()
             .map_err(FailedToBufferBody::from_err)?
-            .to_bytes();
-
-        Ok(bytes)
+        {
+            let Ok(first_chunk) = frame.into_data() else {
+                return Ok(Bytes::new());
+            };
+            first_chunk
+        } else {
+            return Ok(Bytes::new());
+        };
+
+        let mut bytes = if let Some(frame) = body
+            .frame()
+            .await
+            .transpose()
+            .map_err(FailedToBufferBody::from_err)?
+        {
+            let Ok(second_chunk) = frame.into_data() else {
+                return Ok(first_chunk);
+            };
+
+            let cap = first_chunk.remaining()
+                + second_chunk.remaining()
+                + body.size_hint().lower() as usize;
+
+            let mut bytes = BytesMut::with_capacity(cap);
+            bytes.put(first_chunk);
+            bytes.put(second_chunk);
+            bytes
+        } else {
+            return Ok(Bytes::new());
+        };
+
+        body_to_bytes_mut(body, &mut bytes).await?;
+        Ok(bytes.freeze())
     }
 }