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

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
+3 -1
View File
@@ -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.2.8 (10. September, 2022)
+52 -9
View File
@@ -16,8 +16,15 @@ use tower_layer::Layer;
/// [`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;
pub struct DefaultBodyLimit {
kind: DefaultBodyLimitKind,
}
#[derive(Debug, Clone, Copy)]
pub(crate) enum DefaultBodyLimitKind {
Disable,
Limit(usize),
}
impl DefaultBodyLimit {
/// Disable the default request body limit.
@@ -53,7 +60,42 @@ impl DefaultBodyLimit {
/// [`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
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>;
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 +119,7 @@ mod private {
#[derive(Debug, Clone, Copy)]
pub struct DefaultBodyLimitService<S> {
pub(super) inner: S,
pub(super) kind: DefaultBodyLimitKind,
}
impl<B, S> Service<Request<B>> for DefaultBodyLimitService<S>
@@ -94,7 +137,7 @@ mod private {
#[inline]
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)
}
}
+17 -11
View File
@@ -1,6 +1,4 @@
use super::{
default_body_limit::DefaultBodyLimitDisabled, rejection::*, FromRequest, RequestParts,
};
use super::{default_body_limit::DefaultBodyLimitKind, rejection::*, FromRequest, RequestParts};
use crate::BoxError;
use async_trait::async_trait;
use bytes::Bytes;
@@ -100,15 +98,23 @@ where
let body = take_body(req)?;
let bytes = if req.extensions().get::<DefaultBodyLimitDisabled>().is_some() {
crate::body::to_bytes(body)
let limit_kind = req.extensions().get::<DefaultBodyLimitKind>().copied();
let bytes = match limit_kind {
Some(DefaultBodyLimitKind::Disable) => 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)?
.map_err(FailedToBufferBody::from_err)?,
Some(DefaultBodyLimitKind::Limit(limit)) => {
let body = http_body::Limited::new(body, limit);
crate::body::to_bytes(body)
.await
.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)
+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
message points to where the user added the invalid router, rather than
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
[#1397]: https://github.com/tokio-rs/axum/pull/1397
# 0.5.16 (10. September, 2022)
+25
View File
@@ -767,6 +767,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;