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