1use common::{CpuLimit, Error, IoLimit, Limit, MemoryLimit, Result};
2use std::fs;
3use std::path::{Path, PathBuf};
4use std::process::Command;
5
6const CGROUP_ROOT: &str = "/sys/fs/cgroup";
7
8pub const UNLIMIT_CGROUP_NAME: &str = "unlimit";
16
17fn sanitize_cgroup_name(name: &str) -> Result<&str> {
20 if name.is_empty() {
22 return Err(Error::InvalidArgs("cgroup name cannot be empty".into()));
23 }
24
25 if name.contains('/') || name.contains('\\') || name.contains("..") {
27 return Err(Error::InvalidArgs(
28 "cgroup name contains invalid characters".into(),
29 ));
30 }
31
32 if !name
34 .chars()
35 .all(|c| c.is_alphanumeric() || c == '-' || c == '_')
36 {
37 return Err(Error::InvalidArgs(
38 "cgroup name must contain only alphanumeric characters, dashes, or underscores".into(),
39 ));
40 }
41
42 Ok(name)
43}
44
45fn already_limited(pid: u32, cgroup: &str) -> Error {
49 Error::InvalidArgs(format!(
50 "process {pid} is already limited in cgroup '{cgroup}'; run rlm unlimit --cgroup {cgroup} first"
51 ))
52}
53
54fn reject_critical_pid(pid: u32) -> Result<()> {
57 if pid <= 1 {
58 return Err(Error::InvalidArgs(format!(
59 "refusing to limit PID {pid} (init/system critical)"
60 )));
61 }
62 Ok(())
63}
64
65#[derive(Debug, Clone, PartialEq, Eq)]
67pub struct Prepared {
68 pub path: PathBuf,
70 pub created: bool,
72 pub warnings: Vec<String>,
74}
75
76pub fn io_max_lines(devices: &[(u32, u32)], limit: IoLimit) -> Vec<String> {
79 devices
80 .iter()
81 .map(|(major, minor)| {
82 let mut line = format!("{major}:{minor}");
83 if let Some(rbps) = limit.read_bps {
84 line.push_str(&format!(" rbps={rbps}"));
85 }
86 if let Some(wbps) = limit.write_bps {
87 line.push_str(&format!(" wbps={wbps}"));
88 }
89 line
90 })
91 .collect()
92}
93
94pub struct CgroupManager {
95 base_path: PathBuf,
96}
97
98impl CgroupManager {
99 pub fn new() -> Result<Self> {
100 let controllers_path = PathBuf::from(CGROUP_ROOT).join("cgroup.controllers");
102 if !controllers_path.exists() {
103 return Err(Error::CgroupsV2NotAvailable(PathBuf::from(CGROUP_ROOT)));
104 }
105
106 Ok(Self {
107 base_path: Self::default_base_path(),
108 })
109 }
110
111 pub fn at(base_path: PathBuf) -> Self {
114 Self { base_path }
115 }
116
117 pub fn default_base_path() -> PathBuf {
120 let uid = fs::read_to_string("/proc/self/status").ok().and_then(|s| {
125 s.lines()
126 .find(|l| l.starts_with("Uid:"))
127 .and_then(|l| l.split_whitespace().nth(1))
128 .and_then(|u| u.parse::<u32>().ok())
129 });
130
131 if let Some(uid) = uid {
133 let user_slice = PathBuf::from(CGROUP_ROOT).join(format!(
134 "user.slice/user-{uid}.slice/user@{uid}.service/rlm"
135 ));
136
137 if let Some(parent) = user_slice.parent() {
138 if parent.exists() {
139 return user_slice;
140 }
141 }
142 }
143
144 PathBuf::from(CGROUP_ROOT).join("rlm")
146 }
147
148 pub fn base_path(&self) -> &Path {
150 &self.base_path
151 }
152
153 pub fn prepare_cgroup(&self, name: &str, limit: &Limit) -> Result<Prepared> {
158 let safe_name = sanitize_cgroup_name(name)?;
160 let path = self.base_path.join(safe_name);
161 let created = self.create_cgroup(&path)?;
162 match self.set_limits(&path, limit) {
163 Ok(warnings) => Ok(Prepared {
164 path,
165 created,
166 warnings,
167 }),
168 Err(e) => {
169 if created {
170 let _ = self.cleanup_cgroup(safe_name);
171 }
172 Err(e)
173 }
174 }
175 }
176
177 fn set_limits(&self, cgroup_path: &Path, limit: &Limit) -> Result<Vec<String>> {
180 if let Some(mem) = &limit.memory {
181 self.set_memory_limit(cgroup_path, *mem)?;
182 }
183
184 if let Some(cpu) = &limit.cpu {
185 self.set_cpu_limit(cgroup_path, *cpu)?;
186 }
187
188 let mut warnings = Vec::new();
189 if let Some(io) = &limit.io {
190 if !io.is_empty() {
191 warnings = self.set_io_limit(cgroup_path, *io);
192 }
193 }
194
195 Ok(warnings)
196 }
197
198 pub fn placement_command(&self, cgroup_path: &Path, program: &str) -> Command {
213 use std::os::unix::process::CommandExt;
214
215 let mut cmd = Command::new(program);
216 if let Ok(file) = fs::OpenOptions::new()
217 .write(true)
218 .open(cgroup_path.join("cgroup.procs"))
219 {
220 unsafe {
224 cmd.pre_exec(move || {
225 use std::io::Write;
226 let _ = (&file).write_all(b"0");
227 Ok(())
228 });
229 }
230 }
231 cmd
232 }
233
234 pub fn add_to_cgroup(&self, cgroup_path: &Path, pid: u32) -> Result<()> {
236 self.add_process(cgroup_path, pid)?;
237 tracing::info!(pid, ?cgroup_path, "added process to cgroup");
238 Ok(())
239 }
240
241 pub fn find_cgroup_for_pid(&self, pid: u32) -> Option<String> {
245 let entries = fs::read_dir(&self.base_path).ok()?;
246
247 for entry in entries.flatten() {
248 let path = entry.path();
249 if !path.is_dir() || entry.file_name() == UNLIMIT_CGROUP_NAME {
250 continue;
251 }
252
253 let procs_file = path.join("cgroup.procs");
254 if let Ok(content) = fs::read_to_string(&procs_file) {
255 for line in content.lines() {
256 if line.trim().parse::<u32>().ok() == Some(pid) {
257 return path.file_name()?.to_str().map(String::from);
258 }
259 }
260 }
261 }
262 None
263 }
264
265 pub fn apply_limit(&self, pid: u32, limit: &Limit) -> Result<Vec<String>> {
268 reject_critical_pid(pid)?;
269
270 if let Some(existing_cgroup) = self.find_cgroup_for_pid(pid) {
272 if existing_cgroup == format!("pid-{pid}") {
274 let cgroup_path = self.base_path.join(&existing_cgroup);
275 let warnings = self.set_limits(&cgroup_path, limit)?;
276 tracing::info!(pid, "updated existing limits");
277 return Ok(warnings);
278 }
279 return Err(already_limited(pid, &existing_cgroup));
282 }
283
284 let Prepared {
285 path: cgroup_path,
286 warnings,
287 ..
288 } = self.prepare_cgroup(&format!("pid-{pid}"), limit)?;
289
290 if let Err(e) = self.add_process(&cgroup_path, pid) {
293 let _ = self.remove_if_empty(&format!("pid-{pid}"));
295 if !PathBuf::from(format!("/proc/{pid}")).exists() {
297 return Err(Error::ProcessNotFound(pid));
298 }
299 return Err(e);
300 }
301
302 tracing::info!(pid, ?cgroup_path, "applied limits");
303 Ok(warnings)
304 }
305
306 pub fn apply_limit_to_multiple(
311 &self,
312 pids: &[u32],
313 limit: &Limit,
314 cgroup_name: &str,
315 ) -> Result<Vec<String>> {
316 if pids.is_empty() {
317 return Err(Error::InvalidArgs("no processes specified".into()));
318 }
319
320 for pid in pids {
321 reject_critical_pid(*pid)?;
322 }
323
324 let safe_name = sanitize_cgroup_name(cgroup_name)?;
326
327 for pid in pids {
329 if let Some(existing_cgroup) = self.find_cgroup_for_pid(*pid) {
330 if existing_cgroup != safe_name {
332 return Err(already_limited(*pid, &existing_cgroup));
333 }
334 }
335 }
336
337 let Prepared {
339 path: cgroup_path,
340 created,
341 mut warnings,
342 } = self.prepare_cgroup(safe_name, limit)?;
343
344 let mut failed_pids = Vec::new();
346 for pid in pids {
347 if let Err(e) = self.add_process(&cgroup_path, *pid) {
348 tracing::warn!(pid, error = %e, "failed to add process to cgroup");
349 failed_pids.push(*pid);
350 } else {
351 tracing::info!(pid, ?cgroup_path, "added process to shared cgroup");
352 }
353 }
354
355 if failed_pids.len() == pids.len() {
357 if created {
358 let _ = self.cleanup_cgroup(safe_name);
359 }
360 return Err(Error::InvalidArgs(
361 "failed to add any processes to cgroup".into(),
362 ));
363 }
364
365 if !failed_pids.is_empty() {
367 warnings.push(format!(
368 "could not add {} of {} processes: {:?}",
369 failed_pids.len(),
370 pids.len(),
371 failed_pids
372 ));
373 }
374
375 Ok(warnings)
376 }
377
378 pub fn remove_limit(&self, pid: u32) -> Result<()> {
380 self.cleanup_cgroup(&format!("pid-{pid}"))
381 }
382
383 pub fn remove_application_limit(&self, cgroup_name: &str) -> Result<()> {
385 self.cleanup_cgroup(cgroup_name)
386 }
387
388 pub fn cleanup_cgroup(&self, name: &str) -> Result<()> {
390 let safe_name = sanitize_cgroup_name(name)?;
392 let cgroup_path = self.base_path.join(safe_name);
393
394 if !cgroup_path.exists() {
395 return Ok(());
396 }
397
398 if let Ok(content) = fs::read_to_string(cgroup_path.join("cgroup.procs")) {
401 let pids: Vec<u32> = content
402 .lines()
403 .filter_map(|l| l.trim().parse().ok())
404 .collect();
405
406 if !pids.is_empty() {
407 let unlimit_path = self.base_path.join(UNLIMIT_CGROUP_NAME);
409 let _ = fs::create_dir(&unlimit_path);
410 let unlimit_procs = unlimit_path.join("cgroup.procs");
411
412 for pid in pids {
413 if fs::write(&unlimit_procs, pid.to_string()).is_ok() {
414 tracing::debug!(pid, "moved process to unlimit cgroup");
415 }
416 }
417 }
418 }
419
420 for _ in 0..3 {
422 match fs::remove_dir(&cgroup_path) {
423 Ok(()) => {
424 tracing::info!(?cgroup_path, "removed cgroup");
425 return Ok(());
426 }
427 Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(()),
428 Err(_) => std::thread::sleep(std::time::Duration::from_millis(50)),
429 }
430 }
431
432 let still_has_procs = fs::read_to_string(cgroup_path.join("cgroup.procs"))
436 .map(|c| c.lines().any(|l| !l.trim().is_empty()))
437 .unwrap_or(false);
438
439 if still_has_procs {
440 let _ = fs::write(cgroup_path.join("cgroup.freeze"), "0");
443 let _ = fs::write(cgroup_path.join("memory.high"), "max");
444 let _ = fs::write(cgroup_path.join("memory.max"), "max");
445 let _ = fs::write(cgroup_path.join("memory.swap.max"), "max");
446 let _ = fs::write(cgroup_path.join("cpu.max"), "max");
447 let _ = fs::write(cgroup_path.join("io.max"), "");
448 tracing::warn!(
449 ?cgroup_path,
450 "could not remove cgroup (still has live processes); limits reset in place"
451 );
452 return Ok(());
453 }
454
455 Err(Error::Cgroup(format!(
457 "failed to remove cgroup '{safe_name}'"
458 )))
459 }
460
461 pub fn is_populated(&self, name: &str) -> Option<bool> {
464 let content = fs::read_to_string(self.base_path.join(name).join("cgroup.events")).ok()?;
465 crate::guard::cgfs::parse_populated(&content)
466 }
467
468 pub fn oom_kills(&self, name: &str) -> Option<u64> {
470 let content = fs::read_to_string(self.base_path.join(name).join("memory.events")).ok()?;
471 crate::guard::cgfs::parse_events_field(&content, "oom_kill")
472 }
473
474 pub fn memory_max(&self, name: &str) -> Option<u64> {
476 let content = fs::read_to_string(self.base_path.join(name).join("memory.max")).ok()?;
477 content.trim().parse().ok()
478 }
479
480 pub fn remove_if_empty(&self, name: &str) -> Result<bool> {
484 let safe_name = sanitize_cgroup_name(name)?;
485 if self.is_populated(safe_name) == Some(true) {
486 return Ok(false);
487 }
488 let path = self.base_path.join(safe_name);
489 let mut last_err = None;
490 for _ in 0..3 {
491 match fs::remove_dir(&path) {
492 Ok(()) => {
493 tracing::info!(?path, "removed empty cgroup");
494 return Ok(true);
495 }
496 Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(true),
497 Err(e) => {
498 last_err = Some(e);
499 std::thread::sleep(std::time::Duration::from_millis(50));
500 }
501 }
502 }
503 Err(Error::Cgroup(format!(
504 "failed to remove cgroup '{safe_name}': {}",
505 last_err.map(|e| e.to_string()).unwrap_or_default()
506 )))
507 }
508
509 pub fn cgroup_exists(&self, name: &str) -> bool {
511 self.base_path.join(name).is_dir()
512 }
513
514 pub fn pids_in_cgroup(&self, name: &str) -> Vec<u32> {
516 let procs = self.base_path.join(name).join("cgroup.procs");
517 match fs::read_to_string(procs) {
518 Ok(content) => content
519 .lines()
520 .filter_map(|l| l.trim().parse::<u32>().ok())
521 .collect(),
522 Err(_) => Vec::new(),
523 }
524 }
525
526 pub fn sweep_guard_leftovers(&self) -> Result<()> {
536 let Ok(entries) = fs::read_dir(&self.base_path) else {
537 return Ok(());
538 };
539 for entry in entries.flatten() {
540 let Some(name) = entry.file_name().to_str().map(str::to_string) else {
541 continue;
542 };
543 if !name.starts_with("guard-") {
544 continue;
545 }
546 let _ = fs::write(entry.path().join("cgroup.freeze"), "0");
550 let _ = self.cleanup_cgroup(&name);
551 }
552 Ok(())
553 }
554
555 fn create_cgroup(&self, path: &Path) -> Result<bool> {
557 if let Err(e) = fs::create_dir_all(&self.base_path) {
559 if e.kind() == std::io::ErrorKind::PermissionDenied {
560 return Err(Error::PermissionDenied {
561 path: self.base_path.clone(),
562 });
563 } else if e.kind() != std::io::ErrorKind::AlreadyExists {
564 return Err(e.into());
565 }
566 }
567
568 self.enable_controllers(&self.base_path)?;
570
571 match fs::create_dir(path) {
573 Ok(()) => Ok(true),
574 Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => Ok(false),
575 Err(e) if e.kind() == std::io::ErrorKind::PermissionDenied => {
576 Err(Error::PermissionDenied {
577 path: path.to_path_buf(),
578 })
579 }
580 Err(e) => Err(e.into()),
581 }
582 }
583
584 fn enable_controllers(&self, path: &Path) -> Result<()> {
585 let subtree_control = path.join("cgroup.subtree_control");
586
587 let controllers_file = path.join("cgroup.controllers");
589 let available = fs::read_to_string(&controllers_file).unwrap_or_default();
590
591 let mut to_enable = Vec::new();
593 for controller in ["memory", "cpu", "io"] {
594 if available.contains(controller) {
595 to_enable.push(format!("+{controller}"));
596 }
597 }
598
599 if to_enable.is_empty() {
600 return Err(Error::Cgroup(
601 "no controllers available - run as root or configure cgroup delegation".into(),
602 ));
603 }
604
605 fs::write(&subtree_control, to_enable.join(" ")).map_err(|e| {
606 if e.kind() == std::io::ErrorKind::PermissionDenied {
607 Error::Cgroup(
608 "cannot enable cgroup controllers - run as root or configure systemd cgroup delegation".into()
609 )
610 } else {
611 Error::Cgroup(format!("failed to enable controllers: {e}"))
612 }
613 })?;
614
615 Ok(())
616 }
617
618 fn set_memory_limit(&self, cgroup_path: &Path, limit: MemoryLimit) -> Result<()> {
619 let bytes = limit.bytes();
620
621 let high = bytes / 100 * 90;
625 if high > 0 {
626 let _ = fs::write(cgroup_path.join("memory.high"), high.to_string());
627 }
628
629 let memory_max = cgroup_path.join("memory.max");
631 fs::write(&memory_max, bytes.to_string())
632 .map_err(|e| Error::Cgroup(format!("failed to set memory.max: {e}")))?;
633
634 let _ = fs::write(cgroup_path.join("memory.swap.max"), "0");
638
639 Ok(())
640 }
641
642 fn set_cpu_limit(&self, cgroup_path: &Path, limit: CpuLimit) -> Result<()> {
643 let period: u64 = 100_000; let quota = u64::from(limit.percent())
648 .checked_mul(period)
649 .map(|v| v / 100)
650 .ok_or_else(|| Error::InvalidCpu("CPU percentage too large".into()))?;
651
652 let cpu_max = cgroup_path.join("cpu.max");
653 fs::write(&cpu_max, format!("{quota} {period}"))
654 .map_err(|e| Error::Cgroup(format!("failed to set cpu.max: {e}")))?;
655 Ok(())
656 }
657
658 fn add_process(&self, cgroup_path: &Path, pid: u32) -> Result<()> {
659 let procs = cgroup_path.join("cgroup.procs");
660 fs::write(&procs, pid.to_string())
661 .map_err(|e| Error::Cgroup(format!("failed to add process {pid}: {e}")))?;
662 Ok(())
663 }
664
665 fn set_io_limit(&self, cgroup_path: &Path, limit: IoLimit) -> Vec<String> {
669 let devices = match Self::get_real_block_devices() {
670 Ok(d) => d,
671 Err(e) => {
672 return vec![format!(
673 "I/O limits were not applied (could not list block devices: {e}); memory and CPU limits still apply"
674 )]
675 }
676 };
677 if devices.is_empty() {
678 return vec![
679 "I/O limits were not applied (no eligible block devices found); memory and CPU limits still apply"
680 .to_string(),
681 ];
682 }
683
684 let io_max = cgroup_path.join("io.max");
685 let mut warnings = Vec::new();
686 let mut first_err = None;
687 let mut applied = 0;
688 for line in io_max_lines(&devices, limit) {
689 match fs::write(&io_max, &line) {
690 Ok(()) => applied += 1,
691 Err(e) => {
692 let dev = line.split_whitespace().next().unwrap_or_default();
693 warnings.push(format!("I/O limit not applied to device {dev}: {e}"));
694 first_err.get_or_insert(e);
695 }
696 }
697 }
698 if applied == 0 {
699 if let Some(e) = first_err {
700 warnings = vec![format!(
701 "I/O limits were not applied to any device ({e}); memory and CPU limits still apply"
702 )];
703 }
704 }
705 for w in &warnings {
706 tracing::warn!("{w}");
707 }
708 warnings
709 }
710
711 fn get_real_block_devices() -> Result<Vec<(u32, u32)>> {
718 let mut devices = Vec::new();
719
720 let sys_block = Path::new("/sys/block");
721 if !sys_block.exists() {
722 return Ok(devices);
723 }
724
725 for entry in fs::read_dir(sys_block)? {
726 let entry = entry?;
727 let name = entry.file_name();
728 let name_str = name.to_string_lossy();
729
730 if name_str.starts_with("loop")
732 || name_str.starts_with("ram")
733 || name_str.starts_with("nbd")
734 || name_str.starts_with("zram")
735 {
736 continue;
737 }
738
739 let dev_file = entry.path().join("dev");
740 if let Ok(content) = fs::read_to_string(&dev_file) {
741 if let Some((major, minor)) = content.trim().split_once(':') {
742 if let (Ok(major), Ok(minor)) = (major.parse(), minor.parse()) {
743 devices.push((major, minor));
744 }
745 }
746 }
747 }
748
749 Ok(devices)
750 }
751}
752
753#[cfg(test)]
754mod tests {
755 use super::*;
756
757 #[test]
758 fn rejects_init_and_kernel_pids() {
759 assert!(reject_critical_pid(0).is_err()); assert!(reject_critical_pid(1).is_err()); }
762
763 #[test]
764 fn already_limited_hint_names_the_cgroup() {
765 let msg = already_limited(1234, "app-firefox").to_string();
766 assert!(
767 msg.contains("rlm unlimit --cgroup app-firefox"),
768 "hint must name the shared cgroup: {msg}"
769 );
770 assert!(!msg.contains("--pid"), "--pid does not work here: {msg}");
771 let run = already_limited(99, "run-7").to_string();
772 assert!(run.contains("rlm unlimit --cgroup run-7"), "{run}");
773 }
774
775 #[test]
776 fn allows_normal_pids() {
777 assert!(reject_critical_pid(2).is_ok());
778 assert!(reject_critical_pid(1234).is_ok());
779 }
780
781 #[test]
782 fn sanitize_rejects_traversal_and_separators() {
783 assert!(sanitize_cgroup_name("../etc").is_err());
784 assert!(sanitize_cgroup_name("a/b").is_err());
785 assert!(sanitize_cgroup_name("a\\b").is_err());
786 assert!(sanitize_cgroup_name("").is_err());
787 assert!(sanitize_cgroup_name("bad name").is_err()); }
789
790 #[test]
791 fn sanitize_accepts_valid_names() {
792 assert_eq!(sanitize_cgroup_name("pid-1234").unwrap(), "pid-1234");
793 assert_eq!(sanitize_cgroup_name("app_firefox").unwrap(), "app_firefox");
794 assert_eq!(sanitize_cgroup_name("run-42-99").unwrap(), "run-42-99");
795 }
796
797 #[test]
798 fn io_max_is_one_line_per_device_without_newlines() {
799 let l = io_max_lines(
800 &[(8, 0), (259, 0)],
801 IoLimit {
802 read_bps: Some(5_242_880),
803 write_bps: Some(1_048_576),
804 },
805 );
806 assert_eq!(
807 l,
808 vec![
809 "8:0 rbps=5242880 wbps=1048576".to_string(),
810 "259:0 rbps=5242880 wbps=1048576".to_string()
811 ]
812 );
813 }
814
815 #[test]
816 fn released_processes_can_be_limited_again() {
817 let dir = tempfile::tempdir().unwrap();
818 for (name, procs) in [("unlimit", "4242\n"), ("pid-7", "7\n")] {
819 std::fs::create_dir(dir.path().join(name)).unwrap();
820 std::fs::write(dir.path().join(name).join("cgroup.procs"), procs).unwrap();
821 }
822 let m = CgroupManager::at(dir.path().to_path_buf());
823 assert_eq!(
824 m.find_cgroup_for_pid(4242),
825 None,
826 "the unlimit bucket is not a managed cgroup"
827 );
828 assert_eq!(m.find_cgroup_for_pid(7).as_deref(), Some("pid-7"));
829 }
830
831 #[test]
832 fn populated_cgroup_is_not_removed() {
833 let dir = tempfile::tempdir().unwrap();
834 let cg = dir.path().join("run-1-2");
835 std::fs::create_dir(&cg).unwrap();
836 std::fs::write(cg.join("cgroup.events"), "populated 1\nfrozen 0\n").unwrap();
837 let m = CgroupManager::at(dir.path().to_path_buf());
838 assert_eq!(m.is_populated("run-1-2"), Some(true));
839 assert!(!m.remove_if_empty("run-1-2").unwrap());
840 assert!(cg.exists());
841 }
842
843 #[test]
844 fn oom_kills_and_memory_max_are_read() {
845 let dir = tempfile::tempdir().unwrap();
846 let cg = dir.path().join("run-1-2");
847 std::fs::create_dir(&cg).unwrap();
848 std::fs::write(
849 cg.join("memory.events"),
850 "low 0\nhigh 3\nmax 10\noom 1\noom_kill 2\n",
851 )
852 .unwrap();
853 std::fs::write(cg.join("memory.max"), "157286400\n").unwrap();
854 let m = CgroupManager::at(dir.path().to_path_buf());
855 assert_eq!(m.oom_kills("run-1-2"), Some(2));
856 assert_eq!(m.memory_max("run-1-2"), Some(157_286_400));
857 }
858}