sandbox_cgroup/
cgroup.rs

1//! Cgroup v2 management for resource limits
2
3use nix::unistd::Pid;
4use sandbox_core::{Result, SandboxError};
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    pub memory_limit: Option<u64>,
15    pub cpu_weight: Option<u32>,
16    pub cpu_quota: Option<u64>,
17    pub cpu_period: Option<u64>,
18    pub max_pids: Option<u32>,
19}
20
21impl CgroupConfig {
22    pub fn with_memory(limit: u64) -> Self {
23        Self {
24            memory_limit: Some(limit),
25            ..Default::default()
26        }
27    }
28
29    pub fn with_cpu_quota(quota: u64, period: u64) -> Self {
30        Self {
31            cpu_quota: Some(quota),
32            cpu_period: Some(period),
33            ..Default::default()
34        }
35    }
36
37    pub fn validate(&self) -> Result<()> {
38        if let Some(limit) = self.memory_limit
39            && limit == 0
40        {
41            return Err(SandboxError::InvalidConfig(
42                "Memory limit must be greater than 0".to_string(),
43            ));
44        }
45        if let Some(weight) = self.cpu_weight
46            && (!(100..=10000).contains(&weight))
47        {
48            return Err(SandboxError::InvalidConfig(
49                "CPU weight must be between 100-10000".to_string(),
50            ));
51        }
52        Ok(())
53    }
54}
55
56/// Cgroup v2 interface
57pub struct Cgroup {
58    path: PathBuf,
59    pid: Pid,
60}
61
62fn cgroup_root_path() -> PathBuf {
63    std::env::var("SANDBOX_CGROUP_ROOT")
64        .map(PathBuf::from)
65        .unwrap_or_else(|_| PathBuf::from(CGROUP_V2_ROOT))
66}
67
68/// Try to find a delegated cgroup for the current unprivileged user
69pub fn find_delegated_cgroup() -> Option<PathBuf> {
70    let uid = unsafe { libc::geteuid() };
71    if uid == 0 {
72        return Some(PathBuf::from(CGROUP_V2_ROOT));
73    }
74
75    let user_slice = format!("/sys/fs/cgroup/user.slice/user-{}.slice", uid);
76    let path = PathBuf::from(&user_slice);
77
78    if !path.exists() {
79        return None;
80    }
81
82    // Check if we can write to the cgroup directory
83    let test_path = path.join("sandbox-cgroup-probe");
84    match std::fs::create_dir(&test_path) {
85        Ok(()) => {
86            let _ = std::fs::remove_dir(&test_path);
87            Some(path)
88        }
89        Err(_) => None,
90    }
91}
92
93impl Cgroup {
94    pub fn new(name: &str, pid: Pid) -> Result<Self> {
95        let cgroup_path = cgroup_root_path().join(name);
96        fs::create_dir_all(&cgroup_path).map_err(|e| {
97            SandboxError::Cgroup(format!(
98                "Failed to create cgroup directory {}: {}",
99                cgroup_path.display(),
100                e
101            ))
102        })?;
103        Ok(Self {
104            path: cgroup_path,
105            pid,
106        })
107    }
108
109    pub fn apply_config(&self, config: &CgroupConfig) -> Result<()> {
110        config.validate()?;
111        if let Some(memory) = config.memory_limit {
112            self.set_memory_limit(memory)?;
113        }
114        if let Some(weight) = config.cpu_weight {
115            self.set_cpu_weight(weight)?;
116        }
117        if let Some(quota) = config.cpu_quota {
118            let period = config.cpu_period.unwrap_or(100000);
119            self.set_cpu_quota(quota, period)?;
120        }
121        if let Some(max_pids) = config.max_pids {
122            self.set_max_pids(max_pids)?;
123        }
124        Ok(())
125    }
126
127    pub fn add_process(&self, pid: Pid) -> Result<()> {
128        let procs_file = self.path.join("cgroup.procs");
129        self.write_file(&procs_file, &pid.as_raw().to_string())
130    }
131
132    fn set_memory_limit(&self, limit: u64) -> Result<()> {
133        self.write_file(&self.path.join("memory.max"), &limit.to_string())
134    }
135
136    fn set_cpu_weight(&self, weight: u32) -> Result<()> {
137        self.write_file(&self.path.join("cpu.weight"), &weight.to_string())
138    }
139
140    fn set_cpu_quota(&self, quota: u64, period: u64) -> Result<()> {
141        let quota_str = if quota == u64::MAX {
142            "max".to_string()
143        } else {
144            format!("{} {}", quota, period)
145        };
146        self.write_file(&self.path.join("cpu.max"), &quota_str)
147    }
148
149    fn set_max_pids(&self, max_pids: u32) -> Result<()> {
150        self.write_file(&self.path.join("pids.max"), &max_pids.to_string())
151    }
152
153    pub fn get_memory_usage(&self) -> Result<u64> {
154        self.read_file_u64(&self.path.join("memory.current"))
155    }
156
157    pub fn get_memory_limit(&self) -> Result<u64> {
158        self.read_file_u64(&self.path.join("memory.max"))
159    }
160
161    pub fn get_cpu_usage(&self) -> Result<u64> {
162        let cpu_file = self.path.join("cpu.stat");
163        let content = fs::read_to_string(&cpu_file).map_err(|e| {
164            SandboxError::Cgroup(format!("Failed to read {}: {}", cpu_file.display(), e))
165        })?;
166        for line in content.lines() {
167            if line.starts_with("usage_usec") {
168                let parts: Vec<&str> = line.split_whitespace().collect();
169                if parts.len() >= 2 {
170                    return parts[1].parse::<u64>().map_err(|e| {
171                        SandboxError::Cgroup(format!("Failed to parse CPU usage: {}", e))
172                    });
173                }
174            }
175        }
176        Ok(0)
177    }
178
179    pub fn exists(&self) -> bool {
180        self.path.exists()
181    }
182    pub fn pid(&self) -> Pid {
183        self.pid
184    }
185
186    pub fn delete(&self) -> Result<()> {
187        match fs::remove_dir(&self.path) {
188            Ok(()) => Ok(()),
189            Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
190            Err(e) => Err(SandboxError::Cgroup(format!(
191                "Failed to delete cgroup {}: {}",
192                self.path.display(),
193                e
194            ))),
195        }
196    }
197
198    fn write_file(&self, path: &Path, content: &str) -> Result<()> {
199        let mut file = fs::OpenOptions::new().write(true).open(path).map_err(|e| {
200            SandboxError::Cgroup(format!("Failed to open {}: {}", path.display(), e))
201        })?;
202        write!(file, "{}", content).map_err(|e| {
203            SandboxError::Cgroup(format!("Failed to write to {}: {}", path.display(), e))
204        })?;
205        Ok(())
206    }
207
208    fn read_file_u64(&self, path: &Path) -> Result<u64> {
209        let content = fs::read_to_string(path).map_err(|e| {
210            SandboxError::Cgroup(format!("Failed to read {}: {}", path.display(), e))
211        })?;
212        content
213            .trim()
214            .parse::<u64>()
215            .map_err(|e| SandboxError::Cgroup(format!("Failed to parse value: {}", e)))
216    }
217
218    /// Create a Cgroup backed by an arbitrary directory path (for testing)
219    #[doc(hidden)]
220    pub fn for_testing(path: PathBuf) -> Self {
221        Self {
222            path,
223            pid: Pid::from_raw(0),
224        }
225    }
226}
227
228impl Drop for Cgroup {
229    fn drop(&mut self) {
230        let _ = self.delete();
231    }
232}
233
234#[cfg(test)]
235mod tests {
236    use super::*;
237    use tempfile::tempdir;
238
239    fn prepare_cgroup_dir() -> (tempfile::TempDir, PathBuf) {
240        let tmp = tempdir().unwrap();
241        let path = tmp.path().join("cgroup-test");
242        fs::create_dir_all(&path).unwrap();
243        for file in &[
244            "memory.max",
245            "memory.current",
246            "cpu.weight",
247            "cpu.max",
248            "cpu.stat",
249            "pids.max",
250            "cgroup.procs",
251        ] {
252            fs::write(path.join(file), "0").unwrap();
253        }
254        fs::write(path.join("cpu.stat"), "usage_usec 0\n").unwrap();
255        fs::write(path.join("memory.current"), "0\n").unwrap();
256        (tmp, path)
257    }
258
259    #[test]
260    fn test_cgroup_config_default() {
261        let config = CgroupConfig::default();
262        assert!(config.memory_limit.is_none());
263    }
264
265    #[test]
266    fn test_cgroup_config_validate() {
267        assert!(CgroupConfig::default().validate().is_ok());
268        assert!(
269            CgroupConfig {
270                memory_limit: Some(0),
271                ..Default::default()
272            }
273            .validate()
274            .is_err()
275        );
276        assert!(
277            CgroupConfig {
278                cpu_weight: Some(50),
279                ..Default::default()
280            }
281            .validate()
282            .is_err()
283        );
284        assert!(
285            CgroupConfig {
286                cpu_weight: Some(100),
287                ..Default::default()
288            }
289            .validate()
290            .is_ok()
291        );
292    }
293
294    #[test]
295    fn test_cgroup_apply_config_writes_files() {
296        let (_tmp, path) = prepare_cgroup_dir();
297        let cgroup = Cgroup::for_testing(path.clone());
298        let config = CgroupConfig {
299            memory_limit: Some(2048),
300            cpu_weight: Some(500),
301            cpu_quota: Some(50_000),
302            cpu_period: Some(100_000),
303            max_pids: Some(32),
304        };
305        cgroup.apply_config(&config).unwrap();
306        assert_eq!(
307            fs::read_to_string(path.join("memory.max")).unwrap().trim(),
308            "2048"
309        );
310        assert_eq!(
311            fs::read_to_string(path.join("cpu.weight")).unwrap().trim(),
312            "500"
313        );
314        assert_eq!(
315            fs::read_to_string(path.join("cpu.max")).unwrap().trim(),
316            "50000 100000"
317        );
318        assert_eq!(
319            fs::read_to_string(path.join("pids.max")).unwrap().trim(),
320            "32"
321        );
322    }
323
324    #[test]
325    fn test_cgroup_resource_readers() {
326        let (_tmp, path) = prepare_cgroup_dir();
327        fs::write(path.join("memory.current"), "4096").unwrap();
328        fs::write(path.join("cpu.stat"), "usage_usec 900\n").unwrap();
329        let cgroup = Cgroup::for_testing(path);
330        assert_eq!(cgroup.get_memory_usage().unwrap(), 4096);
331        assert_eq!(cgroup.get_cpu_usage().unwrap(), 900);
332    }
333}
sandbox_cgroup/cgroup.rs

sandbox_cgroup/
cgroup.rs
