Limit size of request bodies in Bytes extractor (#1362)

* Limit size of request bodies in `Bytes` extractor (#1346) * Apply default limit to request body size * Support disabling the default limit * docs * changelog * fix doc test * fix docs links * Avoid unhelpful compiler suggestion (#1251) Co-authored-by: Jonas Platte <jplatte+git@posteo.de>
tokio-rs · Sep 10, 2022 · 95e21c1 · 95e21c1
1 parent 3990c3a
commit 95e21c1
Show file tree

Hide file tree

Showing 12 changed files with 245 additions and 39 deletions.
diff --git a/axum-core/CHANGELOG.md b/axum-core/CHANGELOG.md
@@ -7,7 +7,20 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 # Unreleased
 
-- None.
+- **breaking:** Added default limit to how much data `Bytes::from_request` will
+  consume. Previously it would attempt to consume the entire request body
+  without checking its length. This meant if a malicious peer sent an large (or
+  infinite) request body your server might run out of memory and crash.
+
+  The default limit is at 2 MB and can be disabled by adding the new
+  `DefaultBodyLimit::disable()` middleware. See its documentation for more
+  details.
+
+  This also applies to `String` which used `Bytes::from_request` internally.
+
+  ([#1346])
+
+[#1346]: https://github.com/tokio-rs/axum/pull/1346
 
 # 0.2.7 (10. July, 2022)
 

diff --git a/axum-core/Cargo.toml b/axum-core/Cargo.toml
@@ -17,9 +17,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.5" }
 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.5/axum/struct.Json.html
+/// [`Form`]: https://docs.rs/axum/0.5/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.5/axum/struct.Json.html
+    /// [`Form`]: https://docs.rs/axum/0.5/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
@@ -12,9 +12,12 @@ use std::convert::Infallible;
 
 pub mod rejection;
 
+mod default_body_limit;
 mod request_parts;
 mod tuple;
 
+pub use self::default_body_limit::DefaultBodyLimit;
+
 /// Types that can be created from requests.
 ///
 /// See [`axum::extract`] for more details.

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, RequestParts};
+use super::{
+    default_body_limit::DefaultBodyLimitDisabled, rejection::*, FromRequest, RequestParts,
+};
 use crate::BoxError;
 use async_trait::async_trait;
 use bytes::Bytes;
@@ -92,11 +94,22 @@ where
     type Rejection = BytesRejection;
 
     async fn from_request(req: &mut RequestParts<B>) -> 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 body = take_body(req)?;
 
-        let bytes = crate::body::to_bytes(body)
-            .await
-            .map_err(FailedToBufferBody::from_err)?;
+        let bytes = if req.extensions().get::<DefaultBodyLimitDisabled>().is_some() {
+            crate::body::to_bytes(body)
+                .await
+                .map_err(FailedToBufferBody::from_err)?
+        } else {
+            let body = http_body::Limited::new(body, DEFAULT_LIMIT);
+            crate::body::to_bytes(body)
+                .await
+                .map_err(FailedToBufferBody::from_err)?
+        };
 
         Ok(bytes)
     }
@@ -112,14 +125,16 @@ where
     type Rejection = StringRejection;
 
     async fn from_request(req: &mut RequestParts<B>) -> Result<Self, Self::Rejection> {
-        let body = take_body(req)?;
-
-        let bytes = crate::body::to_bytes(body)
-            .await
-            .map_err(FailedToBufferBody::from_err)?
-            .to_vec();
-
-        let string = String::from_utf8(bytes).map_err(InvalidUtf8::from_err)?;
+        let bytes = Bytes::from_request(req).await.map_err(|err| match err {
+            BytesRejection::FailedToBufferBody(inner) => StringRejection::FailedToBufferBody(inner),
+            BytesRejection::BodyAlreadyExtracted(inner) => {
+                StringRejection::BodyAlreadyExtracted(inner)
+            }
+        })?;
+
+        let string = std::str::from_utf8(&bytes)
+            .map_err(InvalidUtf8::from_err)?
+            .to_owned();
 
         Ok(string)
     }

diff --git a/axum-macros/src/debug_handler.rs b/axum-macros/src/debug_handler.rs
@@ -7,20 +7,35 @@ pub(crate) fn expand(attr: Attrs, item_fn: ItemFn) -> TokenStream {
     let check_request_last_extractor = check_request_last_extractor(&item_fn);
     let check_path_extractor = check_path_extractor(&item_fn);
     let check_multiple_body_extractors = check_multiple_body_extractors(&item_fn);
-
-    let check_inputs_impls_from_request = check_inputs_impls_from_request(&item_fn, &attr.body_ty);
     let check_output_impls_into_response = check_output_impls_into_response(&item_fn);
-    let check_future_send = check_future_send(&item_fn);
+
+    // If the function is generic, we can't reliably check its inputs or whether the future it
+    // returns is `Send`. Skip those checks to avoid unhelpful additional compiler errors.
+    let check_inputs_and_future_send = if item_fn.sig.generics.params.is_empty() {
+        let check_inputs_impls_from_request =
+            check_inputs_impls_from_request(&item_fn, &attr.body_ty);
+        let check_future_send = check_future_send(&item_fn);
+
+        quote! {
+            #check_inputs_impls_from_request
+            #check_future_send
+        }
+    } else {
+        syn::Error::new_spanned(
+            &item_fn.sig.generics,
+            "`#[axum_macros::debug_handler]` doesn't support generic functions",
+        )
+        .into_compile_error()
+    };
 
     quote! {
         #item_fn
         #check_extractor_count
         #check_request_last_extractor
         #check_path_extractor
         #check_multiple_body_extractors
-        #check_inputs_impls_from_request
         #check_output_impls_into_response
-        #check_future_send
+        #check_inputs_and_future_send
     }
 }
 
@@ -153,14 +168,6 @@ fn check_multiple_body_extractors(item_fn: &ItemFn) -> TokenStream {
 }
 
 fn check_inputs_impls_from_request(item_fn: &ItemFn, body_ty: &Type) -> TokenStream {
-    if !item_fn.sig.generics.params.is_empty() {
-        return syn::Error::new_spanned(
-            &item_fn.sig.generics,
-            "`#[axum_macros::debug_handler]` doesn't support generic functions",
-        )
-        .into_compile_error();
-    }
-
     item_fn
         .sig
         .inputs

diff --git a/axum-macros/tests/debug_handler/fail/generics.rs b/axum-macros/tests/debug_handler/fail/generics.rs
@@ -1,6 +1,6 @@
 use axum_macros::debug_handler;
 
 #[debug_handler]
-async fn handler<T>() {}
+async fn handler<T>(extract: T) {}
 
 fn main() {}
diff --git a/axum-macros/tests/debug_handler/fail/generics.stderr b/axum-macros/tests/debug_handler/fail/generics.stderr
@@ -1,13 +1,5 @@
 error: `#[axum_macros::debug_handler]` doesn't support generic functions
  --> tests/debug_handler/fail/generics.rs:4:17
   |
-4 | async fn handler<T>() {}
+4 | async fn handler<T>(extract: T) {}
   |                 ^^^
-
-error[E0282]: type annotations needed
- --> tests/debug_handler/fail/generics.rs:4:10
-  |
-4 | async fn handler<T>() {}
-  | -----    ^^^^^^^ cannot infer type for type parameter `T` declared on the function `handler`
-  | |
-  | consider giving `future` a type
diff --git a/axum/CHANGELOG.md b/axum/CHANGELOG.md
@@ -7,7 +7,24 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 # Unreleased
 
-- None.
+## Security
+
+- **breaking:** Added default limit to how much data `Bytes::from_request` will
+  consume. Previously it would attempt to consume the entire request body
+  without checking its length. This meant if a malicious peer sent an large (or
+  infinite) request body your server might run out of memory and crash.
+
+  The default limit is at 2 MB and can be disabled by adding the new
+  `DefaultBodyLimit::disable()` middleware. See its documentation for more
+  details.
+
+  This also applies to these extractors which used `Bytes::from_request`
+  internally:
+  - `Form`
+  - `Json`
+  - `String`
+
+  ([#1346])
 
 # 0.5.15 (9. August, 2022)
 

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` implementations](#accessing-other-extractors-in-fromrequest-implementations)
+- [Request body limits](#request-body-limits)
 - [Request body extractors](#request-body-extractors)
 - [Running extractors from middleware](#running-extractors-from-middleware)
 
@@ -505,6 +506,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
@@ -637,6 +646,7 @@ let app = Router::new().layer(middleware::from_fn(auth_middleware));
 ```
 
 [`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
@@ -17,7 +17,7 @@ mod raw_query;
 mod request_parts;
 
 #[doc(inline)]
-pub use axum_core::extract::{FromRequest, RequestParts};
+pub use axum_core::extract::{DefaultBodyLimit, FromRequest, RequestParts};
 
 #[doc(inline)]
 #[allow(deprecated)]