/home/noah/src/trueno/src/brick/rate_limit.rs

Source
//! DoS Prevention Rate Limiting
//!
//! AWP-15: Request validation and rate limiting for DoS prevention.

use std::fmt;
use std::time::Duration;

// ----------------------------------------------------------------------------
// AWP-15: DoS Prevention Limits
// ----------------------------------------------------------------------------

/// DoS prevention limits for serving.
///
/// # Example
/// ```rust
/// use trueno::brick::ServeLimits;
///
/// let limits = ServeLimits::default();
/// assert!(limits.validate_request(50, 1024).is_ok());
/// assert!(limits.validate_request(200, 1024).is_err());  // Too many headers
/// ```
#[derive(Debug, Clone)]
pub struct ServeLimits {
    /// Maximum request body size (bytes).
    pub max_request_size: usize,
    /// Maximum number of headers.
    pub max_headers: usize,
    /// Maximum header size (bytes).
    pub max_header_size: usize,
    /// Keep-alive timeout.
    pub keep_alive_timeout: Duration,
    /// Client request timeout.
    pub client_timeout: Duration,
    /// Maximum pipelined requests.
    pub max_pipelined: usize,
    /// Maximum concurrent connections.
    pub max_connections: usize,
}

impl Default for ServeLimits {
    fn default() -> Self {
        Self {
            max_request_size: 2 * 1024 * 1024, // 2MB
            max_headers: 100,
            max_header_size: 8 * 1024, // 8KB
            keep_alive_timeout: Duration::from_secs(5),
            client_timeout: Duration::from_secs(5),
            max_pipelined: 16,
            max_connections: 1024,
        }
    }
}

impl ServeLimits {
    /// Create new limits with custom values.
    pub fn new() -> Self {
        Self::default()
    }

    /// Builder: set max request size.
    #[must_use]
    pub fn with_max_request_size(mut self, size: usize) -> Self {
        self.max_request_size = size;
        self
    }

    /// Builder: set max headers.
    #[must_use]
    pub fn with_max_headers(mut self, count: usize) -> Self {
        self.max_headers = count;
        self
    }

    /// Builder: set max connections.
    #[must_use]
    pub fn with_max_connections(mut self, count: usize) -> Self {
        self.max_connections = count;
        self
    }

    /// Validate incoming request against limits.
    pub fn validate_request(
        &self,
        headers_count: usize,
        body_size: usize,
    ) -> Result<(), LimitError> {
        if headers_count > self.max_headers {
            return Err(LimitError::TooManyHeaders {
                count: headers_count,
                max: self.max_headers,
            });
        }
        if body_size > self.max_request_size {
            return Err(LimitError::BodyTooLarge {
                size: body_size,
                max: self.max_request_size,
            });
        }
        Ok(())
    }

    /// Validate header size.
    pub fn validate_header_size(&self, size: usize) -> Result<(), LimitError> {
        if size > self.max_header_size {
            return Err(LimitError::HeaderTooLarge {
                size,
                max: self.max_header_size,
            });
        }
        Ok(())
    }

    /// Validate pipelined request count.
    pub fn validate_pipelined(&self, count: usize) -> Result<(), LimitError> {
        if count > self.max_pipelined {
            return Err(LimitError::TooManyPipelined {
                count,
                max: self.max_pipelined,
            });
        }
        Ok(())
    }

    /// Validate connection count.
    pub fn validate_connections(&self, current: usize) -> Result<(), LimitError> {
        if current >= self.max_connections {
            return Err(LimitError::ConnectionLimitReached {
                current,
                max: self.max_connections,
            });
        }
        Ok(())
    }
}

/// Error when a limit is exceeded.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum LimitError {
    /// Too many headers in request.
    TooManyHeaders {
        /// Actual count.
        count: usize,
        /// Maximum allowed.
        max: usize,
    },
    /// Request body too large.
    BodyTooLarge {
        /// Actual size.
        size: usize,
        /// Maximum allowed.
        max: usize,
    },
    /// Header too large.
    HeaderTooLarge {
        /// Actual size.
        size: usize,
        /// Maximum allowed.
        max: usize,
    },
    /// Too many pipelined requests.
    TooManyPipelined {
        /// Actual count.
        count: usize,
        /// Maximum allowed.
        max: usize,
    },
    /// Connection limit reached.
    ConnectionLimitReached {
        /// Current connections.
        current: usize,
        /// Maximum allowed.
        max: usize,
    },
}

impl fmt::Display for LimitError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            LimitError::TooManyHeaders { count, max } => {
                write!(f, "too many headers: {} (max {})", count, max)
            }
            LimitError::BodyTooLarge { size, max } => {
                write!(f, "body too large: {} bytes (max {})", size, max)
            }
            LimitError::HeaderTooLarge { size, max } => {
                write!(f, "header too large: {} bytes (max {})", size, max)
            }
            LimitError::TooManyPipelined { count, max } => {
                write!(f, "too many pipelined requests: {} (max {})", count, max)
            }
            LimitError::ConnectionLimitReached { current, max } => {
                write!(f, "connection limit reached: {} (max {})", current, max)
            }
        }
    }
}

