1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

use crate::response::IntoResponse;

use super::{rejection::*, FromRequest, RequestParts};
use async_trait::async_trait;
use std::ops::Deref;

/// Extractor that will reject requests with a body larger than some size.
///
/// # Example
///
/// ```rust,no_run
/// use axum::{
///     extract::ContentLengthLimit,
///     routing::post,
///     Router,
/// };
///
/// async fn handler(body: ContentLengthLimit<String, 1024>) {
///     // ...
/// }
///
/// let app = Router::new().route("/", post(handler));
/// # async {
/// # axum::Server::bind(&"".parse().unwrap()).serve(app.into_make_service()).await.unwrap();
/// # };
/// ```
///
/// This requires the request to have a `Content-Length` header.
#[derive(Debug, Clone)]
pub struct ContentLengthLimit<T, const N: u64>(pub T);

#[async_trait]
impl<T, B, const N: u64> FromRequest<B> for ContentLengthLimit<T, N>
where
    T: FromRequest<B>,
    T::Rejection: IntoResponse,
    B: Send,
{
    type Rejection = ContentLengthLimitRejection<T::Rejection>;

    async fn from_request(req: &mut RequestParts<B>) -> Result<Self, Self::Rejection> {
        let content_length = req
            .headers()
            .ok_or(ContentLengthLimitRejection::HeadersAlreadyExtracted(
                HeadersAlreadyExtracted,
            ))?
            .get(http::header::CONTENT_LENGTH);

        let content_length =
            content_length.and_then(|value| value.to_str().ok()?.parse::<u64>().ok());

        if let Some(length) = content_length {
            if length > N {
                return Err(ContentLengthLimitRejection::PayloadTooLarge(
                    PayloadTooLarge,
                ));
            }
        } else {
            return Err(ContentLengthLimitRejection::LengthRequired(LengthRequired));
        };

        let value = T::from_request(req)
            .await
            .map_err(ContentLengthLimitRejection::Inner)?;

        Ok(Self(value))
    }
}

impl<T, const N: u64> Deref for ContentLengthLimit<T, N> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{routing::post, test_helpers::*, Router};
    use bytes::Bytes;
    use http::StatusCode;
    use serde::Deserialize;

    #[tokio::test]
    async fn body_with_length_limit() {
        use std::iter::repeat;

        #[derive(Debug, Deserialize)]
        #[allow(dead_code)]
        struct Input {
            foo: String,
        }

        const LIMIT: u64 = 8;

        let app = Router::new().route(
            "/",
            post(|_body: ContentLengthLimit<Bytes, LIMIT>| async {}),
        );

        let client = TestClient::new(app);
        let res = client
            .post("/")
            .body(repeat(0_u8).take((LIMIT - 1) as usize).collect::<Vec<_>>())
            .send()
            .await;
        assert_eq!(res.status(), StatusCode::OK);

        let res = client
            .post("/")
            .body(repeat(0_u8).take(LIMIT as usize).collect::<Vec<_>>())
            .send()
            .await;
        assert_eq!(res.status(), StatusCode::OK);

        let res = client
            .post("/")
            .body(repeat(0_u8).take((LIMIT + 1) as usize).collect::<Vec<_>>())
            .send()
            .await;
        assert_eq!(res.status(), StatusCode::PAYLOAD_TOO_LARGE);

        let res = client
            .post("/")
            .body(reqwest::Body::wrap_stream(futures_util::stream::iter(
                vec![Ok::<_, std::io::Error>(bytes::Bytes::new())],
            )))
            .send()
            .await;
        assert_eq!(res.status(), StatusCode::LENGTH_REQUIRED);
    }
}