assay_runner_linux/
cgroup.rs

1use anyhow::{anyhow, Context, Result};
2use std::fs;
3#[cfg(unix)]
4use std::os::unix::fs::MetadataExt;
5use std::path::{Path, PathBuf};
6use std::time::{SystemTime, UNIX_EPOCH};
7
8/// Manages Cgroup V2 operations for Assay.
9pub struct CgroupManager {
10    root_path: PathBuf,
11}
12
13/// Represents an active ephemeral Cgroup session.
14pub struct SessionCgroup {
15    path: PathBuf,
16    id: u64, // Inode number (cgroup ID)
17}
18
19impl CgroupManager {
20    /// Checks if Cgroup V2 is available and resolves the correct nesting root.
21    /// SOTA Hardening: Detects current cgroup scope to support systemd slices.
22    pub fn new() -> Result<Self> {
23        let mount_point = PathBuf::from("/sys/fs/cgroup");
24
25        if !mount_point.exists() || !mount_point.is_dir() {
26            return Err(anyhow!("Cgroup V2 mount not found"));
27        }
28
29        // P0 Fix: Robust parsing of /proc/self/cgroup
30        // We look for the line starting with "0::" (Unified hierarchy)
31        let content =
32            fs::read_to_string("/proc/self/cgroup").context("Failed to read /proc/self/cgroup")?;
33
34        let self_cgroup_line = content
35            .lines()
36            .find(|line| line.starts_with("0::"))
37            .ok_or_else(|| {
38                anyhow!("Could not find Unified Hierarchy (0::) in /proc/self/cgroup")
39            })?;
40
41        let self_cgroup_path = self_cgroup_line
42            .split("::")
43            .nth(1)
44            .ok_or_else(|| anyhow!("Invalid cgroup line format"))?;
45
46        // Handle root case "0::/" -> "" (empty relative path)
47        let relative_path = if self_cgroup_path == "/" {
48            Path::new("")
49        } else {
50            self_cgroup_path
51                .strip_prefix('/')
52                .unwrap_or(self_cgroup_path)
53                .as_ref()
54        };
55
56        let root_path = mount_point.join(relative_path);
57
58        if !root_path.exists() {
59            return Err(anyhow!("Could not verify own cgroup path: {:?}", root_path));
60        }
61
62        let root_path = Self::nearest_domain_root(&mount_point, root_path)?;
63
64        Ok(Self { root_path })
65    }
66
67    fn nearest_domain_root(mount_point: &Path, mut path: PathBuf) -> Result<PathBuf> {
68        loop {
69            let cgroup_type = fs::read_to_string(path.join("cgroup.type"))
70                .with_context(|| format!("Failed to read cgroup type for {}", path.display()))?;
71            if cgroup_type.trim() == "domain" && !Self::is_systemd_leaf_unit(&path) {
72                return Ok(path);
73            }
74
75            if path == mount_point {
76                return Err(anyhow!(
77                    "Could not find domain cgroup root at or above {}",
78                    path.display()
79                ));
80            }
81
82            path = path
83                .parent()
84                .ok_or_else(|| anyhow!("Could not ascend from cgroup path {}", path.display()))?
85                .to_path_buf();
86        }
87    }
88
89    fn is_systemd_leaf_unit(path: &Path) -> bool {
90        path.file_name()
91            .and_then(|name| name.to_str())
92            .is_some_and(|name| name.ends_with(".scope") || name.ends_with(".service"))
93    }
94
95    /// Creates a new ephemeral cgroup.
96    pub fn create_session(&self) -> Result<SessionCgroup> {
97        let timestamp = SystemTime::now().duration_since(UNIX_EPOCH)?.as_millis();
98        let name = format!("assay-session-{}", timestamp);
99        let path = self.root_path.join(&name);
100
101        // P0: Best-effort remove if collision (unlikely with timestamp)
102        if path.exists() {
103            let _ = fs::remove_dir(&path);
104        }
105
106        fs::create_dir(&path).context("Failed to create cgroup session dir")?;
107
108        // Best effort enable pids controller
109        let subtree = self.root_path.join("cgroup.subtree_control");
110        if subtree.exists() {
111            // Ignoring errors as we might lack delegation
112            let _ = fs::write(&subtree, "+pids");
113        }
114
115        let meta = fs::metadata(&path)?;
116        #[cfg(unix)]
117        let id = meta.ino();
118        #[cfg(not(unix))]
119        let id = {
120            let _ = meta;
121            0
122        }; // Stub for non-Unix
123        Ok(SessionCgroup { path, id })
124    }
125}
126
127impl SessionCgroup {
128    pub fn id(&self) -> u64 {
129        self.id
130    }
131
132    pub fn path(&self) -> &Path {
133        &self.path
134    }
135
136    pub fn procs_path(&self) -> PathBuf {
137        self.path.join("cgroup.procs")
138    }
139
140    pub fn add_process(&self, pid: u32) -> Result<()> {
141        let procs_path = self.procs_path();
142        fs::write(&procs_path, pid.to_string()).context("Failed to add PID")?;
143        Ok(())
144    }
145
146    pub fn freeze(&self) -> Result<()> {
147        let p = self.path.join("cgroup.freeze");
148        if p.exists() {
149            fs::write(p, "1")?;
150        }
151        Ok(())
152    }
153
154    pub fn thaw(&self) -> Result<()> {
155        let p = self.path.join("cgroup.freeze");
156        if p.exists() {
157            fs::write(p, "0")?;
158        }
159        Ok(())
160    }
161
162    pub fn kill(&self) -> Result<()> {
163        let p = self.path.join("cgroup.kill");
164        if p.exists() {
165            fs::write(p, "1")?;
166        } else {
167            return Err(anyhow!("cgroup.kill missing"));
168        }
169        Ok(())
170    }
171
172    /// P0 Hardening: Graceful kill using pidfd to avoid PID reuse races.
173    pub fn kill_graceful(&self, grace_ms: u64) -> Result<()> {
174        let _ = self.freeze();
175
176        let procs = fs::read_to_string(self.path.join("cgroup.procs"))?;
177        let mut pids: Vec<i32> = procs
178            .lines()
179            .filter_map(|l| l.trim().parse::<i32>().ok())
180            .collect();
181        // SOTA: Dedupe to prevent double-signaling (which is harmless but sloppy)
182        pids.sort_unstable();
183        pids.dedup();
184
185        // Send SIGTERM via pidfd if possible
186        for &pid in &pids {
187            if pid <= 0 {
188                continue;
189            }
190
191            #[cfg(target_os = "linux")]
192            // SAFETY: pidfd_open and pidfd_send_signal are raw libc syscalls.
193            // `pid` is validated above, the returned fd is checked before use,
194            // and the only pointer passed to the kernel is a null siginfo
195            // pointer. Failures fall back to classic kill; nothing is
196            // dereferenced through a Rust reference.
197            unsafe {
198                // syscall(SYS_pidfd_open, pid, 0)
199                let fd = libc::syscall(libc::SYS_pidfd_open, pid, 0) as i32;
200                if fd >= 0 {
201                    // syscall(SYS_pidfd_send_signal, fd, SIGTERM, NULL, 0)
202                    let _ = libc::syscall(
203                        libc::SYS_pidfd_send_signal,
204                        fd,
205                        libc::SIGTERM,
206                        std::ptr::null::<libc::siginfo_t>(),
207                        0,
208                    );
209                    libc::close(fd);
210                } else {
211                    // Fallback to classic kill if pidfd fails
212                    libc::kill(pid, libc::SIGTERM);
213                }
214            }
215
216            #[cfg(all(unix, not(target_os = "linux")))]
217            {
218                use nix::sys::signal::{kill, Signal};
219                use nix::unistd::Pid;
220                let _ = kill(Pid::from_raw(pid), Signal::SIGTERM);
221            }
222            #[cfg(not(unix))]
223            {
224                // Windows/Other: No-op or standard process kill if we had handle
225                // Since this uses PIDs from cgroup.procs (Linux concept), this code is unreachable logic-wise
226                // but must compile.
227                let _ = pid;
228            }
229        }
230
231        let _ = self.thaw();
232        std::thread::sleep(std::time::Duration::from_millis(grace_ms));
233        self.kill()
234    }
235
236    pub fn set_pids_max(&self, max: u32) -> Result<()> {
237        let p = self.path.join("pids.max");
238        if p.exists() {
239            fs::write(p, max.to_string())?;
240        }
241        Ok(())
242    }
243
244    /// P0 Hardening: Reliable cleanup
245    pub fn remove(&self) -> Result<()> {
246        if !self.path.exists() {
247            return Ok(());
248        }
249
250        // Retry loop for cgroup removal (busy/not empty)
251        for _ in 0..3 {
252            match fs::remove_dir(&self.path) {
253                Ok(_) => return Ok(()),
254                Err(_) => {
255                    // Try to kill remainders if still exists
256                    let _ = self.kill();
257                    std::thread::sleep(std::time::Duration::from_millis(50));
258                }
259            }
260        }
261        // Final attempt
262        fs::remove_dir(&self.path).context("Failed to remove cgroup after retries")
263    }
264}
265
266impl Drop for SessionCgroup {
267    fn drop(&mut self) {
268        if let Err(e) = self.remove() {
269            eprintln!("Failed to clean up cgroup session: {:?}", e);
270        }
271    }
272}
273
274#[cfg(test)]
275mod tests {
276    use super::CgroupManager;
277    use std::fs;
278    use std::path::PathBuf;
279    use std::time::{SystemTime, UNIX_EPOCH};
280
281    fn temp_cgroup_tree(name: &str) -> PathBuf {
282        let stamp = SystemTime::now()
283            .duration_since(UNIX_EPOCH)
284            .expect("clock before unix epoch")
285            .as_nanos();
286        let path = std::env::temp_dir().join(format!("assay-cgroup-{name}-{stamp}"));
287        fs::create_dir_all(&path).expect("create temp cgroup tree");
288        path
289    }
290
291    #[test]
292    fn nearest_domain_root_accepts_domain_cgroup() {
293        let root = temp_cgroup_tree("domain");
294        fs::write(root.join("cgroup.type"), "domain\n").expect("write cgroup type");
295
296        let resolved =
297            CgroupManager::nearest_domain_root(&root, root.clone()).expect("resolve domain root");
298
299        assert_eq!(resolved, root);
300        fs::remove_dir_all(resolved).expect("remove temp cgroup tree");
301    }
302
303    #[test]
304    fn nearest_domain_root_ascends_from_domain_threaded_service() {
305        let root = temp_cgroup_tree("domain-threaded");
306        let system_slice = root.join("system.slice");
307        let service = system_slice.join("actions.runner.example.service");
308        fs::create_dir_all(&service).expect("create fake service cgroup");
309        fs::write(root.join("cgroup.type"), "domain\n").expect("write root cgroup type");
310        fs::write(system_slice.join("cgroup.type"), "domain\n")
311            .expect("write system slice cgroup type");
312        fs::write(service.join("cgroup.type"), "domain threaded\n")
313            .expect("write service cgroup type");
314
315        let resolved =
316            CgroupManager::nearest_domain_root(&root, service).expect("resolve domain root");
317
318        assert_eq!(resolved, system_slice);
319        fs::remove_dir_all(root).expect("remove temp cgroup tree");
320    }
321
322    #[test]
323    fn nearest_domain_root_ascends_from_systemd_service_unit() {
324        let root = temp_cgroup_tree("service-unit");
325        let system_slice = root.join("system.slice");
326        let service = system_slice.join("actions.runner.example.service");
327        fs::create_dir_all(&service).expect("create fake service cgroup");
328        fs::write(root.join("cgroup.type"), "domain\n").expect("write root cgroup type");
329        fs::write(system_slice.join("cgroup.type"), "domain\n")
330            .expect("write system slice cgroup type");
331        fs::write(service.join("cgroup.type"), "domain\n").expect("write service cgroup type");
332
333        let resolved =
334            CgroupManager::nearest_domain_root(&root, service).expect("resolve domain root");
335
336        assert_eq!(resolved, system_slice);
337        fs::remove_dir_all(root).expect("remove temp cgroup tree");
338    }
339
340    #[test]
341    fn nearest_domain_root_ascends_from_systemd_session_scope() {
342        let root = temp_cgroup_tree("session-scope");
343        let user_slice = root.join("user.slice");
344        let user_id_slice = user_slice.join("user-1000.slice");
345        let session_scope = user_id_slice.join("session-263.scope");
346        fs::create_dir_all(&session_scope).expect("create fake session cgroup");
347        fs::write(root.join("cgroup.type"), "domain\n").expect("write root cgroup type");
348        fs::write(user_slice.join("cgroup.type"), "domain\n")
349            .expect("write user slice cgroup type");
350        fs::write(user_id_slice.join("cgroup.type"), "domain\n")
351            .expect("write user id slice cgroup type");
352        fs::write(session_scope.join("cgroup.type"), "domain\n")
353            .expect("write session scope cgroup type");
354
355        let resolved =
356            CgroupManager::nearest_domain_root(&root, session_scope).expect("resolve domain root");
357
358        assert_eq!(resolved, user_id_slice);
359        fs::remove_dir_all(root).expect("remove temp cgroup tree");
360    }
361}
assay_runner_linux/cgroup.rs

assay_runner_linux/
cgroup.rs
