cli_testing_specialist/utils/
resource_limits.rs

1use crate::error::{CliTestError, Result};
2use std::time::Duration;
3
4/// Resource limits for DOS attack prevention
5///
6/// Enforces strict limits on resources that can be consumed during analysis:
7/// - Memory usage (prevents memory exhaustion)
8/// - File descriptors (prevents FD exhaustion)
9/// - Process count (prevents fork bombs)
10/// - Execution timeout (prevents infinite loops)
11#[derive(Debug, Clone)]
12pub struct ResourceLimits {
13    /// Maximum memory usage in bytes (default: 500MB)
14    pub max_memory_bytes: u64,
15
16    /// Maximum number of file descriptors (default: 1024)
17    pub max_file_descriptors: u64,
18
19    /// Maximum number of processes (default: 100)
20    pub max_processes: u64,
21
22    /// Maximum execution time (default: 300s)
23    pub execution_timeout: Duration,
24}
25
26impl Default for ResourceLimits {
27    fn default() -> Self {
28        Self {
29            max_memory_bytes: 500 * 1024 * 1024, // 500MB
30            max_file_descriptors: 1024,
31            max_processes: 100,
32            execution_timeout: Duration::from_secs(300), // 5 minutes
33        }
34    }
35}
36
37impl ResourceLimits {
38    /// Create new resource limits with custom values
39    pub fn new(
40        max_memory_bytes: u64,
41        max_file_descriptors: u64,
42        max_processes: u64,
43        execution_timeout: Duration,
44    ) -> Self {
45        Self {
46            max_memory_bytes,
47            max_file_descriptors,
48            max_processes,
49            execution_timeout,
50        }
51    }
52
53    /// Apply resource limits to the current process (Unix only)
54    ///
55    /// This method uses `setrlimit` to enforce hard limits on resources.
56    /// On non-Unix platforms, this is a no-op.
57    #[cfg(unix)]
58    pub fn apply(&self) -> Result<()> {
59        use libc::{rlimit, setrlimit, RLIMIT_AS, RLIMIT_NOFILE, RLIMIT_NPROC};
60
61        // Set memory limit (address space)
62        let mem_limit = rlimit {
63            rlim_cur: self.max_memory_bytes,
64            rlim_max: self.max_memory_bytes,
65        };
66
67        unsafe {
68            if setrlimit(RLIMIT_AS, &mem_limit) != 0 {
69                return Err(CliTestError::ExecutionFailed(
70                    "Failed to set memory limit".to_string(),
71                ));
72            }
73        }
74
75        // Set file descriptor limit
76        let fd_limit = rlimit {
77            rlim_cur: self.max_file_descriptors,
78            rlim_max: self.max_file_descriptors,
79        };
80
81        unsafe {
82            if setrlimit(RLIMIT_NOFILE, &fd_limit) != 0 {
83                return Err(CliTestError::ExecutionFailed(
84                    "Failed to set file descriptor limit".to_string(),
85                ));
86            }
87        }
88
89        // Set process limit
90        let proc_limit = rlimit {
91            rlim_cur: self.max_processes,
92            rlim_max: self.max_processes,
93        };
94
95        unsafe {
96            if setrlimit(RLIMIT_NPROC, &proc_limit) != 0 {
97                return Err(CliTestError::ExecutionFailed(
98                    "Failed to set process limit".to_string(),
99                ));
100            }
101        }
102
103        Ok(())
104    }
105
106    /// Apply resource limits (Windows stub)
107    ///
108    /// Windows does not support setrlimit. Job Objects could be used
109    /// but are significantly more complex. This is marked as future work.
110    #[cfg(not(unix))]
111    pub fn apply(&self) -> Result<()> {
112        log::warn!("Resource limits not supported on this platform");
113        Ok(())
114    }
115
116    /// Get timeout duration
117    pub fn timeout(&self) -> Duration {
118        self.execution_timeout
119    }
120}
121
122#[cfg(test)]
123mod tests {
124    use super::*;
125
126    #[test]
127    fn test_default_limits() {
128        let limits = ResourceLimits::default();
129
130        assert_eq!(limits.max_memory_bytes, 500 * 1024 * 1024);
131        assert_eq!(limits.max_file_descriptors, 1024);
132        assert_eq!(limits.max_processes, 100);
133        assert_eq!(limits.execution_timeout, Duration::from_secs(300));
134    }
135
136    #[test]
137    fn test_custom_limits() {
138        let limits = ResourceLimits::new(100 * 1024 * 1024, 512, 50, Duration::from_secs(60));
139
140        assert_eq!(limits.max_memory_bytes, 100 * 1024 * 1024);
141        assert_eq!(limits.max_file_descriptors, 512);
142        assert_eq!(limits.max_processes, 50);
143        assert_eq!(limits.execution_timeout, Duration::from_secs(60));
144    }
145
146    #[test]
147    fn test_timeout_accessor() {
148        let limits = ResourceLimits::default();
149        assert_eq!(limits.timeout(), Duration::from_secs(300));
150    }
151
152    #[cfg(unix)]
153    #[test]
154    fn test_apply_limits_unix() {
155        // Note: This test may fail if the process doesn't have permission
156        // to set resource limits. It's primarily for compilation checking.
157        let limits = ResourceLimits::new(100 * 1024 * 1024, 512, 50, Duration::from_secs(60));
158
159        // This may fail in restricted environments, so we don't assert success
160        let _ = limits.apply();
161    }
162}
cli_testing_specialist/utils/resource_limits.rs

cli_testing_specialist/utils/
resource_limits.rs
