Skip to main content

yarp_cli/rules/
registry.rs

1use std::borrow::Cow;
2use std::collections::{BTreeMap, BTreeSet};
3use std::ops::Deref;
4use std::path::{Path, PathBuf};
5
6use yarp_rule_pack::{Action, CompiledPack, Rule};
7
8use crate::config;
9
10mod generated {
11    include!(concat!(env!("OUT_DIR"), "/builtin_rules.rs"));
12}
13
14pub use generated::{BUILTIN_PACK_ID, BUILTIN_SOURCE_DIGEST};
15
16#[derive(Clone, Debug, Eq, PartialEq)]
17pub struct PackRequest {
18    pub path: PathBuf,
19    pub expected_digest: Option<[u8; 32]>,
20    pub expected_compiled_digest: Option<[u8; 32]>,
21}
22
23#[derive(Clone, Debug, Eq, PartialEq)]
24pub struct PackReference {
25    pub path: PathBuf,
26    pub source_digest: [u8; 32],
27    pub compiled_digest: [u8; 32],
28}
29
30#[derive(Clone, Debug, Eq, PartialEq)]
31pub struct SelectedRule {
32    pub pack_id: Cow<'static, str>,
33    pub source_digest: [u8; 32],
34    pub rule: SelectedRuleData,
35}
36
37#[derive(Clone, Debug, Eq, PartialEq)]
38pub enum SelectedRuleData {
39    Builtin(&'static Rule),
40    External(Box<Rule>),
41}
42
43impl Deref for SelectedRuleData {
44    type Target = Rule;
45
46    fn deref(&self) -> &Self::Target {
47        match self {
48            Self::Builtin(rule) => rule,
49            Self::External(rule) => rule,
50        }
51    }
52}
53
54#[derive(Clone, Debug, Eq, PartialEq)]
55pub enum Selection {
56    Reduce(SelectedRule),
57    Transform(SelectedRule),
58    Passthrough(Vec<String>),
59    Ambiguous(Vec<String>),
60    Unsupported,
61}
62
63#[derive(Clone, Debug, Eq, PartialEq)]
64pub struct RuleSummary {
65    pub pack_id: String,
66    pub rule: Rule,
67}
68
69struct ExternalPack {
70    pack: CompiledPack,
71    disabled: bool,
72    recheck_digest_after_selection: bool,
73}
74
75pub struct Registry {
76    external: Vec<ExternalPack>,
77    references: Vec<PackReference>,
78    diagnostics: Vec<String>,
79}
80
81impl Registry {
82    /// Load explicit compiled packs and disable deterministic ID conflicts.
83    ///
84    /// # Errors
85    ///
86    /// Returns an error when a requested pack cannot be opened or validated.
87    pub fn load(requests: &[PackRequest]) -> Result<Self, String> {
88        let mut external = Vec::<ExternalPack>::new();
89        let mut seen_paths = BTreeSet::new();
90        for request in requests {
91            let pack = CompiledPack::open(
92                &request.path,
93                request.expected_digest,
94                request.expected_compiled_digest,
95            )?;
96            if !seen_paths.insert(pack.path.clone()) {
97                if request.expected_compiled_digest.is_some()
98                    && let Some(existing) = external
99                        .iter_mut()
100                        .find(|existing| existing.pack.path == pack.path)
101                {
102                    existing.recheck_digest_after_selection = true;
103                }
104                continue;
105            }
106            external.push(ExternalPack {
107                pack,
108                disabled: false,
109                recheck_digest_after_selection: request.expected_compiled_digest.is_some(),
110            });
111        }
112        let references = external
113            .iter()
114            .map(|external| PackReference {
115                path: external.pack.path.clone(),
116                source_digest: external.pack.source_digest,
117                compiled_digest: external.pack.compiled_digest,
118            })
119            .collect();
120        let diagnostics = disable_conflicts(&mut external);
121        Ok(Self {
122            external,
123            references,
124            diagnostics,
125        })
126    }
127
128    #[must_use]
129    pub fn builtins_only() -> Self {
130        Self {
131            external: Vec::new(),
132            references: Vec::new(),
133            diagnostics: Vec::new(),
134        }
135    }
136
137    /// Select one action for an already parsed child argument vector.
138    ///
139    /// # Errors
140    ///
141    /// Returns an error when an indexed external candidate record is corrupt or invalid.
142    pub fn select(&mut self, arguments: &[String]) -> Result<Selection, String> {
143        let Some(program) = arguments.first() else {
144            return Ok(Selection::Unsupported);
145        };
146        let mut matches = MatchAccumulator::default();
147        if let Ok(index) =
148            generated::BUILTIN_PROGRAM_INDEX.binary_search_by(|entry| entry.0.cmp(program.as_str()))
149        {
150            for rule_index in generated::BUILTIN_PROGRAM_INDEX[index].1 {
151                let rule = &generated::BUILTIN_RULES[*rule_index];
152                if rule.matcher.matches(arguments) {
153                    matches.push(SelectedRule {
154                        pack_id: Cow::Borrowed(BUILTIN_PACK_ID),
155                        source_digest: BUILTIN_SOURCE_DIGEST,
156                        rule: SelectedRuleData::Builtin(rule),
157                    });
158                }
159            }
160        }
161        for external in &mut self.external {
162            if external.disabled {
163                continue;
164            }
165            let candidates = external.pack.candidate_indices(program).to_vec();
166            for candidate in candidates {
167                let rule = external.pack.read_rule(candidate)?;
168                if rule.matcher.matches(arguments) {
169                    matches.push(SelectedRule {
170                        pack_id: Cow::Owned(external.pack.id.clone()),
171                        source_digest: external.pack.source_digest,
172                        rule: SelectedRuleData::External(Box::new(rule)),
173                    });
174                }
175            }
176            if external.recheck_digest_after_selection {
177                external.pack.verify_compiled_digest()?;
178            }
179        }
180        Ok(resolve(matches))
181    }
182
183    #[must_use]
184    pub fn references(&self) -> &[PackReference] {
185        &self.references
186    }
187
188    #[must_use]
189    pub fn diagnostics(&self) -> &[String] {
190        &self.diagnostics
191    }
192
193    /// Read summaries for every enabled built-in and external rule.
194    ///
195    /// # Errors
196    ///
197    /// Returns an error when an enabled external record is corrupt or invalid.
198    pub fn summaries(&mut self) -> Result<Vec<RuleSummary>, String> {
199        let mut summaries = generated::BUILTIN_RULES
200            .iter()
201            .cloned()
202            .map(|rule| RuleSummary {
203                pack_id: BUILTIN_PACK_ID.to_owned(),
204                rule,
205            })
206            .collect::<Vec<_>>();
207        for external in &mut self.external {
208            if external.disabled {
209                continue;
210            }
211            for index in 0..external.pack.rules.len() {
212                let rule = external.pack.read_rule(
213                    u32::try_from(index)
214                        .map_err(|_| "external rule index does not fit u32".to_owned())?,
215                )?;
216                summaries.push(RuleSummary {
217                    pack_id: external.pack.id.clone(),
218                    rule,
219                });
220            }
221        }
222        summaries.sort_by(|left, right| {
223            (&left.pack_id, &left.rule.id).cmp(&(&right.pack_id, &right.rule.id))
224        });
225        Ok(summaries)
226    }
227}
228
229#[must_use]
230pub fn requests_from_paths(paths: &[PathBuf]) -> Vec<PackRequest> {
231    paths
232        .iter()
233        .map(|path| PackRequest {
234            path: path.clone(),
235            expected_digest: None,
236            expected_compiled_digest: None,
237        })
238        .collect()
239}
240
241/// Read explicit pack paths from the YARP configuration.
242///
243/// # Errors
244///
245/// Returns an error when configuration loading or pack validation fails.
246pub fn requests_from_config() -> Result<Vec<PackRequest>, String> {
247    config::load().map(|resolved| requests_from_paths(&resolved.rules.packs))
248}
249
250#[must_use]
251pub fn digest_hex(digest: &[u8; 32]) -> String {
252    const HEX: &[u8; 16] = b"0123456789abcdef";
253    let mut output = String::with_capacity(64);
254    for byte in digest {
255        output.push(char::from(HEX[usize::from(byte >> 4)]));
256        output.push(char::from(HEX[usize::from(byte & 0x0f)]));
257    }
258    output
259}
260
261/// Parse one exact SHA-256 digest.
262///
263/// # Errors
264///
265/// Returns an error when the value is not 64 hexadecimal digits.
266pub fn parse_digest(value: &str) -> Result<[u8; 32], String> {
267    if value.len() != 64 {
268        return Err("rule pack digest must contain 32 bytes".to_owned());
269    }
270    let mut output = [0_u8; 32];
271    for (index, pair) in value.as_bytes().chunks_exact(2).enumerate() {
272        let high = hex_digit(pair[0])?;
273        let low = hex_digit(pair[1])?;
274        output[index] = (high << 4) | low;
275    }
276    Ok(output)
277}
278
279fn hex_digit(value: u8) -> Result<u8, String> {
280    match value {
281        b'0'..=b'9' => Ok(value - b'0'),
282        b'a'..=b'f' => Ok(value - b'a' + 10),
283        b'A'..=b'F' => Ok(value - b'A' + 10),
284        _ => Err("rule pack digest must be hexadecimal".to_owned()),
285    }
286}
287
288/// Resolve a regular compiled pack without crossing a symlink outside a trusted project.
289///
290/// # Errors
291///
292/// Returns an error when either path cannot be resolved or any component is unsafe.
293pub fn canonical_project_pack(project_root: &Path, candidate: &Path) -> Result<PathBuf, String> {
294    let root = std::fs::canonicalize(project_root)
295        .map_err(|error| format!("could not resolve project root: {error}"))?;
296    let candidate = if candidate.is_absolute() {
297        candidate.to_path_buf()
298    } else {
299        project_root.join(candidate)
300    };
301    let relative = candidate
302        .strip_prefix(project_root)
303        .or_else(|_| candidate.strip_prefix(&root))
304        .map_err(|_| "project rule pack is not beneath the trusted project root".to_owned())?;
305    let mut current = root.clone();
306    for component in relative.components() {
307        current.push(component.as_os_str());
308        let metadata = std::fs::symlink_metadata(&current)
309            .map_err(|error| format!("could not inspect project rule pack: {error}"))?;
310        if metadata.file_type().is_symlink() {
311            return Err("project rule pack path must not contain symlinks".to_owned());
312        }
313    }
314    let metadata = std::fs::metadata(&candidate)
315        .map_err(|error| format!("could not inspect project rule pack: {error}"))?;
316    if !metadata.is_file() {
317        return Err("project rule pack must be a regular file".to_owned());
318    }
319    let resolved = std::fs::canonicalize(&candidate)
320        .map_err(|error| format!("could not resolve project rule pack: {error}"))?;
321    if !resolved.starts_with(root) {
322        return Err("project rule pack escapes the trusted project root".to_owned());
323    }
324    Ok(resolved)
325}
326
327#[derive(Default)]
328struct MatchAccumulator {
329    passthrough: Vec<String>,
330    actions: ActionMatches,
331}
332
333#[derive(Default)]
334enum ActionMatches {
335    #[default]
336    None,
337    One(SelectedRule),
338    Many(Vec<String>),
339}
340
341impl MatchAccumulator {
342    fn push(&mut self, selected: SelectedRule) {
343        if selected.rule.action == Action::Passthrough {
344            self.passthrough.push(qualified_id(&selected));
345            return;
346        }
347        self.actions = match std::mem::take(&mut self.actions) {
348            ActionMatches::None => ActionMatches::One(selected),
349            ActionMatches::One(first) => {
350                ActionMatches::Many(vec![qualified_id(&first), qualified_id(&selected)])
351            }
352            ActionMatches::Many(mut ids) => {
353                ids.push(qualified_id(&selected));
354                ActionMatches::Many(ids)
355            }
356        };
357    }
358}
359
360fn resolve(mut matches: MatchAccumulator) -> Selection {
361    if !matches.passthrough.is_empty() {
362        matches.passthrough.sort();
363        return Selection::Passthrough(matches.passthrough);
364    }
365    match matches.actions {
366        ActionMatches::None => Selection::Unsupported,
367        ActionMatches::One(selected) => match selected.rule.action {
368            Action::Reduce => Selection::Reduce(selected),
369            Action::Transform => Selection::Transform(selected),
370            Action::Passthrough => {
371                unreachable!("passthrough matches return before action resolution")
372            }
373        },
374        ActionMatches::Many(mut ids) => {
375            ids.sort();
376            Selection::Ambiguous(ids)
377        }
378    }
379}
380
381fn qualified_id(rule: &SelectedRule) -> String {
382    format!("{}/{}", rule.pack_id, rule.rule.id)
383}
384
385fn disable_conflicts(external: &mut [ExternalPack]) -> Vec<String> {
386    let builtin_ids: BTreeSet<&str> = generated::BUILTIN_RULES
387        .iter()
388        .map(|rule| rule.id.as_str())
389        .collect();
390    let mut pack_owners = BTreeMap::<String, Vec<usize>>::new();
391    let mut rule_owners = BTreeMap::<String, Vec<usize>>::new();
392    for (index, candidate) in external.iter().enumerate() {
393        pack_owners
394            .entry(candidate.pack.id.clone())
395            .or_default()
396            .push(index);
397        for rule in &candidate.pack.rules {
398            rule_owners.entry(rule.id.clone()).or_default().push(index);
399        }
400    }
401    let mut disabled = BTreeSet::new();
402    let mut diagnostics = BTreeSet::new();
403    if let Some(owners) = pack_owners.get(BUILTIN_PACK_ID) {
404        disabled.extend(owners.iter().copied());
405        diagnostics.insert(format!(
406            "external pack id conflicts with the built-in pack: {BUILTIN_PACK_ID}"
407        ));
408    }
409    for (pack_id, owners) in pack_owners.iter().filter(|(_, owners)| owners.len() > 1) {
410        disabled.extend(owners.iter().copied());
411        diagnostics.insert(format!(
412            "duplicate external pack id disables all copies: {pack_id}"
413        ));
414    }
415    for (rule_id, owners) in rule_owners {
416        if builtin_ids.contains(rule_id.as_str()) {
417            disabled.extend(owners.iter().copied());
418            diagnostics.insert(format!(
419                "external rule conflicts with a built-in rule: {rule_id}"
420            ));
421        } else if owners.len() > 1 {
422            disabled.extend(owners.iter().copied());
423            diagnostics.insert(format!(
424                "duplicate external rule id disables its packs: {rule_id}"
425            ));
426        }
427    }
428    for index in disabled {
429        external[index].disabled = true;
430    }
431    diagnostics.into_iter().collect()
432}
433
434#[cfg(test)]
435mod tests {
436    use std::io::{Seek as _, SeekFrom, Write as _};
437
438    use serde::Deserialize;
439    use sha2::{Digest as _, Sha256};
440    use tempfile::TempDir;
441    use yarp_rule_pack::SourcePack;
442
443    use super::*;
444
445    #[derive(Deserialize)]
446    #[serde(deny_unknown_fields)]
447    struct FixtureFile {
448        schema_version: u32,
449        cases: Vec<MatchFixture>,
450    }
451
452    #[derive(Deserialize)]
453    #[serde(deny_unknown_fields)]
454    struct MatchFixture {
455        id: String,
456        rule_id: String,
457        argv: Vec<String>,
458        expected_action: String,
459    }
460
461    fn strings(values: &[&str]) -> Vec<String> {
462        values.iter().map(ToString::to_string).collect()
463    }
464
465    #[test]
466    fn builtins_use_indexed_matching_and_guards() {
467        let mut registry = Registry::builtins_only();
468        let selection = registry
469            .select(&strings(&["cargo", "test", "--workspace"]))
470            .expect("selection");
471        assert!(matches!(selection, Selection::Reduce(_)));
472        let selection = registry
473            .select(&strings(&["cargo", "test", "--message-format=json"]))
474            .expect("selection");
475        assert!(matches!(selection, Selection::Passthrough(_)));
476        let selection = registry
477            .select(&strings(&["head", "-50"]))
478            .expect("selection");
479        assert!(matches!(selection, Selection::Transform(_)));
480        let selection = registry
481            .select(&strings(&["head", "--bytes=50"]))
482            .expect("selection");
483        assert!(matches!(selection, Selection::Passthrough(_)));
484        assert_eq!(
485            registry
486                .select(&strings(&["git", "push"]))
487                .expect("selection"),
488            Selection::Unsupported
489        );
490    }
491
492    #[test]
493    fn every_builtin_rule_matches_its_synthetic_fixture() {
494        let fixtures: FixtureFile =
495            serde_json::from_str(include_str!("../../rules/fixtures/builtin.json"))
496                .expect("fixtures");
497        assert_eq!(fixtures.schema_version, 1);
498        let mut covered = BTreeSet::new();
499        for fixture in fixtures.cases {
500            let mut registry = Registry::builtins_only();
501            let selection = registry.select(&fixture.argv).expect("selection");
502            match (fixture.expected_action.as_str(), selection) {
503                ("reduce", Selection::Reduce(selected))
504                | ("transform", Selection::Transform(selected)) => {
505                    assert_eq!(selected.rule.id, fixture.rule_id, "{}", fixture.id);
506                }
507                ("passthrough", Selection::Passthrough(ids)) => {
508                    assert!(
509                        ids.iter().any(|id| id.ends_with(&fixture.rule_id)),
510                        "{}: {ids:?}",
511                        fixture.id
512                    );
513                }
514                (expected, actual) => panic!("{}: expected {expected}, got {actual:?}", fixture.id),
515            }
516            covered.insert(fixture.rule_id);
517        }
518        assert_eq!(covered.len(), generated::BUILTIN_RULES.len());
519    }
520
521    #[cfg(unix)]
522    #[test]
523    fn expected_pack_digest_is_rechecked_after_selection() {
524        let directory = TempDir::new().expect("temp directory");
525        std::fs::create_dir(directory.path().join("rules")).expect("rules directory");
526        std::fs::write(
527            directory.path().join("pack.json"),
528            r#"{"schema_version":1,"id":"external-pack","rules":["rules/test.json"]}"#,
529        )
530        .expect("manifest");
531        std::fs::write(
532            directory.path().join("rules/test.json"),
533            r#"{"id":"tests/external","match":{"program":["external-tool"],"argv_prefix":["run"]},"action":"reduce","reducer":{"kind":"test_summary"},"success":{"max_line_bytes":16384,"max_output_bytes":32768,"min_savings_bytes":120,"min_savings_basis_points":1000},"failure":{"max_line_bytes":16384,"max_output_bytes":65536,"min_savings_bytes":120,"min_savings_basis_points":500}}"#,
534        )
535        .expect("rule");
536        let source = SourcePack::load(directory.path()).expect("source pack");
537        let compiled = yarp_rule_pack::compile(&source).expect("compiled pack");
538        let compiled_digest = Sha256::digest(&compiled).into();
539        let pack_path = directory.path().join("pack.yrp");
540        std::fs::write(&pack_path, &compiled).expect("write pack");
541        let mut registry = Registry::load(&[PackRequest {
542            path: pack_path.clone(),
543            expected_digest: Some(source.source_digest),
544            expected_compiled_digest: Some(compiled_digest),
545        }])
546        .expect("registry");
547
548        let mut file = std::fs::OpenOptions::new()
549            .write(true)
550            .open(pack_path)
551            .expect("open pack for mutation");
552        file.seek(SeekFrom::Start(0)).expect("seek");
553        file.write_all(&[compiled[0] ^ 1]).expect("mutate header");
554        file.flush().expect("flush mutation");
555
556        assert!(
557            registry
558                .select(&strings(&["external-tool", "run"]))
559                .expect_err("changed pack")
560                .contains("changed while loading")
561        );
562    }
563
564    #[test]
565    fn project_pack_rejects_symlinked_path_components() {
566        use std::os::unix::fs::symlink;
567
568        let directory = tempfile::TempDir::new().expect("temp directory");
569        let real = directory.path().join("real");
570        std::fs::create_dir(&real).expect("real directory");
571        std::fs::write(real.join("rules.yrp"), b"pack").expect("pack");
572        symlink(&real, directory.path().join(".yarp")).expect("symlink");
573        let error =
574            canonical_project_pack(directory.path(), &directory.path().join(".yarp/rules.yrp"))
575                .expect_err("symlink rejection");
576        assert!(error.contains("symlink"));
577    }
578
579    #[test]
580    fn digest_encoding_is_strict() {
581        let digest = [7_u8; 32];
582        assert_eq!(parse_digest(&digest_hex(&digest)).expect("digest"), digest);
583        assert!(parse_digest("07").is_err());
584        assert!(parse_digest("not-hex").is_err());
585    }
586}