elif_http/middleware/utils/
body_limit.rs

1//! # Body Limit Middleware
2//!
3//! Framework middleware for request body size limiting using V2 system.
4//! Replaces tower-http RequestBodyLimitLayer with framework-native implementation.
5
6use tracing::warn;
7
8use crate::{
9    middleware::v2::{Middleware, Next, NextFuture},
10    request::ElifRequest,
11    response::{ElifResponse, ElifStatusCode},
12};
13
14/// Configuration for body limit middleware
15#[derive(Debug, Clone)]
16pub struct BodyLimitConfig {
17    /// Maximum allowed body size in bytes
18    pub max_size: usize,
19    /// Whether to log oversized requests
20    pub log_oversized: bool,
21    /// Custom error message for oversized requests
22    pub error_message: String,
23    /// Whether to include Content-Length header in error response
24    pub include_headers: bool,
25}
26
27impl Default for BodyLimitConfig {
28    fn default() -> Self {
29        Self {
30            max_size: 2 * 1024 * 1024, // 2MB default
31            log_oversized: true,
32            error_message: "Request body too large".to_string(),
33            include_headers: true,
34        }
35    }
36}
37
38impl BodyLimitConfig {
39    /// Create new body limit configuration
40    pub fn new(max_size: usize) -> Self {
41        Self {
42            max_size,
43            ..Default::default()
44        }
45    }
46
47    /// Set maximum body size
48    pub fn with_max_size(mut self, max_size: usize) -> Self {
49        self.max_size = max_size;
50        self
51    }
52
53    /// Enable or disable logging of oversized requests
54    pub fn with_logging(mut self, log_oversized: bool) -> Self {
55        self.log_oversized = log_oversized;
56        self
57    }
58
59    /// Set custom error message
60    pub fn with_message<S: Into<String>>(mut self, message: S) -> Self {
61        self.error_message = message.into();
62        self
63    }
64
65    /// Include helpful headers in error response
66    pub fn with_headers(mut self, include_headers: bool) -> Self {
67        self.include_headers = include_headers;
68        self
69    }
70}
71
72/// Framework body limit middleware for HTTP requests
73#[derive(Debug)]
74pub struct BodyLimitMiddleware {
75    config: BodyLimitConfig,
76}
77
78impl BodyLimitMiddleware {
79    /// Create new body limit middleware with default 2MB limit
80    pub fn new() -> Self {
81        Self {
82            config: BodyLimitConfig::default(),
83        }
84    }
85
86    /// Create body limit middleware with specific size limit
87    pub fn with_limit(max_size: usize) -> Self {
88        Self {
89            config: BodyLimitConfig::new(max_size),
90        }
91    }
92
93    /// Create body limit middleware with custom configuration
94    pub fn with_config(config: BodyLimitConfig) -> Self {
95        Self { config }
96    }
97
98    /// Set maximum body size (builder pattern)
99    pub fn max_size(mut self, size: usize) -> Self {
100        self.config = self.config.with_max_size(size);
101        self
102    }
103
104    /// Enable or disable logging (builder pattern)
105    pub fn logging(mut self, enabled: bool) -> Self {
106        self.config = self.config.with_logging(enabled);
107        self
108    }
109
110    /// Set custom error message (builder pattern)
111    pub fn message<S: Into<String>>(mut self, message: S) -> Self {
112        self.config = self.config.with_message(message);
113        self
114    }
115
116    /// Get configured max size
117    pub fn limit(&self) -> usize {
118        self.config.max_size
119    }
120}
121
122impl Default for BodyLimitMiddleware {
123    fn default() -> Self {
124        Self::new()
125    }
126}
127
128impl Middleware for BodyLimitMiddleware {
129    fn handle(&self, request: ElifRequest, next: Next) -> NextFuture<'static> {
130        let config = self.config.clone();
131        Box::pin(async move {
132            // First, check Content-Length header if present
133            let _content_length = {
134                if let Some(content_length) = request.headers.get_str("content-length") {
135                    if let Ok(content_length_str) = content_length.to_str() {
136                        if let Ok(content_length) = content_length_str.parse::<usize>() {
137                            if content_length > config.max_size {
138                                if config.log_oversized {
139                                    warn!(
140                                        "Request body size {} bytes exceeds limit of {} bytes (Content-Length check)",
141                                        content_length,
142                                        config.max_size
143                                    );
144                                }
145
146                                let mut response =
147                                    ElifResponse::with_status(ElifStatusCode::PAYLOAD_TOO_LARGE)
148                                        .text(format!(
149                                            "Request body size {} bytes exceeds limit of {} bytes",
150                                            content_length, config.max_size
151                                        ));
152
153                                if config.include_headers {
154                                    if let Err(e) = response
155                                        .add_header("X-Max-Body-Size", config.max_size.to_string())
156                                    {
157                                        warn!("Failed to add X-Max-Body-Size header: {}", e);
158                                    }
159                                }
160
161                                return response;
162                            }
163                            Some(content_length)
164                        } else {
165                            None
166                        }
167                    } else {
168                        None
169                    }
170                } else {
171                    None
172                }
173            };
174
175            // For streaming bodies or cases where Content-Length is not reliable,
176            // we need to check the actual body size during consumption.
177            // This is typically handled by the framework's built-in body limiting or
178            // custom extractors that check size during body reading.
179
180            // Continue to next middleware/handler
181            let response = next.run(request).await;
182
183            // Log if we're returning a payload too large error
184            if response.status_code() == ElifStatusCode::PAYLOAD_TOO_LARGE && config.log_oversized {
185                warn!("Returned 413 Payload Too Large response due to body size limit");
186            }
187
188            response
189        })
190    }
191
192    fn name(&self) -> &'static str {
193        "BodyLimitMiddleware"
194    }
195}
196
197/// Body limit information for tracking
198#[derive(Debug, Clone)]
199pub struct BodyLimitInfo {
200    pub max_size: usize,
201    pub content_length: Option<usize>,
202    pub error_message: String,
203}
204
205/// Utility functions for common body size limits
206pub mod limits {
207    /// 1KB limit
208    pub const KB: usize = 1024;
209
210    /// 1MB limit
211    pub const MB: usize = 1024 * 1024;
212
213    /// 10MB limit
214    pub const MB_10: usize = 10 * MB;
215
216    /// 100MB limit
217    pub const MB_100: usize = 100 * MB;
218
219    /// 1GB limit (use with caution)
220    pub const GB: usize = 1024 * MB;
221
222    /// Create body limit middleware with common sizes
223    pub mod presets {
224        use super::super::BodyLimitMiddleware;
225        use super::*;
226
227        /// Small API requests (1MB)
228        pub fn small_api() -> BodyLimitMiddleware {
229            BodyLimitMiddleware::with_limit(MB).message("API request body too large (1MB limit)")
230        }
231
232        /// File uploads (10MB)
233        pub fn file_upload() -> BodyLimitMiddleware {
234            BodyLimitMiddleware::with_limit(MB_10).message("File upload too large (10MB limit)")
235        }
236
237        /// Large file uploads (100MB)
238        pub fn large_upload() -> BodyLimitMiddleware {
239            BodyLimitMiddleware::with_limit(MB_100)
240                .message("Large file upload too large (100MB limit)")
241        }
242
243        /// Tiny requests (64KB)
244        pub fn tiny() -> BodyLimitMiddleware {
245            BodyLimitMiddleware::with_limit(64 * KB).message("Request body too large (64KB limit)")
246        }
247    }
248}
249
250#[cfg(test)]
251mod tests {
252    use super::*;
253    use crate::middleware::v2::MiddlewarePipelineV2;
254    use crate::request::{ElifMethod, ElifRequest};
255    use crate::response::headers::ElifHeaderMap;
256    use crate::response::{ElifResponse, ElifStatusCode};
257
258    #[tokio::test]
259    async fn test_body_limit_middleware_v2() {
260        let middleware = BodyLimitMiddleware::new();
261        let pipeline = MiddlewarePipelineV2::new().add(middleware);
262
263        let headers = ElifHeaderMap::new();
264        let request = ElifRequest::new(ElifMethod::POST, "/test".parse().unwrap(), headers);
265
266        let response = pipeline
267            .execute(request, |_req| {
268                Box::pin(async move { ElifResponse::ok().text("Success") })
269            })
270            .await;
271
272        assert_eq!(response.status_code(), ElifStatusCode::OK);
273    }
274
275    #[tokio::test]
276    async fn test_body_limit_middleware_custom_limit() {
277        let middleware = BodyLimitMiddleware::with_limit(1024); // 1KB
278        assert_eq!(middleware.limit(), 1024);
279    }
280
281    #[tokio::test]
282    async fn test_body_limit_middleware_builder() {
283        let middleware = BodyLimitMiddleware::new()
284            .max_size(512)
285            .logging(false)
286            .message("Too big!");
287
288        assert_eq!(middleware.config.max_size, 512);
289        assert!(!middleware.config.log_oversized);
290        assert_eq!(middleware.config.error_message, "Too big!");
291    }
292
293    #[tokio::test]
294    async fn test_content_length_check_within_limit() {
295        let middleware = BodyLimitMiddleware::with_limit(1000);
296        let pipeline = MiddlewarePipelineV2::new().add(middleware);
297
298        let mut headers = ElifHeaderMap::new();
299        headers.insert("content-length".parse().unwrap(), "500".parse().unwrap());
300
301        let request = ElifRequest::new(ElifMethod::POST, "/test".parse().unwrap(), headers);
302
303        let response = pipeline
304            .execute(request, |_req| {
305                Box::pin(async move { ElifResponse::ok().text("Success") })
306            })
307            .await;
308
309        assert_eq!(response.status_code(), ElifStatusCode::OK);
310    }
311
312    #[tokio::test]
313    async fn test_content_length_check_exceeds_limit() {
314        let middleware = BodyLimitMiddleware::with_limit(100);
315        let pipeline = MiddlewarePipelineV2::new().add(middleware);
316
317        let mut headers = ElifHeaderMap::new();
318        headers.insert("content-length".parse().unwrap(), "200".parse().unwrap());
319
320        let request = ElifRequest::new(ElifMethod::POST, "/test".parse().unwrap(), headers);
321
322        let response = pipeline
323            .execute(request, |_req| {
324                Box::pin(async move { ElifResponse::ok().text("Should not reach here") })
325            })
326            .await;
327
328        assert_eq!(response.status_code(), ElifStatusCode::PAYLOAD_TOO_LARGE);
329        assert!(response.has_header("X-Max-Body-Size"));
330    }
331
332    #[tokio::test]
333    async fn test_body_limit_config() {
334        let config = BodyLimitConfig::new(512)
335            .with_logging(false)
336            .with_message("Custom message")
337            .with_headers(false);
338
339        let middleware = BodyLimitMiddleware::with_config(config);
340
341        assert_eq!(middleware.config.max_size, 512);
342        assert!(!middleware.config.log_oversized);
343        assert_eq!(middleware.config.error_message, "Custom message");
344        assert!(!middleware.config.include_headers);
345    }
346
347    #[tokio::test]
348    async fn test_body_limit_middleware_name() {
349        let middleware = BodyLimitMiddleware::new();
350        assert_eq!(middleware.name(), "BodyLimitMiddleware");
351    }
352
353    #[tokio::test]
354    async fn test_body_limit_presets() {
355        let small = limits::presets::small_api();
356        assert_eq!(small.limit(), limits::MB);
357
358        let upload = limits::presets::file_upload();
359        assert_eq!(upload.limit(), limits::MB_10);
360
361        let large = limits::presets::large_upload();
362        assert_eq!(large.limit(), limits::MB_100);
363
364        let tiny = limits::presets::tiny();
365        assert_eq!(tiny.limit(), 64 * limits::KB);
366    }
367
368    #[tokio::test]
369    async fn test_body_limit_constants() {
370        assert_eq!(limits::KB, 1024);
371        assert_eq!(limits::MB, 1024 * 1024);
372        assert_eq!(limits::MB_10, 10 * 1024 * 1024);
373        assert_eq!(limits::MB_100, 100 * 1024 * 1024);
374        assert_eq!(limits::GB, 1024 * 1024 * 1024);
375    }
376
377    #[tokio::test]
378    async fn test_invalid_content_length_header() {
379        let middleware = BodyLimitMiddleware::with_limit(1000);
380        let pipeline = MiddlewarePipelineV2::new().add(middleware);
381
382        let mut headers = ElifHeaderMap::new();
383        headers.insert(
384            "content-length".parse().unwrap(),
385            "not-a-number".parse().unwrap(),
386        );
387
388        let request = ElifRequest::new(ElifMethod::POST, "/test".parse().unwrap(), headers);
389
390        // Should not error on invalid content-length, just ignore it
391        let response = pipeline
392            .execute(request, |_req| {
393                Box::pin(async move { ElifResponse::ok().text("Success") })
394            })
395            .await;
396
397        assert_eq!(response.status_code(), ElifStatusCode::OK);
398    }
399}
elif_http/middleware/utils/body_limit.rs

elif_http/middleware/utils/
body_limit.rs
