1use std::collections::BTreeSet;
2use std::fs;
3#[cfg(unix)]
4use std::fs::File;
5use std::io::Write as _;
6use std::path::{Path, PathBuf};
7
8use serde::{Deserialize, Serialize};
9use tempfile::NamedTempFile;
10use toml_edit::{Array, DocumentMut, Item, Table, value};
11use yarp_rule_pack::CompiledPack;
12
13pub const CONFIG_VERSION: u32 = 1;
14pub const DEFAULT_OUTPUT_CAP_BYTES: usize = 5 * 1024;
15pub const MIN_OUTPUT_CAP_BYTES: usize = 1024;
16pub const MAX_OUTPUT_CAP_BYTES: usize = 16 * 1024 * 1024;
17pub const DEFAULT_RECOVERY_CAP_BYTES: usize = 32 * 1024;
18pub const MIN_RECOVERY_CAP_BYTES: usize = 1024;
19pub const MAX_RECOVERY_CAP_BYTES: usize = 48 * 1024;
20pub const DEFAULT_RECOVERY_CAP_LINES: usize = 1900;
21pub const MIN_RECOVERY_CAP_LINES: usize = 1;
22pub const MAX_RECOVERY_CAP_LINES: usize = 1900;
23
24const DEFAULT_DOCUMENT: &str = r"# YARP configuration
25version = 1
26
27[pruning]
28enabled = true
29
30[output]
31cap_bytes = 5120
32recovery_cap_bytes = 32768
33recovery_cap_lines = 1900
34
35[archive]
36enabled = true
37# Omit path to use the XDG data directory.
38
39[rules]
40packs = []
41";
42
43#[derive(Clone, Debug, Deserialize)]
44#[serde(deny_unknown_fields)]
45struct FileConfig {
46 version: u32,
47 #[serde(default)]
48 pruning: FilePruning,
49 #[serde(default)]
50 output: FileOutput,
51 #[serde(default)]
52 archive: FileArchive,
53 #[serde(default)]
54 rules: FileRules,
55}
56
57#[derive(Clone, Debug, Default, Deserialize)]
58#[serde(deny_unknown_fields)]
59struct FilePruning {
60 enabled: Option<bool>,
61}
62
63#[derive(Clone, Debug, Default, Deserialize)]
64#[serde(deny_unknown_fields)]
65struct FileOutput {
66 cap_bytes: Option<usize>,
67 recovery_cap_bytes: Option<usize>,
68 recovery_cap_lines: Option<usize>,
69}
70
71#[derive(Clone, Debug, Default, Deserialize)]
72#[serde(deny_unknown_fields)]
73struct FileArchive {
74 enabled: Option<bool>,
75 path: Option<String>,
76}
77
78#[derive(Clone, Debug, Default, Deserialize)]
79#[serde(deny_unknown_fields)]
80struct FileRules {
81 #[serde(default)]
82 packs: Option<Vec<String>>,
83}
84
85#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
86pub struct ResolvedConfig {
87 pub version: u32,
88 pub pruning: PruningConfig,
89 pub output: OutputConfig,
90 pub archive: ArchiveConfig,
91 pub rules: RulesConfig,
92}
93
94#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
95pub struct PruningConfig {
96 pub enabled: bool,
97}
98
99#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
100pub struct OutputConfig {
101 pub cap_bytes: usize,
102 pub recovery_cap_bytes: usize,
103 pub recovery_cap_lines: usize,
104}
105
106#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
107pub struct ArchiveConfig {
108 pub enabled: bool,
109 pub path: PathBuf,
110 #[serde(skip)]
111 pub is_default_path: bool,
112}
113
114#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
115pub struct RulesConfig {
116 pub packs: Vec<PathBuf>,
117}
118
119pub fn path() -> Result<PathBuf, String> {
125 if let Some(home) = std::env::var_os("XDG_CONFIG_HOME") {
126 return absolute_directory(&home, "XDG_CONFIG_HOME")
127 .map(|directory| directory.join("yarp/config.toml"));
128 }
129 let home = std::env::var_os("HOME")
130 .ok_or_else(|| "could not resolve YARP configuration: HOME is not set".to_owned())?;
131 absolute_directory(&home, "HOME").map(|directory| directory.join(".config/yarp/config.toml"))
132}
133
134pub fn load() -> Result<ResolvedConfig, String> {
140 load_from(&path()?)
141}
142
143pub fn load_from(config_path: &Path) -> Result<ResolvedConfig, String> {
149 let body = match fs::read_to_string(config_path) {
150 Ok(body) => body,
151 Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
152 return defaults();
153 }
154 Err(error) => {
155 return Err(format!("could not read {}: {error}", config_path.display()));
156 }
157 };
158 let document = body
159 .parse::<DocumentMut>()
160 .map_err(|error| format!("invalid {}: {error}", config_path.display()))?;
161 resolve_document(&document, config_path)
162}
163
164pub fn defaults() -> Result<ResolvedConfig, String> {
170 Ok(ResolvedConfig {
171 version: CONFIG_VERSION,
172 pruning: PruningConfig { enabled: true },
173 output: OutputConfig {
174 cap_bytes: DEFAULT_OUTPUT_CAP_BYTES,
175 recovery_cap_bytes: DEFAULT_RECOVERY_CAP_BYTES,
176 recovery_cap_lines: DEFAULT_RECOVERY_CAP_LINES,
177 },
178 archive: ArchiveConfig {
179 enabled: true,
180 path: default_archive_path()?,
181 is_default_path: true,
182 },
183 rules: RulesConfig { packs: Vec::new() },
184 })
185}
186
187pub fn show_toml(config: &ResolvedConfig) -> Result<String, String> {
193 toml_edit::ser::to_string_pretty(config)
194 .map_err(|error| format!("could not render resolved configuration: {error}"))
195}
196
197pub fn show_json(config: &ResolvedConfig) -> Result<String, String> {
203 serde_json::to_string_pretty(config)
204 .map(|mut text| {
205 text.push('\n');
206 text
207 })
208 .map_err(|error| format!("could not render resolved configuration: {error}"))
209}
210
211pub fn get(config: &ResolvedConfig, key: &str) -> Result<String, String> {
217 match key {
218 "version" => Ok(config.version.to_string()),
219 "pruning.enabled" => Ok(config.pruning.enabled.to_string()),
220 "output.cap_bytes" => Ok(config.output.cap_bytes.to_string()),
221 "output.recovery_cap_bytes" => Ok(config.output.recovery_cap_bytes.to_string()),
222 "output.recovery_cap_lines" => Ok(config.output.recovery_cap_lines.to_string()),
223 "archive.enabled" => Ok(config.archive.enabled.to_string()),
224 "archive.path" => path_text(&config.archive.path),
225 "rules.packs" => serde_json::to_string(&config.rules.packs)
226 .map_err(|error| format!("could not render rules.packs: {error}")),
227 _ => Err(format!("unknown configuration field {key}")),
228 }
229}
230
231pub fn init() -> Result<(), String> {
237 let config_path = path()?;
238 let parent = config_parent(&config_path)?;
239 prepare_directory(parent)?;
240 if fs::symlink_metadata(&config_path).is_ok() {
241 return Err(format!("{} already exists", config_path.display()));
242 }
243 let mut temporary = private_temporary(parent)?;
244 temporary
245 .write_all(DEFAULT_DOCUMENT.as_bytes())
246 .map_err(|error| format!("could not write temporary configuration: {error}"))?;
247 sync_temporary(&temporary)?;
248 temporary.persist_noclobber(&config_path).map_err(|error| {
249 format!(
250 "could not create {}: {}",
251 config_path.display(),
252 error.error
253 )
254 })?;
255 sync_directory(parent)?;
256 Ok(())
257}
258
259pub fn set(key: &str, arguments: &[String]) -> Result<ResolvedConfig, String> {
265 let config_path = path()?;
266 let mut document = editable_document(&config_path)?;
267 set_document_value(&mut document, key, arguments)?;
268 let resolved = resolve_document(&document, &config_path)?;
269 write_document(&config_path, &document)?;
270 Ok(resolved)
271}
272
273pub fn unset(key: &str) -> Result<ResolvedConfig, String> {
279 if key == "version" {
280 return Err("version cannot be unset".to_owned());
281 }
282 let config_path = path()?;
283 let mut document = editable_document(&config_path)?;
284 let (section, field) = split_key(key)?;
285 let remove_section = {
286 let item = document
287 .get_mut(section)
288 .ok_or_else(|| format!("configuration field {key} is not set"))?;
289 if let Some(table) = item.as_table_mut() {
290 if table.remove(field).is_none() {
291 return Err(format!("configuration field {key} is not set"));
292 }
293 table.is_empty()
294 } else if let Some(table) = item.as_inline_table_mut() {
295 if table.remove(field).is_none() {
296 return Err(format!("configuration field {key} is not set"));
297 }
298 table.is_empty()
299 } else {
300 return Err(format!("configuration field {key} is not set"));
301 }
302 };
303 if remove_section {
304 document.remove(section);
305 }
306 let resolved = resolve_document(&document, &config_path)?;
307 write_document(&config_path, &document)?;
308 Ok(resolved)
309}
310
311fn resolve_document(document: &DocumentMut, config_path: &Path) -> Result<ResolvedConfig, String> {
312 let parsed: FileConfig = toml_edit::de::from_document(document.clone())
313 .map_err(|error| format!("invalid {}: {error}", config_path.display()))?;
314 if parsed.version != CONFIG_VERSION {
315 return Err(format!(
316 "unsupported configuration version {}; expected {CONFIG_VERSION}",
317 parsed.version
318 ));
319 }
320 let cap_bytes = parsed.output.cap_bytes.unwrap_or(DEFAULT_OUTPUT_CAP_BYTES);
321 if cap_bytes != 0 && !(MIN_OUTPUT_CAP_BYTES..=MAX_OUTPUT_CAP_BYTES).contains(&cap_bytes) {
322 return Err(format!(
323 "output.cap_bytes must be 0 or {MIN_OUTPUT_CAP_BYTES} through {MAX_OUTPUT_CAP_BYTES}"
324 ));
325 }
326 let recovery_cap_bytes = parsed
327 .output
328 .recovery_cap_bytes
329 .unwrap_or(DEFAULT_RECOVERY_CAP_BYTES);
330 if !(MIN_RECOVERY_CAP_BYTES..=MAX_RECOVERY_CAP_BYTES).contains(&recovery_cap_bytes) {
331 return Err(format!(
332 "output.recovery_cap_bytes must be {MIN_RECOVERY_CAP_BYTES} through {MAX_RECOVERY_CAP_BYTES}"
333 ));
334 }
335 let recovery_cap_lines = parsed
336 .output
337 .recovery_cap_lines
338 .unwrap_or(DEFAULT_RECOVERY_CAP_LINES);
339 if !(MIN_RECOVERY_CAP_LINES..=MAX_RECOVERY_CAP_LINES).contains(&recovery_cap_lines) {
340 return Err(format!(
341 "output.recovery_cap_lines must be {MIN_RECOVERY_CAP_LINES} through {MAX_RECOVERY_CAP_LINES}"
342 ));
343 }
344 let config_directory = config_parent(config_path)?;
345 let (archive_path, is_default_path) = match parsed.archive.path {
346 Some(path) => (
347 resolve_path(config_directory, &path, "archive.path")?,
348 false,
349 ),
350 None => (default_archive_path()?, true),
351 };
352 let mut packs = Vec::new();
353 let mut pack_ids = BTreeSet::new();
354 for path in parsed.rules.packs.unwrap_or_default() {
355 let resolved = resolve_path(config_directory, &path, "rules.packs")?;
356 let pack = CompiledPack::open(&resolved, None, None)?;
357 if !pack_ids.insert(pack.id.clone()) {
358 return Err(format!("configured rule pack id {} is duplicated", pack.id));
359 }
360 packs.push(pack.path);
361 }
362 Ok(ResolvedConfig {
363 version: CONFIG_VERSION,
364 pruning: PruningConfig {
365 enabled: parsed.pruning.enabled.unwrap_or(true),
366 },
367 output: OutputConfig {
368 cap_bytes,
369 recovery_cap_bytes,
370 recovery_cap_lines,
371 },
372 archive: ArchiveConfig {
373 enabled: parsed.archive.enabled.unwrap_or(true),
374 path: archive_path,
375 is_default_path,
376 },
377 rules: RulesConfig { packs },
378 })
379}
380
381fn default_archive_path() -> Result<PathBuf, String> {
382 if let Some(home) = std::env::var_os("XDG_DATA_HOME") {
383 return absolute_directory(&home, "XDG_DATA_HOME")
384 .map(|directory| directory.join("yarp/tool-calls.sqlite3"));
385 }
386 let home = std::env::var_os("HOME")
387 .ok_or_else(|| "could not resolve YARP archive: HOME is not set".to_owned())?;
388 absolute_directory(&home, "HOME")
389 .map(|directory| directory.join(".local/share/yarp/tool-calls.sqlite3"))
390}
391
392fn absolute_directory(value: &std::ffi::OsStr, name: &str) -> Result<PathBuf, String> {
393 let path = PathBuf::from(value);
394 if !path.is_absolute() {
395 return Err(format!("{name} must be an absolute path"));
396 }
397 Ok(path)
398}
399
400fn resolve_path(base: &Path, value: &str, field: &str) -> Result<PathBuf, String> {
401 if value.is_empty() {
402 return Err(format!("{field} must not be empty"));
403 }
404 let path = PathBuf::from(value);
405 let joined = if path.is_absolute() {
406 path
407 } else {
408 base.join(path)
409 };
410 std::path::absolute(&joined)
411 .map_err(|error| format!("could not resolve {field} {}: {error}", joined.display()))
412}
413
414fn path_text(path: &Path) -> Result<String, String> {
415 path.to_str()
416 .map(ToOwned::to_owned)
417 .ok_or_else(|| format!("path is not valid UTF-8: {}", path.display()))
418}
419
420fn editable_document(config_path: &Path) -> Result<DocumentMut, String> {
421 match fs::read_to_string(config_path) {
422 Ok(body) => body
423 .parse::<DocumentMut>()
424 .map_err(|error| format!("invalid {}: {error}", config_path.display())),
425 Err(error) if error.kind() == std::io::ErrorKind::NotFound => "version = 1\n"
426 .parse::<DocumentMut>()
427 .map_err(|error| format!("could not create default configuration: {error}")),
428 Err(error) => Err(format!("could not read {}: {error}", config_path.display())),
429 }
430}
431
432fn set_document_value(
433 document: &mut DocumentMut,
434 key: &str,
435 arguments: &[String],
436) -> Result<(), String> {
437 if key == "version" {
438 require_count(key, arguments, 1)?;
439 let version = parse_integer(key, &arguments[0])?;
440 document["version"] = value(version);
441 return Ok(());
442 }
443 let (section, field) = split_key(key)?;
444 ensure_table(document, section)?;
445 let item = match key {
446 "pruning.enabled" | "archive.enabled" => {
447 require_count(key, arguments, 1)?;
448 value(parse_boolean(key, &arguments[0])?)
449 }
450 "output.cap_bytes" | "output.recovery_cap_bytes" | "output.recovery_cap_lines" => {
451 require_count(key, arguments, 1)?;
452 value(parse_integer(key, &arguments[0])?)
453 }
454 "archive.path" => {
455 require_count(key, arguments, 1)?;
456 value(arguments[0].as_str())
457 }
458 "rules.packs" => {
459 let mut array = Array::new();
460 for argument in arguments {
461 array.push(argument.as_str());
462 }
463 value(array)
464 }
465 _ => return Err(format!("unknown configuration field {key}")),
466 };
467 if let Some(table) = document[section].as_table_mut() {
468 table.insert(field, item);
469 return Ok(());
470 }
471 let item = item
472 .into_value()
473 .map_err(|_| format!("configuration value {key} is not a TOML value"))?;
474 document[section]
475 .as_inline_table_mut()
476 .ok_or_else(|| format!("configuration section {section} is not a table"))?
477 .insert(field, item);
478 Ok(())
479}
480
481fn split_key(key: &str) -> Result<(&str, &str), String> {
482 match key {
483 "pruning.enabled" => Ok(("pruning", "enabled")),
484 "output.cap_bytes" => Ok(("output", "cap_bytes")),
485 "output.recovery_cap_bytes" => Ok(("output", "recovery_cap_bytes")),
486 "output.recovery_cap_lines" => Ok(("output", "recovery_cap_lines")),
487 "archive.enabled" => Ok(("archive", "enabled")),
488 "archive.path" => Ok(("archive", "path")),
489 "rules.packs" => Ok(("rules", "packs")),
490 _ => Err(format!("unknown configuration field {key}")),
491 }
492}
493
494fn ensure_table(document: &mut DocumentMut, section: &str) -> Result<(), String> {
495 if !document.contains_key(section) {
496 document[section] = Item::Table(Table::new());
497 }
498 if !document[section].is_table() && document[section].as_inline_table().is_none() {
499 return Err(format!("configuration section {section} is not a table"));
500 }
501 Ok(())
502}
503
504fn require_count(key: &str, arguments: &[String], expected: usize) -> Result<(), String> {
505 if arguments.len() != expected {
506 return Err(format!("{key} requires {expected} value(s)"));
507 }
508 Ok(())
509}
510
511fn parse_boolean(key: &str, value: &str) -> Result<bool, String> {
512 match value {
513 "true" => Ok(true),
514 "false" => Ok(false),
515 _ => Err(format!("{key} must be true or false")),
516 }
517}
518
519fn parse_integer(key: &str, value: &str) -> Result<i64, String> {
520 if value.is_empty()
521 || (value.len() > 1 && value.starts_with('0'))
522 || !value.bytes().all(|byte| byte.is_ascii_digit())
523 {
524 return Err(format!("{key} must be a canonical unsigned integer"));
525 }
526 value
527 .parse::<i64>()
528 .map_err(|_| format!("{key} integer is too large"))
529}
530
531fn config_parent(config_path: &Path) -> Result<&Path, String> {
532 config_path.parent().ok_or_else(|| {
533 format!(
534 "configuration path has no parent: {}",
535 config_path.display()
536 )
537 })
538}
539
540fn prepare_directory(path: &Path) -> Result<(), String> {
541 fs::create_dir_all(path)
542 .map_err(|error| format!("could not create {}: {error}", path.display()))?;
543 #[cfg(unix)]
544 {
545 use std::os::unix::fs::PermissionsExt as _;
546 fs::set_permissions(path, fs::Permissions::from_mode(0o700))
547 .map_err(|error| format!("could not protect {}: {error}", path.display()))?;
548 }
549 Ok(())
550}
551
552fn private_temporary(directory: &Path) -> Result<NamedTempFile, String> {
553 let file = NamedTempFile::new_in(directory)
554 .map_err(|error| format!("could not create temporary configuration: {error}"))?;
555 #[cfg(unix)]
556 {
557 use std::os::unix::fs::PermissionsExt as _;
558 file.as_file()
559 .set_permissions(fs::Permissions::from_mode(0o600))
560 .map_err(|error| format!("could not protect temporary configuration: {error}"))?;
561 }
562 Ok(file)
563}
564
565fn sync_temporary(file: &NamedTempFile) -> Result<(), String> {
566 file.as_file()
567 .sync_all()
568 .map_err(|error| format!("could not flush temporary configuration: {error}"))
569}
570
571fn write_document(config_path: &Path, document: &DocumentMut) -> Result<(), String> {
572 let parent = config_parent(config_path)?;
573 prepare_directory(parent)?;
574 if fs::symlink_metadata(config_path).is_ok_and(|metadata| metadata.file_type().is_symlink()) {
575 return Err(format!(
576 "refusing to replace symlink {}",
577 config_path.display()
578 ));
579 }
580 let mut temporary = private_temporary(parent)?;
581 temporary
582 .write_all(document.to_string().as_bytes())
583 .map_err(|error| format!("could not write temporary configuration: {error}"))?;
584 sync_temporary(&temporary)?;
585 temporary.persist(config_path).map_err(|error| {
586 format!(
587 "could not replace {}: {}",
588 config_path.display(),
589 error.error
590 )
591 })?;
592 sync_directory(parent)
593}
594
595#[cfg(unix)]
596fn sync_directory(path: &Path) -> Result<(), String> {
597 File::open(path)
598 .and_then(|directory| directory.sync_all())
599 .map_err(|error| format!("could not flush {}: {error}", path.display()))
600}
601
602#[cfg(not(unix))]
603fn sync_directory(_path: &Path) -> Result<(), String> {
604 Ok(())
605}
606
607#[cfg(test)]
608mod tests {
609 use super::*;
610 use tempfile::TempDir;
611
612 #[test]
613 fn missing_file_uses_defaults() {
614 let directory = TempDir::new().expect("temp directory");
615 let config = load_from(&directory.path().join("config.toml")).expect("defaults");
616 assert!(config.pruning.enabled);
617 assert!(config.archive.enabled);
618 assert_eq!(config.output.cap_bytes, DEFAULT_OUTPUT_CAP_BYTES);
619 assert_eq!(config.output.recovery_cap_bytes, DEFAULT_RECOVERY_CAP_BYTES);
620 assert_eq!(config.output.recovery_cap_lines, DEFAULT_RECOVERY_CAP_LINES);
621 }
622
623 #[test]
624 fn resolves_relative_paths_and_rejects_unknown_fields() {
625 let directory = TempDir::new().expect("temp directory");
626 let path = directory.path().join("config.toml");
627 fs::write(
628 &path,
629 "version = 1\n[archive]\npath = 'data/archive.sqlite3'\n[output]\nunknown = 1\n",
630 )
631 .expect("write config");
632 assert!(load_from(&path).is_err());
633 fs::write(
634 &path,
635 "version = 1\n[archive]\npath = 'data/archive.sqlite3'\n",
636 )
637 .expect("write config");
638 let config = load_from(&path).expect("valid config");
639 assert_eq!(
640 config.archive.path,
641 directory.path().join("data/archive.sqlite3")
642 );
643 }
644
645 #[test]
646 fn rejects_invalid_versions_and_ranges() {
647 let directory = TempDir::new().expect("temp directory");
648 let path = directory.path().join("config.toml");
649 for body in [
650 "version = 2\n",
651 "version = 1\n[output]\ncap_bytes = 1\n",
652 "version = 1\n[output]\nrecovery_cap_bytes = 50000\n",
653 "version = 1\n[output]\nrecovery_cap_lines = 1901\n",
654 ] {
655 fs::write(&path, body).expect("write config");
656 assert!(load_from(&path).is_err(), "accepted {body:?}");
657 }
658 }
659}