Add RequestExt::{with_limited_body, into_limited_body} (#1420)

* Move RequestExt and RequestPartsExt into axum-core

* Add RequestExt::into_limited_body

… and use it for Bytes extraction.

* Add RequestExt::with_limited_body

… and use it for Multipart extraction.

Co-authored-by: David Pedersen <[email protected]>

axum-core/src/ext_traits/mod.rs

@@ -0,0 +1,50 @@
+pub(crate) mod request;
+pub(crate) mod request_parts;
+
+#[cfg(test)]
+mod tests {
+    use std::convert::Infallible;
+
+    use crate::extract::{FromRef, FromRequestParts};
+    use async_trait::async_trait;
+    use http::request::Parts;
+
+    #[derive(Debug, Default, Clone, Copy)]
+    pub(crate) struct State<S>(pub(crate) S);
+
+    #[async_trait]
+    impl<OuterState, InnerState> FromRequestParts<OuterState> for State<InnerState>
+    where
+        InnerState: FromRef<OuterState>,
+        OuterState: Send + Sync,
+    {
+        type Rejection = Infallible;
+
+        async fn from_request_parts(
+            _parts: &mut Parts,
+            state: &OuterState,
+        ) -> Result<Self, Self::Rejection> {
+            let inner_state = InnerState::from_ref(state);
+            Ok(Self(inner_state))
+        }
+    }
+
+    // some extractor that requires the state, such as `SignedCookieJar`
+    pub(crate) struct RequiresState(pub(crate) String);
+
+    #[async_trait]
+    impl<S> FromRequestParts<S> for RequiresState
+    where
+        S: Send + Sync,
+        String: FromRef<S>,
+    {
+        type Rejection = Infallible;
+
+        async fn from_request_parts(
+            _parts: &mut Parts,
+            state: &S,
+        ) -> Result<Self, Self::Rejection> {
+            Ok(Self(String::from_ref(state)))
+        }
+    }
+}

axum-core/src/ext_traits/request.rs

@@ -0,0 +1,233 @@
+use crate::extract::{DefaultBodyLimitKind, FromRequest, FromRequestParts};
+use futures_util::future::BoxFuture;
+use http::Request;
+use http_body::Limited;
+
+mod sealed {
+    pub trait Sealed<B> {}
+    impl<B> Sealed<B> for http::Request<B> {}
+}
+
+/// Extension trait that adds additional methods to [`Request`].
+pub trait RequestExt<B>: sealed::Sealed<B> + Sized {
+    /// Apply an extractor to this `Request`.
+    ///
+    /// This is just a convenience for `E::from_request(req, &())`.
+    ///
+    /// Note this consumes the request. Use [`RequestExt::extract_parts`] if you're not extracting
+    /// the body and don't want to consume the request.
+    fn extract<E, M>(self) -> BoxFuture<'static, Result<E, E::Rejection>>
+    where
+        E: FromRequest<(), B, M> + 'static,
+        M: 'static;
+
+    /// Apply an extractor that requires some state to this `Request`.
+    ///
+    /// This is just a convenience for `E::from_request(req, state)`.
+    ///
+    /// Note this consumes the request. Use [`RequestExt::extract_parts_with_state`] if you're not
+    /// extracting the body and don't want to consume the request.
+    fn extract_with_state<E, S, M>(self, state: &S) -> BoxFuture<'_, Result<E, E::Rejection>>
+    where
+        E: FromRequest<S, B, M> + 'static,
+        S: Send + Sync;
+
+    /// Apply a parts extractor to this `Request`.
+    ///
+    /// This is just a convenience for `E::from_request_parts(parts, state)`.
+    fn extract_parts<E>(&mut self) -> BoxFuture<'_, Result<E, E::Rejection>>
+    where
+        E: FromRequestParts<()> + 'static;
+
+    /// Apply a parts extractor that requires some state to this `Request`.
+    ///
+    /// This is just a convenience for `E::from_request_parts(parts, state)`.
+    fn extract_parts_with_state<'a, E, S>(
+        &'a mut self,
+        state: &'a S,
+    ) -> BoxFuture<'a, Result<E, E::Rejection>>
+    where
+        E: FromRequestParts<S> + 'static,
+        S: Send + Sync;
+
+    /// Apply the [default body limit](crate::extract::DefaultBodyLimit).
+    ///
+    /// If it is disabled, return the request as-is in `Err`.
+    fn with_limited_body(self) -> Result<Request<Limited<B>>, Request<B>>;
+
+    /// Consumes the request, returning the body wrapped in [`Limited`] if a
+    /// [default limit](crate::extract::DefaultBodyLimit) is in place, or not wrapped if the
+    /// default limit is disabled.
+    fn into_limited_body(self) -> Result<Limited<B>, B>;
+}
+
+impl<B> RequestExt<B> for Request<B>
+where
+    B: Send + 'static,
+{
+    fn extract<E, M>(self) -> BoxFuture<'static, Result<E, E::Rejection>>
+    where
+        E: FromRequest<(), B, M> + 'static,
+        M: 'static,
+    {
+        self.extract_with_state(&())
+    }
+
+    fn extract_with_state<E, S, M>(self, state: &S) -> BoxFuture<'_, Result<E, E::Rejection>>
+    where
+        E: FromRequest<S, B, M> + 'static,
+        S: Send + Sync,
+    {
+        E::from_request(self, state)
+    }
+
+    fn extract_parts<E>(&mut self) -> BoxFuture<'_, Result<E, E::Rejection>>
+    where
+        E: FromRequestParts<()> + 'static,
+    {
+        self.extract_parts_with_state(&())
+    }
+
+    fn extract_parts_with_state<'a, E, S>(
+        &'a mut self,
+        state: &'a S,
+    ) -> BoxFuture<'a, Result<E, E::Rejection>>
+    where
+        E: FromRequestParts<S> + 'static,
+        S: Send + Sync,
+    {
+        let mut req = Request::new(());
+        *req.version_mut() = self.version();
+        *req.method_mut() = self.method().clone();
+        *req.uri_mut() = self.uri().clone();
+        *req.headers_mut() = std::mem::take(self.headers_mut());
+        *req.extensions_mut() = std::mem::take(self.extensions_mut());
+        let (mut parts, _) = req.into_parts();
+
+        Box::pin(async move {
+            let result = E::from_request_parts(&mut parts, state).await;
+
+            *self.version_mut() = parts.version;
+            *self.method_mut() = parts.method.clone();
+            *self.uri_mut() = parts.uri.clone();
+            *self.headers_mut() = std::mem::take(&mut parts.headers);
+            *self.extensions_mut() = std::mem::take(&mut parts.extensions);
+
+            result
+        })
+    }
+
+    fn with_limited_body(self) -> Result<Request<Limited<B>>, Request<B>> {
+        // update docs in `axum-core/src/extract/default_body_limit.rs` and
+        // `axum/src/docs/extract.md` if this changes
+        const DEFAULT_LIMIT: usize = 2_097_152; // 2 mb
+
+        match self.extensions().get::<DefaultBodyLimitKind>().copied() {
+            Some(DefaultBodyLimitKind::Disable) => Err(self),
+            Some(DefaultBodyLimitKind::Limit(limit)) => {
+                Ok(self.map(|b| http_body::Limited::new(b, limit)))
+            }
+            None => Ok(self.map(|b| http_body::Limited::new(b, DEFAULT_LIMIT))),
+        }
+    }
+
+    fn into_limited_body(self) -> Result<Limited<B>, B> {
+        self.with_limited_body()
+            .map(Request::into_body)
+            .map_err(Request::into_body)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{
+        ext_traits::tests::{RequiresState, State},
+        extract::FromRef,
+    };
+    use async_trait::async_trait;
+    use http::Method;
+    use hyper::Body;
+
+    #[tokio::test]
+    async fn extract_without_state() {
+        let req = Request::new(());
+
+        let method: Method = req.extract().await.unwrap();
+
+        assert_eq!(method, Method::GET);
+    }
+
+    #[tokio::test]
+    async fn extract_body_without_state() {
+        let req = Request::new(Body::from("foobar"));
+
+        let body: String = req.extract().await.unwrap();
+
+        assert_eq!(body, "foobar");
+    }
+
+    #[tokio::test]
+    async fn extract_with_state() {
+        let req = Request::new(());
+
+        let state = "state".to_owned();
+
+        let State(extracted_state): State<String> = req.extract_with_state(&state).await.unwrap();
+
+        assert_eq!(extracted_state, state);
+    }
+
+    #[tokio::test]
+    async fn extract_parts_without_state() {
+        let mut req = Request::builder().header("x-foo", "foo").body(()).unwrap();
+
+        let method: Method = req.extract_parts().await.unwrap();
+
+        assert_eq!(method, Method::GET);
+        assert_eq!(req.headers()["x-foo"], "foo");
+    }
+
+    #[tokio::test]
+    async fn extract_parts_with_state() {
+        let mut req = Request::builder().header("x-foo", "foo").body(()).unwrap();
+
+        let state = "state".to_owned();
+
+        let State(extracted_state): State<String> =
+            req.extract_parts_with_state(&state).await.unwrap();
+
+        assert_eq!(extracted_state, state);
+        assert_eq!(req.headers()["x-foo"], "foo");
+    }
+
+    // this stuff just needs to compile
+    #[allow(dead_code)]
+    struct WorksForCustomExtractor {
+        method: Method,
+        from_state: String,
+        body: String,
+    }
+
+    #[async_trait]
+    impl<S, B> FromRequest<S, B> for WorksForCustomExtractor
+    where
+        S: Send + Sync,
+        B: Send + 'static,
+        String: FromRef<S> + FromRequest<(), B>,
+    {
+        type Rejection = <String as FromRequest<(), B>>::Rejection;
+
+        async fn from_request(mut req: Request<B>, state: &S) -> Result<Self, Self::Rejection> {
+            let RequiresState(from_state) = req.extract_parts_with_state(state).await.unwrap();
+            let method = req.extract_parts().await.unwrap();
+            let body = req.extract().await?;
+
+            Ok(Self {
+                method,
+                from_state,
+                body,
+            })
+        }
+    }
+}

