1#[cfg(test)]
2use crate::rule_pack as yarp_rule_pack;
3use crate::rule_pack::{OutputPolicy, Rule, Transform};
4use crate::rules::{PackRequest, Registry, Selection, digest_hex};
5use crate::shell::{self, Connector, ShellItem, SimpleCommand};
6use serde::Serialize;
7
8#[derive(Clone, Copy, Debug)]
10pub struct ArchiveCommandRef<'a> {
11 pub agent: &'a str,
12 pub account: &'a str,
13 pub session_id: &'a str,
14 pub call_id: &'a str,
15}
16
17#[must_use]
19pub fn rewrite(command: &str) -> Option<String> {
20 rewrite_with_archive(command, None)
21}
22
23#[must_use]
25pub fn rewrite_with_archive(
26 command: &str,
27 archive: Option<ArchiveCommandRef<'_>>,
28) -> Option<String> {
29 rewrite_with_options(command, archive, &[]).ok().flatten()
30}
31
32#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
33pub struct ShellPlan {
34 pub version: u32,
35 pub execution: ExecutionPlan,
36 pub result: ResultPolicy,
37}
38
39#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
40#[serde(tag = "kind", rename_all = "snake_case")]
41pub enum ExecutionPlan {
42 Original,
43 Rewrite { command: String },
44}
45
46#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
47#[serde(tag = "kind", rename_all = "snake_case")]
48pub enum ResultPolicy {
49 Ordinary,
50 Recovery,
51}
52
53pub fn plan_with_options(
59 command: &str,
60 archive: Option<ArchiveCommandRef<'_>>,
61 packs: &[PackRequest],
62) -> Result<ShellPlan, String> {
63 let command = command.trim();
64 let Some(words) = shell::parse_simple_words(command).ok() else {
65 return Ok(original_plan(ResultPolicy::Ordinary));
66 };
67 if is_recovery_command(&words) {
68 return Ok(original_plan(ResultPolicy::Recovery));
69 }
70 let mut registry = Registry::load(packs)?;
71 let Selection::Reduce(selected) = registry.select(&words)? else {
72 return Ok(original_plan(ResultPolicy::Ordinary));
73 };
74 let wrapper = selected_wrapper(command, archive, ®istry, &selected)?;
75 Ok(ShellPlan {
76 version: 1,
77 execution: ExecutionPlan::Rewrite { command: wrapper },
78 result: ResultPolicy::Ordinary,
79 })
80}
81
82pub fn rewrite_with_options(
89 command: &str,
90 archive: Option<ArchiveCommandRef<'_>>,
91 packs: &[PackRequest],
92) -> Result<Option<String>, String> {
93 plan_with_options(command, archive, packs).map(|plan| match plan.execution {
94 ExecutionPlan::Original => None,
95 ExecutionPlan::Rewrite { command } => Some(command),
96 })
97}
98
99const fn original_plan(result: ResultPolicy) -> ShellPlan {
100 ShellPlan {
101 version: 1,
102 execution: ExecutionPlan::Original,
103 result,
104 }
105}
106
107fn is_recovery_command(words: &[String]) -> bool {
108 matches!(
109 words,
110 [program, subcommand, arguments @ ..]
111 if program == "yarp"
112 && matches!(subcommand.as_str(), "search" | "read")
113 && !arguments.iter().any(|argument| matches!(argument.as_str(), "--help" | "-h"))
114 )
115}
116
117fn selected_wrapper(
118 command: &str,
119 archive: Option<ArchiveCommandRef<'_>>,
120 registry: &Registry,
121 selected: &crate::rules::SelectedRule,
122) -> Result<String, String> {
123 let mut wrapper = format!(
124 "yarp run --selected-pack {} --selected-rule {} --selected-digest {}",
125 shell_quote(&selected.pack_id),
126 shell_quote(&selected.rule.id),
127 shell_quote(&digest_hex(&selected.source_digest)),
128 );
129 for reference in registry.references() {
130 let path = reference
131 .path
132 .to_str()
133 .ok_or_else(|| "rule pack path is not valid UTF-8".to_owned())?;
134 wrapper.push_str(" --rule-pack ");
135 wrapper.push_str(&shell_quote(path));
136 wrapper.push_str(" --rule-pack-digest ");
137 wrapper.push_str(&shell_quote(&digest_hex(&reference.source_digest)));
138 wrapper.push_str(" --rule-pack-compiled-digest ");
139 wrapper.push_str(&shell_quote(&digest_hex(&reference.compiled_digest)));
140 }
141 if let Some(reference) = archive {
142 wrapper.push_str(" --archive-agent ");
143 wrapper.push_str(&shell_quote(reference.agent));
144 wrapper.push_str(" --archive-account ");
145 wrapper.push_str(&shell_quote(reference.account));
146 wrapper.push_str(" --archive-session ");
147 wrapper.push_str(&shell_quote(reference.session_id));
148 wrapper.push_str(" --archive-call ");
149 wrapper.push_str(&shell_quote(reference.call_id));
150 }
151 wrapper.push_str(" -- ");
152 wrapper.push_str(command);
153 Ok(wrapper)
154}
155
156fn shell_quote(value: &str) -> String {
157 format!("'{}'", value.replace('\'', "'\\''"))
158}
159
160pub fn select_builtin_command(command: &str) -> Result<(Vec<String>, Selection), String> {
166 let words = shell::parse_simple_words(command)?;
167 let selection = Registry::builtins_only().select(&words)?;
168 Ok((words, selection))
169}
170
171pub fn select_builtin_argv(arguments: &[String]) -> Result<Selection, String> {
177 Registry::builtins_only().select(arguments)
178}
179
180#[derive(Clone, Copy, Debug, Eq, PartialEq)]
181pub enum StatusConfidence {
182 Complete,
183 FinalStageOnly,
184 Conditional,
185}
186
187const TRANSFORM_DIAGNOSTIC_PREFIXES: [(&str, &[u8]); 6] = [
188 ("cat", b"cat:"),
189 ("tee", b"tee:"),
190 ("head", b"head:"),
191 ("tail", b"tail:"),
192 ("sort", b"sort:"),
193 ("uniq", b"uniq:"),
194];
195
196#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
197pub struct TransformDiagnostics(u8);
198
199impl TransformDiagnostics {
200 fn insert(&mut self, program: &str) {
201 if let Some(index) = TRANSFORM_DIAGNOSTIC_PREFIXES
202 .iter()
203 .position(|(candidate, _)| *candidate == program)
204 {
205 self.0 |= 1_u8 << index;
206 }
207 }
208
209 fn merge(&mut self, other: Self) {
210 self.0 |= other.0;
211 }
212
213 #[must_use]
214 pub fn matches_line(self, line: &[u8]) -> bool {
215 TRANSFORM_DIAGNOSTIC_PREFIXES
216 .iter()
217 .enumerate()
218 .any(|(index, (_, prefix))| self.0 & (1_u8 << index) != 0 && line.starts_with(prefix))
219 }
220}
221
222#[derive(Clone, Debug, Eq, PartialEq)]
223pub struct ResultPlan {
224 pub rule: Rule,
225 pub status_confidence: StatusConfidence,
226 pub transform_diagnostics: TransformDiagnostics,
227 pub fail_open_setup_diagnostics: bool,
228}
229
230pub fn select_result_plan(command: &str) -> Result<ResultPlan, String> {
237 let program = shell::parse(command)?;
238 let mut selected: Option<Rule> = None;
239 let mut confidence = StatusConfidence::Complete;
240 let mut transform_diagnostics = TransformDiagnostics::default();
241 let mut fail_open_setup_diagnostics = false;
242 let mut pipefail = Some(false);
243 let mut previous_had_output = false;
244
245 for (index, item) in program.items.iter().enumerate() {
246 let connector = index
247 .checked_sub(1)
248 .and_then(|connector| program.connectors.get(connector))
249 .copied();
250 if let ShellItem::Simple(command) = item
251 && let Some(value) = pipefail_setting(&command.words)
252 {
253 pipefail = if matches!(connector, None | Some(Connector::Sequence)) {
254 Some(value)
255 } else {
256 None
257 };
258 }
259 let (candidate, setup) = match item {
260 ShellItem::Simple(command) => select_simple(command)?,
261 ShellItem::Pipeline(stages) => {
262 if pipefail != Some(true) {
263 confidence = merge_confidence(confidence, StatusConfidence::FinalStageOnly);
264 }
265 let (rule, diagnostics) = select_pipeline(stages)?;
266 transform_diagnostics.merge(diagnostics);
267 (Some(rule), false)
268 }
269 };
270 let has_output = candidate.is_some();
271 let may_have_setup_diagnostics =
272 setup && matches!(item, ShellItem::Simple(command) if setup_may_emit(&command.words));
273 fail_open_setup_diagnostics |= may_have_setup_diagnostics;
274 let may_have_visible_output = has_output || may_have_setup_diagnostics;
275 match connector {
276 Some(Connector::Sequence) if previous_had_output => {
277 confidence = merge_confidence(confidence, StatusConfidence::FinalStageOnly);
278 }
279 Some(Connector::Or) if previous_had_output || may_have_visible_output => {
280 confidence = merge_confidence(confidence, StatusConfidence::Conditional);
281 }
282 None if index > 0 => {
283 return Err("compound command is missing a connector".to_owned());
284 }
285 Some(Connector::And | Connector::Sequence | Connector::Or) | None => {}
286 }
287 if let Some(candidate) = candidate {
288 merge_rule(&mut selected, candidate)?;
289 }
290 previous_had_output = previous_had_output || may_have_visible_output;
291 }
292
293 let rule =
294 selected.ok_or_else(|| "compound command has no supported output command".to_owned())?;
295 Ok(ResultPlan {
296 rule,
297 status_confidence: confidence,
298 transform_diagnostics,
299 fail_open_setup_diagnostics,
300 })
301}
302
303pub fn select_result_rule(command: &str) -> Result<Rule, String> {
309 select_result_plan(command).map(|plan| plan.rule)
310}
311
312fn select_simple(command: &SimpleCommand) -> Result<(Option<Rule>, bool), String> {
313 if is_setup_command(&command.words) || command.words.iter().all(|word| is_assignment(word)) {
314 return Ok((None, true));
315 }
316 let Some(words) = normalize_result_words(command.words.clone())? else {
317 return Ok((None, true));
318 };
319 let mut registry = Registry::builtins_only();
320 match registry
321 .select(&words)
322 .map_err(|error| format!("could not classify compound command: {error}"))?
323 {
324 Selection::Reduce(candidate) => Ok((Some((*candidate.rule).clone()), false)),
325 Selection::Transform(_) => Err("line transform has no pipeline input".to_owned()),
326 Selection::Passthrough(_) | Selection::Ambiguous(_) | Selection::Unsupported => {
327 Err("compound command contains an unsupported or guarded command".to_owned())
328 }
329 }
330}
331
332fn select_pipeline(stages: &[SimpleCommand]) -> Result<(Rule, TransformDiagnostics), String> {
333 let mut selected = None;
334 let mut diagnostics = TransformDiagnostics::default();
335 for stage in stages {
336 if is_setup_command(&stage.words) || stage.words.iter().all(|word| is_assignment(word)) {
337 return Err("pipeline contains a setup command".to_owned());
338 }
339 let words = normalize_result_words(stage.words.clone())?
340 .ok_or_else(|| "pipeline stage has no output command".to_owned())?;
341 let mut registry = Registry::builtins_only();
342 match registry
343 .select(&words)
344 .map_err(|error| format!("could not classify pipeline stage: {error}"))?
345 {
346 Selection::Reduce(candidate) => merge_rule(&mut selected, (*candidate.rule).clone())?,
347 Selection::Transform(candidate) => {
348 if candidate.rule.transform != Some(Transform::LinePreserving) {
349 return Err("pipeline transform is unsupported".to_owned());
350 }
351 validate_line_preserving(&words)?;
352 if selected.is_none() {
353 return Err("pipeline transform has no typed input".to_owned());
354 }
355 diagnostics.insert(&words[0]);
356 }
357 Selection::Passthrough(_) | Selection::Ambiguous(_) | Selection::Unsupported => {
358 return Err("pipeline contains an unsupported or guarded command".to_owned());
359 }
360 }
361 }
362 selected
363 .map(|rule| (rule, diagnostics))
364 .ok_or_else(|| "pipeline has no supported output command".to_owned())
365}
366
367fn validate_line_preserving(words: &[String]) -> Result<(), String> {
368 let Some((program, arguments)) = words.split_first() else {
369 return Err("pipeline transform has no program".to_owned());
370 };
371 match program.as_str() {
372 "cat" => validate_stdin_only(arguments),
373 "tee" => validate_tee(arguments),
374 "head" | "tail" => validate_line_selector(program, arguments),
375 "sort" => validate_sort(arguments),
376 "uniq" => validate_uniq(arguments),
377 _ => Err("pipeline transform program is unsupported".to_owned()),
378 }
379}
380
381fn validate_stdin_only(arguments: &[String]) -> Result<(), String> {
382 let mut operands = false;
383 for argument in arguments {
384 if argument == "--" && !operands {
385 operands = true;
386 } else if argument != "-" {
387 return Err("pipeline transform must read only standard input".to_owned());
388 }
389 }
390 Ok(())
391}
392
393fn validate_tee(arguments: &[String]) -> Result<(), String> {
394 let mut operands = false;
395 for argument in arguments {
396 if operands {
397 continue;
398 }
399 match argument.as_str() {
400 "--" => operands = true,
401 "-a" | "-i" | "-p" | "--append" | "--ignore-interrupts" | "--output-error" => {}
402 value if value.starts_with("--output-error=") => {}
403 value if !value.starts_with('-') || value == "-" => operands = true,
404 _ => return Err("tee transform uses an unsupported option".to_owned()),
405 }
406 }
407 Ok(())
408}
409
410fn validate_line_selector(program: &str, arguments: &[String]) -> Result<(), String> {
411 let mut index = 0;
412 let mut operands = false;
413 while index < arguments.len() {
414 let argument = &arguments[index];
415 if operands {
416 if argument != "-" {
417 return Err(format!("{program} transform must read only standard input"));
418 }
419 index += 1;
420 continue;
421 }
422 if argument == "--" {
423 operands = true;
424 index += 1;
425 continue;
426 }
427 if argument == "-" {
428 operands = true;
429 index += 1;
430 continue;
431 }
432 if matches!(argument.as_str(), "-q" | "--quiet" | "--silent") {
433 index += 1;
434 continue;
435 }
436 if argument == "-n" || argument == "--lines" {
437 let count = arguments
438 .get(index + 1)
439 .ok_or_else(|| format!("{program} line option is missing its value"))?;
440 validate_line_count(count, program)?;
441 index += 2;
442 continue;
443 }
444 if let Some(count) = argument.strip_prefix("--lines=") {
445 validate_line_count(count, program)?;
446 index += 1;
447 continue;
448 }
449 if let Some(count) = argument.strip_prefix("-n")
450 && !count.is_empty()
451 {
452 validate_line_count(count, program)?;
453 index += 1;
454 continue;
455 }
456 if argument.strip_prefix('-').is_some_and(|count| {
457 !count.is_empty() && count.bytes().all(|byte| byte.is_ascii_digit())
458 }) {
459 index += 1;
460 continue;
461 }
462 return Err(format!(
463 "{program} transform uses an unsupported option or operand"
464 ));
465 }
466 Ok(())
467}
468
469fn validate_line_count(value: &str, program: &str) -> Result<(), String> {
470 let digits = value.strip_prefix(['+', '-']).unwrap_or(value).as_bytes();
471 if digits.is_empty() || !digits.iter().all(u8::is_ascii_digit) {
472 return Err(format!("{program} line count is not a literal integer"));
473 }
474 Ok(())
475}
476
477fn validate_sort(arguments: &[String]) -> Result<(), String> {
478 let mut index = 0;
479 let mut operands = false;
480 while index < arguments.len() {
481 let argument = &arguments[index];
482 if operands || argument == "-" || !argument.starts_with('-') {
483 if argument != "-" {
484 return Err("sort transform must read only standard input".to_owned());
485 }
486 operands = true;
487 index += 1;
488 continue;
489 }
490 if argument == "--" {
491 operands = true;
492 index += 1;
493 continue;
494 }
495 if sort_option_takes_value(argument) {
496 if sort_option_has_attached_value(argument) {
497 index += 1;
498 } else {
499 index = index
500 .checked_add(2)
501 .filter(|next| *next <= arguments.len())
502 .ok_or_else(|| "sort option is missing its value".to_owned())?;
503 }
504 continue;
505 }
506 if is_safe_sort_flag(argument) {
507 index += 1;
508 continue;
509 }
510 return Err("sort transform uses an unsupported option".to_owned());
511 }
512 Ok(())
513}
514
515fn sort_option_takes_value(value: &str) -> bool {
516 [
517 "-k",
518 "-t",
519 "-S",
520 "-T",
521 "--key",
522 "--field-separator",
523 "--buffer-size",
524 "--temporary-directory",
525 "--batch-size",
526 "--random-source",
527 ]
528 .iter()
529 .any(|option| value == *option || value.starts_with(&format!("{option}=")))
530 || ["-k", "-t", "-S", "-T"]
531 .iter()
532 .any(|option| value.starts_with(option) && value.len() > option.len())
533}
534
535fn sort_option_has_attached_value(value: &str) -> bool {
536 value.contains('=')
537 || ["-k", "-t", "-S", "-T"]
538 .iter()
539 .any(|option| value.starts_with(option) && value.len() > option.len())
540}
541
542fn is_safe_sort_flag(value: &str) -> bool {
543 if value.starts_with("--") {
544 return matches!(
545 value,
546 "--stable"
547 | "--unique"
548 | "--reverse"
549 | "--numeric-sort"
550 | "--general-numeric-sort"
551 | "--human-numeric-sort"
552 | "--version-sort"
553 | "--month-sort"
554 | "--random-sort"
555 | "--ignore-case"
556 | "--ignore-leading-blanks"
557 | "--dictionary-order"
558 | "--ignore-nonprinting"
559 );
560 }
561 value.len() > 1
562 && value[1..].bytes().all(|byte| {
563 matches!(
564 byte,
565 b'b' | b'd'
566 | b'f'
567 | b'g'
568 | b'h'
569 | b'i'
570 | b'M'
571 | b'n'
572 | b'R'
573 | b'r'
574 | b's'
575 | b'u'
576 | b'V'
577 )
578 })
579}
580
581fn validate_uniq(arguments: &[String]) -> Result<(), String> {
582 let mut index = 0;
583 let mut operands = false;
584 let mut operand_count = 0_usize;
585 while index < arguments.len() {
586 let argument = &arguments[index];
587 if operands || argument == "-" || !argument.starts_with('-') {
588 if argument != "-" || operand_count > 0 {
589 return Err("uniq transform must read only standard input".to_owned());
590 }
591 operands = true;
592 operand_count += 1;
593 index += 1;
594 continue;
595 }
596 if argument == "--" {
597 operands = true;
598 index += 1;
599 continue;
600 }
601 if uniq_option_takes_value(argument) {
602 if uniq_option_has_attached_value(argument) {
603 index += 1;
604 } else {
605 index = index
606 .checked_add(2)
607 .filter(|next| *next <= arguments.len())
608 .ok_or_else(|| "uniq option is missing its value".to_owned())?;
609 }
610 continue;
611 }
612 if is_safe_uniq_flag(argument) {
613 index += 1;
614 continue;
615 }
616 return Err("uniq transform uses an unsupported option".to_owned());
617 }
618 Ok(())
619}
620
621fn uniq_option_takes_value(value: &str) -> bool {
622 [
623 "-f",
624 "-s",
625 "-w",
626 "--skip-fields",
627 "--skip-chars",
628 "--check-chars",
629 ]
630 .iter()
631 .any(|option| value == *option || value.starts_with(&format!("{option}=")))
632 || ["-f", "-s", "-w"]
633 .iter()
634 .any(|option| value.starts_with(option) && value.len() > option.len())
635}
636
637fn uniq_option_has_attached_value(value: &str) -> bool {
638 value.contains('=')
639 || ["-f", "-s", "-w"]
640 .iter()
641 .any(|option| value.starts_with(option) && value.len() > option.len())
642}
643
644fn is_safe_uniq_flag(value: &str) -> bool {
645 if value.starts_with("--") {
646 return matches!(
647 value,
648 "--all-repeated" | "--repeated" | "--unique" | "--ignore-case"
649 );
650 }
651 value.len() > 1
652 && value[1..]
653 .bytes()
654 .all(|byte| matches!(byte, b'D' | b'd' | b'u' | b'i'))
655}
656
657fn merge_rule(selected: &mut Option<Rule>, candidate: Rule) -> Result<(), String> {
658 let Some(existing) = selected else {
659 *selected = Some(candidate);
660 return Ok(());
661 };
662 if existing.reducer != candidate.reducer {
663 return Err("compound command mixes reducer families".to_owned());
664 }
665 let existing_success = existing
666 .success
667 .ok_or_else(|| "selected rule has no success policy".to_owned())?;
668 let candidate_success = candidate
669 .success
670 .ok_or_else(|| "candidate rule has no success policy".to_owned())?;
671 let existing_failure = existing
672 .failure
673 .ok_or_else(|| "selected rule has no failure policy".to_owned())?;
674 let candidate_failure = candidate
675 .failure
676 .ok_or_else(|| "candidate rule has no failure policy".to_owned())?;
677 existing.success = Some(merge_policy(existing_success, candidate_success));
678 existing.failure = Some(merge_policy(existing_failure, candidate_failure));
679 Ok(())
680}
681
682const fn merge_confidence(left: StatusConfidence, right: StatusConfidence) -> StatusConfidence {
683 match (left, right) {
684 (StatusConfidence::Conditional, _) | (_, StatusConfidence::Conditional) => {
685 StatusConfidence::Conditional
686 }
687 (StatusConfidence::FinalStageOnly, _) | (_, StatusConfidence::FinalStageOnly) => {
688 StatusConfidence::FinalStageOnly
689 }
690 (StatusConfidence::Complete, StatusConfidence::Complete) => StatusConfidence::Complete,
691 }
692}
693
694fn pipefail_setting(words: &[String]) -> Option<bool> {
695 match words {
696 [program, option, value] if program == "set" && value == "pipefail" && option == "-o" => {
697 Some(true)
698 }
699 [program, option, value] if program == "set" && value == "pipefail" && option == "+o" => {
700 Some(false)
701 }
702 _ => None,
703 }
704}
705
706fn merge_policy(left: OutputPolicy, right: OutputPolicy) -> OutputPolicy {
707 OutputPolicy {
708 max_line_bytes: left.max_line_bytes.max(right.max_line_bytes),
709 max_output_bytes: left.max_output_bytes.max(right.max_output_bytes),
710 min_savings_bytes: left.min_savings_bytes.max(right.min_savings_bytes),
711 min_savings_basis_points: left
712 .min_savings_basis_points
713 .max(right.min_savings_basis_points),
714 }
715}
716
717fn normalize_result_words(mut words: Vec<String>) -> Result<Option<Vec<String>>, String> {
718 if is_setup_command(&words) {
719 return Ok(None);
720 }
721 while words.first().is_some_and(|word| is_assignment(word)) {
722 words.remove(0);
723 }
724 let Some(program) = words.first().map(String::as_str) else {
725 return Ok(None);
726 };
727 match program {
728 "env" => {
729 words.remove(0);
730 if words.first().is_some_and(|word| word == "--") {
731 words.remove(0);
732 }
733 while words.first().is_some_and(|word| is_assignment(word)) {
734 words.remove(0);
735 }
736 }
737 "command" | "exec" => {
738 words.remove(0);
739 if words.first().is_some_and(|word| word == "--") {
740 words.remove(0);
741 }
742 }
743 "time" if words.get(1).is_some_and(|word| !word.starts_with('-')) => {
744 words.remove(0);
745 }
746 "timeout"
747 if words.len() >= 3 && words.get(1).is_some_and(|word| !word.starts_with('-')) =>
748 {
749 words.drain(0..2);
750 }
751 _ => {}
752 }
753 if words.is_empty() {
754 return Ok(None);
755 }
756 if words.first().is_some_and(|word| {
757 matches!(
758 word.as_str(),
759 "env" | "command" | "exec" | "time" | "timeout"
760 )
761 }) {
762 return Err("compound command uses an unsupported command wrapper".to_owned());
763 }
764 Ok(Some(words))
765}
766
767fn is_setup_command(words: &[String]) -> bool {
768 match words {
769 [program] if matches!(program.as_str(), ":" | "true" | "false") => true,
770 [program, path] if program == "cd" && !path.starts_with('-') => true,
771 [program, flag] if program == "set" && is_quiet_set_flag(flag) => true,
772 [program, flag, value]
773 if program == "set" && matches!(flag.as_str(), "-o" | "+o") && value == "pipefail" =>
774 {
775 true
776 }
777 [program, values @ ..] if program == "export" && !values.is_empty() => values
778 .iter()
779 .all(|value| is_assignment(value) || is_identifier(value)),
780 [program, value] if program == "umask" && !value.starts_with('-') => true,
781 _ => false,
782 }
783}
784
785#[must_use]
787pub fn contains_setup_diagnostic(input: &[u8]) -> bool {
788 [
789 b"cd:".as_slice(),
790 b"umask:".as_slice(),
791 b"export:".as_slice(),
792 b"readonly variable".as_slice(),
793 b"not a valid identifier".as_slice(),
794 ]
795 .iter()
796 .any(|pattern| {
797 input.windows(pattern.len()).any(|window| {
798 window
799 .iter()
800 .zip(*pattern)
801 .all(|(left, right)| left.eq_ignore_ascii_case(right))
802 })
803 })
804}
805
806fn setup_may_emit(words: &[String]) -> bool {
807 words.iter().all(|word| is_assignment(word))
808 || words
809 .first()
810 .is_some_and(|program| matches!(program.as_str(), "cd" | "export" | "umask"))
811}
812
813fn is_quiet_set_flag(value: &str) -> bool {
814 matches!(
815 value,
816 "-e" | "+e"
817 | "-u"
818 | "+u"
819 | "-f"
820 | "+f"
821 | "-C"
822 | "+C"
823 | "-m"
824 | "+m"
825 | "-b"
826 | "+b"
827 | "-n"
828 | "+n"
829 | "-E"
830 | "+E"
831 | "-T"
832 | "+T"
833 | "-P"
834 | "+P"
835 | "-h"
836 | "+h"
837 | "-k"
838 | "+k"
839 | "-p"
840 | "+p"
841 )
842}
843
844fn is_identifier(value: &str) -> bool {
845 let mut characters = value.chars();
846 characters
847 .next()
848 .is_some_and(|character| character == '_' || character.is_ascii_alphabetic())
849 && characters.all(|character| character == '_' || character.is_ascii_alphanumeric())
850}
851
852fn is_assignment(word: &str) -> bool {
853 word.split_once('=')
854 .is_some_and(|(name, _)| is_identifier(name))
855}
856
857#[cfg(test)]
858mod tests {
859 use super::*;
860
861 #[test]
862 fn rewrites_selected_commands_without_changing_the_original_text() {
863 let expected = format!(
864 "yarp run --selected-pack 'yarp-builtins' --selected-rule 'git/diff' --selected-digest '{}' -- git diff -- 'file name'",
865 digest_hex(&crate::rules::BUILTIN_SOURCE_DIGEST)
866 );
867 assert_eq!(
868 rewrite(" git diff -- 'file name' ").as_deref(),
869 Some(expected.as_str())
870 );
871 assert!(
872 rewrite("npm run lint -- --fix")
873 .is_some_and(|value| value.ends_with("-- npm run lint -- --fix"))
874 );
875 assert!(rewrite("pytest -q").is_some());
876 assert!(rewrite("rg -n needle .").is_some());
877 assert!(rewrite("egrep -n needle file").is_some());
878 assert!(rewrite("fgrep needle file").is_some());
879 assert!(rewrite("podman ps").is_some());
880 assert!(rewrite("git reflog --date=iso").is_some());
881 assert!(rewrite("git remote show origin").is_some());
882 assert!(rewrite("gh repo view osolmaz/yarp").is_some());
883 assert!(rewrite("pnpm -s build").is_some());
884 assert!(rewrite("git -C /repo diff --check").is_some());
885 assert!(rewrite("pnpm -C frontend test").is_some());
886 assert!(rewrite("pnpm --filter web -r lint").is_some());
887 assert!(rewrite("ls -la").is_some());
888 assert!(rewrite("find src -name '*.rs'").is_some());
889 assert!(rewrite("ps aux").is_some());
890 }
891
892 #[test]
893 fn adds_archive_identifiers_with_shell_quoting() {
894 let rewritten = rewrite_with_archive(
895 "git status --short",
896 Some(ArchiveCommandRef {
897 agent: "pi",
898 account: "o'nur",
899 session_id: "session-1",
900 call_id: "call-1",
901 }),
902 )
903 .expect("rewrite");
904 assert_eq!(
905 rewritten,
906 format!(
907 "yarp run --selected-pack 'yarp-builtins' --selected-rule 'git/status' --selected-digest '{}' --archive-agent 'pi' --archive-account 'o'\\''nur' --archive-session 'session-1' --archive-call 'call-1' -- git status --short",
908 digest_hex(&crate::rules::BUILTIN_SOURCE_DIGEST)
909 )
910 );
911 }
912
913 #[test]
914 fn classifies_only_direct_archive_queries_as_recovery() {
915 for command in [
916 "yarp search yr_0123456789abcdef0123456789abcdef error",
917 "yarp read yr_0123456789abcdef0123456789abcdef stdout 1:20",
918 "yarp search yr_0123456789abcdef0123456789abcdef 'error|warning' -C 2",
919 ] {
920 let plan = plan_with_options(command, None, &[]).expect("recovery plan");
921 assert_eq!(plan.execution, ExecutionPlan::Original);
922 assert_eq!(plan.result, ResultPolicy::Recovery);
923 }
924 for command in [
925 "command yarp search ref error",
926 "env DEBUG=1 yarp search ref error",
927 "/usr/bin/yarp read ref 1:20",
928 "yarp search ref error | head",
929 "yarp read ref 1:20 > output",
930 "yarp search ref error; printf extra",
931 "yarp search $(printf ref) error",
932 "yarp search ref error &",
933 "yarp search --help",
934 ] {
935 let plan = plan_with_options(command, None, &[]).expect("ordinary plan");
936 assert_eq!(
937 plan.result,
938 ResultPolicy::Ordinary,
939 "classified {command:?}"
940 );
941 }
942 }
943
944 #[test]
945 fn leaves_unknown_and_guarded_commands_unchanged() {
946 for command in [
947 "",
948 "cat .env",
949 "curl https://example.com",
950 "git push",
951 "npm install",
952 "git show HEAD:file",
953 "rg --json needle .",
954 "rg --files",
955 "rg -l needle .",
956 "rg -0l needle .",
957 "grep -l needle file",
958 "grep -lZ needle file",
959 "egrep -l needle file",
960 "fgrep --count needle file",
961 "git status -sbz",
962 "git grep -l needle",
963 "git grep -lz needle",
964 "git log --oneline -z",
965 "ls --zero",
966 "find . -print0",
967 "find . -exec echo {} +",
968 "fd -0 needle .",
969 "fd -x echo {}",
970 "tree -J .",
971 "du --null .",
972 "df --output=source,target",
973 "lsof -F pcfn",
974 "free --seconds 1",
975 "ps aux -o pid=",
976 "cmake -S . -B build --trace",
977 "git worktree list --porcelain",
978 "git tag --list -z",
979 "git reflog delete main@{0}",
980 "git reflog --format=%H",
981 "gh workflow view ci.yml --yaml",
982 "docker ps --format json",
983 "docker images --quiet",
984 "podman ps --format json",
985 "podman logs --follow container",
986 "journalctl --output=json",
987 "systemctl status -o json sshd",
988 "ninja -t compdb",
989 "tsc --showConfig",
990 "eslint --format json .",
991 "npx eslint --format=json .",
992 "pyright --outputjson",
993 "ruff check --output-format json .",
994 "uv run ruff --show-settings",
995 "mypy --output=json .",
996 "pytest --json-report",
997 "python3 -m pytest --collect-only",
998 "go vet -json ./...",
999 "go test -json ./...",
1000 "cargo test --message-format=json",
1001 "cargo audit --json",
1002 "pnpm exec oxlint --format json",
1003 "gh pr view 1 --json=number,title",
1004 "git diff --stat",
1005 "kubectl get pods -ojson",
1006 "npm test -- --reporter=json",
1007 "codex review --json --base main",
1008 "yarp run -- git status",
1009 ] {
1010 assert_eq!(rewrite(command), None, "accepted {command:?}");
1011 }
1012 }
1013
1014 #[test]
1015 fn rejects_shell_control_syntax_and_expansion() {
1016 for command in [
1017 "git status | cat",
1018 "git status > result",
1019 "git status && echo done",
1020 "git status; echo done",
1021 "git status $(touch bad)",
1022 "git status `touch bad`",
1023 "git status # comment",
1024 "git status\nwhoami",
1025 "git status $HOME",
1026 ] {
1027 assert_eq!(rewrite(command), None, "accepted {command:?}");
1028 }
1029 }
1030
1031 #[test]
1032 fn classifies_only_compatible_compound_results() {
1033 let rule = select_result_rule("cd repo && cargo check; cargo build")
1034 .expect("compatible build result");
1035 assert_eq!(rule.action, yarp_rule_pack::Action::Reduce);
1036 assert!(matches!(
1037 rule.reducer,
1038 Some(yarp_rule_pack::Reducer::BuildSummary)
1039 ));
1040 assert!(select_result_rule("cargo test && cargo build").is_err());
1041 assert!(matches!(
1042 select_result_rule("GOWORK=off go test ./... 2>&1 || true")
1043 .expect("assignment and stream merge"),
1044 Rule {
1045 reducer: Some(yarp_rule_pack::Reducer::TestSummary),
1046 ..
1047 }
1048 ));
1049 assert!(matches!(
1050 select_result_rule("env CI=1 pnpm build && :").expect("environment wrapper"),
1051 Rule {
1052 reducer: Some(yarp_rule_pack::Reducer::BuildSummary),
1053 ..
1054 }
1055 ));
1056 assert!(matches!(
1057 select_result_rule("cargo test | cat").expect("line-preserving pipeline"),
1058 Rule {
1059 reducer: Some(yarp_rule_pack::Reducer::TestSummary),
1060 ..
1061 }
1062 ));
1063 assert!(select_result_rule("cargo test > result.log").is_err());
1064 assert!(select_result_rule("cargo test && echo done").is_err());
1065 assert!(select_result_rule("set && cargo test").is_err());
1066 assert!(select_result_rule("yarp search ref error").is_err());
1067 }
1068
1069 #[test]
1070 fn plans_safe_composite_results_and_rejects_uncertain_stages() {
1071 let search = select_result_plan("rg TODO . | sort | head -50").expect("search pipeline");
1072 assert!(matches!(
1073 search.rule.reducer,
1074 Some(yarp_rule_pack::Reducer::SearchSummary)
1075 ));
1076 assert_eq!(search.status_confidence, StatusConfidence::FinalStageOnly);
1077
1078 let list = select_result_plan("find . -type f | sort | uniq").expect("list pipeline");
1079 assert!(matches!(
1080 list.rule.reducer,
1081 Some(yarp_rule_pack::Reducer::ListSummary)
1082 ));
1083 assert_eq!(list.status_confidence, StatusConfidence::FinalStageOnly);
1084
1085 let test = select_result_plan("cargo test | tee test.log").expect("test pipeline");
1086 assert!(matches!(
1087 test.rule.reducer,
1088 Some(yarp_rule_pack::Reducer::TestSummary)
1089 ));
1090 assert_eq!(test.status_confidence, StatusConfidence::FinalStageOnly);
1091
1092 let pipefail = select_result_plan("set -o pipefail && cargo test | head -100")
1093 .expect("pipefail pipeline");
1094 assert_eq!(pipefail.status_confidence, StatusConfidence::Complete);
1095 let skipped_pipefail = select_result_plan("false && set -o pipefail; cargo test | head -1")
1096 .expect("conditionally skipped pipefail");
1097 assert_eq!(
1098 skipped_pipefail.status_confidence,
1099 StatusConfidence::FinalStageOnly
1100 );
1101 let conditional_disable =
1102 select_result_plan("set -o pipefail; false && set +o pipefail; cargo test | head -1")
1103 .expect("conditional pipefail disable");
1104 assert_eq!(
1105 conditional_disable.status_confidence,
1106 StatusConfidence::FinalStageOnly
1107 );
1108
1109 let sequence = select_result_plan("cargo test; cargo test").expect("test sequence");
1110 assert_eq!(sequence.status_confidence, StatusConfidence::FinalStageOnly);
1111 for command in [
1112 "cd missing; cargo test",
1113 "export CI=1; cargo test",
1114 "VALUE=1; cargo test",
1115 "umask invalid; cargo test",
1116 ] {
1117 let plan = select_result_plan(command).expect("fallible setup sequence");
1118 assert_eq!(
1119 plan.status_confidence,
1120 StatusConfidence::FinalStageOnly,
1121 "{command:?}"
1122 );
1123 }
1124 let conjunction = select_result_plan("cargo test && cargo test").expect("test conjunction");
1125 assert_eq!(conjunction.status_confidence, StatusConfidence::Complete);
1126 let multiline = select_result_plan("cargo test\ncargo test").expect("multiline tests");
1127 assert_eq!(
1128 multiline.status_confidence,
1129 StatusConfidence::FinalStageOnly
1130 );
1131 let masked = select_result_plan("cargo test; true").expect("masked test status");
1132 assert_eq!(masked.status_confidence, StatusConfidence::FinalStageOnly);
1133 assert!(select_result_plan("export CI=1; umask 077; cargo test").is_ok());
1134
1135 for command in [
1136 "cargo test | head source.txt",
1137 "find . -type f | sort existing.txt",
1138 "find . -type f | sort --compress-program=cat",
1139 "find . -type f | sort --compress-program cat",
1140 "ls -D",
1141 "git tag --list --format=%(refname)",
1142 "git stash list --pretty=format:%H",
1143 "rg --json TODO . | jq .",
1144 "find . -print0 | xargs -0 echo",
1145 "cat source.rs | sed 's/x/y/'",
1146 "set -o; cargo test",
1147 "set -x; cargo test",
1148 "set -v; cargo test",
1149 "export -p; cargo test",
1150 "umask -S; cargo test",
1151 "echo \"$VALUE\" | rg x",
1152 "cargo test | unknown-filter",
1153 "cargo test |& head",
1154 "cat source.rs | head",
1155 ] {
1156 assert!(select_result_plan(command).is_err(), "planned {command:?}");
1157 }
1158 }
1159
1160 #[test]
1161 fn handles_quotes_and_escapes_conservatively() {
1162 assert!(rewrite("git diff file\\ name").is_some());
1163 assert!(rewrite("'git' \"status\"").is_some());
1164 assert!(rewrite("git diff '$HOME'").is_some());
1165 assert_eq!(rewrite("git diff \"$HOME\""), None);
1166 assert_eq!(rewrite("git diff 'unterminated"), None);
1167 assert_eq!(rewrite("git diff trailing\\"), None);
1168 }
1169
1170 #[test]
1171 fn validates_parsed_child_arguments() {
1172 assert!(matches!(
1173 select_builtin_argv(&strings(&["cargo", "test", "--workspace"])),
1174 Ok(Selection::Reduce(_))
1175 ));
1176 assert!(matches!(
1177 select_builtin_argv(&strings(&["go", "test", "./..."])),
1178 Ok(Selection::Reduce(_))
1179 ));
1180 assert!(matches!(
1181 select_builtin_argv(&strings(&["git", "show", "HEAD:file"])),
1182 Ok(Selection::Passthrough(_))
1183 ));
1184 assert_eq!(
1185 select_builtin_argv(&strings(&["git", "push"])),
1186 Ok(Selection::Unsupported)
1187 );
1188 }
1189
1190 #[test]
1191 fn classifies_reviewed_commands_and_guards_streaming_or_structured_output() {
1192 for arguments in [
1193 &["hf", "jobs", "logs", "job-id"][..],
1194 &["hf", "spaces", "logs", "owner/space", "--build"][..],
1195 ] {
1196 let Ok(Selection::Reduce(selected)) = select_builtin_argv(&strings(arguments)) else {
1197 panic!("did not reduce {arguments:?}");
1198 };
1199 assert!(matches!(
1200 selected.rule.reducer,
1201 Some(yarp_rule_pack::Reducer::LogSummary)
1202 ));
1203 }
1204 for arguments in [
1205 &["herdr", "pane", "list"][..],
1206 &["herdr", "tab", "list", "--workspace", "w1"][..],
1207 &["herdr", "workspace", "list"][..],
1208 ] {
1209 let Ok(Selection::Reduce(selected)) = select_builtin_argv(&strings(arguments)) else {
1210 panic!("did not reduce {arguments:?}");
1211 };
1212 assert!(matches!(
1213 selected.rule.reducer,
1214 Some(yarp_rule_pack::Reducer::ListSummary)
1215 ));
1216 }
1217 let arguments = ["pnpm", "vitest", "run", "src/example.test.ts"];
1218 let Ok(Selection::Reduce(selected)) = select_builtin_argv(&strings(&arguments)) else {
1219 panic!("did not reduce {arguments:?}");
1220 };
1221 assert!(matches!(
1222 selected.rule.reducer,
1223 Some(yarp_rule_pack::Reducer::TestSummary)
1224 ));
1225 for arguments in [
1226 &["hf", "jobs", "logs", "job-id", "--follow"][..],
1227 &["hf", "jobs", "logs", "job-id", "-f"][..],
1228 &["hf", "jobs", "logs", "job-id", "--json"][..],
1229 &["hf", "jobs", "logs", "job-id", "--format=json"][..],
1230 &["hf", "jobs", "logs", "job-id", "--format", "quiet"][..],
1231 &["pnpm", "vitest", "run", "--watch"][..],
1232 ] {
1233 assert!(matches!(
1234 select_builtin_argv(&strings(arguments)),
1235 Ok(Selection::Passthrough(_))
1236 ));
1237 }
1238 for arguments in [&["pnpm", "vitest"][..], &["pnpm", "vitest", "watch"][..]] {
1239 assert!(!matches!(
1240 select_builtin_argv(&strings(arguments)),
1241 Ok(Selection::Reduce(_))
1242 ));
1243 }
1244 }
1245
1246 fn strings(values: &[&str]) -> Vec<String> {
1247 values.iter().map(ToString::to_string).collect()
1248 }
1249}