rustapi_core/middleware/
body_limit.rs

1//! Body size limit middleware for RustAPI
2//!
3//! This module provides middleware to enforce request body size limits,
4//! protecting against denial-of-service attacks via large payloads.
5//!
6//! # Example
7//!
8//! ```rust,ignore
9//! use rustapi_rs::prelude::*;
10//! use rustapi_core::middleware::BodyLimitLayer;
11//!
12//! RustApi::new()
13//!     .layer(BodyLimitLayer::new(1024 * 1024)) // 1MB limit
14//!     .route("/upload", post(upload_handler))
15//!     .run("127.0.0.1:8080")
16//!     .await
17//! ```
18
19use 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
27/// Default body size limit: 1MB
28pub const DEFAULT_BODY_LIMIT: usize = 1024 * 1024;
29
30/// Body size limit middleware layer
31///
32/// Enforces a maximum size for request bodies. When a request body exceeds
33/// the configured limit, a 413 Payload Too Large response is returned.
34#[derive(Clone)]
35pub struct BodyLimitLayer {
36    limit: usize,
37}
38
39impl BodyLimitLayer {
40    /// Create a new body limit layer with the specified limit in bytes
41    ///
42    /// # Arguments
43    ///
44    /// * `limit` - Maximum body size in bytes
45    ///
46    /// # Example
47    ///
48    /// ```rust,ignore
49    /// // 2MB limit
50    /// let layer = BodyLimitLayer::new(2 * 1024 * 1024);
51    /// ```
52    pub fn new(limit: usize) -> Self {
53        Self { limit }
54    }
55
56    /// Create a body limit layer with the default limit (1MB)
57    pub fn default_limit() -> Self {
58        Self::new(DEFAULT_BODY_LIMIT)
59    }
60
61    /// Get the configured limit
62    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            // Check Content-Length header first if available
83            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            // Also check actual body size (for cases without Content-Length or streaming)
99            // The body has already been read at this point in the pipeline
100            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            // Body is within limits, continue to next middleware/handler
112            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    /// Create a test request with the given body
132    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        // Set Content-Length header
137        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    /// Create a test request without Content-Length header
146    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    /// Create a simple handler that returns 200 OK
157    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    // **Feature: phase4-ergonomics-v1, Property 3: Body Size Limit Enforcement**
170    //
171    // For any configured body size limit L and any request body B where size(B) > L,
172    // the system should return a 413 Payload Too Large response.
173    //
174    // **Validates: Requirements 2.2, 2.3, 2.4, 2.5**
175    proptest! {
176        #![proptest_config(ProptestConfig::with_cases(100))]
177
178        #[test]
179        fn prop_body_size_limit_enforcement(
180            // Generate limit between 1 and 10KB for testing
181            limit in 1usize..10240usize,
182            // Generate body size relative to limit
183            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                    // Body exceeds limit - should return 413
198                    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                    // Body within limit - should return 200
207                    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                // Create request without Content-Length header
230                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                    // Body exceeds limit - should return 413
239                    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                    // Body within limit - should return 200
248                    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}