Skip to main content

rustlavel_http/
body_limit.rs

1//! A per-route ceiling on the request body.
2//!
3//! The server has a limit of its own ([`crate::Limits`]) that protects the
4//! process: a body over it is refused before it is buffered. That limit has to
5//! be as large as the biggest upload the application accepts anywhere, which
6//! makes it useless as policy — a JSON endpoint that expects a kilobyte should
7//! not accept the fifty megabytes the avatar upload needs. This middleware is
8//! the policy: tighter, and per group.
9//!
10//! ```ignore
11//! r.group("/api", |api| {
12//!     api.middleware(BodyLimit::kilobytes(64));
13//!     …
14//! });
15//! ```
16//!
17//! By the time middleware runs the body has been read, so this does not save
18//! memory — the server's limit does that. It saves the handler from parsing
19//! something it was never meant to receive, and tells the client why.
20
21use crate::handler::BoxFuture;
22use crate::middleware::{Middleware, Next};
23use crate::request::Request;
24use crate::response::Response;
25use crate::status::Status;
26use rustlavel_core::Json;
27
28#[derive(Debug, Clone, Copy)]
29pub struct BodyLimit {
30    max: usize,
31}
32
33impl BodyLimit {
34    pub fn bytes(max: usize) -> Self {
35        BodyLimit { max }
36    }
37
38    pub fn kilobytes(max: usize) -> Self {
39        BodyLimit::bytes(max * 1024)
40    }
41
42    pub fn megabytes(max: usize) -> Self {
43        BodyLimit::bytes(max * 1024 * 1024)
44    }
45}
46
47impl Middleware for BodyLimit {
48    fn handle(&self, request: Request, next: Next) -> BoxFuture<Response> {
49        // The declared length is checked as well as the actual one, so a
50        // request that lies in either direction is still caught.
51        let declared = request.headers().content_length().unwrap_or(0);
52        let actual = request.body().len();
53        let size = declared.max(actual);
54
55        if size <= self.max {
56            return next.run(request);
57        }
58
59        let max = self.max;
60        let wants_json = request.wants_json();
61        Box::pin(async move {
62            let message = format!("The request body is {size} bytes; this endpoint accepts at most {max}.");
63            let response = Response::new(Status::PAYLOAD_TOO_LARGE);
64            if wants_json {
65                response.with_json(Json::object([("message", Json::from(message))]))
66            } else {
67                response.with_text(message)
68            }
69        })
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76    use crate::method::Method;
77    use crate::router::Router;
78    use crate::testing::TestClient;
79
80    fn client(limit: BodyLimit) -> TestClient {
81        let mut router = Router::new();
82        router.middleware(limit);
83        router.post("/notes", |req: Request| async move {
84            Response::text(format!("{} bytes", req.body().len()))
85        });
86        TestClient::new(router)
87    }
88
89    #[tokio::test]
90    async fn a_body_under_the_limit_reaches_the_handler() {
91        let request = Request::new(Method::Post, "/notes").with_body(vec![b'x'; 100]);
92        let response = client(BodyLimit::bytes(100)).send(request).await;
93        let response = response.assert_ok();
94        assert_eq!(response.body(), "100 bytes");
95    }
96
97    #[tokio::test]
98    async fn a_body_over_the_limit_is_a_413_with_the_numbers() {
99        let request = Request::new(Method::Post, "/notes")
100            .with_body(vec![b'x'; 101])
101            .with_header("accept", "application/json");
102        let response = client(BodyLimit::bytes(100)).send(request).await;
103        let response = response.assert_status(413);
104        let message = response.json().get("message").and_then(Json::as_str).unwrap().to_string();
105        assert!(message.contains("101 bytes") && message.contains("at most 100"), "{message}");
106    }
107
108    #[tokio::test]
109    async fn a_declared_length_over_the_limit_is_refused_too() {
110        let request = Request::new(Method::Post, "/notes").with_header("content-length", "5000000");
111        client(BodyLimit::kilobytes(64)).send(request).await.assert_status(413);
112    }
113
114    #[test]
115    fn the_unit_helpers_multiply_correctly() {
116        assert_eq!(BodyLimit::kilobytes(2).max, 2048);
117        assert_eq!(BodyLimit::megabytes(1).max, 1_048_576);
118    }
119}