Skip to main content

safe_chains/registry/
dispatch.rs

1use crate::parse::Token;
2use crate::policy::FlagSet;
3use crate::verdict::{SafetyLevel, Verdict};
4
5use super::policy::check_owned;
6use super::types::*;
7use super::{CMD_HANDLERS, SUB_HANDLERS};
8
9type HandlerMap = std::collections::HashMap<&'static str, super::HandlerFn>;
10
11fn short_flag_char(s: &str) -> Option<char> {
12    let bytes = s.as_bytes();
13    if bytes.len() == 2 && bytes[0] == b'-' && bytes[1] != b'-' {
14        s.chars().nth(1)
15    } else {
16        None
17    }
18}
19
20fn is_combined_short(s: &str) -> bool {
21    let bytes = s.as_bytes();
22    bytes.len() > 2 && bytes[0] == b'-' && bytes[1] != b'-'
23}
24
25fn dispatch_first_arg(tokens: &[Token], patterns: &[String], level: SafetyLevel) -> Verdict {
26    if tokens.len() == 2 && (tokens[1] == "--help" || tokens[1] == "-h") {
27        return Verdict::Allowed(SafetyLevel::Inert);
28    }
29    let Some(arg) = tokens.get(1) else {
30        return Verdict::Denied;
31    };
32    let arg_str = arg.as_str();
33    let matches = patterns.iter().any(|p| {
34        if let Some(prefix) = p.strip_suffix('*') {
35            arg_str.starts_with(prefix)
36        } else {
37            arg_str == p
38        }
39    });
40    if matches { Verdict::Allowed(level) } else { Verdict::Denied }
41}
42
43fn dispatch_require_any(
44    tokens: &[Token],
45    require_any: &[String],
46    policy: &OwnedPolicy,
47    level: SafetyLevel,
48    accept_bare_help: bool,
49) -> Verdict {
50    if tokens.len() == 2 {
51        let t = tokens[1].as_str();
52        if t == "--help" || t == "-h" || (accept_bare_help && t == "help") {
53            return Verdict::Allowed(SafetyLevel::Inert);
54        }
55    }
56    let has_required = tokens[1..].iter().any(|t| {
57        require_any.iter().any(|r| {
58            let t_str = t.as_str();
59            if t_str == r.as_str() {
60                return true;
61            }
62            if r.starts_with("--") && t_str.starts_with(&format!("{r}=")) {
63                return true;
64            }
65            if let Some(short_char) = short_flag_char(r)
66                && is_combined_short(t_str)
67                && t_str[1..].contains(short_char)
68            {
69                return true;
70            }
71            false
72        })
73    });
74    if has_required && check_owned(tokens, policy) {
75        Verdict::Allowed(level)
76    } else {
77        Verdict::Denied
78    }
79}
80
81fn skip_pre_flags(
82    tokens: &[Token],
83    pre_standalone: &[String],
84    pre_valued: &[String],
85    start: usize,
86) -> usize {
87    let mut i = start;
88    while i < tokens.len() {
89        let t = &tokens[i];
90        let s = t.as_str();
91        if !s.starts_with('-') {
92            break;
93        }
94        if pre_valued.contains_flag(s) {
95            i += 2;
96            continue;
97        }
98        if let Some((flag, _)) = s.split_once('=')
99            && pre_valued.contains_flag(flag)
100        {
101            i += 1;
102            continue;
103        }
104        if pre_standalone.contains_flag(s) {
105            i += 1;
106            continue;
107        }
108        // POSIX-style short-flag cluster (`-vv`, `-vy`): every byte after
109        // the dash must be a known standalone short. Mirrors the same
110        // logic in policy::check_flags for non-wrapper subs.
111        let bytes = s.as_bytes();
112        if bytes.len() > 2
113            && bytes[1] != b'-'
114            && bytes[1..].iter().all(|&b| pre_standalone.contains_short(b))
115        {
116            i += 1;
117            continue;
118        }
119        break;
120    }
121    i
122}
123
124fn dispatch_branching(
125    tokens: &[Token],
126    subs: &[SubSpec],
127    bare_flags: &[String],
128    bare_ok: bool,
129    pre_flags: (&[String], &[String]),
130    first_arg: &[String],
131    first_arg_level: SafetyLevel,
132    credential_first_arg: &[String],
133) -> Verdict {
134    let (pre_standalone, pre_valued) = pre_flags;
135    let start = skip_pre_flags(tokens, pre_standalone, pre_valued, 1);
136    if start >= tokens.len() {
137        return if bare_ok { Verdict::Allowed(SafetyLevel::Inert) } else { Verdict::Denied };
138    }
139    let arg = tokens[start].as_str();
140    let is_bare_flag = bare_flags.iter().any(|f| f == arg)
141        || (bare_flags.is_empty() && matches!(arg, "--help" | "-h"));
142    if is_bare_flag {
143        let after = skip_pre_flags(tokens, pre_standalone, pre_valued, start + 1);
144        if after >= tokens.len() {
145            return Verdict::Allowed(SafetyLevel::Inert);
146        }
147        if bare_flags.is_empty() {
148            return Verdict::Denied;
149        }
150    }
151    if let Some(sub) = subs.iter().find(|s| s.name == arg) {
152        return dispatch_kind(&tokens[start..], &sub.kind, &SUB_HANDLERS);
153    }
154    let glob_match = |p: &str| match p.strip_suffix('*') {
155        Some(prefix) => arg.starts_with(prefix),
156        None => arg == p,
157    };
158    // A first positional naming secret material (`kubectl get secret`/`secret/x`, `aws configure get
159    // aws_secret_access_key`) is a CREDENTIAL-READ — deny before the allow-glob admits it. The
160    // value-dependent complement to `profile=credential-read`; `arg` is already flag-aware.
161    if credential_first_arg.iter().any(|p| glob_match(p)) {
162        return Verdict::Denied;
163    }
164    if !first_arg.is_empty() && first_arg.iter().any(|p| glob_match(p)) {
165        return Verdict::Allowed(first_arg_level);
166    }
167    Verdict::Denied
168}
169
170fn dispatch_wrapper(
171    tokens: &[Token],
172    standalone: &[String],
173    valued: &[String],
174    positional_skip: usize,
175    separator: Option<&str>,
176    bare_ok: bool,
177) -> Verdict {
178    let mut i = 1;
179    while i < tokens.len() {
180        let t = &tokens[i];
181        if let Some(sep) = separator
182            && t == sep
183        {
184            i += 1;
185            break;
186        }
187        if !t.starts_with('-') {
188            break;
189        }
190        if valued.iter().any(|f| t == f.as_str()) {
191            i += 2;
192            continue;
193        }
194        if valued.iter().any(|f| t.as_str().starts_with(&format!("{f}="))) {
195            i += 1;
196            continue;
197        }
198        if standalone.iter().any(|f| t == f.as_str()) {
199            i += 1;
200            continue;
201        }
202        return Verdict::Denied;
203    }
204    for _ in 0..positional_skip {
205        if i >= tokens.len() {
206            return if bare_ok {
207                Verdict::Allowed(SafetyLevel::Inert)
208            } else {
209                Verdict::Denied
210            };
211        }
212        i += 1;
213    }
214    if i >= tokens.len() {
215        return if bare_ok {
216            Verdict::Allowed(SafetyLevel::Inert)
217        } else {
218            Verdict::Denied
219        };
220    }
221    let inner = shell_words::join(tokens[i..].iter().map(|t| t.as_str()));
222    crate::command_verdict(&inner)
223}
224
225fn dispatch_kind(tokens: &[Token], kind: &DispatchKind, handlers: &HandlerMap) -> Verdict {
226    match kind {
227        DispatchKind::Policy { policy, level } => {
228            if check_owned(tokens, policy) {
229                Verdict::Allowed(*level)
230            } else {
231                Verdict::Denied
232            }
233        }
234        DispatchKind::FirstArg { patterns, level } => {
235            dispatch_first_arg(tokens, patterns, *level)
236        }
237        DispatchKind::RequireAny { require_any, policy, level, accept_bare_help } => {
238            dispatch_require_any(tokens, require_any, policy, *level, *accept_bare_help)
239        }
240        DispatchKind::Branching {
241            subs, bare_flags, bare_ok, pre_standalone, pre_valued, first_arg, first_arg_level,
242            credential_first_arg,
243        } => {
244            dispatch_branching(
245                tokens, subs, bare_flags, *bare_ok, (pre_standalone, pre_valued),
246                first_arg, *first_arg_level, credential_first_arg,
247            )
248        }
249        DispatchKind::WriteFlagged { policy, base_level, write_flags } => {
250            if !check_owned(tokens, policy) {
251                return Verdict::Denied;
252            }
253            let has_write = tokens[1..].iter().any(|t| {
254                write_flags.iter().any(|f| t == f.as_str() || t.as_str().starts_with(&format!("{f}=")))
255            });
256            if has_write {
257                Verdict::Allowed(SafetyLevel::SafeWrite)
258            } else {
259                Verdict::Allowed(*base_level)
260            }
261        }
262        DispatchKind::DelegateAfterSeparator { separator } => {
263            let sep_pos = tokens[1..].iter().position(|t| t == separator.as_str());
264            let Some(pos) = sep_pos else {
265                return Verdict::Denied;
266            };
267            let inner_start = pos + 2;
268            if inner_start >= tokens.len() {
269                return Verdict::Denied;
270            }
271            let inner = shell_words::join(tokens[inner_start..].iter().map(|t| t.as_str()));
272            crate::command_verdict(&inner)
273        }
274        DispatchKind::DelegateSkip { skip } => {
275            if tokens.len() <= *skip {
276                return Verdict::Denied;
277            }
278            let inner = shell_words::join(tokens[*skip..].iter().map(|t| t.as_str()));
279            crate::command_verdict(&inner)
280        }
281        DispatchKind::Wrapper {
282            standalone, valued, positional_skip, separator, bare_ok,
283        } => {
284            dispatch_wrapper(tokens, standalone, valued, *positional_skip, separator.as_deref(), *bare_ok)
285        }
286        DispatchKind::VerbChain(spec) => dispatch_verb_chain(tokens, spec),
287        DispatchKind::Executor { policy, level, kind, redirect_flag, shape } => {
288            dispatch_executor(tokens, policy, *kind, *level, redirect_flag.as_deref(), *shape)
289        }
290        DispatchKind::Custom { handler_name, .. } => {
291            handlers
292                .get(handler_name.as_str())
293                .map(|f| f(tokens))
294                .unwrap_or(Verdict::Denied)
295        }
296    }
297}
298
299pub fn dispatch_spec(tokens: &[Token], spec: &CommandSpec) -> Verdict {
300    dispatch_kind(tokens, &spec.kind, &CMD_HANDLERS)
301}
302
303/// Dispatches a sub's kind directly, used by `registry::try_sub_dispatch`
304/// when a handler-using command consults its TOML-declared subs.
305pub(super) fn dispatch_sub_kind(tokens: &[Token], kind: &DispatchKind) -> Verdict {
306    dispatch_kind(tokens, kind, &SUB_HANDLERS)
307}
308
309pub(super) fn check_handler_policy_owned(tokens: &[Token], policy: &OwnedPolicy) -> bool {
310    check_owned(tokens, policy)
311}
312
313pub(super) fn dispatch_matrix_action(
314    tokens: &[Token],
315    policy: &OwnedPolicy,
316    level: SafetyLevel,
317) -> Verdict {
318    if check_owned(tokens, policy) {
319        Verdict::Allowed(level)
320    } else {
321        Verdict::Denied
322    }
323}
324
325/// A `verb-chain` grammar (`mlr`): `CMD [main-flags…] verb [args…] then verb …`. The main-flag
326/// region is a STRICT allowlist (standalone / valued / variadic) — an unlisted flag denies, so a
327/// mutating flag like mlr's `-I`/`--in-place` (omitted) is caught by omission. The region ends at
328/// `--` or the first non-flag token, which opens the verb chain: every `then`-segment's first token
329/// (the verb NAME) must be on the `verbs` allowlist; verb ARGS are open-ended (a pure verb has no
330/// shell/file escape). See `types::VerbChainSpec`.
331pub(super) fn dispatch_verb_chain(tokens: &[Token], spec: &VerbChainSpec) -> Verdict {
332    if tokens.len() == 1 {
333        return Verdict::Denied;
334    }
335    let has = |set: &[String], s: &str| set.iter().any(|f| f == s);
336    let mut i = 1;
337    while i < tokens.len() {
338        let s = tokens[i].as_str();
339        // `--` terminates the main region; the verb chain follows it.
340        if s == "--" {
341            return verb_chain_tail(&tokens[i + 1..], spec);
342        }
343        // The first non-flag token opens the verb chain.
344        if !tokens[i].starts_with('-') {
345            return verb_chain_tail(&tokens[i..], spec);
346        }
347        if has(&spec.main_standalone, s) {
348            i += 1;
349        } else if has(&spec.main_variadic, s) {
350            // Consume input files up to the `--` terminator (or end); the `--` is left for the
351            // loop to treat as the main/verb separator next iteration.
352            i += 1;
353            while i < tokens.len() && tokens[i].as_str() != "--" {
354                i += 1;
355            }
356        } else if has(&spec.main_valued, s) {
357            i += if crate::policy::consumes_next_value(tokens.get(i + 1)) { 2 } else { 1 };
358        } else if s.split_once('=').is_some_and(|(f, _)| has(&spec.main_valued, f)) {
359            i += 1;
360        } else {
361            return Verdict::Denied;
362        }
363    }
364    // Only allowlisted main flags, no verb (`mlr --version`, `mlr --csv`): nothing runs.
365    Verdict::Allowed(spec.level)
366}
367
368/// The verb region: a `then`-chain where every segment's first token must be an allowlisted verb.
369fn verb_chain_tail(region: &[Token], spec: &VerbChainSpec) -> Verdict {
370    let mut expect_verb = true;
371    for t in region {
372        if t.as_str() == spec.separator {
373            expect_verb = true;
374            continue;
375        }
376        if expect_verb {
377            if !spec.verbs.contains(t.as_str()) {
378                return Verdict::Denied;
379            }
380            expect_verb = false;
381        }
382    }
383    // An empty region or a dangling separator leaves `expect_verb` set — no valid verb ran.
384    if expect_verb { Verdict::Denied } else { Verdict::Allowed(spec.level) }
385}
386
387/// Gate a code-execution command through the execution-origin engine (worktree code allows,
388/// foreign denies) instead of a flat level. Shared by executor SUBS (`go run`, `cargo run`)
389/// and executor FALLBACKS (interpreters). See docs/design/behavioral-taxonomy-execution-origin.md.
390///
391/// - `File`: the first positional is the executor path — `bash x.sh`, `go run ./cmd`. Tokens
392///   AFTER it are the script's own arguments (data, not validated). With no positional, it's a
393///   flag-only invocation (`python3 --version`) validated by `policy` (denies a bare REPL).
394/// - `Project`: the current project is the executor — `cargo run`. A `redirect_flag` value
395///   (`--manifest-path DIR/Cargo.toml`) moves the executor OUT of the project and is locus-gated
396///   like a file executor; without it, the project (worktree) is the executor.
397pub(super) fn dispatch_executor(
398    tokens: &[Token],
399    policy: &OwnedPolicy,
400    kind: ExecutorKind,
401    level: SafetyLevel,
402    redirect_flag: Option<&str>,
403    shape: Option<crate::policy::PositionalShape>,
404) -> Verdict {
405    match kind {
406        ExecutorKind::File => match super::policy::first_positional(tokens, policy) {
407            // A declared shape the executor path must satisfy (`go run` → `go-package`):
408            // a remote import path is not a worktree executor, so it denies here.
409            Some(first) if shape.is_some_and(|s| !s.matches(first)) => Verdict::Denied,
410            Some(first) => crate::engine::resolve::execute_file_verdict(first),
411            None if check_owned(tokens, policy) => Verdict::Allowed(level),
412            None => Verdict::Denied,
413        },
414        ExecutorKind::Project => {
415            if !check_owned(tokens, policy) {
416                return Verdict::Denied;
417            }
418            match redirect_flag.and_then(|f| flag_value(tokens, f)) {
419                Some(redirected) => crate::engine::resolve::execute_file_verdict(redirected),
420                None => crate::engine::resolve::execute_project_verdict(),
421            }
422        }
423    }
424}
425
426/// The value of valued flag `flag` in `tokens` — the space form (`--manifest-path P`) or the
427/// glued `=` form (`--manifest-path=P`). Scans only up to a `--` terminator, so a program
428/// argument after `--` can't be mistaken for the flag's value. `None` if the flag is absent.
429fn flag_value<'a>(tokens: &'a [Token], flag: &str) -> Option<&'a str> {
430    let mut i = 1;
431    while i < tokens.len() {
432        let t = tokens[i].as_str();
433        if t == "--" {
434            return None;
435        }
436        if t == flag {
437            return tokens.get(i + 1).map(Token::as_str);
438        }
439        if let Some(rest) = t.strip_prefix(flag).and_then(|r| r.strip_prefix('=')) {
440            return Some(rest);
441        }
442        i += 1;
443    }
444    None
445}
446
447/// Applies a TOML-declared fallback grammar. Used by
448/// `registry::try_fallback_grammar()`.
449pub(super) fn dispatch_fallback(tokens: &[Token], spec: &FallbackSpec) -> Verdict {
450    if let Some(kind) = spec.executor {
451        return dispatch_executor(
452            tokens, &spec.policy, kind, spec.level,
453            spec.executor_redirect_flag.as_deref(), spec.positional_shape,
454        );
455    }
456    if let Some(shape) = spec.positional_shape
457        && let Some(first) = super::policy::first_positional(tokens, &spec.policy)
458        && !shape.matches(first)
459    {
460        return Verdict::Denied;
461    }
462    if !check_owned(tokens, &spec.policy) {
463        return Verdict::Denied;
464    }
465    Verdict::Allowed(spec.level)
466}