rustapi_core/middleware/
body_limit.rs1use crate::error::ApiError;
20use crate::request::Request;
21use crate::response::{IntoResponse, Response};
22use super::{BoxedNext, MiddlewareLayer};
23use http::StatusCode;
24use std::future::Future;
25use std::pin::Pin;
26
27pub const DEFAULT_BODY_LIMIT: usize = 1024 * 1024;
29
30#[derive(Clone)]
35pub struct BodyLimitLayer {
36 limit: usize,
37}
38
39impl BodyLimitLayer {
40 pub fn new(limit: usize) -> Self {
53 Self { limit }
54 }
55
56 pub fn default_limit() -> Self {
58 Self::new(DEFAULT_BODY_LIMIT)
59 }
60
61 pub fn limit(&self) -> usize {
63 self.limit
64 }
65}
66
67impl Default for BodyLimitLayer {
68 fn default() -> Self {
69 Self::default_limit()
70 }
71}
72
73impl MiddlewareLayer for BodyLimitLayer {
74 fn call(
75 &self,
76 req: Request,
77 next: BoxedNext,
78 ) -> Pin<Box<dyn Future<Output = Response> + Send + 'static>> {
79 let limit = self.limit;
80
81 Box::pin(async move {
82 if let Some(content_length) = req.headers().get(http::header::CONTENT_LENGTH) {
84 if let Ok(length_str) = content_length.to_str() {
85 if let Ok(length) = length_str.parse::<usize>() {
86 if length > limit {
87 return ApiError::new(
88 StatusCode::PAYLOAD_TOO_LARGE,
89 "payload_too_large",
90 format!("Request body exceeds limit of {} bytes", limit),
91 )
92 .into_response();
93 }
94 }
95 }
96 }
97
98 if let Some(body) = &req.body {
101 if body.len() > limit {
102 return ApiError::new(
103 StatusCode::PAYLOAD_TOO_LARGE,
104 "payload_too_large",
105 format!("Request body exceeds limit of {} bytes", limit),
106 )
107 .into_response();
108 }
109 }
110
111 next(req).await
113 })
114 }
115
116 fn clone_box(&self) -> Box<dyn MiddlewareLayer> {
117 Box::new(self.clone())
118 }
119}
120
121#[cfg(test)]
122mod tests {
123 use super::*;
124 use crate::request::Request;
125 use bytes::Bytes;
126 use http::{Extensions, Method};
127 use proptest::prelude::*;
128 use std::collections::HashMap;
129 use std::sync::Arc;
130
131 fn create_test_request_with_body(body: Bytes) -> Request {
133 let uri: http::Uri = "/test".parse().unwrap();
134 let mut builder = http::Request::builder().method(Method::POST).uri(uri);
135
136 builder = builder.header(http::header::CONTENT_LENGTH, body.len().to_string());
138
139 let req = builder.body(()).unwrap();
140 let (parts, _) = req.into_parts();
141
142 Request::new(parts, body, Arc::new(Extensions::new()), HashMap::new())
143 }
144
145 fn create_test_request_without_content_length(body: Bytes) -> Request {
147 let uri: http::Uri = "/test".parse().unwrap();
148 let builder = http::Request::builder().method(Method::POST).uri(uri);
149
150 let req = builder.body(()).unwrap();
151 let (parts, _) = req.into_parts();
152
153 Request::new(parts, body, Arc::new(Extensions::new()), HashMap::new())
154 }
155
156 fn ok_handler() -> BoxedNext {
158 Arc::new(|_req: Request| {
159 Box::pin(async {
160 http::Response::builder()
161 .status(StatusCode::OK)
162 .body(http_body_util::Full::new(Bytes::from("ok")))
163 .unwrap()
164 }) as Pin<Box<dyn Future<Output = Response> + Send + 'static>>
165 })
166 }
167
168
169 proptest! {
176 #![proptest_config(ProptestConfig::with_cases(100))]
177
178 #[test]
179 fn prop_body_size_limit_enforcement(
180 limit in 1usize..10240usize,
182 body_size_factor in 0.5f64..2.0f64,
184 ) {
185 let rt = tokio::runtime::Runtime::new().unwrap();
186 rt.block_on(async {
187 let body_size = ((limit as f64) * body_size_factor) as usize;
188 let body = Bytes::from(vec![b'x'; body_size]);
189 let request = create_test_request_with_body(body.clone());
190
191 let layer = BodyLimitLayer::new(limit);
192 let handler = ok_handler();
193
194 let response = layer.call(request, handler).await;
195
196 if body_size > limit {
197 prop_assert_eq!(
199 response.status(),
200 StatusCode::PAYLOAD_TOO_LARGE,
201 "Expected 413 for body size {} > limit {}",
202 body_size,
203 limit
204 );
205 } else {
206 prop_assert_eq!(
208 response.status(),
209 StatusCode::OK,
210 "Expected 200 for body size {} <= limit {}",
211 body_size,
212 limit
213 );
214 }
215
216 Ok(())
217 })?;
218 }
219
220 #[test]
221 fn prop_body_limit_without_content_length_header(
222 limit in 1usize..10240usize,
223 body_size_factor in 0.5f64..2.0f64,
224 ) {
225 let rt = tokio::runtime::Runtime::new().unwrap();
226 rt.block_on(async {
227 let body_size = ((limit as f64) * body_size_factor) as usize;
228 let body = Bytes::from(vec![b'x'; body_size]);
229 let request = create_test_request_without_content_length(body.clone());
231
232 let layer = BodyLimitLayer::new(limit);
233 let handler = ok_handler();
234
235 let response = layer.call(request, handler).await;
236
237 if body_size > limit {
238 prop_assert_eq!(
240 response.status(),
241 StatusCode::PAYLOAD_TOO_LARGE,
242 "Expected 413 for body size {} > limit {} (no Content-Length)",
243 body_size,
244 limit
245 );
246 } else {
247 prop_assert_eq!(
249 response.status(),
250 StatusCode::OK,
251 "Expected 200 for body size {} <= limit {} (no Content-Length)",
252 body_size,
253 limit
254 );
255 }
256
257 Ok(())
258 })?;
259 }
260 }
261
262 #[tokio::test]
263 async fn test_body_at_exact_limit() {
264 let limit = 100;
265 let body = Bytes::from(vec![b'x'; limit]);
266 let request = create_test_request_with_body(body);
267
268 let layer = BodyLimitLayer::new(limit);
269 let handler = ok_handler();
270
271 let response = layer.call(request, handler).await;
272 assert_eq!(response.status(), StatusCode::OK);
273 }
274
275 #[tokio::test]
276 async fn test_body_one_byte_over_limit() {
277 let limit = 100;
278 let body = Bytes::from(vec![b'x'; limit + 1]);
279 let request = create_test_request_with_body(body);
280
281 let layer = BodyLimitLayer::new(limit);
282 let handler = ok_handler();
283
284 let response = layer.call(request, handler).await;
285 assert_eq!(response.status(), StatusCode::PAYLOAD_TOO_LARGE);
286 }
287
288 #[tokio::test]
289 async fn test_body_one_byte_under_limit() {
290 let limit = 100;
291 let body = Bytes::from(vec![b'x'; limit - 1]);
292 let request = create_test_request_with_body(body);
293
294 let layer = BodyLimitLayer::new(limit);
295 let handler = ok_handler();
296
297 let response = layer.call(request, handler).await;
298 assert_eq!(response.status(), StatusCode::OK);
299 }
300
301 #[tokio::test]
302 async fn test_empty_body() {
303 let limit = 100;
304 let body = Bytes::new();
305 let request = create_test_request_with_body(body);
306
307 let layer = BodyLimitLayer::new(limit);
308 let handler = ok_handler();
309
310 let response = layer.call(request, handler).await;
311 assert_eq!(response.status(), StatusCode::OK);
312 }
313
314 #[tokio::test]
315 async fn test_default_limit() {
316 let layer = BodyLimitLayer::default();
317 assert_eq!(layer.limit(), DEFAULT_BODY_LIMIT);
318 }
319
320 #[test]
321 fn test_clone() {
322 let layer = BodyLimitLayer::new(1024);
323 let cloned = layer.clone();
324 assert_eq!(layer.limit(), cloned.limit());
325 }
326}