impl std::error::Error for LimitError {}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_serve_limits_default() {
        let limits = ServeLimits::default();
        assert_eq!(limits.max_request_size, 2 * 1024 * 1024);
        assert_eq!(limits.max_headers, 100);
        assert_eq!(limits.max_header_size, 8 * 1024);
        assert_eq!(limits.max_pipelined, 16);
        assert_eq!(limits.max_connections, 1024);
    }

    #[test]
    fn test_serve_limits_builder() {
        let limits = ServeLimits::new()
            .with_max_request_size(1024)
            .with_max_headers(50)
            .with_max_connections(100);

        assert_eq!(limits.max_request_size, 1024);
        assert_eq!(limits.max_headers, 50);
        assert_eq!(limits.max_connections, 100);
    }

    #[test]
    fn test_validate_request_ok() {
        let limits = ServeLimits::default();
        assert!(limits.validate_request(50, 1024).is_ok());
    }

    #[test]
    fn test_validate_request_too_many_headers() {
        let limits = ServeLimits::default();
        let result = limits.validate_request(200, 1024);
        assert!(matches!(result, Err(LimitError::TooManyHeaders { .. })));
    }

    #[test]
    fn test_validate_request_body_too_large() {
        let limits = ServeLimits::default();
        let result = limits.validate_request(50, 10 * 1024 * 1024);
        assert!(matches!(result, Err(LimitError::BodyTooLarge { .. })));
    }

    #[test]
    fn test_validate_header_size_ok() {
        let limits = ServeLimits::default();
        assert!(limits.validate_header_size(1024).is_ok());
    }

    #[test]
    fn test_validate_header_size_too_large() {
        let limits = ServeLimits::default();
        let result = limits.validate_header_size(16 * 1024);
        assert!(matches!(result, Err(LimitError::HeaderTooLarge { .. })));
    }

    #[test]
    fn test_validate_pipelined_ok() {
        let limits = ServeLimits::default();
        assert!(limits.validate_pipelined(10).is_ok());
    }

    #[test]
    fn test_validate_pipelined_too_many() {
        let limits = ServeLimits::default();
        let result = limits.validate_pipelined(20);
        assert!(matches!(result, Err(LimitError::TooManyPipelined { .. })));
    }

    #[test]
    fn test_validate_connections_ok() {
        let limits = ServeLimits::default();
        assert!(limits.validate_connections(500).is_ok());
    }

    #[test]
    fn test_validate_connections_limit_reached() {
        let limits = ServeLimits::default();
        let result = limits.validate_connections(1024);
        assert!(matches!(
            result,
            Err(LimitError::ConnectionLimitReached { .. })
        ));
    }

    #[test]
    fn test_limit_error_display() {
        let err = LimitError::TooManyHeaders { count: 150, max: 100 };
        assert_eq!(format!("{}", err), "too many headers: 150 (max 100)");

        let err = LimitError::BodyTooLarge {
            size: 5000000,
            max: 2097152,
        };
        assert_eq!(
            format!("{}", err),
            "body too large: 5000000 bytes (max 2097152)"
        );
    }

    #[test]
    fn test_limit_error_eq() {
        let err1 = LimitError::TooManyHeaders { count: 150, max: 100 };
        let err2 = LimitError::TooManyHeaders { count: 150, max: 100 };
        let err3 = LimitError::TooManyHeaders { count: 200, max: 100 };

        assert_eq!(err1, err2);
        assert_ne!(err1, err3);
    }

    /// FALSIFICATION TEST: Verify boundary conditions
    ///
    /// Limits must reject exactly at the boundary, not before or after.
    #[test]
    fn test_falsify_exact_boundaries() {
        let limits = ServeLimits::new()
            .with_max_headers(100)
            .with_max_request_size(1000)
            .with_max_connections(10);

        // Headers: 100 is the max, so 100 should be OK but 101 should fail
        assert!(limits.validate_request(100, 0).is_ok());
        assert!(limits.validate_request(101, 0).is_err());

        // Body: 1000 is the max, so 1000 should be OK but 1001 should fail
        assert!(limits.validate_request(0, 1000).is_ok());
        assert!(limits.validate_request(0, 1001).is_err());

        // Connections: Using >=, so 10 should fail but 9 should be OK
        assert!(limits.validate_connections(9).is_ok());
        assert!(limits.validate_connections(10).is_err());
    }

    /// FALSIFICATION TEST: Verify all validation passes for zero values
    #[test]
    fn test_falsify_zero_values_pass() {
        let limits = ServeLimits::default();

        // Zero headers and zero body should always pass
        assert!(
            limits.validate_request(0, 0).is_ok(),
            "FALSIFICATION FAILED: Zero values should always pass"
        );
        assert!(
            limits.validate_header_size(0).is_ok(),
            "FALSIFICATION FAILED: Zero header size should pass"
        );
        assert!(
            limits.validate_pipelined(0).is_ok(),
            "FALSIFICATION FAILED: Zero pipelined should pass"
        );
        assert!(
            limits.validate_connections(0).is_ok(),
            "FALSIFICATION FAILED: Zero connections should pass"
        );
    }
}

Coverage Report

Created: 2026-01-25 15:05

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		}