Skip to main content

safe_chains/
command.rs

1use crate::parse::{has_flag, Segment, Token};
2use crate::policy::{self, FlagPolicy};
3
4pub type CheckFn = fn(&[Token], &dyn Fn(&Segment) -> bool) -> bool;
5
6pub enum SubDef {
7    Policy {
8        name: &'static str,
9        policy: &'static FlagPolicy,
10    },
11    Nested {
12        name: &'static str,
13        subs: &'static [SubDef],
14    },
15    Guarded {
16        name: &'static str,
17        guard_short: Option<&'static str>,
18        guard_long: &'static str,
19        policy: &'static FlagPolicy,
20    },
21    Custom {
22        name: &'static str,
23        check: CheckFn,
24        doc: &'static str,
25        test_suffix: Option<&'static str>,
26    },
27    Delegation {
28        name: &'static str,
29        skip: usize,
30        doc: &'static str,
31    },
32}
33
34pub struct CommandDef {
35    pub name: &'static str,
36    pub subs: &'static [SubDef],
37    pub bare_flags: &'static [&'static str],
38    pub help_eligible: bool,
39}
40
41impl SubDef {
42    pub fn name(&self) -> &'static str {
43        match self {
44            Self::Policy { name, .. }
45            | Self::Nested { name, .. }
46            | Self::Guarded { name, .. }
47            | Self::Custom { name, .. }
48            | Self::Delegation { name, .. } => name,
49        }
50    }
51
52    pub fn check(&self, tokens: &[Token], is_safe: &dyn Fn(&Segment) -> bool) -> bool {
53        match self {
54            Self::Policy { policy, .. } => policy::check(tokens, policy),
55            Self::Nested { subs, .. } => {
56                if tokens.len() < 2 {
57                    return false;
58                }
59                let sub = tokens[1].as_str();
60                subs.iter()
61                    .any(|s| s.name() == sub && s.check(&tokens[1..], is_safe))
62            }
63            Self::Guarded {
64                guard_short,
65                guard_long,
66                policy,
67                ..
68            } => {
69                has_flag(tokens, *guard_short, Some(guard_long))
70                    && policy::check(tokens, policy)
71            }
72            Self::Custom { check: f, .. } => f(tokens, is_safe),
73            Self::Delegation { skip, .. } => {
74                if tokens.len() <= *skip {
75                    return false;
76                }
77                let inner = Token::join(&tokens[*skip..]);
78                is_safe(&inner)
79            }
80        }
81    }
82}
83
84impl CommandDef {
85    pub fn check(&self, tokens: &[Token], is_safe: &dyn Fn(&Segment) -> bool) -> bool {
86        if tokens.len() < 2 {
87            return false;
88        }
89        let arg = tokens[1].as_str();
90        if tokens.len() == 2 && self.bare_flags.contains(&arg) {
91            return true;
92        }
93        self.subs
94            .iter()
95            .find(|s| s.name() == arg)
96            .is_some_and(|s| s.check(&tokens[1..], is_safe))
97    }
98
99    pub fn dispatch(
100        &self,
101        cmd: &str,
102        tokens: &[Token],
103        is_safe: &dyn Fn(&Segment) -> bool,
104    ) -> Option<bool> {
105        if cmd == self.name {
106            Some(self.check(tokens, is_safe))
107        } else {
108            None
109        }
110    }
111
112    pub fn to_doc(&self) -> crate::docs::CommandDoc {
113        let mut parts = Vec::new();
114
115        let mut policy_names: Vec<&str> = Vec::new();
116        let mut nested_names: Vec<String> = Vec::new();
117        let mut guarded_descs: Vec<String> = Vec::new();
118        let mut extra_docs: Vec<&str> = Vec::new();
119
120        for sub in self.subs {
121            match sub {
122                SubDef::Policy { name, .. } => {
123                    policy_names.push(name);
124                }
125                SubDef::Nested { name, subs } => {
126                    let visible: Vec<_> = subs.iter()
127                        .filter(|s| !s.name().starts_with('-'))
128                        .collect();
129                    if visible.len() <= 5 {
130                        for s in &visible {
131                            nested_names.push(format!("{name} {}", s.name()));
132                        }
133                    } else {
134                        nested_names.push((*name).to_string());
135                    }
136                }
137                SubDef::Guarded { name, guard_long, .. } => {
138                    guarded_descs.push(format!("{name} (requires {guard_long})"));
139                }
140                SubDef::Custom { name, doc, .. } => {
141                    if doc.is_empty() {
142                        policy_names.push(name);
143                    } else if !doc.trim().is_empty() {
144                        extra_docs.push(doc);
145                    }
146                }
147                SubDef::Delegation { doc, .. } => {
148                    if !doc.is_empty() {
149                        extra_docs.push(doc);
150                    }
151                }
152            }
153        }
154
155        if !policy_names.is_empty() {
156            policy_names.sort();
157            parts.push(format!("Subcommands: {}.", policy_names.join(", ")));
158        }
159
160        if !nested_names.is_empty() {
161            nested_names.sort();
162            parts.push(format!("Multi-level: {}.", nested_names.join(", ")));
163        }
164
165        if !self.bare_flags.is_empty() {
166            parts.push(format!("Info flags: {}.", self.bare_flags.join(", ")));
167        }
168
169        if !guarded_descs.is_empty() {
170            parts.push(format!("{}.", guarded_descs.join(", ")));
171        }
172
173        for doc in extra_docs {
174            parts.push(doc.to_string());
175        }
176
177        crate::docs::CommandDoc::handler(self.name, parts.join(" "))
178    }
179}
180
181#[cfg(test)]
182impl CommandDef {
183    pub fn auto_test_reject_unknown(&self) {
184        let mut failures = Vec::new();
185
186        assert!(
187            !crate::is_safe_command(self.name),
188            "{}: accepted bare invocation",
189            self.name,
190        );
191
192        let test = format!("{} xyzzy-unknown-42", self.name);
193        assert!(
194            !crate::is_safe_command(&test),
195            "{}: accepted unknown subcommand: {test}",
196            self.name,
197        );
198
199        for sub in self.subs {
200            auto_test_sub(self.name, sub, &mut failures);
201        }
202        assert!(
203            failures.is_empty(),
204            "{}: unknown flags/subcommands accepted:\n{}",
205            self.name,
206            failures.join("\n"),
207        );
208    }
209}
210
211#[cfg(test)]
212fn auto_test_sub(prefix: &str, sub: &SubDef, failures: &mut Vec<String>) {
213    const UNKNOWN: &str = "--xyzzy-unknown-42";
214
215    match sub {
216        SubDef::Policy { name, .. } => {
217            let test = format!("{prefix} {name} {UNKNOWN}");
218            if crate::is_safe_command(&test) {
219                failures.push(format!("{prefix} {name}: accepted unknown flag: {test}"));
220            }
221        }
222        SubDef::Nested { name, subs } => {
223            let path = format!("{prefix} {name}");
224            let test = format!("{path} xyzzy-unknown-42");
225            if crate::is_safe_command(&test) {
226                failures.push(format!("{path}: accepted unknown subcommand: {test}"));
227            }
228            for s in *subs {
229                auto_test_sub(&path, s, failures);
230            }
231        }
232        SubDef::Guarded {
233            name, guard_long, ..
234        } => {
235            let test = format!("{prefix} {name} {guard_long} {UNKNOWN}");
236            if crate::is_safe_command(&test) {
237                failures.push(format!("{prefix} {name}: accepted unknown flag: {test}"));
238            }
239        }
240        SubDef::Custom {
241            name, test_suffix, ..
242        } => {
243            if let Some(suffix) = test_suffix {
244                let test = format!("{prefix} {name} {suffix} {UNKNOWN}");
245                if crate::is_safe_command(&test) {
246                    failures.push(format!(
247                        "{prefix} {name}: accepted unknown flag: {test}"
248                    ));
249                }
250            }
251        }
252        SubDef::Delegation { .. } => {}
253    }
254}
255
256#[cfg(test)]
257mod tests {
258    use super::*;
259    use crate::parse::WordSet;
260    use crate::policy::FlagStyle;
261
262    fn toks(words: &[&str]) -> Vec<Token> {
263        words.iter().map(|s| Token::from_test(s)).collect()
264    }
265
266    fn no_safe(_: &Segment) -> bool {
267        false
268    }
269
270    static TEST_POLICY: FlagPolicy = FlagPolicy {
271        standalone: WordSet::new(&["--verbose"]),
272        standalone_short: b"v",
273        valued: WordSet::new(&["--output"]),
274        valued_short: b"o",
275        bare: true,
276        max_positional: None,
277        flag_style: FlagStyle::Strict,
278    };
279
280    static SIMPLE_CMD: CommandDef = CommandDef {
281        name: "mycmd",
282        subs: &[SubDef::Policy {
283            name: "build",
284            policy: &TEST_POLICY,
285        }],
286        bare_flags: &["--info"],
287        help_eligible: true,
288    };
289
290    #[test]
291    fn bare_rejected() {
292        assert!(!SIMPLE_CMD.check(&toks(&["mycmd"]), &no_safe));
293    }
294
295    #[test]
296    fn bare_flag_accepted() {
297        assert!(SIMPLE_CMD.check(&toks(&["mycmd", "--info"]), &no_safe));
298    }
299
300    #[test]
301    fn bare_flag_with_extra_rejected() {
302        assert!(!SIMPLE_CMD.check(&toks(&["mycmd", "--info", "extra"]), &no_safe));
303    }
304
305    #[test]
306    fn policy_sub_bare() {
307        assert!(SIMPLE_CMD.check(&toks(&["mycmd", "build"]), &no_safe));
308    }
309
310    #[test]
311    fn policy_sub_with_flag() {
312        assert!(SIMPLE_CMD.check(&toks(&["mycmd", "build", "--verbose"]), &no_safe));
313    }
314
315    #[test]
316    fn policy_sub_unknown_flag() {
317        assert!(!SIMPLE_CMD.check(&toks(&["mycmd", "build", "--bad"]), &no_safe));
318    }
319
320    #[test]
321    fn unknown_sub_rejected() {
322        assert!(!SIMPLE_CMD.check(&toks(&["mycmd", "deploy"]), &no_safe));
323    }
324
325    #[test]
326    fn dispatch_matches() {
327        assert_eq!(
328            SIMPLE_CMD.dispatch("mycmd", &toks(&["mycmd", "build"]), &no_safe),
329            Some(true)
330        );
331    }
332
333    #[test]
334    fn dispatch_no_match() {
335        assert_eq!(
336            SIMPLE_CMD.dispatch("other", &toks(&["other", "build"]), &no_safe),
337            None
338        );
339    }
340
341    static NESTED_CMD: CommandDef = CommandDef {
342        name: "nested",
343        subs: &[SubDef::Nested {
344            name: "package",
345            subs: &[SubDef::Policy {
346                name: "describe",
347                policy: &TEST_POLICY,
348            }],
349        }],
350        bare_flags: &[],
351        help_eligible: false,
352    };
353
354    #[test]
355    fn nested_sub() {
356        assert!(NESTED_CMD.check(&toks(&["nested", "package", "describe"]), &no_safe));
357    }
358
359    #[test]
360    fn nested_sub_with_flag() {
361        assert!(NESTED_CMD.check(
362            &toks(&["nested", "package", "describe", "--verbose"]),
363            &no_safe,
364        ));
365    }
366
367    #[test]
368    fn nested_bare_rejected() {
369        assert!(!NESTED_CMD.check(&toks(&["nested", "package"]), &no_safe));
370    }
371
372    #[test]
373    fn nested_unknown_sub_rejected() {
374        assert!(!NESTED_CMD.check(&toks(&["nested", "package", "deploy"]), &no_safe));
375    }
376
377    static GUARDED_POLICY: FlagPolicy = FlagPolicy {
378        standalone: WordSet::new(&["--all", "--check"]),
379        standalone_short: b"",
380        valued: WordSet::new(&[]),
381        valued_short: b"",
382        bare: false,
383        max_positional: None,
384        flag_style: FlagStyle::Strict,
385    };
386
387    static GUARDED_CMD: CommandDef = CommandDef {
388        name: "guarded",
389        subs: &[SubDef::Guarded {
390            name: "fmt",
391            guard_short: None,
392            guard_long: "--check",
393            policy: &GUARDED_POLICY,
394        }],
395        bare_flags: &[],
396        help_eligible: false,
397    };
398
399    #[test]
400    fn guarded_with_guard() {
401        assert!(GUARDED_CMD.check(&toks(&["guarded", "fmt", "--check"]), &no_safe));
402    }
403
404    #[test]
405    fn guarded_without_guard() {
406        assert!(!GUARDED_CMD.check(&toks(&["guarded", "fmt"]), &no_safe));
407    }
408
409    #[test]
410    fn guarded_with_guard_and_flag() {
411        assert!(GUARDED_CMD.check(
412            &toks(&["guarded", "fmt", "--check", "--all"]),
413            &no_safe,
414        ));
415    }
416
417    fn safe_echo(seg: &Segment) -> bool {
418        seg.as_str() == "echo hello"
419    }
420
421    static DELEGATION_CMD: CommandDef = CommandDef {
422        name: "runner",
423        subs: &[SubDef::Delegation {
424            name: "run",
425            skip: 2,
426            doc: "run delegates to inner command.",
427        }],
428        bare_flags: &[],
429        help_eligible: false,
430    };
431
432    #[test]
433    fn delegation_safe_inner() {
434        assert!(DELEGATION_CMD.check(
435            &toks(&["runner", "run", "stable", "echo", "hello"]),
436            &safe_echo,
437        ));
438    }
439
440    #[test]
441    fn delegation_unsafe_inner() {
442        assert!(!DELEGATION_CMD.check(
443            &toks(&["runner", "run", "stable", "rm", "-rf"]),
444            &no_safe,
445        ));
446    }
447
448    #[test]
449    fn delegation_no_inner() {
450        assert!(!DELEGATION_CMD.check(
451            &toks(&["runner", "run", "stable"]),
452            &no_safe,
453        ));
454    }
455
456    fn custom_check(tokens: &[Token], _is_safe: &dyn Fn(&Segment) -> bool) -> bool {
457        tokens.len() >= 2 && tokens[1] == "safe"
458    }
459
460    static CUSTOM_CMD: CommandDef = CommandDef {
461        name: "custom",
462        subs: &[SubDef::Custom {
463            name: "special",
464            check: custom_check,
465            doc: "special (safe only).",
466            test_suffix: Some("safe"),
467        }],
468        bare_flags: &[],
469        help_eligible: false,
470    };
471
472    #[test]
473    fn custom_passes() {
474        assert!(CUSTOM_CMD.check(&toks(&["custom", "special", "safe"]), &no_safe));
475    }
476
477    #[test]
478    fn custom_fails() {
479        assert!(!CUSTOM_CMD.check(&toks(&["custom", "special", "bad"]), &no_safe));
480    }
481
482    #[test]
483    fn doc_simple() {
484        let doc = SIMPLE_CMD.to_doc();
485        assert_eq!(doc.name, "mycmd");
486        assert_eq!(doc.description, "Subcommands: build. Info flags: --info.");
487    }
488
489    #[test]
490    fn doc_nested() {
491        let doc = NESTED_CMD.to_doc();
492        assert_eq!(doc.description, "Multi-level: package describe.");
493    }
494
495    #[test]
496    fn doc_guarded() {
497        let doc = GUARDED_CMD.to_doc();
498        assert_eq!(doc.description, "fmt (requires --check).");
499    }
500
501    #[test]
502    fn doc_delegation() {
503        let doc = DELEGATION_CMD.to_doc();
504        assert_eq!(doc.description, "run delegates to inner command.");
505    }
506
507    #[test]
508    fn doc_custom() {
509        let doc = CUSTOM_CMD.to_doc();
510        assert_eq!(doc.description, "special (safe only).");
511    }
512}