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

Parameterize ContentLengthLimit

Browse Source
This commit is contained in:
David Pedersen
2021-06-09 08:14:20 +02:00
parent 09f76f3c87
commit 90c3e5ba74
6 changed files with 21 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/extract/mod.rs
+10 -8
View File
@@ -388,14 +388,14 @@ impl FromRequest for Body {
}
}
/// Extractor that will buffer request bodies up to a certain size.
/// Extractor that will reject requests with a body larger than some size.
///
/// # Example
///
/// ```rust,no_run
/// use tower_web::prelude::*;
///
/// async fn handler(req: Request<Body>, body: extract::BytesMaxLength<1024>) {
/// async fn handler(req: Request<Body>, body: extract::ContentLengthLimit<String, 1024>) {
/// // ...
/// }
///
@@ -404,15 +404,17 @@ impl FromRequest for Body {
///
/// This requires the request to have a `Content-Length` header.
#[derive(Debug, Clone)]
pub struct BytesMaxLength<const N: u64>(pub Bytes);
pub struct ContentLengthLimit<T, const N: u64>(pub T);
#[async_trait]
impl<const N: u64> FromRequest for BytesMaxLength<N> {
impl<T, const N: u64> FromRequest for ContentLengthLimit<T, N>
where
T: FromRequest,
{
type Rejection = Response<Body>;
async fn from_request(req: &mut Request<Body>) -> Result<Self, Self::Rejection> {
let content_length = req.headers().get(http::header::CONTENT_LENGTH).cloned();
let body = take_body(req).map_err(|reject| reject.into_response())?;
let content_length =
content_length.and_then(|value| value.to_str().ok()?.parse::<u64>().ok());
@@ -425,11 +427,11 @@ impl<const N: u64> FromRequest for BytesMaxLength<N> {
return Err(LengthRequired.into_response());
};
let bytes = hyper::body::to_bytes(body)
let value = T::from_request(req)
.await
.map_err(|e| FailedToBufferBody::from_err(e).into_response())?;
.map_err(IntoResponse::into_response)?;
Ok(BytesMaxLength(bytes))
Ok(Self(value))
}
}
src/extract/rejection.rs
+4 -4
View File
@@ -103,16 +103,16 @@ define_rejection! {
define_rejection! {
#[status = PAYLOAD_TOO_LARGE]
#[body = "Request payload is too large"]
/// Rejection type for [`BytesMaxLength`](super::BytesMaxLength) if the
/// request body is too large.
/// Rejection type for [`ContentLengthLimit`](super::ContentLengthLimit) if
/// the request body is too large.
pub struct PayloadTooLarge;
}
define_rejection! {
#[status = LENGTH_REQUIRED]
#[body = "Content length header is required"]
/// Rejection type for [`BytesMaxLength`](super::BytesMaxLength) if the
/// request is missing the `Content-Length` header or it is invalid.
/// Rejection type for [`ContentLengthLimit`](super::ContentLengthLimit) if
/// the request is missing the `Content-Length` header or it is invalid.
pub struct LengthRequired;
}