From b8b4506ad38245609ed265e5750a200c44262c64 Mon Sep 17 00:00:00 2001
From: David Pedersen <david.pdrsn@gmail.com>
Date: Sun, 22 May 2022 10:43:35 +0200
Subject: [PATCH] Automatically handle `http_body::LengthLimitError`

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 https://github.com/tower-rs/tower-http/pull/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));
```
---
 axum-core/CHANGELOG.md             |  6 ++-
 axum-core/Cargo.toml               |  2 +-
 axum-core/src/extract/rejection.rs | 73 +++++++++++++++++++++++++++---
 3 files changed, 73 insertions(+), 8 deletions(-)

diff --git a/axum-core/CHANGELOG.md b/axum-core/CHANGELOG.md
index d6d0ea07b3..b18c2faf9a 100644
--- a/axum-core/CHANGELOG.md
+++ b/axum-core/CHANGELOG.md
@@ -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)
 
diff --git a/axum-core/Cargo.toml b/axum-core/Cargo.toml
index 9ce2a6423a..2b817302c7 100644
--- a/axum-core/Cargo.toml
+++ b/axum-core/Cargo.toml
@@ -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]
diff --git a/axum-core/src/extract/rejection.rs b/axum-core/src/extract/rejection.rs
index fc81eda7db..068799e554 100644
--- a/axum-core/src/extract/rejection.rs
+++ b/axum-core/src/extract/rejection.rs
@@ -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
@@ -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! {