/home/noah/src/trueno/src/brick/rate_limit.rs
Line | Count | Source |
1 | | //! DoS Prevention Rate Limiting |
2 | | //! |
3 | | //! AWP-15: Request validation and rate limiting for DoS prevention. |
4 | | |
5 | | use std::fmt; |
6 | | use std::time::Duration; |
7 | | |
8 | | // ---------------------------------------------------------------------------- |
9 | | // AWP-15: DoS Prevention Limits |
10 | | // ---------------------------------------------------------------------------- |
11 | | |
12 | | /// DoS prevention limits for serving. |
13 | | /// |
14 | | /// # Example |
15 | | /// ```rust |
16 | | /// use trueno::brick::ServeLimits; |
17 | | /// |
18 | | /// let limits = ServeLimits::default(); |
19 | | /// assert!(limits.validate_request(50, 1024).is_ok()); |
20 | | /// assert!(limits.validate_request(200, 1024).is_err()); // Too many headers |
21 | | /// ``` |
22 | | #[derive(Debug, Clone)] |
23 | | pub struct ServeLimits { |
24 | | /// Maximum request body size (bytes). |
25 | | pub max_request_size: usize, |
26 | | /// Maximum number of headers. |
27 | | pub max_headers: usize, |
28 | | /// Maximum header size (bytes). |
29 | | pub max_header_size: usize, |
30 | | /// Keep-alive timeout. |
31 | | pub keep_alive_timeout: Duration, |
32 | | /// Client request timeout. |
33 | | pub client_timeout: Duration, |
34 | | /// Maximum pipelined requests. |
35 | | pub max_pipelined: usize, |
36 | | /// Maximum concurrent connections. |
37 | | pub max_connections: usize, |
38 | | } |
39 | | |
40 | | impl Default for ServeLimits { |
41 | 0 | fn default() -> Self { |
42 | 0 | Self { |
43 | 0 | max_request_size: 2 * 1024 * 1024, // 2MB |
44 | 0 | max_headers: 100, |
45 | 0 | max_header_size: 8 * 1024, // 8KB |
46 | 0 | keep_alive_timeout: Duration::from_secs(5), |
47 | 0 | client_timeout: Duration::from_secs(5), |
48 | 0 | max_pipelined: 16, |
49 | 0 | max_connections: 1024, |
50 | 0 | } |
51 | 0 | } |
52 | | } |
53 | | |
54 | | impl ServeLimits { |
55 | | /// Create new limits with custom values. |
56 | 0 | pub fn new() -> Self { |
57 | 0 | Self::default() |
58 | 0 | } |
59 | | |
60 | | /// Builder: set max request size. |
61 | | #[must_use] |
62 | 0 | pub fn with_max_request_size(mut self, size: usize) -> Self { |
63 | 0 | self.max_request_size = size; |
64 | 0 | self |
65 | 0 | } |
66 | | |
67 | | /// Builder: set max headers. |
68 | | #[must_use] |
69 | 0 | pub fn with_max_headers(mut self, count: usize) -> Self { |
70 | 0 | self.max_headers = count; |
71 | 0 | self |
72 | 0 | } |
73 | | |
74 | | /// Builder: set max connections. |
75 | | #[must_use] |
76 | 0 | pub fn with_max_connections(mut self, count: usize) -> Self { |
77 | 0 | self.max_connections = count; |
78 | 0 | self |
79 | 0 | } |
80 | | |
81 | | /// Validate incoming request against limits. |
82 | 0 | pub fn validate_request( |
83 | 0 | &self, |
84 | 0 | headers_count: usize, |
85 | 0 | body_size: usize, |
86 | 0 | ) -> Result<(), LimitError> { |
87 | 0 | if headers_count > self.max_headers { |
88 | 0 | return Err(LimitError::TooManyHeaders { |
89 | 0 | count: headers_count, |
90 | 0 | max: self.max_headers, |
91 | 0 | }); |
92 | 0 | } |
93 | 0 | if body_size > self.max_request_size { |
94 | 0 | return Err(LimitError::BodyTooLarge { |
95 | 0 | size: body_size, |
96 | 0 | max: self.max_request_size, |
97 | 0 | }); |
98 | 0 | } |
99 | 0 | Ok(()) |
100 | 0 | } |
101 | | |
102 | | /// Validate header size. |
103 | 0 | pub fn validate_header_size(&self, size: usize) -> Result<(), LimitError> { |
104 | 0 | if size > self.max_header_size { |
105 | 0 | return Err(LimitError::HeaderTooLarge { |
106 | 0 | size, |
107 | 0 | max: self.max_header_size, |
108 | 0 | }); |
109 | 0 | } |
110 | 0 | Ok(()) |
111 | 0 | } |
112 | | |
113 | | /// Validate pipelined request count. |
114 | 0 | pub fn validate_pipelined(&self, count: usize) -> Result<(), LimitError> { |
115 | 0 | if count > self.max_pipelined { |
116 | 0 | return Err(LimitError::TooManyPipelined { |
117 | 0 | count, |
118 | 0 | max: self.max_pipelined, |
119 | 0 | }); |
120 | 0 | } |
121 | 0 | Ok(()) |
122 | 0 | } |
123 | | |
124 | | /// Validate connection count. |
125 | 0 | pub fn validate_connections(&self, current: usize) -> Result<(), LimitError> { |
126 | 0 | if current >= self.max_connections { |
127 | 0 | return Err(LimitError::ConnectionLimitReached { |
128 | 0 | current, |
129 | 0 | max: self.max_connections, |
130 | 0 | }); |
131 | 0 | } |
132 | 0 | Ok(()) |
133 | 0 | } |
134 | | } |
135 | | |
136 | | /// Error when a limit is exceeded. |
137 | | #[derive(Debug, Clone, PartialEq, Eq)] |
138 | | pub enum LimitError { |
139 | | /// Too many headers in request. |
140 | | TooManyHeaders { |
141 | | /// Actual count. |
142 | | count: usize, |
143 | | /// Maximum allowed. |
144 | | max: usize, |
145 | | }, |
146 | | /// Request body too large. |
147 | | BodyTooLarge { |
148 | | /// Actual size. |
149 | | size: usize, |
150 | | /// Maximum allowed. |
151 | | max: usize, |
152 | | }, |
153 | | /// Header too large. |
154 | | HeaderTooLarge { |
155 | | /// Actual size. |
156 | | size: usize, |
157 | | /// Maximum allowed. |
158 | | max: usize, |
159 | | }, |
160 | | /// Too many pipelined requests. |
161 | | TooManyPipelined { |
162 | | /// Actual count. |
163 | | count: usize, |
164 | | /// Maximum allowed. |
165 | | max: usize, |
166 | | }, |
167 | | /// Connection limit reached. |
168 | | ConnectionLimitReached { |
169 | | /// Current connections. |
170 | | current: usize, |
171 | | /// Maximum allowed. |
172 | | max: usize, |
173 | | }, |
174 | | } |
175 | | |
176 | | impl fmt::Display for LimitError { |
177 | 0 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
178 | 0 | match self { |
179 | 0 | LimitError::TooManyHeaders { count, max } => { |
180 | 0 | write!(f, "too many headers: {} (max {})", count, max) |
181 | | } |
182 | 0 | LimitError::BodyTooLarge { size, max } => { |
183 | 0 | write!(f, "body too large: {} bytes (max {})", size, max) |
184 | | } |
185 | 0 | LimitError::HeaderTooLarge { size, max } => { |
186 | 0 | write!(f, "header too large: {} bytes (max {})", size, max) |
187 | | } |
188 | 0 | LimitError::TooManyPipelined { count, max } => { |
189 | 0 | write!(f, "too many pipelined requests: {} (max {})", count, max) |
190 | | } |
191 | 0 | LimitError::ConnectionLimitReached { current, max } => { |
192 | 0 | write!(f, "connection limit reached: {} (max {})", current, max) |
193 | | } |
194 | | } |
195 | 0 | } |
196 | | } |
197 | | |
198 | | impl std::error::Error for LimitError {} |
199 | | |
200 | | #[cfg(test)] |
201 | | mod tests { |
202 | | use super::*; |
203 | | |
204 | | #[test] |
205 | | fn test_serve_limits_default() { |
206 | | let limits = ServeLimits::default(); |
207 | | assert_eq!(limits.max_request_size, 2 * 1024 * 1024); |
208 | | assert_eq!(limits.max_headers, 100); |
209 | | assert_eq!(limits.max_header_size, 8 * 1024); |
210 | | assert_eq!(limits.max_pipelined, 16); |
211 | | assert_eq!(limits.max_connections, 1024); |
212 | | } |
213 | | |
214 | | #[test] |
215 | | fn test_serve_limits_builder() { |
216 | | let limits = ServeLimits::new() |
217 | | .with_max_request_size(1024) |
218 | | .with_max_headers(50) |
219 | | .with_max_connections(100); |
220 | | |
221 | | assert_eq!(limits.max_request_size, 1024); |
222 | | assert_eq!(limits.max_headers, 50); |
223 | | assert_eq!(limits.max_connections, 100); |
224 | | } |
225 | | |
226 | | #[test] |
227 | | fn test_validate_request_ok() { |
228 | | let limits = ServeLimits::default(); |
229 | | assert!(limits.validate_request(50, 1024).is_ok()); |
230 | | } |
231 | | |
232 | | #[test] |
233 | | fn test_validate_request_too_many_headers() { |
234 | | let limits = ServeLimits::default(); |
235 | | let result = limits.validate_request(200, 1024); |
236 | | assert!(matches!(result, Err(LimitError::TooManyHeaders { .. }))); |
237 | | } |
238 | | |
239 | | #[test] |
240 | | fn test_validate_request_body_too_large() { |
241 | | let limits = ServeLimits::default(); |
242 | | let result = limits.validate_request(50, 10 * 1024 * 1024); |
243 | | assert!(matches!(result, Err(LimitError::BodyTooLarge { .. }))); |
244 | | } |
245 | | |
246 | | #[test] |
247 | | fn test_validate_header_size_ok() { |
248 | | let limits = ServeLimits::default(); |
249 | | assert!(limits.validate_header_size(1024).is_ok()); |
250 | | } |
251 | | |
252 | | #[test] |
253 | | fn test_validate_header_size_too_large() { |
254 | | let limits = ServeLimits::default(); |
255 | | let result = limits.validate_header_size(16 * 1024); |
256 | | assert!(matches!(result, Err(LimitError::HeaderTooLarge { .. }))); |
257 | | } |
258 | | |
259 | | #[test] |
260 | | fn test_validate_pipelined_ok() { |
261 | | let limits = ServeLimits::default(); |
262 | | assert!(limits.validate_pipelined(10).is_ok()); |
263 | | } |
264 | | |
265 | | #[test] |
266 | | fn test_validate_pipelined_too_many() { |
267 | | let limits = ServeLimits::default(); |
268 | | let result = limits.validate_pipelined(20); |
269 | | assert!(matches!(result, Err(LimitError::TooManyPipelined { .. }))); |
270 | | } |
271 | | |
272 | | #[test] |
273 | | fn test_validate_connections_ok() { |
274 | | let limits = ServeLimits::default(); |
275 | | assert!(limits.validate_connections(500).is_ok()); |
276 | | } |
277 | | |
278 | | #[test] |
279 | | fn test_validate_connections_limit_reached() { |
280 | | let limits = ServeLimits::default(); |
281 | | let result = limits.validate_connections(1024); |
282 | | assert!(matches!( |
283 | | result, |
284 | | Err(LimitError::ConnectionLimitReached { .. }) |
285 | | )); |
286 | | } |
287 | | |
288 | | #[test] |
289 | | fn test_limit_error_display() { |
290 | | let err = LimitError::TooManyHeaders { count: 150, max: 100 }; |
291 | | assert_eq!(format!("{}", err), "too many headers: 150 (max 100)"); |
292 | | |
293 | | let err = LimitError::BodyTooLarge { |
294 | | size: 5000000, |
295 | | max: 2097152, |
296 | | }; |
297 | | assert_eq!( |
298 | | format!("{}", err), |
299 | | "body too large: 5000000 bytes (max 2097152)" |
300 | | ); |
301 | | } |
302 | | |
303 | | #[test] |
304 | | fn test_limit_error_eq() { |
305 | | let err1 = LimitError::TooManyHeaders { count: 150, max: 100 }; |
306 | | let err2 = LimitError::TooManyHeaders { count: 150, max: 100 }; |
307 | | let err3 = LimitError::TooManyHeaders { count: 200, max: 100 }; |
308 | | |
309 | | assert_eq!(err1, err2); |
310 | | assert_ne!(err1, err3); |
311 | | } |
312 | | |
313 | | /// FALSIFICATION TEST: Verify boundary conditions |
314 | | /// |
315 | | /// Limits must reject exactly at the boundary, not before or after. |
316 | | #[test] |
317 | | fn test_falsify_exact_boundaries() { |
318 | | let limits = ServeLimits::new() |
319 | | .with_max_headers(100) |
320 | | .with_max_request_size(1000) |
321 | | .with_max_connections(10); |
322 | | |
323 | | // Headers: 100 is the max, so 100 should be OK but 101 should fail |
324 | | assert!(limits.validate_request(100, 0).is_ok()); |
325 | | assert!(limits.validate_request(101, 0).is_err()); |
326 | | |
327 | | // Body: 1000 is the max, so 1000 should be OK but 1001 should fail |
328 | | assert!(limits.validate_request(0, 1000).is_ok()); |
329 | | assert!(limits.validate_request(0, 1001).is_err()); |
330 | | |
331 | | // Connections: Using >=, so 10 should fail but 9 should be OK |
332 | | assert!(limits.validate_connections(9).is_ok()); |
333 | | assert!(limits.validate_connections(10).is_err()); |
334 | | } |
335 | | |
336 | | /// FALSIFICATION TEST: Verify all validation passes for zero values |
337 | | #[test] |
338 | | fn test_falsify_zero_values_pass() { |
339 | | let limits = ServeLimits::default(); |
340 | | |
341 | | // Zero headers and zero body should always pass |
342 | | assert!( |
343 | | limits.validate_request(0, 0).is_ok(), |
344 | | "FALSIFICATION FAILED: Zero values should always pass" |
345 | | ); |
346 | | assert!( |
347 | | limits.validate_header_size(0).is_ok(), |
348 | | "FALSIFICATION FAILED: Zero header size should pass" |
349 | | ); |
350 | | assert!( |
351 | | limits.validate_pipelined(0).is_ok(), |
352 | | "FALSIFICATION FAILED: Zero pipelined should pass" |
353 | | ); |
354 | | assert!( |
355 | | limits.validate_connections(0).is_ok(), |
356 | | "FALSIFICATION FAILED: Zero connections should pass" |
357 | | ); |
358 | | } |
359 | | } |