1use std::collections::{HashMap, HashSet};
10use std::os::unix::fs::MetadataExt;
11
12use crate::guard::cgfs;
13use crate::process::ProcessInfo;
14use crate::CgroupManager;
15use common::{AppRule, Config, Limit};
16
17pub struct CompiledRule {
19 pub name: String,
20 pub match_exe: Vec<String>,
21 pub limit: Limit,
22 pub cgroup: String,
24}
25
26#[derive(Debug, Clone, PartialEq, Eq)]
28pub enum RuleAction {
29 EnsureCgroup { rule: String },
31 AddPid { rule: String, pid: u32 },
33 TeardownEmpty { rule: String },
35}
36
37pub fn cgroup_name_for(rule_name: &str) -> String {
40 format!("app-{}", rule_name.replace(['/', ' '], "_"))
41}
42
43impl CompiledRule {
44 fn compile(name: &str, rule: &AppRule) -> Option<Self> {
45 match rule.to_limit() {
46 Ok(limit) => Some(CompiledRule {
47 name: name.to_string(),
48 match_exe: rule.match_exe.clone(),
49 limit,
50 cgroup: cgroup_name_for(name),
51 }),
52 Err(e) => {
53 tracing::warn!(rule = name, error = %e, "skipping rule with invalid limits");
54 None
55 }
56 }
57 }
58
59 fn matches(&self, proc: &ProcessInfo) -> bool {
60 self.match_exe.iter().any(|want| {
61 proc.name == *want
62 || proc
63 .executable
64 .as_ref()
65 .and_then(|exe| exe.file_name())
66 .and_then(|n| n.to_str())
67 .map(|n| n == want)
68 .unwrap_or(false)
69 })
70 }
71}
72
73pub fn plan(
93 rule: &CompiledRule,
94 procs: &[ProcessInfo],
95 already_placed: &[u32],
96 cgroup_exists: bool,
97 held: bool,
98) -> Vec<RuleAction> {
99 if held {
100 return Vec::new();
101 }
102
103 let matches: Vec<&ProcessInfo> = procs.iter().filter(|p| rule.matches(p)).collect();
104
105 if matches.is_empty() {
106 return if cgroup_exists && already_placed.is_empty() {
110 vec![RuleAction::TeardownEmpty {
111 rule: rule.name.clone(),
112 }]
113 } else {
114 Vec::new()
115 };
116 }
117
118 let mut actions = vec![RuleAction::EnsureCgroup {
119 rule: rule.name.clone(),
120 }];
121 for p in matches {
122 if !already_placed.contains(&p.pid) {
123 actions.push(RuleAction::AddPid {
124 rule: rule.name.clone(),
125 pid: p.pid,
126 });
127 }
128 }
129 actions
130}
131
132pub fn needs_ensure(recorded_inode: Option<u64>, current_inode: Option<u64>) -> bool {
140 current_inode.is_none() || recorded_inode != current_inode
141}
142
143pub struct RulesEnforcer {
145 rules: Vec<CompiledRule>,
146 ensured: HashMap<String, u64>,
148}
149
150impl RulesEnforcer {
151 pub fn new(cfg: &Config) -> Self {
154 let rules = cfg
155 .rules
156 .iter()
157 .filter_map(|(name, rule)| CompiledRule::compile(name, rule))
158 .collect();
159 Self {
160 rules,
161 ensured: HashMap::new(),
162 }
163 }
164
165 pub fn rule_count(&self) -> usize {
166 self.rules.len()
167 }
168
169 pub fn reconcile(
186 &mut self,
187 mgr: &CgroupManager,
188 procs: &[ProcessInfo],
189 held_cgroups: &[String],
190 ) -> Vec<RuleAction> {
191 let rlm_rel = crate::guard::sampler::strip_cgroup_root(mgr.base_path());
197 let held: HashSet<&str> = held_cgroups.iter().map(String::as_str).collect();
198
199 let mut applied = Vec::new();
200 for rule in &self.rules {
201 let mut blocked = false;
202 if let Some(rel) = &rlm_rel {
203 let cg_path = format!("{rel}/{}", rule.cgroup);
204
205 if held.contains(cg_path.as_str()) {
209 blocked = true;
210 }
211
212 if !blocked && cgfs::read_frozen(&cg_path) == Some(true) {
221 blocked = true;
222 }
223 }
224
225 let placed = mgr.pids_in_cgroup(&rule.cgroup);
227 let exists = !placed.is_empty() || mgr.cgroup_exists(&rule.cgroup);
228
229 for action in plan(rule, procs, &placed, exists, blocked) {
230 match apply(mgr, rule, &action, &mut self.ensured) {
231 Ok(true) => applied.push(action),
232 Ok(false) => {}
233 Err(e) => tracing::warn!(?action, error = %e, "rules: action failed"),
234 }
235 }
236 }
237 applied
238 }
239}
240
241fn cgroup_inode(mgr: &CgroupManager, cgroup: &str) -> Option<u64> {
243 std::fs::metadata(mgr.base_path().join(cgroup))
244 .ok()
245 .map(|m| m.ino())
246}
247
248fn apply(
250 mgr: &CgroupManager,
251 rule: &CompiledRule,
252 action: &RuleAction,
253 ensured: &mut HashMap<String, u64>,
254) -> common::Result<bool> {
255 match action {
256 RuleAction::EnsureCgroup { .. } => {
257 let current = cgroup_inode(mgr, &rule.cgroup);
258 if !needs_ensure(ensured.get(&rule.cgroup).copied(), current) {
259 return Ok(false);
260 }
261 let prepared = mgr.prepare_cgroup(&rule.cgroup, &rule.limit)?;
263 for w in &prepared.warnings {
264 tracing::warn!(cgroup = %rule.cgroup, "{w}");
265 }
266 match cgroup_inode(mgr, &rule.cgroup) {
267 Some(ino) => ensured.insert(rule.cgroup.clone(), ino),
268 None => ensured.remove(&rule.cgroup),
269 };
270 Ok(true)
271 }
272 RuleAction::AddPid { pid, .. } => {
273 let path = mgr.base_path().join(&rule.cgroup);
274 mgr.add_to_cgroup(&path, *pid).map(|()| true)
275 }
276 RuleAction::TeardownEmpty { .. } => {
277 ensured.remove(&rule.cgroup);
278 mgr.cleanup_cgroup(&rule.cgroup).map(|()| true)
279 }
280 }
281}
282
283#[cfg(test)]
284mod tests {
285 use super::*;
286 use std::path::PathBuf;
287
288 fn rule(name: &str, exes: &[&str]) -> CompiledRule {
289 CompiledRule {
290 name: name.to_string(),
291 match_exe: exes.iter().map(|s| s.to_string()).collect(),
292 limit: Limit::default(),
293 cgroup: cgroup_name_for(name),
294 }
295 }
296
297 fn proc(pid: u32, name: &str, exe: Option<&str>) -> ProcessInfo {
298 ProcessInfo {
299 pid,
300 name: name.to_string(),
301 executable: exe.map(PathBuf::from),
302 ..Default::default()
303 }
304 }
305
306 #[test]
307 fn ensure_only_when_new_or_recreated() {
308 assert!(needs_ensure(None, None), "missing cgroup");
309 assert!(needs_ensure(None, Some(7)), "never written");
310 assert!(needs_ensure(Some(7), Some(9)), "recreated with a new inode");
311 assert!(needs_ensure(Some(7), None), "removed since");
312 assert!(!needs_ensure(Some(7), Some(7)), "unchanged: no writes");
313 }
314
315 #[test]
316 fn cgroup_name_matches_cli_scheme() {
317 assert_eq!(cgroup_name_for("firefox"), "app-firefox");
318 assert_eq!(cgroup_name_for("my app/x"), "app-my_app_x");
319 }
320
321 #[test]
322 fn matches_by_comm_or_exe_basename() {
323 let r = rule("firefox", &["firefox"]);
324 assert!(r.matches(&proc(1, "firefox", None)));
325 assert!(r.matches(&proc(2, "Web Content", Some("/usr/lib/firefox/firefox"))));
326 assert!(!r.matches(&proc(3, "code", Some("/usr/bin/code"))));
327 }
328
329 #[test]
330 fn plan_ensures_and_adds_unplaced_matches() {
331 let r = rule("firefox", &["firefox"]);
332 let procs = vec![proc(10, "firefox", None), proc(11, "firefox", None)];
333 let actions = plan(&r, &procs, &[], false, false);
334 assert_eq!(
335 actions[0],
336 RuleAction::EnsureCgroup {
337 rule: "firefox".into()
338 }
339 );
340 assert!(actions.contains(&RuleAction::AddPid {
341 rule: "firefox".into(),
342 pid: 10
343 }));
344 assert!(actions.contains(&RuleAction::AddPid {
345 rule: "firefox".into(),
346 pid: 11
347 }));
348 }
349
350 #[test]
351 fn plan_is_idempotent_when_all_placed() {
352 let r = rule("firefox", &["firefox"]);
353 let procs = vec![proc(10, "firefox", None)];
354 let actions = plan(&r, &procs, &[10], true, false);
355 assert_eq!(
357 actions,
358 vec![RuleAction::EnsureCgroup {
359 rule: "firefox".into()
360 }]
361 );
362 }
363
364 #[test]
365 fn plan_adds_only_new_pid() {
366 let r = rule("firefox", &["firefox"]);
367 let procs = vec![proc(10, "firefox", None), proc(12, "firefox", None)];
368 let actions = plan(&r, &procs, &[10], true, false);
369 assert_eq!(
370 actions,
371 vec![
372 RuleAction::EnsureCgroup {
373 rule: "firefox".into()
374 },
375 RuleAction::AddPid {
376 rule: "firefox".into(),
377 pid: 12
378 },
379 ]
380 );
381 }
382
383 #[test]
384 fn plan_teardown_only_when_empty_and_present() {
385 let r = rule("firefox", &["firefox"]);
386 let actions = plan(&r, &[proc(1, "code", None)], &[], true, false);
388 assert_eq!(
389 actions,
390 vec![RuleAction::TeardownEmpty {
391 rule: "firefox".into()
392 }]
393 );
394 }
395
396 #[test]
397 fn plan_does_not_evict_occupied_cgroup_with_no_matches() {
398 let r = rule("firefox", &["firefox"]);
402 let actions = plan(&r, &[proc(1, "code", None)], &[999], true, false);
403 assert!(
404 actions.is_empty(),
405 "must not evict an occupied cgroup: {actions:?}"
406 );
407 }
408
409 #[test]
410 fn plan_noop_when_no_matches_and_no_cgroup() {
411 let r = rule("firefox", &["firefox"]);
412 let actions = plan(&r, &[proc(1, "code", None)], &[], false, false);
413 assert!(actions.is_empty());
414 }
415
416 #[test]
419 fn plan_produces_no_actions_when_rule_cgroup_is_guard_held() {
420 let r = rule("firefox", &["firefox"]);
425 let procs = vec![proc(10, "firefox", None), proc(11, "firefox", None)];
426 let actions = plan(&r, &procs, &[10], true, true);
427 assert!(
428 actions.is_empty(),
429 "a guard-held rule cgroup must get no actions: {actions:?}"
430 );
431 }
432
433 #[test]
434 fn plan_unheld_rule_is_unaffected() {
435 let r = rule("firefox", &["firefox"]);
437 let procs = vec![proc(10, "firefox", None), proc(11, "firefox", None)];
438 let actions = plan(&r, &procs, &[10], true, false);
439 assert!(actions.contains(&RuleAction::EnsureCgroup {
440 rule: "firefox".into()
441 }));
442 assert!(actions.contains(&RuleAction::AddPid {
443 rule: "firefox".into(),
444 pid: 11
445 }));
446 }
447}