From 022de38eb72a038171c3db5252f7b4eba19ae206 Mon Sep 17 00:00:00 2001 From: David Pedersen Date: Mon, 19 Sep 2022 21:31:17 +0200 Subject: [PATCH 1/4] convert default body limit into an enum --- axum-core/src/extract/default_body_limit.rs | 27 ++++++++++++++------- axum-core/src/extract/request_parts.rs | 20 ++++++++------- 2 files changed, 29 insertions(+), 18 deletions(-) diff --git a/axum-core/src/extract/default_body_limit.rs b/axum-core/src/extract/default_body_limit.rs index fcf44b790f..9bfe5c7706 100644 --- a/axum-core/src/extract/default_body_limit.rs +++ b/axum-core/src/extract/default_body_limit.rs @@ -16,8 +16,14 @@ use tower_layer::Layer; /// [`Json`]: https://docs.rs/axum/0.6.0-rc.2/axum/struct.Json.html /// [`Form`]: https://docs.rs/axum/0.6.0-rc.2/axum/struct.Form.html #[derive(Debug, Clone)] -#[non_exhaustive] -pub struct DefaultBodyLimit; +pub struct DefaultBodyLimit { + kind: DefaultBodyLimitKind, +} + +#[derive(Debug, Clone, Copy)] +pub(crate) enum DefaultBodyLimitKind { + Disable, +} impl DefaultBodyLimit { /// Disable the default request body limit. @@ -53,7 +59,9 @@ impl DefaultBodyLimit { /// [`Json`]: https://docs.rs/axum/0.6.0-rc.2/axum/struct.Json.html /// [`Form`]: https://docs.rs/axum/0.6.0-rc.2/axum/struct.Form.html pub fn disable() -> Self { - Self + Self { + kind: DefaultBodyLimitKind::Disable, + } } } @@ -61,15 +69,15 @@ impl Layer for DefaultBodyLimit { type Service = DefaultBodyLimitService; fn layer(&self, inner: S) -> Self::Service { - DefaultBodyLimitService { inner } + DefaultBodyLimitService { + inner, + kind: self.kind, + } } } -#[derive(Copy, Clone, Debug)] -pub(crate) struct DefaultBodyLimitDisabled; - mod private { - use super::DefaultBodyLimitDisabled; + use super::DefaultBodyLimitKind; use http::Request; use std::task::Context; use tower_service::Service; @@ -77,6 +85,7 @@ mod private { #[derive(Debug, Clone, Copy)] pub struct DefaultBodyLimitService { pub(super) inner: S, + pub(super) kind: DefaultBodyLimitKind, } impl Service> for DefaultBodyLimitService @@ -94,7 +103,7 @@ mod private { #[inline] fn call(&mut self, mut req: Request) -> Self::Future { - req.extensions_mut().insert(DefaultBodyLimitDisabled); + req.extensions_mut().insert(self.kind); self.inner.call(req) } } diff --git a/axum-core/src/extract/request_parts.rs b/axum-core/src/extract/request_parts.rs index 60ef1b79ad..aa88ecf939 100644 --- a/axum-core/src/extract/request_parts.rs +++ b/axum-core/src/extract/request_parts.rs @@ -1,5 +1,5 @@ use super::{ - default_body_limit::DefaultBodyLimitDisabled, rejection::*, FromRequest, FromRequestParts, + default_body_limit::DefaultBodyLimitKind, rejection::*, FromRequest, FromRequestParts, }; use crate::BoxError; use async_trait::async_trait; @@ -88,15 +88,17 @@ where // `axum/src/docs/extract.md` if this changes const DEFAULT_LIMIT: usize = 2_097_152; // 2 mb - let bytes = if req.extensions().get::().is_some() { - crate::body::to_bytes(req.into_body()) + let limit_kind = req.extensions().get::(); + let bytes = match limit_kind { + Some(DefaultBodyLimitKind::Disable) => 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)? + .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) From 983945232eced9c81214a565a8564e22b8a27865 Mon Sep 17 00:00:00 2001 From: David Pedersen Date: Mon, 19 Sep 2022 21:40:18 +0200 Subject: [PATCH 2/4] Support changing the default body limit --- axum-core/src/extract/default_body_limit.rs | 30 +++++++++++++++++++++ axum-core/src/extract/request_parts.rs | 8 +++++- axum/src/routing/tests/mod.rs | 25 +++++++++++++++++ 3 files changed, 62 insertions(+), 1 deletion(-) diff --git a/axum-core/src/extract/default_body_limit.rs b/axum-core/src/extract/default_body_limit.rs index 9bfe5c7706..cc07a7bbea 100644 --- a/axum-core/src/extract/default_body_limit.rs +++ b/axum-core/src/extract/default_body_limit.rs @@ -23,6 +23,7 @@ pub struct DefaultBodyLimit { #[derive(Debug, Clone, Copy)] pub(crate) enum DefaultBodyLimitKind { Disable, + Limit(usize), } impl DefaultBodyLimit { @@ -63,6 +64,35 @@ impl DefaultBodyLimit { kind: DefaultBodyLimitKind::Disable, } } + + /// Set the default request body limit. + /// + /// By default the limit of request body sizes that [`Bytes::from_request`] (and other + /// extractors built on top of it such as `String`, [`Json`], and [`Form`]) is 2MB. This method + /// can be used to change that 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> = Router::new() + /// .route("/", get(|body: Bytes| async {})) + /// // Replace the default of 2MB with 1024 bytes. + /// .layer(DefaultBodyLimit::max(1024)); + /// ``` + pub fn max(limit: usize) -> Self { + Self { + kind: DefaultBodyLimitKind::Limit(limit), + } + } } impl Layer for DefaultBodyLimit { diff --git a/axum-core/src/extract/request_parts.rs b/axum-core/src/extract/request_parts.rs index aa88ecf939..fc151a2da0 100644 --- a/axum-core/src/extract/request_parts.rs +++ b/axum-core/src/extract/request_parts.rs @@ -88,11 +88,17 @@ where // `axum/src/docs/extract.md` if this changes const DEFAULT_LIMIT: usize = 2_097_152; // 2 mb - let limit_kind = req.extensions().get::(); + let limit_kind = req.extensions().get::().copied(); let bytes = match limit_kind { Some(DefaultBodyLimitKind::Disable) => crate::body::to_bytes(req.into_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) diff --git a/axum/src/routing/tests/mod.rs b/axum/src/routing/tests/mod.rs index 08efd28f07..4d4c58cd38 100644 --- a/axum/src/routing/tests/mod.rs +++ b/axum/src/routing/tests/mod.rs @@ -671,6 +671,31 @@ async fn limited_body_with_content_length() { assert_eq!(res.status(), StatusCode::PAYLOAD_TOO_LARGE); } +#[tokio::test] +async fn changing_the_default_limit() { + let new_limit = 2; + + let app = Router::new() + .route("/", post(|_: Bytes| async {})) + .layer(DefaultBodyLimit::max(new_limit)); + + let client = TestClient::new(app); + + let res = client + .post("/") + .body(Body::from("a".repeat(new_limit))) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + + let res = client + .post("/") + .body(Body::from("a".repeat(new_limit + 1))) + .send() + .await; + assert_eq!(res.status(), StatusCode::PAYLOAD_TOO_LARGE); +} + #[tokio::test] async fn limited_body_with_streaming_body() { const LIMIT: usize = 3; From db64908049e65b925070d6e1ebb4cf889815a8c3 Mon Sep 17 00:00:00 2001 From: David Pedersen Date: Mon, 19 Sep 2022 21:46:06 +0200 Subject: [PATCH 3/4] changelog --- axum-core/CHANGELOG.md | 4 +++- axum/CHANGELOG.md | 2 ++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/axum-core/CHANGELOG.md b/axum-core/CHANGELOG.md index 71cccf5df7..fea1a3ed42 100644 --- a/axum-core/CHANGELOG.md +++ b/axum-core/CHANGELOG.md @@ -7,7 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 # Unreleased -- None. +- **added:** Add `DefaultBodyLimit::max` for changing the default body limit ([#1397]) + +[#1397]: https://github.com/tokio-rs/axum/pull/1397 # 0.3.0-rc.2 (10. September, 2022) diff --git a/axum/CHANGELOG.md b/axum/CHANGELOG.md index 1bedc2ac4a..9405265397 100644 --- a/axum/CHANGELOG.md +++ b/axum/CHANGELOG.md @@ -15,10 +15,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - **fixed:** Support streaming/chunked requests in `ContentLengthLimit` ([#1389]) - **fixed:** Used `400 Bad Request` for `FailedToDeserializeQueryString` rejections, instead of `422 Unprocessable Entity` ([#1387]) +- **added:** Add `DefaultBodyLimit::max` for changing the default body limit ([#1397]) [#1371]: https://github.com/tokio-rs/axum/pull/1371 [#1387]: https://github.com/tokio-rs/axum/pull/1387 [#1389]: https://github.com/tokio-rs/axum/pull/1389 +[#1397]: https://github.com/tokio-rs/axum/pull/1397 # 0.6.0-rc.2 (10. September, 2022) From 0d89b49dc20ade6ae3f4f28eeef1a8b76dc70cd7 Mon Sep 17 00:00:00 2001 From: David Pedersen Date: Mon, 19 Sep 2022 22:02:05 +0200 Subject: [PATCH 4/4] fix docs --- axum-core/src/extract/default_body_limit.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/axum-core/src/extract/default_body_limit.rs b/axum-core/src/extract/default_body_limit.rs index cc07a7bbea..21616287c6 100644 --- a/axum-core/src/extract/default_body_limit.rs +++ b/axum-core/src/extract/default_body_limit.rs @@ -88,6 +88,10 @@ impl DefaultBodyLimit { /// // Replace the default of 2MB with 1024 bytes. /// .layer(DefaultBodyLimit::max(1024)); /// ``` + /// + /// [`Bytes::from_request`]: bytes::Bytes + /// [`Json`]: https://docs.rs/axum/0.6.0-rc.2/axum/struct.Json.html + /// [`Form`]: https://docs.rs/axum/0.6.0-rc.2/axum/struct.Form.html pub fn max(limit: usize) -> Self { Self { kind: DefaultBodyLimitKind::Limit(limit),