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" {
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    /// Creates a new ephemeral cgroup.
90    pub fn create_session(&self) -> Result<SessionCgroup> {
91        let timestamp = SystemTime::now().duration_since(UNIX_EPOCH)?.as_millis();
92        let name = format!("assay-session-{}", timestamp);
93        let path = self.root_path.join(&name);
94
95        // P0: Best-effort remove if collision (unlikely with timestamp)
96        if path.exists() {
97            let _ = fs::remove_dir(&path);
98        }
99
100        fs::create_dir(&path).context("Failed to create cgroup session dir")?;
101
102        // Best effort enable pids controller
103        let subtree = self.root_path.join("cgroup.subtree_control");
104        if subtree.exists() {
105            // Ignoring errors as we might lack delegation
106            let _ = fs::write(&subtree, "+pids");
107        }
108
109        let meta = fs::metadata(&path)?;
110        #[cfg(unix)]
111        let id = meta.ino();
112        #[cfg(not(unix))]
113        let id = {
114            let _ = meta;
115            0
116        }; // Stub for non-Unix
117        Ok(SessionCgroup { path, id })
118    }
119}
120
121impl SessionCgroup {
122    pub fn id(&self) -> u64 {
123        self.id
124    }
125
126    pub fn path(&self) -> &Path {
127        &self.path
128    }
129
130    pub fn procs_path(&self) -> PathBuf {
131        self.path.join("cgroup.procs")
132    }
133
134    pub fn add_process(&self, pid: u32) -> Result<()> {
135        let procs_path = self.procs_path();
136        fs::write(&procs_path, pid.to_string()).context("Failed to add PID")?;
137        Ok(())
138    }
139
140    pub fn freeze(&self) -> Result<()> {
141        let p = self.path.join("cgroup.freeze");
142        if p.exists() {
143            fs::write(p, "1")?;
144        }
145        Ok(())
146    }
147
148    pub fn thaw(&self) -> Result<()> {
149        let p = self.path.join("cgroup.freeze");
150        if p.exists() {
151            fs::write(p, "0")?;
152        }
153        Ok(())
154    }
155
156    pub fn kill(&self) -> Result<()> {
157        let p = self.path.join("cgroup.kill");
158        if p.exists() {
159            fs::write(p, "1")?;
160        } else {
161            return Err(anyhow!("cgroup.kill missing"));
162        }
163        Ok(())
164    }
165
166    /// P0 Hardening: Graceful kill using pidfd to avoid PID reuse races.
167    pub fn kill_graceful(&self, grace_ms: u64) -> Result<()> {
168        let _ = self.freeze();
169
170        let procs = fs::read_to_string(self.path.join("cgroup.procs"))?;
171        let mut pids: Vec<i32> = procs
172            .lines()
173            .filter_map(|l| l.trim().parse::<i32>().ok())
174            .collect();
175        // SOTA: Dedupe to prevent double-signaling (which is harmless but sloppy)
176        pids.sort_unstable();
177        pids.dedup();
178
179        // Send SIGTERM via pidfd if possible
180        for &pid in &pids {
181            if pid <= 0 {
182                continue;
183            }
184
185            #[cfg(target_os = "linux")]
186            unsafe {
187                // syscall(SYS_pidfd_open, pid, 0)
188                let fd = libc::syscall(libc::SYS_pidfd_open, pid, 0) as i32;
189                if fd >= 0 {
190                    // syscall(SYS_pidfd_send_signal, fd, SIGTERM, NULL, 0)
191                    let _ = libc::syscall(
192                        libc::SYS_pidfd_send_signal,
193                        fd,
194                        libc::SIGTERM,
195                        std::ptr::null::<libc::siginfo_t>(),
196                        0,
197                    );
198                    libc::close(fd);
199                } else {
200                    // Fallback to classic kill if pidfd fails
201                    libc::kill(pid, libc::SIGTERM);
202                }
203            }
204
205            #[cfg(all(unix, not(target_os = "linux")))]
206            {
207                use nix::sys::signal::{kill, Signal};
208                use nix::unistd::Pid;
209                let _ = kill(Pid::from_raw(pid), Signal::SIGTERM);
210            }
211            #[cfg(not(unix))]
212            {
213                // Windows/Other: No-op or standard process kill if we had handle
214                // Since this uses PIDs from cgroup.procs (Linux concept), this code is unreachable logic-wise
215                // but must compile.
216                let _ = pid;
217            }
218        }
219
220        let _ = self.thaw();
221        std::thread::sleep(std::time::Duration::from_millis(grace_ms));
222        self.kill()
223    }
224
225    pub fn set_pids_max(&self, max: u32) -> Result<()> {
226        let p = self.path.join("pids.max");
227        if p.exists() {
228            fs::write(p, max.to_string())?;
229        }
230        Ok(())
231    }
232
233    /// P0 Hardening: Reliable cleanup
234    pub fn remove(&self) -> Result<()> {
235        if !self.path.exists() {
236            return Ok(());
237        }
238
239        // Retry loop for cgroup removal (busy/not empty)
240        for _ in 0..3 {
241            match fs::remove_dir(&self.path) {
242                Ok(_) => return Ok(()),
243                Err(_) => {
244                    // Try to kill remainders if still exists
245                    let _ = self.kill();
246                    std::thread::sleep(std::time::Duration::from_millis(50));
247                }
248            }
249        }
250        // Final attempt
251        fs::remove_dir(&self.path).context("Failed to remove cgroup after retries")
252    }
253}
254
255impl Drop for SessionCgroup {
256    fn drop(&mut self) {
257        if let Err(e) = self.remove() {
258            eprintln!("Failed to clean up cgroup session: {:?}", e);
259        }
260    }
261}
262
263#[cfg(test)]
264mod tests {
265    use super::CgroupManager;
266    use std::fs;
267    use std::path::PathBuf;
268    use std::time::{SystemTime, UNIX_EPOCH};
269
270    fn temp_cgroup_tree(name: &str) -> PathBuf {
271        let stamp = SystemTime::now()
272            .duration_since(UNIX_EPOCH)
273            .expect("clock before unix epoch")
274            .as_nanos();
275        let path = std::env::temp_dir().join(format!("assay-cgroup-{name}-{stamp}"));
276        fs::create_dir_all(&path).expect("create temp cgroup tree");
277        path
278    }
279
280    #[test]
281    fn nearest_domain_root_accepts_domain_cgroup() {
282        let root = temp_cgroup_tree("domain");
283        fs::write(root.join("cgroup.type"), "domain\n").expect("write cgroup type");
284
285        let resolved =
286            CgroupManager::nearest_domain_root(&root, root.clone()).expect("resolve domain root");
287
288        assert_eq!(resolved, root);
289        fs::remove_dir_all(resolved).expect("remove temp cgroup tree");
290    }
291
292    #[test]
293    fn nearest_domain_root_ascends_from_domain_threaded_service() {
294        let root = temp_cgroup_tree("domain-threaded");
295        let system_slice = root.join("system.slice");
296        let service = system_slice.join("actions.runner.example.service");
297        fs::create_dir_all(&service).expect("create fake service cgroup");
298        fs::write(root.join("cgroup.type"), "domain\n").expect("write root cgroup type");
299        fs::write(system_slice.join("cgroup.type"), "domain\n")
300            .expect("write system slice cgroup type");
301        fs::write(service.join("cgroup.type"), "domain threaded\n")
302            .expect("write service cgroup type");
303
304        let resolved =
305            CgroupManager::nearest_domain_root(&root, service).expect("resolve domain root");
306
307        assert_eq!(resolved, system_slice);
308        fs::remove_dir_all(root).expect("remove temp cgroup tree");
309    }
310}
assay_runner_linux/cgroup.rs

assay_runner_linux/
cgroup.rs
