1use crate::{Error, Limit, Result};
2use serde::{Deserialize, Serialize};
3use std::collections::{HashMap, HashSet};
4use std::fs;
5use std::path::{Path, PathBuf};
6
7const MAX_CONFIG_SIZE: u64 = 1_048_576;
9
10#[derive(Debug, Default, Serialize, Deserialize)]
11#[serde(deny_unknown_fields)]
12pub struct Config {
13 #[serde(default)]
14 pub profiles: HashMap<String, Profile>,
15
16 #[serde(default, skip_serializing_if = "GuardConfig::is_default")]
19 pub guard: GuardConfig,
20
21 #[serde(default, skip_serializing_if = "HashMap::is_empty")]
25 pub rules: HashMap<String, AppRule>,
26}
27
28#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
32#[serde(deny_unknown_fields)]
33pub struct AppRule {
34 #[serde(default, skip_serializing_if = "Vec::is_empty")]
36 pub match_exe: Vec<String>,
37
38 #[serde(skip_serializing_if = "Option::is_none")]
40 pub memory: Option<String>,
41
42 #[serde(skip_serializing_if = "Option::is_none")]
44 pub cpu: Option<String>,
45
46 #[serde(skip_serializing_if = "Option::is_none")]
48 pub io_read: Option<String>,
49
50 #[serde(skip_serializing_if = "Option::is_none")]
52 pub io_write: Option<String>,
53}
54
55impl AppRule {
56 pub fn to_limit(&self) -> Result<Limit> {
57 use crate::{CpuLimit, IoLimit, MemoryLimit};
58
59 let read_bps = self
60 .io_read
61 .as_ref()
62 .map(|s| IoLimit::parse_bps(s))
63 .transpose()?;
64 let write_bps = self
65 .io_write
66 .as_ref()
67 .map(|s| IoLimit::parse_bps(s))
68 .transpose()?;
69 let io = if read_bps.is_some() || write_bps.is_some() {
70 Some(IoLimit {
71 read_bps,
72 write_bps,
73 })
74 } else {
75 None
76 };
77
78 Ok(Limit {
79 memory: self
80 .memory
81 .as_ref()
82 .map(|s| MemoryLimit::parse(s))
83 .transpose()?,
84 cpu: self.cpu.as_ref().map(|s| CpuLimit::parse(s)).transpose()?,
85 io,
86 })
87 }
88}
89
90#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
93#[serde(default, deny_unknown_fields)]
94pub struct GuardConfig {
95 pub enabled: bool,
96 pub trigger: GuardTrigger,
97 pub timing: GuardTiming,
98 pub selection: GuardSelection,
99 pub notify: bool,
100}
101
102impl Default for GuardConfig {
103 fn default() -> Self {
104 Self {
105 enabled: true,
106 trigger: GuardTrigger::default(),
107 timing: GuardTiming::default(),
108 selection: GuardSelection::default(),
109 notify: true,
110 }
111 }
112}
113
114impl GuardConfig {
115 pub fn is_default(&self) -> bool {
116 *self == GuardConfig::default()
117 }
118
119 pub fn validate(&self) -> Result<()> {
122 let bad = |m: &str| Err(Error::Config(format!("guard: {m}")));
123 let t = &self.trigger;
124 let pct = |v: f64| v > 0.0 && v <= 100.0;
125 if !pct(t.psi_some_warn) || !pct(t.psi_some_high) || !pct(t.psi_full_critical) {
126 return bad("trigger PSI thresholds must be between 0 (exclusive) and 100");
127 }
128 if t.psi_some_warn >= t.psi_some_high {
129 return bad("trigger.psi_some_warn must be below trigger.psi_some_high");
130 }
131 if !(1..=100).contains(&t.act_below_available_pct) {
132 return bad("trigger.act_below_available_pct must be between 1 and 100");
133 }
134 let tm = &self.timing;
135 if !(100..=60_000).contains(&tm.sample_interval_ms) {
136 return bad("timing.sample_interval_ms must be between 100 and 60000");
137 }
138 if !(1..=60).contains(&tm.freeze_hold_secs) {
139 return bad("timing.freeze_hold_secs must be between 1 and 60");
140 }
141 if tm.calm_hold_secs == 0 {
142 return bad("timing.calm_hold_secs must be at least 1");
143 }
144 if tm.freeze_cooldown_secs < tm.freeze_hold_secs {
145 return bad("timing.freeze_cooldown_secs must be at least timing.freeze_hold_secs");
146 }
147 if self.selection.protect.iter().any(|p| p.trim().is_empty()) {
148 return bad("selection.protect must not contain empty names");
149 }
150 Ok(())
151 }
152}
153
154#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
156#[serde(default, deny_unknown_fields)]
157pub struct GuardTrigger {
158 pub psi_some_warn: f64,
160 pub psi_some_high: f64,
162 pub psi_full_critical: f64,
164 pub mem_available_floor_mb: u64,
166 pub act_below_available_pct: u64,
168}
169
170impl Default for GuardTrigger {
171 fn default() -> Self {
172 Self {
173 psi_some_warn: 10.0,
174 psi_some_high: 30.0,
175 psi_full_critical: 10.0,
176 mem_available_floor_mb: 400,
177 act_below_available_pct: 20,
178 }
179 }
180}
181
182#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
184#[serde(default, deny_unknown_fields)]
185pub struct GuardTiming {
186 pub freeze_hold_secs: u64,
188 pub calm_hold_secs: u64,
190 pub freeze_cooldown_secs: u64,
192 pub sample_interval_ms: u64,
194}
195
196impl Default for GuardTiming {
197 fn default() -> Self {
198 Self {
199 freeze_hold_secs: 5,
200 calm_hold_secs: 30,
201 freeze_cooldown_secs: 60,
202 sample_interval_ms: 1000,
203 }
204 }
205}
206
207#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
209#[serde(default, deny_unknown_fields)]
210pub struct GuardSelection {
211 pub min_rss_mb: u64,
213 pub protect: Vec<String>,
215}
216
217impl Default for GuardSelection {
218 fn default() -> Self {
219 Self {
220 min_rss_mb: 200,
221 protect: Vec::new(),
222 }
223 }
224}
225
226pub const BUILTIN_PROTECT: &[&str] = &[
228 "gnome-shell",
229 "kwin_wayland",
230 "kwin_x11",
231 "plasmashell",
232 "sway",
233 "Hyprland",
234 "Xwayland",
235 "Xorg",
236 "sshd",
237 "systemd",
238 "dbus-daemon",
239 "pipewire",
240 "wireplumber",
241 "pulseaudio",
242 "rlm-guard",
243 "bash",
244 "zsh",
245 "fish",
246];
247
248pub fn protect_set(extra: &[String]) -> HashSet<String> {
250 BUILTIN_PROTECT
251 .iter()
252 .map(|s| (*s).to_string())
253 .chain(extra.iter().cloned())
254 .collect()
255}
256
257pub fn is_protected(set: &HashSet<String>, comm: &str, exe: Option<&str>) -> bool {
261 exe.is_some_and(|e| set.contains(e)) || set.contains(comm)
262}
263
264#[derive(Debug, Clone, Default, Serialize, Deserialize)]
265#[serde(deny_unknown_fields)]
266pub struct Profile {
267 #[serde(default, skip_serializing_if = "Vec::is_empty")]
269 pub match_exe: Vec<String>,
270
271 #[serde(skip_serializing_if = "Option::is_none")]
273 pub memory: Option<String>,
274
275 #[serde(skip_serializing_if = "Option::is_none")]
277 pub cpu: Option<String>,
278
279 #[serde(skip_serializing_if = "Option::is_none")]
281 pub io_read: Option<String>,
282
283 #[serde(skip_serializing_if = "Option::is_none")]
285 pub io_write: Option<String>,
286}
287
288impl Profile {
289 pub fn validate(&self) -> Result<()> {
292 let l = self.to_limit()?;
293 if l.is_empty() {
294 return Err(Error::Config("profile sets no limits".into()));
295 }
296 Ok(())
297 }
298
299 pub fn to_limit(&self) -> Result<Limit> {
300 use crate::{CpuLimit, IoLimit, MemoryLimit};
301
302 let read_bps = self
303 .io_read
304 .as_ref()
305 .map(|s| IoLimit::parse_bps(s))
306 .transpose()?;
307 let write_bps = self
308 .io_write
309 .as_ref()
310 .map(|s| IoLimit::parse_bps(s))
311 .transpose()?;
312 let io = if read_bps.is_some() || write_bps.is_some() {
313 Some(IoLimit {
314 read_bps,
315 write_bps,
316 })
317 } else {
318 None
319 };
320
321 Ok(Limit {
322 memory: self
323 .memory
324 .as_ref()
325 .map(|s| MemoryLimit::parse(s))
326 .transpose()?,
327 cpu: self.cpu.as_ref().map(|s| CpuLimit::parse(s)).transpose()?,
328 io,
329 })
330 }
331}
332
333pub fn builtin_presets() -> HashMap<String, Profile> {
335 let mut presets = HashMap::new();
336
337 presets.insert(
338 "Light".to_string(),
339 Profile {
340 match_exe: Vec::new(),
341 memory: Some("512M".to_string()),
342 cpu: Some("25%".to_string()),
343 io_read: None,
344 io_write: None,
345 },
346 );
347
348 presets.insert(
349 "Medium".to_string(),
350 Profile {
351 match_exe: Vec::new(),
352 memory: Some("2G".to_string()),
353 cpu: Some("50%".to_string()),
354 io_read: Some("50M".to_string()),
355 io_write: Some("25M".to_string()),
356 },
357 );
358
359 presets.insert(
360 "Heavy".to_string(),
361 Profile {
362 match_exe: Vec::new(),
363 memory: Some("4G".to_string()),
364 cpu: Some("100%".to_string()),
365 io_read: Some("100M".to_string()),
366 io_write: Some("50M".to_string()),
367 },
368 );
369
370 presets.insert(
371 "Browser".to_string(),
372 Profile {
373 match_exe: vec![
374 "firefox".to_string(),
375 "chrome".to_string(),
376 "chromium".to_string(),
377 ],
378 memory: Some("4G".to_string()),
379 cpu: Some("75%".to_string()),
380 io_read: None,
381 io_write: None,
382 },
383 );
384
385 presets
386}
387
388impl Config {
389 pub fn load() -> Result<Self> {
391 let mut config = Config::default();
392
393 let system_path = PathBuf::from("/etc/rlm/config.yaml");
395 if system_path.exists() {
396 config.merge_from(&system_path)?;
397 }
398
399 if let Some(user_path) = Self::user_config_path() {
401 if user_path.exists() {
402 config.merge_from(&user_path)?;
403 }
404
405 let profiles_dir = user_path
407 .parent()
408 .map(|p| p.join("profiles.d"))
409 .unwrap_or_else(|| PathBuf::from("profiles.d"));
410 if profiles_dir.exists() {
411 config.load_profiles_dir(&profiles_dir)?;
412 }
413 }
414
415 Ok(config)
416 }
417
418 pub fn load_validated() -> Result<Self> {
420 let c = Self::load()?;
421 c.guard.validate()?;
422 Ok(c)
423 }
424
425 pub fn load_from(path: &Path) -> Result<Self> {
427 let metadata = fs::metadata(path)?;
429 if metadata.len() > MAX_CONFIG_SIZE {
430 return Err(Error::Config(format!(
431 "config file {} exceeds maximum size of 1MB",
432 path.display()
433 )));
434 }
435
436 let content = fs::read_to_string(path)?;
437 serde_yaml_ng::from_str(&content)
438 .map_err(|e| Error::Config(format!("failed to parse {}: {e}", path.display())))
439 }
440
441 fn merge_from(&mut self, path: &Path) -> Result<()> {
442 let other = Self::load_from(path)?;
443 self.profiles.extend(other.profiles);
444 self.rules.extend(other.rules);
445 if !other.guard.is_default() {
447 self.guard = other.guard;
448 }
449 Ok(())
450 }
451
452 fn load_profiles_dir(&mut self, dir: &Path) -> Result<()> {
453 for entry in fs::read_dir(dir)? {
454 let entry = entry?;
455 let path = entry.path();
456 if path.extension().is_some_and(|e| e == "yaml" || e == "yml") {
457 self.merge_from(&path)?;
458 }
459 }
460 Ok(())
461 }
462
463 fn user_config_path() -> Option<PathBuf> {
464 dirs::config_dir().map(|d| d.join("rlm").join("config.yaml"))
465 }
466
467 pub fn get_profile(&self, name: &str) -> Option<Profile> {
471 let resolved = self.resolve_profile_name(name)?;
472 self.all_profiles().get(&resolved).cloned()
473 }
474
475 pub fn resolve_profile_name(&self, name: &str) -> Option<String> {
479 let all = self.all_profiles();
480 if all.contains_key(name) {
481 return Some(name.to_string());
482 }
483 let mut matches = all.keys().filter(|k| k.eq_ignore_ascii_case(name));
484 let first = matches.next()?.clone();
485 if matches.next().is_some() {
486 None
487 } else {
488 Some(first)
489 }
490 }
491
492 pub fn profile_names(&self) -> Vec<String> {
495 let mut names: Vec<String> = self.all_profiles().into_keys().collect();
496 names.sort_by(|a, b| {
497 a.to_lowercase()
498 .cmp(&b.to_lowercase())
499 .then_with(|| a.cmp(b))
500 });
501 names
502 }
503
504 pub fn all_profiles(&self) -> HashMap<String, Profile> {
506 let mut all = builtin_presets();
507 for (name, profile) in &self.profiles {
509 all.insert(name.clone(), profile.clone());
510 }
511 all
512 }
513
514 pub fn add_rule(&mut self, name: impl Into<String>, rule: AppRule) {
516 self.rules.insert(name.into(), rule);
517 }
518
519 pub fn remove_rule(&mut self, name: &str) -> bool {
521 self.rules.remove(name).is_some()
522 }
523
524 pub fn save(&self) -> Result<()> {
526 let path = Self::user_config_path()
527 .ok_or_else(|| Error::Config("No config directory found".into()))?;
528
529 if let Some(parent) = path.parent() {
530 fs::create_dir_all(parent)?;
531 }
532
533 let yaml = serde_yaml_ng::to_string(self)
534 .map_err(|e| Error::Config(format!("Failed to serialize config: {e}")))?;
535
536 let tmp_path = path.with_extension("yaml.tmp");
538 fs::write(&tmp_path, &yaml)?;
539 fs::rename(&tmp_path, &path)?;
540 Ok(())
541 }
542}
543
544#[cfg(test)]
545mod tests {
546 use super::*;
547
548 #[test]
549 fn app_rule_to_limit_parses_fields() {
550 let rule = AppRule {
551 match_exe: vec!["firefox".into()],
552 memory: Some("4G".into()),
553 cpu: Some("75%".into()),
554 io_read: None,
555 io_write: None,
556 };
557 let limit = rule.to_limit().unwrap();
558 assert_eq!(limit.memory.unwrap().bytes(), 4 * 1024 * 1024 * 1024);
559 assert_eq!(limit.cpu.unwrap().percent(), 75);
560 assert!(limit.io.is_none());
561 }
562
563 #[test]
564 fn app_rule_invalid_limit_errors() {
565 let rule = AppRule {
566 match_exe: vec!["x".into()],
567 memory: Some("notasize".into()),
568 ..Default::default()
569 };
570 assert!(rule.to_limit().is_err());
571 }
572
573 #[test]
574 fn empty_rules_omitted_from_yaml() {
575 let cfg = Config::default();
576 let yaml = serde_yaml_ng::to_string(&cfg).unwrap();
577 assert!(
578 !yaml.contains("rules:"),
579 "empty rules must be omitted: {yaml}"
580 );
581 }
582
583 #[test]
584 fn rules_round_trip_through_yaml() {
585 let mut cfg = Config::default();
586 cfg.add_rule(
587 "firefox",
588 AppRule {
589 match_exe: vec!["firefox".into()],
590 memory: Some("4G".into()),
591 cpu: Some("75%".into()),
592 io_read: None,
593 io_write: None,
594 },
595 );
596 let yaml = serde_yaml_ng::to_string(&cfg).unwrap();
597 assert!(yaml.contains("rules:"));
598 let back: Config = serde_yaml_ng::from_str(&yaml).unwrap();
599 let r = back.rules.get("firefox").expect("rule present");
600 assert_eq!(r.match_exe, vec!["firefox".to_string()]);
601 assert_eq!(r.memory.as_deref(), Some("4G"));
602 }
603
604 #[test]
605 fn add_and_remove_rule() {
606 let mut cfg = Config::default();
607 cfg.add_rule("code", AppRule::default());
608 assert!(cfg.rules.contains_key("code"));
609 assert!(cfg.remove_rule("code"));
610 assert!(!cfg.remove_rule("code"));
611 assert!(cfg.rules.is_empty());
612 }
613
614 #[test]
615 fn protect_set_merges_builtin_and_extra() {
616 let s = protect_set(&["gnome-control-center".into()]);
617 assert!(s.contains("gnome-shell"));
618 assert!(s.contains("gnome-control-center"));
619 }
620
621 #[test]
622 fn is_protected_prefers_full_exe_name_over_truncated_comm() {
623 let s = protect_set(&["gnome-control-center".into()]);
624 assert!(is_protected(
625 &s,
626 "gnome-control-c",
627 Some("gnome-control-center")
628 ));
629 assert!(
630 !is_protected(&s, "gnome-control-c", None),
631 "truncated comm alone cannot match"
632 );
633 assert!(is_protected(&s, "bash", None));
634 assert!(!is_protected(&s, "firefox", Some("firefox")));
635 }
636
637 #[test]
638 fn readme_guard_example_parses_and_validates() {
639 let yaml = "guard:\n enabled: true\n trigger: { psi_some_warn: 10, psi_some_high: 30, psi_full_critical: 10, mem_available_floor_mb: 400 }\n timing: { freeze_hold_secs: 5, calm_hold_secs: 30, freeze_cooldown_secs: 60, sample_interval_ms: 1000 }\n selection: { min_rss_mb: 200, protect: [] }\n notify: true\n";
640 let cfg: Config = serde_yaml_ng::from_str(yaml).unwrap();
641 cfg.guard.validate().unwrap();
642 assert_eq!(cfg.guard.trigger.act_below_available_pct, 20);
643 }
644
645 #[test]
646 fn unknown_guard_key_is_an_error() {
647 let err = serde_yaml_ng::from_str::<Config>("guard:\n selection: { min_rss: 100 }\n")
648 .unwrap_err()
649 .to_string();
650 assert!(err.contains("min_rss"), "{err}");
651 }
652
653 #[test]
654 fn unknown_top_level_and_profile_keys_are_errors() {
655 assert!(serde_yaml_ng::from_str::<Config>("gaurd:\n enabled: false\n").is_err());
656 assert!(serde_yaml_ng::from_str::<Config>("profiles:\n a: { memroy: 2G }\n").is_err());
657 }
658
659 #[test]
660 fn claude_md_profile_example_still_parses() {
661 let yaml = "profiles:\n browser:\n match_exe: [firefox, chrome]\n memory: \"4G\"\n cpu: \"75%\"\n io_read: \"100M\"\n io_write: \"50M\"\n";
662 let cfg: Config = serde_yaml_ng::from_str(yaml).unwrap();
663 assert_eq!(cfg.profiles["browser"].memory.as_deref(), Some("4G"));
664 }
665
666 #[test]
667 fn default_guard_config_validates() {
668 GuardConfig::default().validate().unwrap();
669 }
670
671 #[test]
672 #[allow(clippy::type_complexity)]
673 fn validate_rejects_bad_values() {
674 let bad: Vec<Box<dyn Fn(&mut GuardConfig)>> = vec![
675 Box::new(|c| {
676 c.trigger.psi_some_warn = 40.0;
677 c.trigger.psi_some_high = 30.0
678 }),
679 Box::new(|c| c.trigger.psi_full_critical = 0.0),
680 Box::new(|c| c.trigger.psi_some_high = f64::NAN),
681 Box::new(|c| c.trigger.act_below_available_pct = 0),
682 Box::new(|c| c.trigger.act_below_available_pct = 101),
683 Box::new(|c| c.timing.sample_interval_ms = 0),
684 Box::new(|c| c.timing.freeze_hold_secs = 0),
685 Box::new(|c| c.timing.calm_hold_secs = 0),
686 Box::new(|c| c.timing.freeze_cooldown_secs = 1),
687 Box::new(|c| c.selection.protect = vec![" ".into()]),
688 ];
689 for (i, f) in bad.iter().enumerate() {
690 let mut c = GuardConfig::default();
691 f(&mut c);
692 assert!(c.validate().is_err(), "case {i} should be rejected");
693 }
694 }
695
696 #[test]
697 fn profile_lookup_is_case_insensitive_when_unique() {
698 let cfg = Config::default();
699 assert!(cfg.get_profile("browser").is_some());
700 assert!(cfg.get_profile("MEDIUM").is_some());
701 assert!(cfg.get_profile("nope").is_none());
702 }
703
704 #[test]
705 fn exact_profile_name_wins_and_ambiguity_is_refused() {
706 let mut cfg = Config::default();
707 cfg.profiles.insert(
708 "browser".into(),
709 Profile {
710 memory: Some("1G".into()),
711 ..Default::default()
712 },
713 );
714 assert_eq!(
715 cfg.get_profile("browser").unwrap().memory.as_deref(),
716 Some("1G")
717 );
718 assert_eq!(
719 cfg.get_profile("Browser").unwrap().memory.as_deref(),
720 Some("4G")
721 );
722 assert!(
723 cfg.get_profile("BROWSER").is_none(),
724 "two case-insensitive matches"
725 );
726 }
727
728 #[test]
729 fn profile_names_are_sorted_case_insensitively() {
730 let mut cfg = Config::default();
731 cfg.profiles.insert(
732 "aaa".into(),
733 Profile {
734 cpu: Some("10%".into()),
735 ..Default::default()
736 },
737 );
738 assert_eq!(
739 cfg.profile_names(),
740 vec!["aaa", "Browser", "Heavy", "Light", "Medium"]
741 );
742 }
743
744 #[test]
745 fn profile_validate_rejects_empty_and_invalid() {
746 assert!(Profile::default().validate().is_err());
747 assert!(Profile {
748 memory: Some("lots".into()),
749 ..Default::default()
750 }
751 .validate()
752 .is_err());
753 assert!(Profile {
754 cpu: Some("50%".into()),
755 ..Default::default()
756 }
757 .validate()
758 .is_ok());
759 }
760
761 #[test]
762 fn load_from_names_the_file_on_parse_error() {
763 let dir = tempfile::tempdir().unwrap();
764 let p = dir.path().join("config.yaml");
765 std::fs::write(&p, "profiles: [\n").unwrap();
766 let err = Config::load_from(&p).unwrap_err().to_string();
767 assert!(err.contains("config.yaml"), "{err}");
768 }
769}