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 cpu_time(mut self, seconds: u64) -> Self {
60        self.cpu_time = Some(seconds);
61        self
62    }
63
64    pub fn file_size(mut self, bytes: u64) -> Self {
65        self.file_size = Some(bytes);
66        self
67    }
68
69    pub fn open_files(mut self, count: u64) -> Self {
70        self.open_files = Some(count);
71        self
72    }
73
74    pub fn processes(mut self, count: u64) -> Self {
75        self.processes = Some(count);
76        self
77    }
78
79    pub fn 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        if let Some(nproc) = self.processes {
116            let rlim = Rlimit {
117                current: Some(nproc),
118                maximum: Some(nproc),
119            };
120            setrlimit(Resource::Nproc, rlim)
121                .map_err(|e| SecurityError::ResourceLimit(format!("NPROC: {}", e)))?;
122        }
123
124        #[cfg(target_os = "linux")]
125        if let Some(vmem) = self.virtual_memory {
126            let rlim = Rlimit {
127                current: Some(vmem),
128                maximum: Some(vmem),
129            };
130            setrlimit(Resource::As, rlim)
131                .map_err(|e| SecurityError::ResourceLimit(format!("AS: {}", e)))?;
132        }
133
134        if let Some(data) = self.data_size {
135            let rlim = Rlimit {
136                current: Some(data),
137                maximum: Some(data),
138            };
139            setrlimit(Resource::Data, rlim)
140                .map_err(|e| SecurityError::ResourceLimit(format!("DATA: {}", e)))?;
141        }
142
143        if let Some(stack) = self.stack_size {
144            let rlim = Rlimit {
145                current: Some(stack),
146                maximum: Some(stack),
147            };
148            setrlimit(Resource::Stack, rlim)
149                .map_err(|e| SecurityError::ResourceLimit(format!("STACK: {}", e)))?;
150        }
151
152        Ok(())
153    }
154
155    #[cfg(not(unix))]
156    pub fn apply(&self) -> Result<(), SecurityError> {
157        Ok(())
158    }
159}
160
161#[cfg(test)]
162mod tests {
163    use super::*;
164
165    #[test]
166    fn test_default_limits() {
167        let limits = ResourceLimits::default();
168        assert_eq!(limits.cpu_time, Some(300));
169        assert_eq!(limits.open_files, Some(256));
170    }
171
172    #[test]
173    fn test_strict_limits() {
174        let limits = ResourceLimits::strict();
175        assert_eq!(limits.cpu_time, Some(60));
176        assert_eq!(limits.processes, Some(10));
177    }
178
179    #[test]
180    fn test_none_limits() {
181        let limits = ResourceLimits::none();
182        assert!(limits.cpu_time.is_none());
183        assert!(limits.file_size.is_none());
184    }
185
186    #[test]
187    fn test_builder() {
188        let limits = ResourceLimits::none()
189            .cpu_time(120)
190            .file_size(1024 * 1024)
191            .open_files(128);
192
193        assert_eq!(limits.cpu_time, Some(120));
194        assert_eq!(limits.file_size, Some(1024 * 1024));
195        assert_eq!(limits.open_files, Some(128));
196    }
197}
claude_agent/security/limits/rlimit.rs

claude_agent/security/limits/
rlimit.rs
