1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
| | // SPDX-FileCopyrightText: 2026 Demi Marie Obenour <demiobenour@gmail.com>
// SPDX-License-Identifier: EUPL-1.2+
mod cgroup;
use cgroup::{Cgroup, OpenFlags, openat2_simple, write_value};
use rustix::{
fs::{FlockOperation, Mode, XattrFlags},
io::Errno,
};
use std::{
env::ArgsOs,
os::unix::prelude::*,
path::{Path, PathBuf},
};
fn exec_in_cgroup(
mut args: std::iter::Peekable<ArgsOs>,
cgroup: Option<&dyn AsFd>,
) -> Result<(), String> {
let Some(program_name) = args.next() else {
return Ok(());
};
if let Some(cgroup) = cgroup {
let pid = std::process::id().to_string();
write_value(
cgroup,
Path::new("$inner.service/cgroup.procs"),
pid.as_bytes(),
)
.map_err(|e| format!("Cannot move process to child cgroup: {e}"))?;
}
let e = std::process::Command::new(&program_name).args(args).exec();
Err(format!("Cannot exec child {program_name:?}: {e}"))
}
// Check that the path is canonical,
// then split it into basename and filename.
fn split_canonical_path(path: &Path) -> Result<(&Path, &Path), String> {
let bytes = path.as_os_str().as_bytes();
// Path::components() skips ., so use string manipulation instead.
for component in bytes[path.is_absolute() as usize..].split(|&b| b == b'/') {
if matches!(component, b"" | b"." | b"..") {
return Err(format!("cgroup path {path:?} isn't canonical"));
}
}
Ok((path.parent().unwrap(), Path::new(path.file_name().unwrap())))
}
fn cgroup_setup(args: ArgsOs) -> Result<(), String> {
let mut wait = true;
let mut args = args.peekable();
while let Some(arg) = args.peek() {
if !arg.as_bytes().starts_with(b"-") {
break;
}
let arg = args.next().unwrap();
let Some(option) = arg.as_bytes().strip_prefix(b"--") else {
return Err("takes no short options".to_owned());
};
match option {
b"" => break,
b"no-wait" => wait = false,
_ => return Err(format!("unknown long option {arg:?}")),
}
}
let Some(cgroup_path) = args.next().map(PathBuf::from) else {
return Err("have no positional arguments, expected at least 1".to_owned());
};
let (parent_cgroup_path, child_cgroup_path) = split_canonical_path(&cgroup_path)?;
let cgroup = Cgroup::new(parent_cgroup_path)?;
match rustix::fs::mkdirat(&cgroup, child_cgroup_path, Mode::from_raw_mode(0o755)) {
Ok(()) | Err(Errno::EXIST) => {}
Err(e) => {
return Err(format!(
"Cannot create child cgroup {child_cgroup_path:?}: {e}"
));
}
}
let child = openat2_simple(&cgroup, child_cgroup_path, OpenFlags::Directory)
.map_err(|e| format!("Cannot open child cgroup: {e}"))?;
// While waiting, hold an exclusive lock on the child.
// This avoids two processes both waiting for the same cgroup to become
// empty, then execing processes in the same cgroup.
rustix::fs::flock(&child, FlockOperation::LockExclusive)
.map_err(|e| format!("Cannot take an exclusive lock on child cgroup: {e}"))?;
if wait {
Cgroup::wait_for_empty(&child)
.map_err(|e| format!("Cannot wait for {parent_cgroup_path:?} to be empty: {e}"))?;
}
// systemd-aware programs expect to have user.delegate=1
rustix::fs::fsetxattr(&child, c"user.delegate", b"1", XattrFlags::empty())
.map_err(|e| format!("Cannot enable cgroup delegation: {e}"))?;
// If the child process will need to manage cgroups itself, it will need
// to set up a sub-cgroup due to the "no internal processes" rule. It's
// simplest to just do it automatically. If the cgroup already exists,
// that isn't an error.
match rustix::fs::mkdirat(&child, cgroup::DEFAULT_LEAF, Mode::from_raw_mode(0o755)) {
Ok(()) | Err(Errno::EXIST) => {}
Err(e) => return Err(format!("Cannot create $inner.service cgroup: {e}")),
}
exec_in_cgroup(args, Some(&child))
}
fn cgroup_purge(mut args: ArgsOs) -> Result<(), String> {
if args.len() != 1 {
return Err("usage: cgroup-purge CGROUP_TO_PURGE".to_owned());
}
let arg = args.next().unwrap();
let (parent, child) = split_canonical_path(Path::new(&arg))?;
Cgroup::new(parent)?.purge_child(child)
}
fn run(prog_name: &Path, args: ArgsOs) -> Result<(), String> {
match prog_name.file_name().map(|f| f.as_bytes()) {
Some(b"cgroup-setup") => cgroup_setup(args),
Some(b"cgroup-purge") => cgroup_purge(args),
_ => Err(format!(
"must be invoked as \"cgroup-setup\" or \
\"cgroup-purge\", got {prog_name:?}",
)),
}
}
fn main() {
let mut args = std::env::args_os();
let Some(prog_name) = args.next() else {
eprintln!("No command line arguments (argv[0] is NULL)");
std::process::exit(1);
};
match run(Path::new(&prog_name), args) {
Ok(()) => {}
Err(e) => {
eprintln!("{prog_name:?}: {}", e);
std::process::exit(1);
}
}
}
|