tokio-rs · davidpdrsn · Sep 10, 2022 · Sep 3, 2022 · Sep 3, 2022 · Sep 3, 2022
diff --git a/axum-core/Cargo.toml b/axum-core/Cargo.toml
@@ -18,9 +18,12 @@ futures-util = { version = "0.3", default-features = false, features = ["alloc"]
 http = "0.2.7"
 http-body = "0.4.5"
 mime = "0.3.16"
+tower-layer = "0.3"
+tower-service = "0.3"
 
 [dev-dependencies]
 axum = { path = "../axum", version = "0.6.0-rc.1" }
 futures-util = "0.3"
 hyper = "0.14"
 tokio = { version = "1.0", features = ["macros"] }
+tower-http = { version = "0.3.4", features = ["limit"] }
diff --git a/axum-core/src/extract/default_body_limit.rs b/axum-core/src/extract/default_body_limit.rs
@@ -0,0 +1,101 @@
+use self::private::DefaultBodyLimitService;
+use tower_layer::Layer;
+
+/// Layer for configuring the default request body limit.
+///
+/// For security reasons, [`Bytes`] will, by default, not accept bodies larger than 2MB. This also
+/// applies to extractors that uses [`Bytes`] internally such as `String`, [`Json`], and [`Form`].
+///
+/// This middleware provides ways to configure that.
+///
+/// Note that if an extractor consumes the body directly with [`Body::data`], or similar, the
+/// default limit is _not_ applied.
+///
+/// [`Body::data`]: http_body::Body::data
+/// [`Bytes`]: bytes::Bytes
+/// [`Json`]: https://docs.rs/axum/0.6.0-rc.1/axum/struct.Json.html
+/// [`Form`]: https://docs.rs/axum/0.6.0-rc.1/axum/struct.Form.html
+#[derive(Debug, Clone)]
+#[non_exhaustive]
+pub struct DefaultBodyLimit;
+
+impl DefaultBodyLimit {
+    /// Disable the default request body limit.
+    ///
+    /// This must be used to receive bodies larger than the default limit of 2MB using [`Bytes`] or
+    /// an extractor built on it such as `String`, [`Json`], [`Form`].
+    ///
+    /// Note that if you're accepting data from untrusted remotes it is recommend to add your own
+    /// limit such as [`tower_http::limit`].
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// use axum::{
+    ///     Router,
+    ///     routing::get,
+    ///     body::{Bytes, Body},
+    ///     extract::DefaultBodyLimit,
+    /// };
+    /// use tower_http::limit::RequestBodyLimitLayer;
+    /// use http_body::Limited;
+    ///
+    /// let app: Router<_, Limited<Body>> = Router::new()
+    ///     .route("/", get(|body: Bytes| async {}))
+    ///     // Disable the default limit
+    ///     .layer(DefaultBodyLimit::disable())
+    ///     // Set a different limit
+    ///     .layer(RequestBodyLimitLayer::new(10 * 1000 * 1000));
+    /// ```
+    ///
+    /// [`tower_http::limit`]: https://docs.rs/tower-http/0.3.4/tower_http/limit/index.html
+    /// [`Bytes`]: bytes::Bytes
+    /// [`Json`]: https://docs.rs/axum/0.6.0-rc.1/axum/struct.Json.html
+    /// [`Form`]: https://docs.rs/axum/0.6.0-rc.1/axum/struct.Form.html
+    pub fn disable() -> Self {
+        Self
+    }
+}
+
+impl<S> Layer<S> for DefaultBodyLimit {
+    type Service = DefaultBodyLimitService<S>;
+
+    fn layer(&self, inner: S) -> Self::Service {
+        DefaultBodyLimitService { inner }
+    }
+}
+
+#[derive(Copy, Clone, Debug)]
+pub(crate) struct DefaultBodyLimitDisabled;
+
+mod private {
+    use super::DefaultBodyLimitDisabled;
+    use http::Request;
+    use std::task::Context;
+    use tower_service::Service;
+
+    #[derive(Debug, Clone, Copy)]
+    pub struct DefaultBodyLimitService<S> {
+        pub(super) inner: S,
+    }
+
+    impl<B, S> Service<Request<B>> for DefaultBodyLimitService<S>
+    where
+        S: Service<Request<B>>,
+    {
+        type Response = S::Response;
+        type Error = S::Error;
+        type Future = S::Future;
+
+        #[inline]
+        fn poll_ready(&mut self, cx: &mut Context<'_>) -> std::task::Poll<Result<(), Self::Error>> {
+            self.inner.poll_ready(cx)
+        }
+
+        #[inline]
+        fn call(&mut self, mut req: Request<B>) -> Self::Future {
+            req.extensions_mut().insert(DefaultBodyLimitDisabled);
+            self.inner.call(req)
+        }
+    }
+}
diff --git a/axum-core/src/extract/mod.rs b/axum-core/src/extract/mod.rs
@@ -11,11 +11,12 @@ use std::convert::Infallible;
 
 pub mod rejection;
 
+mod default_body_limit;
 mod from_ref;
 mod request_parts;
 mod tuple;
 
-pub use self::from_ref::FromRef;
+pub use self::{default_body_limit::DefaultBodyLimit, from_ref::FromRef};
 
 mod private {
     #[derive(Debug, Clone, Copy)]

diff --git a/axum-core/src/extract/request_parts.rs b/axum-core/src/extract/request_parts.rs
@@ -1,4 +1,6 @@
-use super::{rejection::*, FromRequest, FromRequestParts};
+use super::{
+    default_body_limit::DefaultBodyLimitDisabled, rejection::*, FromRequest, FromRequestParts,
+};
 use crate::BoxError;
 use async_trait::async_trait;
 use bytes::Bytes;
@@ -82,11 +84,20 @@ where
     type Rejection = BytesRejection;
 
     async fn from_request(req: Request<B>, _: &S) -> Result<Self, Self::Rejection> {
-        let body = req.into_body();
-
-        let bytes = crate::body::to_bytes(body)
-            .await
-            .map_err(FailedToBufferBody::from_err)?;
+        // 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 bytes = if req.extensions().get::<DefaultBodyLimitDisabled>().is_some() {
+            crate::body::to_bytes(req.into_body())
+                .await
+                .map_err(FailedToBufferBody::from_err)?
+        } else {
+            let body = http_body::Limited::new(req.into_body(), DEFAULT_LIMIT);
+            crate::body::to_bytes(body)
+                .await
+                .map_err(FailedToBufferBody::from_err)?
+        };
 
         Ok(bytes)
     }
@@ -102,15 +113,18 @@ where
 {
     type Rejection = StringRejection;
 
-    async fn from_request(req: Request<B>, _: &S) -> Result<Self, Self::Rejection> {
-        let body = req.into_body();
-
-        let bytes = crate::body::to_bytes(body)
+    async fn from_request(req: Request<B>, state: &S) -> Result<Self, Self::Rejection> {
+        let bytes = Bytes::from_request(req, state)
             .await
-            .map_err(FailedToBufferBody::from_err)?
-            .to_vec();
-
-        let string = String::from_utf8(bytes).map_err(InvalidUtf8::from_err)?;
+            .map_err(|err| match err {
+                BytesRejection::FailedToBufferBody(inner) => {
+                    StringRejection::FailedToBufferBody(inner)
+                }
+            })?;
+
+        let string = std::str::from_utf8(&bytes)
+            .map_err(InvalidUtf8::from_err)?
+            .to_owned();
 
         Ok(string)
     }

diff --git a/axum/src/docs/extract.md b/axum/src/docs/extract.md
@@ -11,6 +11,7 @@ Types and traits for extracting data from requests.
 - [Accessing inner errors](#accessing-inner-errors)
 - [Defining custom extractors](#defining-custom-extractors)
 - [Accessing other extractors in `FromRequest` or `FromRequestParts` implementations](#accessing-other-extractors-in-fromrequest-or-fromrequestparts-implementations)
+- [Request body limits](#request-body-limits)
 - [Request body extractors](#request-body-extractors)
 - [Running extractors from middleware](#running-extractors-from-middleware)
 - [Wrapping extractors](#wrapping-extractors)
@@ -620,6 +621,14 @@ let app = Router::new().route("/", get(handler)).layer(Extension(state));
 # };
 ```
 
+# Request body limits
+
+For security reasons, [`Bytes`] will, by default, not accept bodies larger than
+2MB. This also applies to extractors that uses [`Bytes`] internally such as
+`String`, [`Json`], and [`Form`].
+
+For more details, including how to disable this limit, see [`DefaultBodyLimit`].
+
 # Request body extractors
 
 Most of the time your request body type will be [`body::Body`] (a re-export
@@ -816,6 +825,7 @@ async fn handler(
 ```
 
 [`body::Body`]: crate::body::Body
+[`Bytes`]: crate::body::Bytes
 [customize-extractor-error]: https://github.com/tokio-rs/axum/blob/main/examples/customize-extractor-error/src/main.rs
 [`HeaderMap`]: https://docs.rs/http/latest/http/header/struct.HeaderMap.html
 [`Request`]: https://docs.rs/http/latest/http/struct.Request.html

diff --git a/axum/src/extract/mod.rs b/axum/src/extract/mod.rs
@@ -16,7 +16,7 @@ mod request_parts;
 mod state;
 
 #[doc(inline)]
-pub use axum_core::extract::{FromRef, FromRequest, FromRequestParts};
+pub use axum_core::extract::{DefaultBodyLimit, FromRef, FromRequest, FromRequestParts};
 
 #[doc(inline)]
 #[allow(deprecated)]

diff --git a/axum/src/routing/tests/mod.rs b/axum/src/routing/tests/mod.rs
@@ -1,13 +1,14 @@
 use crate::{
     body::{Bytes, Empty},
     error_handling::HandleErrorLayer,
-    extract::{self, FromRef, Path, State},
+    extract::{self, DefaultBodyLimit, FromRef, Path, State},
     handler::{Handler, HandlerWithoutStateExt},
     response::IntoResponse,
     routing::{delete, get, get_service, on, on_service, patch, patch_service, post, MethodFilter},
     test_helpers::*,
     BoxError, Json, Router,
 };
+use futures_util::stream::StreamExt;
 use http::{header::CONTENT_LENGTH, HeaderMap, Request, Response, StatusCode, Uri};
 use hyper::Body;
 use serde_json::json;
@@ -604,6 +605,50 @@ async fn routes_must_start_with_slash() {
     TestClient::new(app);
 }
 
+#[tokio::test]
+async fn body_limited_by_default() {
+    let app = Router::new()
+        .route("/bytes", post(|_: Bytes| async {}))
+        .route("/string", post(|_: String| async {}))
+        .route("/json", post(|_: Json<serde_json::Value>| async {}));
+
+    let client = TestClient::new(app);
+
+    for uri in ["/bytes", "/string", "/json"] {
+        println!("calling {}", uri);
+
+        let stream = futures_util::stream::repeat("a".repeat(1000)).map(Ok::<_, hyper::Error>);
+        let body = Body::wrap_stream(stream);
+
+        let res_future = client
+            .post(uri)
+            .header("content-type", "application/json")
+            .body(body)
+            .send();
+        let res = tokio::time::timeout(Duration::from_secs(3), res_future)
+            .await
+            .expect("never got response");
+
+        assert_eq!(res.status(), StatusCode::PAYLOAD_TOO_LARGE);
+    }
+}
+
+#[tokio::test]
+async fn disabling_the_default_limit() {
+    let app = Router::new()
+        .route("/", post(|_: Bytes| async {}))
+        .layer(DefaultBodyLimit::disable());
+
+    let client = TestClient::new(app);
+
+    // `DEFAULT_LIMIT` is 2mb so make a body larger than that
+    let body = Body::from("a".repeat(3_000_000));
+
+    let res = client.post("/").body(body).send().await;
+
+    assert_eq!(res.status(), StatusCode::OK);
+}
+
 #[tokio::test]
 async fn limited_body_with_content_length() {
     const LIMIT: usize = 3;