sandbox_rs/resources/
cgroup.rs

1//! Cgroup v2 management for resource limits
2
3use crate::errors::{Result, SandboxError};
4use nix::unistd::Pid;
5use std::fs;
6use std::io::Write;
7use std::path::{Path, PathBuf};
8
9const CGROUP_V2_ROOT: &str = "/sys/fs/cgroup";
10
11/// Cgroup v2 resource limits configuration
12#[derive(Debug, Clone, Default)]
13pub struct CgroupConfig {
14    /// Memory limit in bytes (e.g., 100MB)
15    pub memory_limit: Option<u64>,
16    /// CPU weight (100-10000, default 100)
17    pub cpu_weight: Option<u32>,
18    /// CPU quota in microseconds
19    pub cpu_quota: Option<u64>,
20    /// CPU period in microseconds (default 100000)
21    pub cpu_period: Option<u64>,
22    /// Max PIDs allowed
23    pub max_pids: Option<u32>,
24}
25
26impl CgroupConfig {
27    /// Create cgroup config with memory limit
28    pub fn with_memory(limit: u64) -> Self {
29        Self {
30            memory_limit: Some(limit),
31            ..Default::default()
32        }
33    }
34
35    /// Create cgroup config with CPU quota
36    pub fn with_cpu_quota(quota: u64, period: u64) -> Self {
37        Self {
38            cpu_quota: Some(quota),
39            cpu_period: Some(period),
40            ..Default::default()
41        }
42    }
43
44    /// Validate configuration
45    pub fn validate(&self) -> Result<()> {
46        if let Some(limit) = self.memory_limit
47            && limit == 0
48        {
49            return Err(SandboxError::InvalidConfig(
50                "Memory limit must be greater than 0".to_string(),
51            ));
52        }
53
54        if let Some(weight) = self.cpu_weight
55            && (!(100..=10000).contains(&weight))
56        {
57            return Err(SandboxError::InvalidConfig(
58                "CPU weight must be between 100-10000".to_string(),
59            ));
60        }
61
62        Ok(())
63    }
64}
65
66/// Cgroup v2 interface
67pub struct Cgroup {
68    path: PathBuf,
69    pid: Pid,
70}
71
72fn cgroup_root_path() -> PathBuf {
73    std::env::var("SANDBOX_CGROUP_ROOT")
74        .map(PathBuf::from)
75        .unwrap_or_else(|_| PathBuf::from(CGROUP_V2_ROOT))
76}
77
78impl Cgroup {
79    /// Create new cgroup for a process
80    pub fn new(name: &str, pid: Pid) -> Result<Self> {
81        let cgroup_path = cgroup_root_path().join(name);
82
83        // Create cgroup directory
84        fs::create_dir_all(&cgroup_path).map_err(|e| {
85            SandboxError::Cgroup(format!(
86                "Failed to create cgroup directory {}: {}",
87                cgroup_path.display(),
88                e
89            ))
90        })?;
91
92        ensure_controller_files(&cgroup_path)?;
93
94        Ok(Self {
95            path: cgroup_path,
96            pid,
97        })
98    }
99
100    /// Apply configuration to cgroup
101    pub fn apply_config(&self, config: &CgroupConfig) -> Result<()> {
102        config.validate()?;
103
104        if let Some(memory) = config.memory_limit {
105            self.set_memory_limit(memory)?;
106        }
107
108        if let Some(weight) = config.cpu_weight {
109            self.set_cpu_weight(weight)?;
110        }
111
112        if let Some(quota) = config.cpu_quota {
113            let period = config.cpu_period.unwrap_or(100000);
114            self.set_cpu_quota(quota, period)?;
115        }
116
117        if let Some(max_pids) = config.max_pids {
118            self.set_max_pids(max_pids)?;
119        }
120
121        Ok(())
122    }
123
124    /// Add process to cgroup
125    pub fn add_process(&self, pid: Pid) -> Result<()> {
126        let procs_file = self.path.join("cgroup.procs");
127        let mut file = fs::OpenOptions::new()
128            .write(true)
129            .open(&procs_file)
130            .map_err(|e| {
131                SandboxError::Cgroup(format!("Failed to open {}: {}", procs_file.display(), e))
132            })?;
133
134        write!(file, "{}", pid.as_raw())
135            .map_err(|e| SandboxError::Cgroup(format!("Failed to add process to cgroup: {}", e)))?;
136
137        Ok(())
138    }
139
140    /// Set memory limit
141    fn set_memory_limit(&self, limit: u64) -> Result<()> {
142        let mem_file = self.path.join("memory.max");
143        self.write_file(&mem_file, &limit.to_string())
144    }
145
146    /// Set CPU weight
147    fn set_cpu_weight(&self, weight: u32) -> Result<()> {
148        let cpu_file = self.path.join("cpu.weight");
149        self.write_file(&cpu_file, &weight.to_string())
150    }
151
152    /// Set CPU quota (microseconds)
153    fn set_cpu_quota(&self, quota: u64, period: u64) -> Result<()> {
154        let quota_file = self.path.join("cpu.max");
155        let quota_str = if quota == u64::MAX {
156            "max".to_string()
157        } else {
158            format!("{} {}", quota, period)
159        };
160        self.write_file(&quota_file, &quota_str)
161    }
162
163    /// Set max PIDs
164    fn set_max_pids(&self, max_pids: u32) -> Result<()> {
165        let pids_file = self.path.join("pids.max");
166        self.write_file(&pids_file, &max_pids.to_string())
167    }
168
169    /// Read memory usage
170    pub fn get_memory_usage(&self) -> Result<u64> {
171        let mem_file = self.path.join("memory.current");
172        self.read_file_u64(&mem_file)
173    }
174
175    /// Read memory limit
176    pub fn get_memory_limit(&self) -> Result<u64> {
177        let mem_file = self.path.join("memory.max");
178        self.read_file_u64(&mem_file)
179    }
180
181    /// Read CPU usage in microseconds
182    pub fn get_cpu_usage(&self) -> Result<u64> {
183        let cpu_file = self.path.join("cpu.stat");
184        let content = fs::read_to_string(&cpu_file).map_err(|e| {
185            SandboxError::Cgroup(format!("Failed to read {}: {}", cpu_file.display(), e))
186        })?;
187
188        // Parse "usage_usec 123456"
189        for line in content.lines() {
190            if line.starts_with("usage_usec") {
191                let parts: Vec<&str> = line.split_whitespace().collect();
192                if parts.len() >= 2 {
193                    return parts[1].parse::<u64>().map_err(|e| {
194                        SandboxError::Cgroup(format!("Failed to parse CPU usage: {}", e))
195                    });
196                }
197            }
198        }
199
200        Ok(0)
201    }
202
203    /// Check if cgroup exists
204    pub fn exists(&self) -> bool {
205        self.path.exists()
206    }
207
208    /// Get the PID this cgroup was created for
209    pub fn pid(&self) -> Pid {
210        self.pid
211    }
212
213    /// Delete cgroup
214    pub fn delete(&self) -> Result<()> {
215        if self.exists() {
216            fs::remove_dir(&self.path).map_err(|e| {
217                SandboxError::Cgroup(format!(
218                    "Failed to delete cgroup {}: {}",
219                    self.path.display(),
220                    e
221                ))
222            })?;
223        }
224        Ok(())
225    }
226
227    fn write_file(&self, path: &Path, content: &str) -> Result<()> {
228        let mut file = fs::OpenOptions::new().write(true).open(path).map_err(|e| {
229            SandboxError::Cgroup(format!("Failed to open {}: {}", path.display(), e))
230        })?;
231
232        write!(file, "{}", content).map_err(|e| {
233            SandboxError::Cgroup(format!("Failed to write to {}: {}", path.display(), e))
234        })?;
235
236        Ok(())
237    }
238
239    fn read_file_u64(&self, path: &Path) -> Result<u64> {
240        let content = fs::read_to_string(path).map_err(|e| {
241            SandboxError::Cgroup(format!("Failed to read {}: {}", path.display(), e))
242        })?;
243
244        content
245            .trim()
246            .parse::<u64>()
247            .map_err(|e| SandboxError::Cgroup(format!("Failed to parse value: {}", e)))
248    }
249
250    #[cfg(test)]
251    pub(crate) fn for_testing(path: PathBuf) -> Self {
252        Self {
253            path,
254            pid: Pid::from_raw(0),
255        }
256    }
257}
258
259fn ensure_controller_files(path: &Path) -> Result<()> {
260    let files = [
261        ("memory.max", "max"),
262        ("memory.current", "0"),
263        ("cpu.weight", "100"),
264        ("cpu.max", "max 100000"),
265        ("cpu.stat", "usage_usec 0\n"),
266        ("pids.max", "max"),
267        ("cgroup.procs", ""),
268    ];
269
270    for (name, default_content) in files {
271        let file_path = path.join(name);
272        if !file_path.exists() {
273            fs::write(&file_path, default_content).map_err(|e| {
274                SandboxError::Cgroup(format!("Failed to create {}: {}", file_path.display(), e))
275            })?;
276        }
277    }
278
279    Ok(())
280}
281
282impl Drop for Cgroup {
283    fn drop(&mut self) {
284        // Clean up cgroup on drop (best effort)
285        let _ = self.delete();
286    }
287}
288
289#[cfg(test)]
290mod tests {
291    use super::*;
292    use std::env;
293    use std::fs;
294    use tempfile::tempdir;
295
296    fn prepare_cgroup_dir() -> (tempfile::TempDir, std::path::PathBuf) {
297        let tmp = tempdir().unwrap();
298        let path = tmp.path().join("cgroup-test");
299        fs::create_dir_all(&path).unwrap();
300        for file in &[
301            "memory.max",
302            "memory.current",
303            "cpu.weight",
304            "cpu.max",
305            "cpu.stat",
306            "pids.max",
307            "cgroup.procs",
308        ] {
309            let file_path = path.join(file);
310            if let Some(parent) = file_path.parent() {
311                fs::create_dir_all(parent).unwrap();
312            }
313            fs::write(&file_path, "0").unwrap();
314        }
315        fs::write(path.join("cpu.stat"), "usage_usec 0\n").unwrap();
316        fs::write(path.join("memory.current"), "0\n").unwrap();
317        (tmp, path)
318    }
319
320    #[test]
321    fn test_cgroup_config_default() {
322        let config = CgroupConfig::default();
323        assert!(config.memory_limit.is_none());
324        assert!(config.cpu_weight.is_none());
325    }
326
327    #[test]
328    fn test_cgroup_config_with_memory() {
329        let config = CgroupConfig::with_memory(100 * 1024 * 1024);
330        assert_eq!(config.memory_limit, Some(100 * 1024 * 1024));
331    }
332
333    #[test]
334    fn test_cgroup_config_with_cpu_quota() {
335        let config = CgroupConfig::with_cpu_quota(50000, 100000);
336        assert_eq!(config.cpu_quota, Some(50000));
337        assert_eq!(config.cpu_period, Some(100000));
338    }
339
340    #[test]
341    fn test_cgroup_config_validate() {
342        let config = CgroupConfig::default();
343        assert!(config.validate().is_ok());
344
345        let bad_config = CgroupConfig {
346            memory_limit: Some(0),
347            ..Default::default()
348        };
349        assert!(bad_config.validate().is_err());
350
351        let bad_cpu_config = CgroupConfig {
352            cpu_weight: Some(50),
353            ..Default::default()
354        };
355        assert!(bad_cpu_config.validate().is_err());
356
357        let good_cpu_config = CgroupConfig {
358            cpu_weight: Some(100),
359            ..Default::default()
360        };
361        assert!(good_cpu_config.validate().is_ok());
362    }
363
364    #[test]
365    fn test_cgroup_path_creation() {
366        // This test may only work if running as root and cgroup v2 is available
367        // We'll test the logic without actually creating cgroups
368        let test_path = Path::new(CGROUP_V2_ROOT);
369        if test_path.exists() {
370            // Cgroup v2 is available
371            let result = Cgroup::new(
372                "sandbox-test-delete-me",
373                Pid::from_raw(std::process::id() as i32),
374            );
375            // Don't assert, as it may fail due to permissions
376            let _ = result;
377        }
378    }
379
380    #[test]
381    fn test_cgroup_apply_config_writes_files() {
382        let (_tmp, path) = prepare_cgroup_dir();
383        let cgroup = Cgroup::for_testing(path.clone());
384
385        let config = CgroupConfig {
386            memory_limit: Some(2048),
387            cpu_weight: Some(500),
388            cpu_quota: Some(50_000),
389            cpu_period: Some(100_000),
390            max_pids: Some(32),
391        };
392
393        cgroup.apply_config(&config).unwrap();
394
395        assert_eq!(
396            fs::read_to_string(path.join("memory.max")).unwrap().trim(),
397            "2048"
398        );
399        assert_eq!(
400            fs::read_to_string(path.join("cpu.weight")).unwrap().trim(),
401            "500"
402        );
403        assert_eq!(
404            fs::read_to_string(path.join("cpu.max")).unwrap().trim(),
405            "50000 100000"
406        );
407        assert_eq!(
408            fs::read_to_string(path.join("pids.max")).unwrap().trim(),
409            "32"
410        );
411    }
412
413    #[test]
414    fn test_cgroup_add_process_writes_pid() {
415        let (_tmp, path) = prepare_cgroup_dir();
416        let cgroup = Cgroup::for_testing(path.clone());
417
418        cgroup.add_process(Pid::from_raw(1234)).unwrap();
419        assert_eq!(
420            fs::read_to_string(path.join("cgroup.procs")).unwrap(),
421            "1234"
422        );
423    }
424
425    #[test]
426    fn test_cgroup_resource_readers() {
427        let (_tmp, path) = prepare_cgroup_dir();
428        fs::write(path.join("memory.current"), "4096").unwrap();
429        fs::write(path.join("cpu.stat"), "usage_usec 900\n").unwrap();
430        let cgroup = Cgroup::for_testing(path.clone());
431
432        assert_eq!(cgroup.get_memory_usage().unwrap(), 4096);
433        assert_eq!(cgroup.get_cpu_usage().unwrap(), 900);
434    }
435
436    #[test]
437    fn test_cgroup_delete_removes_directory() {
438        let (tmp, path) = prepare_cgroup_dir();
439        let cgroup = Cgroup::for_testing(path.clone());
440        assert!(path.exists());
441        for entry in fs::read_dir(&path).unwrap() {
442            let entry = entry.unwrap();
443            if entry.path().is_file() {
444                fs::remove_file(entry.path()).unwrap();
445            }
446        }
447        cgroup.delete().unwrap();
448        assert!(!path.exists());
449        drop(tmp);
450    }
451
452    #[test]
453    fn test_cgroup_new_uses_env_override() {
454        let tmp = tempdir().unwrap();
455        let prev = env::var("SANDBOX_CGROUP_ROOT").ok();
456        unsafe {
457            env::set_var("SANDBOX_CGROUP_ROOT", tmp.path());
458        }
459
460        let cg = Cgroup::new("env-test", Pid::from_raw(0)).unwrap();
461        assert!(cg.exists());
462        assert!(tmp.path().join("env-test").exists());
463
464        if let Some(value) = prev {
465            unsafe {
466                env::set_var("SANDBOX_CGROUP_ROOT", value);
467            }
468        } else {
469            unsafe {
470                env::remove_var("SANDBOX_CGROUP_ROOT");
471            }
472        }
473    }
474}
sandbox_rs/resources/cgroup.rs

sandbox_rs/resources/
cgroup.rs
