claude_agent/security/limits/
rlimit.rs

1//! Resource limits using setrlimit.
2
3use crate::security::SecurityError;
4
5const KB: u64 = 1024;
6const MB: u64 = 1024 * KB;
7const GB: u64 = 1024 * MB;
8
9#[derive(Debug, Clone)]
10pub struct ResourceLimits {
11    pub cpu_time: Option<u64>,
12    pub file_size: Option<u64>,
13    pub open_files: Option<u64>,
14    pub processes: Option<u64>,
15    pub virtual_memory: Option<u64>,
16    pub data_size: Option<u64>,
17    pub stack_size: Option<u64>,
18}
19
20impl Default for ResourceLimits {
21    fn default() -> Self {
22        Self {
23            cpu_time: Some(300),       // 5 minutes
24            file_size: Some(100 * MB), // 100 MB
25            open_files: Some(256),
26            processes: Some(32),
27            virtual_memory: Some(2 * GB), // 2 GB
28            data_size: Some(GB),          // 1 GB
29            stack_size: Some(8 * MB),     // 8 MB
30        }
31    }
32}
33
34impl ResourceLimits {
35    pub fn none() -> Self {
36        Self {
37            cpu_time: None,
38            file_size: None,
39            open_files: None,
40            processes: None,
41            virtual_memory: None,
42            data_size: None,
43            stack_size: None,
44        }
45    }
46
47    pub fn strict() -> Self {
48        Self {
49            cpu_time: Some(60),       // 1 minute
50            file_size: Some(10 * MB), // 10 MB
51            open_files: Some(64),
52            processes: Some(10),
53            virtual_memory: Some(512 * MB), // 512 MB
54            data_size: Some(256 * MB),      // 256 MB
55            stack_size: Some(MB),           // 1 MB
56        }
57    }
58
59    pub fn with_cpu_time(mut self, seconds: u64) -> Self {
60        self.cpu_time = Some(seconds);
61        self
62    }
63
64    pub fn with_file_size(mut self, bytes: u64) -> Self {
65        self.file_size = Some(bytes);
66        self
67    }
68
69    pub fn with_open_files(mut self, count: u64) -> Self {
70        self.open_files = Some(count);
71        self
72    }
73
74    pub fn with_processes(mut self, count: u64) -> Self {
75        self.processes = Some(count);
76        self
77    }
78
79    pub fn with_virtual_memory(mut self, bytes: u64) -> Self {
80        self.virtual_memory = Some(bytes);
81        self
82    }
83
84    #[cfg(unix)]
85    pub fn apply(&self) -> Result<(), SecurityError> {
86        use rustix::process::{Resource, Rlimit, setrlimit};
87
88        if let Some(cpu) = self.cpu_time {
89            let rlim = Rlimit {
90                current: Some(cpu),
91                maximum: Some(cpu),
92            };
93            setrlimit(Resource::Cpu, rlim)
94                .map_err(|e| SecurityError::ResourceLimit(format!("CPU: {}", e)))?;
95        }
96
97        if let Some(fsize) = self.file_size {
98            let rlim = Rlimit {
99                current: Some(fsize),
100                maximum: Some(fsize),
101            };
102            setrlimit(Resource::Fsize, rlim)
103                .map_err(|e| SecurityError::ResourceLimit(format!("FSIZE: {}", e)))?;
104        }
105
106        if let Some(nofile) = self.open_files {
107            let rlim = Rlimit {
108                current: Some(nofile),
109                maximum: Some(nofile),
110            };
111            setrlimit(Resource::Nofile, rlim)
112                .map_err(|e| SecurityError::ResourceLimit(format!("NOFILE: {}", e)))?;
113        }
114
115        #[cfg(target_os = "linux")]
116        if let Some(nproc) = self.processes {
117            let rlim = Rlimit {
118                current: Some(nproc),
119                maximum: Some(nproc),
120            };
121            setrlimit(Resource::Nproc, rlim)
122                .map_err(|e| SecurityError::ResourceLimit(format!("NPROC: {}", e)))?;
123        }
124
125        #[cfg(target_os = "linux")]
126        if let Some(vmem) = self.virtual_memory {
127            let rlim = Rlimit {
128                current: Some(vmem),
129                maximum: Some(vmem),
130            };
131            setrlimit(Resource::As, rlim)
132                .map_err(|e| SecurityError::ResourceLimit(format!("AS: {}", e)))?;
133        }
134
135        if let Some(data) = self.data_size {
136            let rlim = Rlimit {
137                current: Some(data),
138                maximum: Some(data),
139            };
140            setrlimit(Resource::Data, rlim)
141                .map_err(|e| SecurityError::ResourceLimit(format!("DATA: {}", e)))?;
142        }
143
144        if let Some(stack) = self.stack_size {
145            let rlim = Rlimit {
146                current: Some(stack),
147                maximum: Some(stack),
148            };
149            setrlimit(Resource::Stack, rlim)
150                .map_err(|e| SecurityError::ResourceLimit(format!("STACK: {}", e)))?;
151        }
152
153        Ok(())
154    }
155
156    #[cfg(not(unix))]
157    pub fn apply(&self) -> Result<(), SecurityError> {
158        Ok(())
159    }
160}
161
162#[cfg(test)]
163mod tests {
164    use super::*;
165
166    #[test]
167    fn test_default_limits() {
168        let limits = ResourceLimits::default();
169        assert_eq!(limits.cpu_time, Some(300));
170        assert_eq!(limits.open_files, Some(256));
171    }
172
173    #[test]
174    fn test_strict_limits() {
175        let limits = ResourceLimits::strict();
176        assert_eq!(limits.cpu_time, Some(60));
177        assert_eq!(limits.processes, Some(10));
178    }
179
180    #[test]
181    fn test_none_limits() {
182        let limits = ResourceLimits::none();
183        assert!(limits.cpu_time.is_none());
184        assert!(limits.file_size.is_none());
185    }
186
187    #[test]
188    fn test_builder() {
189        let limits = ResourceLimits::none()
190            .with_cpu_time(120)
191            .with_file_size(1024 * 1024)
192            .with_open_files(128);
193
194        assert_eq!(limits.cpu_time, Some(120));
195        assert_eq!(limits.file_size, Some(1024 * 1024));
196        assert_eq!(limits.open_files, Some(128));
197    }
198}
claude_agent/security/limits/rlimit.rs

claude_agent/security/limits/
rlimit.rs
