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

yarp_cli/rules/
registry.rs