axum-core/src/ext_traits/request_parts.rs

@@ -0,0 +1,108 @@
+use crate::extract::FromRequestParts;
+use futures_util::future::BoxFuture;
+use http::request::Parts;
+
+mod sealed {
+    pub trait Sealed {}
+    impl Sealed for http::request::Parts {}
+}
+
+/// Extension trait that adds additional methods to [`Parts`].
+pub trait RequestPartsExt: sealed::Sealed + Sized {
+    /// Apply an extractor to this `Parts`.
+    ///
+    /// This is just a convenience for `E::from_request_parts(parts, &())`.
+    fn extract<E>(&mut self) -> BoxFuture<'_, Result<E, E::Rejection>>
+    where
+        E: FromRequestParts<()> + 'static;
+
+    /// Apply an extractor that requires some state to this `Parts`.
+    ///
+    /// This is just a convenience for `E::from_request_parts(parts, state)`.
+    fn extract_with_state<'a, E, S>(
+        &'a mut self,
+        state: &'a S,
+    ) -> BoxFuture<'a, Result<E, E::Rejection>>
+    where
+        E: FromRequestParts<S> + 'static,
+        S: Send + Sync;
+}
+
+impl RequestPartsExt for Parts {
+    fn extract<E>(&mut self) -> BoxFuture<'_, Result<E, E::Rejection>>
+    where
+        E: FromRequestParts<()> + 'static,
+    {
+        self.extract_with_state(&())
+    }
+
+    fn extract_with_state<'a, E, S>(
+        &'a mut self,
+        state: &'a S,
+    ) -> BoxFuture<'a, Result<E, E::Rejection>>
+    where
+        E: FromRequestParts<S> + 'static,
+        S: Send + Sync,
+    {
+        E::from_request_parts(self, state)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::convert::Infallible;
+
+    use super::*;
+    use crate::{
+        ext_traits::tests::{RequiresState, State},
+        extract::FromRef,
+    };
+    use async_trait::async_trait;
+    use http::{Method, Request};
+
+    #[tokio::test]
+    async fn extract_without_state() {
+        let (mut parts, _) = Request::new(()).into_parts();
+
+        let method: Method = parts.extract().await.unwrap();
+
+        assert_eq!(method, Method::GET);
+    }
+
+    #[tokio::test]
+    async fn extract_with_state() {
+        let (mut parts, _) = Request::new(()).into_parts();
+
+        let state = "state".to_owned();
+
+        let State(extracted_state): State<String> = parts
+            .extract_with_state::<State<String>, String>(&state)
+            .await
+            .unwrap();
+
+        assert_eq!(extracted_state, state);
+    }
+
+    // this stuff just needs to compile
+    #[allow(dead_code)]
+    struct WorksForCustomExtractor {
+        method: Method,
+        from_state: String,
+    }
+
+    #[async_trait]
+    impl<S> FromRequestParts<S> for WorksForCustomExtractor
+    where
+        S: Send + Sync,
+        String: FromRef<S>,
+    {
+        type Rejection = Infallible;
+
+        async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
+            let RequiresState(from_state) = parts.extract_with_state(state).await?;
+            let method = parts.extract().await?;
+
+            Ok(Self { method, from_state })
+        }
+    }
+}

