rustlavel_http/
body_limit.rs1use 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 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}