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_scope(&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_scope(path: &Path) -> bool {
90        path.file_name()
91            .and_then(|name| name.to_str())
92            .is_some_and(|name| name.ends_with(".scope"))
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            unsafe {
193                // syscall(SYS_pidfd_open, pid, 0)
194                let fd = libc::syscall(libc::SYS_pidfd_open, pid, 0) as i32;
195                if fd >= 0 {
196                    // syscall(SYS_pidfd_send_signal, fd, SIGTERM, NULL, 0)
197                    let _ = libc::syscall(
198                        libc::SYS_pidfd_send_signal,
199                        fd,
200                        libc::SIGTERM,
201                        std::ptr::null::<libc::siginfo_t>(),
202                        0,
203                    );
204                    libc::close(fd);
205                } else {
206                    // Fallback to classic kill if pidfd fails
207                    libc::kill(pid, libc::SIGTERM);
208                }
209            }
210
211            #[cfg(all(unix, not(target_os = "linux")))]
212            {
213                use nix::sys::signal::{kill, Signal};
214                use nix::unistd::Pid;
215                let _ = kill(Pid::from_raw(pid), Signal::SIGTERM);
216            }
217            #[cfg(not(unix))]
218            {
219                // Windows/Other: No-op or standard process kill if we had handle
220                // Since this uses PIDs from cgroup.procs (Linux concept), this code is unreachable logic-wise
221                // but must compile.
222                let _ = pid;
223            }
224        }
225
226        let _ = self.thaw();
227        std::thread::sleep(std::time::Duration::from_millis(grace_ms));
228        self.kill()
229    }
230
231    pub fn set_pids_max(&self, max: u32) -> Result<()> {
232        let p = self.path.join("pids.max");
233        if p.exists() {
234            fs::write(p, max.to_string())?;
235        }
236        Ok(())
237    }
238
239    /// P0 Hardening: Reliable cleanup
240    pub fn remove(&self) -> Result<()> {
241        if !self.path.exists() {
242            return Ok(());
243        }
244
245        // Retry loop for cgroup removal (busy/not empty)
246        for _ in 0..3 {
247            match fs::remove_dir(&self.path) {
248                Ok(_) => return Ok(()),
249                Err(_) => {
250                    // Try to kill remainders if still exists
251                    let _ = self.kill();
252                    std::thread::sleep(std::time::Duration::from_millis(50));
253                }
254            }
255        }
256        // Final attempt
257        fs::remove_dir(&self.path).context("Failed to remove cgroup after retries")
258    }
259}
260
261impl Drop for SessionCgroup {
262    fn drop(&mut self) {
263        if let Err(e) = self.remove() {
264            eprintln!("Failed to clean up cgroup session: {:?}", e);
265        }
266    }
267}
268
269#[cfg(test)]
270mod tests {
271    use super::CgroupManager;
272    use std::fs;
273    use std::path::PathBuf;
274    use std::time::{SystemTime, UNIX_EPOCH};
275
276    fn temp_cgroup_tree(name: &str) -> PathBuf {
277        let stamp = SystemTime::now()
278            .duration_since(UNIX_EPOCH)
279            .expect("clock before unix epoch")
280            .as_nanos();
281        let path = std::env::temp_dir().join(format!("assay-cgroup-{name}-{stamp}"));
282        fs::create_dir_all(&path).expect("create temp cgroup tree");
283        path
284    }
285
286    #[test]
287    fn nearest_domain_root_accepts_domain_cgroup() {
288        let root = temp_cgroup_tree("domain");
289        fs::write(root.join("cgroup.type"), "domain\n").expect("write cgroup type");
290
291        let resolved =
292            CgroupManager::nearest_domain_root(&root, root.clone()).expect("resolve domain root");
293
294        assert_eq!(resolved, root);
295        fs::remove_dir_all(resolved).expect("remove temp cgroup tree");
296    }
297
298    #[test]
299    fn nearest_domain_root_ascends_from_domain_threaded_service() {
300        let root = temp_cgroup_tree("domain-threaded");
301        let system_slice = root.join("system.slice");
302        let service = system_slice.join("actions.runner.example.service");
303        fs::create_dir_all(&service).expect("create fake service cgroup");
304        fs::write(root.join("cgroup.type"), "domain\n").expect("write root cgroup type");
305        fs::write(system_slice.join("cgroup.type"), "domain\n")
306            .expect("write system slice cgroup type");
307        fs::write(service.join("cgroup.type"), "domain threaded\n")
308            .expect("write service cgroup type");
309
310        let resolved =
311            CgroupManager::nearest_domain_root(&root, service).expect("resolve domain root");
312
313        assert_eq!(resolved, system_slice);
314        fs::remove_dir_all(root).expect("remove temp cgroup tree");
315    }
316
317    #[test]
318    fn nearest_domain_root_ascends_from_systemd_session_scope() {
319        let root = temp_cgroup_tree("session-scope");
320        let user_slice = root.join("user.slice");
321        let user_id_slice = user_slice.join("user-1000.slice");
322        let session_scope = user_id_slice.join("session-263.scope");
323        fs::create_dir_all(&session_scope).expect("create fake session cgroup");
324        fs::write(root.join("cgroup.type"), "domain\n").expect("write root cgroup type");
325        fs::write(user_slice.join("cgroup.type"), "domain\n")
326            .expect("write user slice cgroup type");
327        fs::write(user_id_slice.join("cgroup.type"), "domain\n")
328            .expect("write user id slice cgroup type");
329        fs::write(session_scope.join("cgroup.type"), "domain\n")
330            .expect("write session scope cgroup type");
331
332        let resolved =
333            CgroupManager::nearest_domain_root(&root, session_scope).expect("resolve domain root");
334
335        assert_eq!(resolved, user_id_slice);
336        fs::remove_dir_all(root).expect("remove temp cgroup tree");
337    }
338}
assay_runner_linux/cgroup.rs

assay_runner_linux/
cgroup.rs