axum-core/src/extract/mod.rs

@@ -16,6 +16,7 @@ mod from_ref;
 mod request_parts;
 mod tuple;
 
+pub(crate) use self::default_body_limit::DefaultBodyLimitKind;
 pub use self::{default_body_limit::DefaultBodyLimit, from_ref::FromRef};
 
 mod private {

axum-core/src/extract/request_parts.rs

@@ -1,7 +1,5 @@
-use super::{
-    default_body_limit::DefaultBodyLimitKind, rejection::*, FromRequest, FromRequestParts,
-};
-use crate::BoxError;
+use super::{rejection::*, FromRequest, FromRequestParts};
+use crate::{BoxError, RequestExt};
 use async_trait::async_trait;
 use bytes::Bytes;
 use http::{request::Parts, HeaderMap, Method, Request, Uri, Version};
@@ -84,27 +82,13 @@ where
     type Rejection = BytesRejection;
 
     async fn from_request(req: Request<B>, _: &S) -> Result<Self, Self::Rejection> {
-        // update docs in `axum-core/src/extract/default_body_limit.rs` and
-        // `axum/src/docs/extract.md` if this changes
-        const DEFAULT_LIMIT: usize = 2_097_152; // 2 mb
-
-        let limit_kind = req.extensions().get::<DefaultBodyLimitKind>().copied();
-        let bytes = match limit_kind {
-            Some(DefaultBodyLimitKind::Disable) => crate::body::to_bytes(req.into_body())
+        let bytes = match req.into_limited_body() {
+            Ok(limited_body) => crate::body::to_bytes(limited_body)
+                .await
+                .map_err(FailedToBufferBody::from_err)?,
+            Err(unlimited_body) => crate::body::to_bytes(unlimited_body)
                 .await
                 .map_err(FailedToBufferBody::from_err)?,
-            Some(DefaultBodyLimitKind::Limit(limit)) => {
-                let body = http_body::Limited::new(req.into_body(), limit);
-                crate::body::to_bytes(body)
-                    .await
-                    .map_err(FailedToBufferBody::from_err)?
-            }
-            None => {
-                let body = http_body::Limited::new(req.into_body(), DEFAULT_LIMIT);
-                crate::body::to_bytes(body)
-                    .await
-                    .map_err(FailedToBufferBody::from_err)?
-            }
         };
 
         Ok(bytes)

axum-core/src/lib.rs

@@ -52,6 +52,7 @@
 pub(crate) mod macros;
 
 mod error;
+mod ext_traits;
 pub use self::error::Error;
 
 pub mod body;
@@ -60,3 +61,5 @@ pub mod response;
 
 /// Alias for a type-erased error type.
 pub type BoxError = Box<dyn std::error::Error + Send + Sync>;
+
+pub use self::ext_traits::{request::RequestExt, request_parts::RequestPartsExt};