1use std::time::Duration;
4use thiserror::Error;
5
6#[derive(Error, Debug, Clone, PartialEq)]
8pub enum LimitViolation {
9 #[error("execution time limit exceeded: {limit:?} (actual: {actual:?})")]
11 TimeExceeded {
12 limit: Duration,
14 actual: Duration,
16 },
17
18 #[error("memory limit exceeded: {limit} bytes (actual: {actual} bytes)")]
20 MemoryExceeded {
21 limit: usize,
23 actual: usize,
25 },
26
27 #[error("instruction limit exceeded: {limit} (actual: {actual})")]
29 InstructionsExceeded {
30 limit: u64,
32 actual: u64,
34 },
35
36 #[error("stack depth limit exceeded: {limit} (actual: {actual})")]
38 StackDepthExceeded {
39 limit: usize,
41 actual: usize,
43 },
44
45 #[error("output size limit exceeded: {limit} bytes (actual: {actual} bytes)")]
47 OutputSizeExceeded {
48 limit: usize,
50 actual: usize,
52 },
53
54 #[error("filesystem operation limit exceeded: {limit} operations")]
56 FsOpsExceeded {
57 limit: usize,
59 },
60
61 #[error("network operation limit exceeded: {limit} operations")]
63 NetOpsExceeded {
64 limit: usize,
66 },
67}
68
69#[derive(Debug, Clone, PartialEq)]
74pub struct Limits {
75 pub timeout: Option<Duration>,
77
78 pub memory_bytes: Option<usize>,
80
81 pub max_instructions: Option<u64>,
83
84 pub max_stack_depth: Option<usize>,
86
87 pub max_output_bytes: Option<usize>,
89
90 pub max_fs_ops: Option<usize>,
92
93 pub max_net_ops: Option<usize>,
95
96 pub max_concurrent_tasks: Option<usize>,
98}
99
100impl Default for Limits {
101 fn default() -> Self {
102 Self {
103 timeout: Some(Duration::from_secs(30)),
104 memory_bytes: Some(64 * 1024 * 1024), max_instructions: Some(10_000_000), max_stack_depth: Some(1000),
107 max_output_bytes: Some(1024 * 1024), max_fs_ops: Some(100),
109 max_net_ops: Some(10),
110 max_concurrent_tasks: Some(16),
111 }
112 }
113}
114
115impl Limits {
116 pub fn unlimited() -> Self {
118 Self {
119 timeout: None,
120 memory_bytes: None,
121 max_instructions: None,
122 max_stack_depth: None,
123 max_output_bytes: None,
124 max_fs_ops: None,
125 max_net_ops: None,
126 max_concurrent_tasks: None,
127 }
128 }
129
130 pub fn strict() -> Self {
132 Self {
133 timeout: Some(Duration::from_secs(5)),
134 memory_bytes: Some(16 * 1024 * 1024), max_instructions: Some(1_000_000), max_stack_depth: Some(100),
137 max_output_bytes: Some(64 * 1024), max_fs_ops: Some(0), max_net_ops: Some(0), max_concurrent_tasks: Some(4),
141 }
142 }
143
144 pub fn with_timeout(mut self, timeout: Duration) -> Self {
146 self.timeout = Some(timeout);
147 self
148 }
149
150 pub fn with_memory_bytes(mut self, bytes: usize) -> Self {
152 self.memory_bytes = Some(bytes);
153 self
154 }
155
156 pub fn with_memory_mb(mut self, mb: usize) -> Self {
158 self.memory_bytes = Some(mb * 1024 * 1024);
159 self
160 }
161
162 pub fn with_max_instructions(mut self, count: u64) -> Self {
164 self.max_instructions = Some(count);
165 self
166 }
167
168 pub fn with_max_stack_depth(mut self, depth: usize) -> Self {
170 self.max_stack_depth = Some(depth);
171 self
172 }
173
174 pub fn with_max_output_bytes(mut self, bytes: usize) -> Self {
176 self.max_output_bytes = Some(bytes);
177 self
178 }
179
180 pub fn with_max_fs_ops(mut self, ops: usize) -> Self {
182 self.max_fs_ops = Some(ops);
183 self
184 }
185
186 pub fn with_max_net_ops(mut self, ops: usize) -> Self {
188 self.max_net_ops = Some(ops);
189 self
190 }
191
192 pub fn with_max_concurrent_tasks(mut self, tasks: usize) -> Self {
194 self.max_concurrent_tasks = Some(tasks);
195 self
196 }
197
198 pub fn no_timeout(mut self) -> Self {
200 self.timeout = None;
201 self
202 }
203
204 pub fn check_time(&self, elapsed: Duration) -> Result<(), LimitViolation> {
206 if let Some(limit) = self.timeout {
207 if elapsed > limit {
208 return Err(LimitViolation::TimeExceeded {
209 limit,
210 actual: elapsed,
211 });
212 }
213 }
214 Ok(())
215 }
216
217 pub fn check_memory(&self, used: usize) -> Result<(), LimitViolation> {
219 if let Some(limit) = self.memory_bytes {
220 if used > limit {
221 return Err(LimitViolation::MemoryExceeded {
222 limit,
223 actual: used,
224 });
225 }
226 }
227 Ok(())
228 }
229
230 pub fn check_instructions(&self, count: u64) -> Result<(), LimitViolation> {
232 if let Some(limit) = self.max_instructions {
233 if count > limit {
234 return Err(LimitViolation::InstructionsExceeded {
235 limit,
236 actual: count,
237 });
238 }
239 }
240 Ok(())
241 }
242
243 pub fn check_stack_depth(&self, depth: usize) -> Result<(), LimitViolation> {
245 if let Some(limit) = self.max_stack_depth {
246 if depth > limit {
247 return Err(LimitViolation::StackDepthExceeded {
248 limit,
249 actual: depth,
250 });
251 }
252 }
253 Ok(())
254 }
255}
256
257#[derive(Debug, Clone)]
259pub struct LimitTracker {
260 limits: Limits,
261 start_time: std::time::Instant,
262 instructions_executed: u64,
263 memory_used: usize,
264 current_stack_depth: usize,
265 output_bytes: usize,
266 fs_ops: usize,
267 net_ops: usize,
268}
269
270impl LimitTracker {
271 pub fn new(limits: Limits) -> Self {
273 Self {
274 limits,
275 start_time: std::time::Instant::now(),
276 instructions_executed: 0,
277 memory_used: 0,
278 current_stack_depth: 0,
279 output_bytes: 0,
280 fs_ops: 0,
281 net_ops: 0,
282 }
283 }
284
285 pub fn reset(&mut self) {
287 self.start_time = std::time::Instant::now();
288 self.instructions_executed = 0;
289 self.memory_used = 0;
290 self.current_stack_depth = 0;
291 self.output_bytes = 0;
292 self.fs_ops = 0;
293 self.net_ops = 0;
294 }
295
296 pub fn check_timeout(&self) -> Result<(), LimitViolation> {
298 self.limits.check_time(self.start_time.elapsed())
299 }
300
301 pub fn record_instructions(&mut self, count: u64) -> Result<(), LimitViolation> {
303 self.instructions_executed += count;
304 self.limits.check_instructions(self.instructions_executed)
305 }
306
307 pub fn record_memory(&mut self, bytes: usize) -> Result<(), LimitViolation> {
309 self.memory_used = bytes;
310 self.limits.check_memory(self.memory_used)
311 }
312
313 pub fn push_stack(&mut self) -> Result<(), LimitViolation> {
315 self.current_stack_depth += 1;
316 self.limits.check_stack_depth(self.current_stack_depth)
317 }
318
319 pub fn pop_stack(&mut self) {
321 self.current_stack_depth = self.current_stack_depth.saturating_sub(1);
322 }
323
324 pub fn record_output(&mut self, bytes: usize) -> Result<(), LimitViolation> {
326 self.output_bytes += bytes;
327 if let Some(limit) = self.limits.max_output_bytes {
328 if self.output_bytes > limit {
329 return Err(LimitViolation::OutputSizeExceeded {
330 limit,
331 actual: self.output_bytes,
332 });
333 }
334 }
335 Ok(())
336 }
337
338 pub fn record_fs_op(&mut self) -> Result<(), LimitViolation> {
340 self.fs_ops += 1;
341 if let Some(limit) = self.limits.max_fs_ops {
342 if self.fs_ops > limit {
343 return Err(LimitViolation::FsOpsExceeded { limit });
344 }
345 }
346 Ok(())
347 }
348
349 pub fn record_net_op(&mut self) -> Result<(), LimitViolation> {
351 self.net_ops += 1;
352 if let Some(limit) = self.limits.max_net_ops {
353 if self.net_ops > limit {
354 return Err(LimitViolation::NetOpsExceeded { limit });
355 }
356 }
357 Ok(())
358 }
359
360 pub fn elapsed(&self) -> Duration {
362 self.start_time.elapsed()
363 }
364
365 pub fn memory_used(&self) -> usize {
367 self.memory_used
368 }
369
370 pub fn instructions_executed(&self) -> u64 {
372 self.instructions_executed
373 }
374}
375
376#[cfg(test)]
377mod tests {
378 use super::*;
379
380 #[test]
381 fn test_default_limits() {
382 let limits = Limits::default();
383 assert!(limits.timeout.is_some());
384 assert!(limits.memory_bytes.is_some());
385 assert!(limits.max_instructions.is_some());
386 }
387
388 #[test]
389 fn test_unlimited() {
390 let limits = Limits::unlimited();
391 assert!(limits.timeout.is_none());
392 assert!(limits.memory_bytes.is_none());
393 }
394
395 #[test]
396 fn test_strict_limits() {
397 let limits = Limits::strict();
398 assert_eq!(limits.max_fs_ops, Some(0));
399 assert_eq!(limits.max_net_ops, Some(0));
400 }
401
402 #[test]
403 fn test_builder_pattern() {
404 let limits = Limits::default()
405 .with_timeout(Duration::from_secs(10))
406 .with_memory_mb(32);
407
408 assert_eq!(limits.timeout, Some(Duration::from_secs(10)));
409 assert_eq!(limits.memory_bytes, Some(32 * 1024 * 1024));
410 }
411
412 #[test]
413 fn test_limit_checks() {
414 let limits = Limits::default().with_memory_bytes(1000);
415
416 assert!(limits.check_memory(500).is_ok());
417 assert!(limits.check_memory(1500).is_err());
418 }
419
420 #[test]
421 fn test_limit_tracker() {
422 let limits = Limits::default().with_max_instructions(100);
423 let mut tracker = LimitTracker::new(limits);
424
425 assert!(tracker.record_instructions(50).is_ok());
426 assert!(tracker.record_instructions(60).is_err());
427 }
428}