tokio-rs/axum
1
0
Fork 0
mirror of https://github.com/tokio-rs/axum.git synced 2026-09-07 00:00:12 +02:00
Code Issues Packages Projects Releases Wiki Activity

Add DefaultBodyLimit::max to change the body size limit (#1397)

Browse Source
This commit is contained in:
David Pedersen
2022-10-20 20:02:46 +02:00
parent 9991b85f1f
commit c8693498c1
5 changed files with 99 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
axum-core/CHANGELOG.md
+3 -1
View File
@@ -7,7 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
# Unreleased # Unreleased
- None. - **added:** Add `DefaultBodyLimit::max` for changing the default body limit ([#1397])
[#1397]: https://github.com/tokio-rs/axum/pull/1397
# 0.2.8 (10. September, 2022) # 0.2.8 (10. September, 2022)
axum-core/src/extract/default_body_limit.rs
+52 -9
View File
@@ -16,8 +16,15 @@ use tower_layer::Layer;
/// [`Json`]: https://docs.rs/axum/0.5/axum/struct.Json.html /// [`Json`]: https://docs.rs/axum/0.5/axum/struct.Json.html
/// [`Form`]: https://docs.rs/axum/0.5/axum/struct.Form.html /// [`Form`]: https://docs.rs/axum/0.5/axum/struct.Form.html
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
#[non_exhaustive] pub struct DefaultBodyLimit {
pub struct DefaultBodyLimit; kind: DefaultBodyLimitKind,
}
#[derive(Debug, Clone, Copy)]
pub(crate) enum DefaultBodyLimitKind {
Disable,
Limit(usize),
}
impl DefaultBodyLimit { impl DefaultBodyLimit {
/// Disable the default request body limit. /// Disable the default request body limit.
@@ -53,7 +60,42 @@ impl DefaultBodyLimit {
/// [`Json`]: https://docs.rs/axum/0.5/axum/struct.Json.html /// [`Json`]: https://docs.rs/axum/0.5/axum/struct.Json.html
/// [`Form`]: https://docs.rs/axum/0.5/axum/struct.Form.html /// [`Form`]: https://docs.rs/axum/0.5/axum/struct.Form.html
pub fn disable() -> Self { pub fn disable() -> Self {
Self Self {
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<Body>> = Router::new()
/// .route("/", get(|body: Bytes| async {}))
/// // 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),
}
} }
} }
@@ -61,15 +103,15 @@ impl<S> Layer<S> for DefaultBodyLimit {
type Service = DefaultBodyLimitService<S>; type Service = DefaultBodyLimitService<S>;
fn layer(&self, inner: S) -> Self::Service { fn layer(&self, inner: S) -> Self::Service {
DefaultBodyLimitService { inner } DefaultBodyLimitService {
inner,
kind: self.kind,
}
} }
} }
#[derive(Copy, Clone, Debug)]
pub(crate) struct DefaultBodyLimitDisabled;
mod private { mod private {
use super::DefaultBodyLimitDisabled; use super::DefaultBodyLimitKind;
use http::Request; use http::Request;
use std::task::Context; use std::task::Context;
use tower_service::Service; use tower_service::Service;
@@ -77,6 +119,7 @@ mod private {
#[derive(Debug, Clone, Copy)] #[derive(Debug, Clone, Copy)]
pub struct DefaultBodyLimitService<S> { pub struct DefaultBodyLimitService<S> {
pub(super) inner: S, pub(super) inner: S,
pub(super) kind: DefaultBodyLimitKind,
} }
impl<B, S> Service<Request<B>> for DefaultBodyLimitService<S> impl<B, S> Service<Request<B>> for DefaultBodyLimitService<S>
@@ -94,7 +137,7 @@ mod private {
#[inline] #[inline]
fn call(&mut self, mut req: Request<B>) -> Self::Future { fn call(&mut self, mut req: Request<B>) -> Self::Future {
req.extensions_mut().insert(DefaultBodyLimitDisabled); req.extensions_mut().insert(self.kind);
self.inner.call(req) self.inner.call(req)
} }
} }
axum-core/src/extract/request_parts.rs
+17 -11
View File
@@ -1,6 +1,4 @@
use super::{ use super::{default_body_limit::DefaultBodyLimitKind, rejection::*, FromRequest, RequestParts};
default_body_limit::DefaultBodyLimitDisabled, rejection::*, FromRequest, RequestParts,
};
use crate::BoxError; use crate::BoxError;
use async_trait::async_trait; use async_trait::async_trait;
use bytes::Bytes; use bytes::Bytes;
@@ -100,15 +98,23 @@ where
let body = take_body(req)?; let body = take_body(req)?;
let bytes = if req.extensions().get::<DefaultBodyLimitDisabled>().is_some() { let limit_kind = req.extensions().get::<DefaultBodyLimitKind>().copied();
crate::body::to_bytes(body) let bytes = match limit_kind {
Some(DefaultBodyLimitKind::Disable) => crate::body::to_bytes(body)
.await .await
.map_err(FailedToBufferBody::from_err)? .map_err(FailedToBufferBody::from_err)?,
} else { Some(DefaultBodyLimitKind::Limit(limit)) => {
let body = http_body::Limited::new(body, DEFAULT_LIMIT); let body = http_body::Limited::new(body, limit);
crate::body::to_bytes(body) crate::body::to_bytes(body)
.await .await
.map_err(FailedToBufferBody::from_err)? .map_err(FailedToBufferBody::from_err)?
}
None => {
let body = http_body::Limited::new(body, DEFAULT_LIMIT);
crate::body::to_bytes(body)
.await
.map_err(FailedToBufferBody::from_err)?
}
}; };
Ok(bytes) Ok(bytes)
axum/CHANGELOG.md
+2
View File
@@ -10,8 +10,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
- **fixed:** Annotate panicking functions with `#[track_caller]` so the error - **fixed:** Annotate panicking functions with `#[track_caller]` so the error
message points to where the user added the invalid router, rather than message points to where the user added the invalid router, rather than
somewhere internally in axum ([#1248]) somewhere internally in axum ([#1248])
- **added:** Add `DefaultBodyLimit::max` for changing the default body limit ([#1397])
[#1248]: https://github.com/tokio-rs/axum/pull/1248 [#1248]: https://github.com/tokio-rs/axum/pull/1248
[#1397]: https://github.com/tokio-rs/axum/pull/1397
# 0.5.16 (10. September, 2022) # 0.5.16 (10. September, 2022)
axum/src/routing/tests/mod.rs
+25
View File
@@ -767,6 +767,31 @@ async fn limited_body_with_content_length() {
assert_eq!(res.status(), StatusCode::PAYLOAD_TOO_LARGE); 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] #[tokio::test]
async fn limited_body_with_streaming_body() { async fn limited_body_with_streaming_body() {
const LIMIT: usize = 3; const LIMIT: usize = 3;