Skip to content

Commit

Permalink
Automatically handle http_body::LengthLimitError
Browse files Browse the repository at this point in the history
This adds special handling of `http_body::LengthLimitError` to
`FailedToBufferBody` such that this just works and returns `413 Payload
Too Large` if the limit is exceeded:

```rust
use http_body::Limited;
use tower_http::map_request_body::MapRequestBodyLayer;

Router::new()
    .route("/", post(|body: Bytes| { ... }))
    .layer(MapRequestBodyLayer::new(|body| {
        Limited::new(body, 1024)
    }));
```

This is a draft until tower-rs/tower-http#271 is
merged because I wanna make sure that works as well. The use case I have
in mind is:

```rust
use http_body::Limited;
use tower_http::limit::RequestBodyLimitLayer;

Router::new()
    .route("/", post(|body: Bytes| { ... }))
    .layer(RequestBodyLimitLayer::new(1024));
```
  • Loading branch information
davidpdrsn committed May 22, 2022
1 parent 1d7878c commit b8b4506
Show file tree
Hide file tree
Showing 3 changed files with 73 additions and 8 deletions.
6 changes: 5 additions & 1 deletion axum-core/CHANGELOG.md
Expand Up @@ -7,7 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0

# Unreleased

- None.
- **added:** Automatically handle `http_body::LengthLimitError` in `FailedToBufferBody` and map
such errors to `413 Payload Too Large`
- **added:** `FailedToBufferBody::is_length_limit_error` to check if the underlying error is
`http_body::LengthLimitError`. Its source error can also be downcast to
`http_body::LengthLimitError`

# 0.2.4 (02. May, 2022)

Expand Down
2 changes: 1 addition & 1 deletion axum-core/Cargo.toml
Expand Up @@ -15,7 +15,7 @@ async-trait = "0.1"
bytes = "1.0"
futures-util = { version = "0.3", default-features = false, features = ["alloc"] }
http = "0.2.7"
http-body = "0.4"
http-body = "0.4.5"
mime = "0.3.16"

[dev-dependencies]
Expand Down
73 changes: 67 additions & 6 deletions axum-core/src/extract/rejection.rs
Expand Up @@ -2,6 +2,7 @@

use crate::response::{IntoResponse, Response};
use http::StatusCode;
use http_body::LengthLimitError;
use std::fmt;

/// Rejection type used if you try and extract the request body more than
Expand All @@ -28,12 +29,72 @@ impl fmt::Display for BodyAlreadyExtracted {

impl std::error::Error for BodyAlreadyExtracted {}

define_rejection! {
#[status = BAD_REQUEST]
#[body = "Failed to buffer the request body"]
/// Rejection type for extractors that buffer the request body. Used if the
/// request body cannot be buffered due to an error.
pub struct FailedToBufferBody(Error);
/// Rejection type for extractors that buffer the request body. Used if the
/// request body cannot be buffered due to an error.
// TODO: in next major for axum-core make this a #[non_exhaustive] enum so we don't need the
// additional indirection
#[derive(Debug)]
pub struct FailedToBufferBody(FailedToBufferBodyInner);

impl FailedToBufferBody {
/// Check if the body failed to be buffered because a length limit was hit.
///
/// This can _only_ happen when you're using [`tower_http::limit::RequestBodyLimitLayer`] or
/// otherwise wrapping request bodies in [`http_body::Limited`].
pub fn is_length_limit_error(&self) -> bool {
matches!(self.0, FailedToBufferBodyInner::LengthLimitError(_))
}
}

#[derive(Debug)]
enum FailedToBufferBodyInner {
Unknown(crate::Error),
LengthLimitError(LengthLimitError),
}

impl FailedToBufferBody {
pub(crate) fn from_err<E>(err: E) -> Self
where
E: Into<crate::BoxError>,
{
let err = err.into();
match err.downcast::<LengthLimitError>() {
Ok(err) => Self(FailedToBufferBodyInner::LengthLimitError(*err)),
Err(err) => Self(FailedToBufferBodyInner::Unknown(crate::Error::new(err))),
}
}
}

impl crate::response::IntoResponse for FailedToBufferBody {
fn into_response(self) -> crate::response::Response {
match self.0 {
FailedToBufferBodyInner::Unknown(err) => (
http::StatusCode::BAD_REQUEST,
format!(concat!("Failed to buffer the request body", ": {}"), err),
)
.into_response(),
FailedToBufferBodyInner::LengthLimitError(err) => (
StatusCode::PAYLOAD_TOO_LARGE,
format!(concat!("Failed to buffer the request body", ": {}"), err),
)
.into_response(),
}
}
}

impl std::fmt::Display for FailedToBufferBody {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "Failed to buffer the request body")
}
}

impl std::error::Error for FailedToBufferBody {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match &self.0 {
FailedToBufferBodyInner::Unknown(err) => Some(err),
FailedToBufferBodyInner::LengthLimitError(err) => Some(err),
}
}
}

define_rejection! {
Expand Down

0 comments on commit b8b4506

Please sign in to comment.