Skip to main content

safe_chains/registry/
mod.rs

1mod build;
2mod custom;
3mod dispatch;
4mod docs;
5mod policy;
6pub(crate) mod types;
7
8use std::collections::HashMap;
9use std::sync::LazyLock;
10
11use crate::parse::Token;
12use crate::verdict::Verdict;
13
14pub use build::{build_registry, load_toml};
15pub use dispatch::dispatch_spec;
16pub use types::{CommandSpec, OwnedPolicy};
17
18use types::DispatchKind;
19
20type HandlerFn = fn(&[Token]) -> Verdict;
21
22static CMD_HANDLERS: LazyLock<HashMap<&'static str, HandlerFn>> =
23    LazyLock::new(crate::handlers::custom_cmd_handlers);
24
25static SUB_HANDLERS: LazyLock<HashMap<&'static str, HandlerFn>> =
26    LazyLock::new(crate::handlers::custom_sub_handlers);
27
28static TOML_REGISTRY: LazyLock<HashMap<String, CommandSpec>> = LazyLock::new(||
29    include!(concat!(env!("OUT_DIR"), "/toml_includes.rs"))
30);
31
32static CUSTOM_REGISTRY: LazyLock<HashMap<String, CommandSpec>> = LazyLock::new(|| {
33    let mut map = HashMap::new();
34    custom::apply_custom(&mut map);
35    map
36});
37
38pub fn toml_dispatch(tokens: &[Token]) -> Option<Verdict> {
39    let cmd = tokens[0].command_name();
40    TOML_REGISTRY.get(cmd).map(|spec| dispatch_spec(tokens, spec))
41}
42
43/// Looks up the command in the runtime custom registry (project-local
44/// `.safe-chains.toml`, then user-level `~/.config/safe-chains.toml`).
45/// A match here wins over the built-in hardcoded handlers, which is how
46/// an override of `gh` takes effect.
47pub fn custom_dispatch(tokens: &[Token]) -> Option<Verdict> {
48    let cmd = tokens[0].command_name();
49    CUSTOM_REGISTRY.get(cmd).map(|spec| dispatch_spec(tokens, spec))
50}
51
52pub fn toml_command_names() -> Vec<&'static str> {
53    TOML_REGISTRY
54        .keys()
55        .map(|k| k.as_str())
56        .collect()
57}
58
59/// Look up `cmd_name`'s TOML-declared subs (set via `[[command.sub]]`
60/// blocks alongside `handler = "..."`) and dispatch the one whose name
61/// matches `tokens[1]`. Returns `None` if no sub matched, so the
62/// handler can fall through to its fallback grammar (or deny).
63pub fn try_sub_dispatch(cmd_name: &str, tokens: &[Token]) -> Option<Verdict> {
64    let spec = handler_spec(cmd_name)?;
65    let DispatchKind::Custom { subs, .. } = &spec.kind else {
66        return None;
67    };
68    let arg = tokens.get(1)?.as_str();
69    let sub = subs.iter().find(|s| s.name == arg)?;
70    Some(dispatch::dispatch_sub_kind(&tokens[1..], &sub.kind))
71}
72
73/// Apply `cmd_name`'s TOML-declared `[command.fallback]` grammar.
74/// Returns `None` if no fallback is declared.
75pub fn try_fallback_grammar(cmd_name: &str, tokens: &[Token]) -> Option<Verdict> {
76    let spec = handler_spec(cmd_name)?;
77    let DispatchKind::Custom { fallback, .. } = &spec.kind else {
78        return None;
79    };
80    let f = fallback.as_ref()?;
81    Some(dispatch::dispatch_fallback(tokens, f))
82}
83
84/// Dispatch `tokens` against `cmd_name`'s `[[command.matrix]]`
85/// blocks. Looks at `tokens[1]` (parent) and `tokens[2]` (action),
86/// finds the first matrix whose `parents` contains the parent and
87/// whose `actions` map contains the action, then validates
88/// `tokens[2..]` against the named policy (and a guard flag if the
89/// matrix entry declared one). Returns `None` if no matrix matched —
90/// the handler can then fall through to its remaining special cases
91/// or deny.
92pub fn try_matrix_dispatch(cmd_name: &str, tokens: &[Token]) -> Option<Verdict> {
93    let spec = handler_spec(cmd_name)?;
94    let DispatchKind::Custom { matrices, handler_policies, .. } = &spec.kind else {
95        return None;
96    };
97    let parent = tokens.get(1)?.as_str();
98    let action = tokens.get(2)?.as_str();
99    for matrix in matrices {
100        if !matrix.parents.iter().any(|p| p == parent) {
101            continue;
102        }
103        let Some(action_spec) = matrix.actions.get(action) else { continue; };
104        if let Some(long) = action_spec.guard.as_deref()
105            && !crate::parse::has_flag(&tokens[2..], action_spec.guard_short.as_deref(), Some(long))
106        {
107            return Some(Verdict::Denied);
108        }
109        let Some(policy) = handler_policies.get(&action_spec.policy_key) else {
110            return Some(Verdict::Denied);
111        };
112        return Some(dispatch::dispatch_matrix_action(&tokens[2..], policy, matrix.level));
113    }
114    None
115}
116
117/// Validate `tokens` against `cmd_name`'s named flag policy declared
118/// in a `[command.handler_policy.KEY]` block. Returns `false` if no
119/// such policy is declared or the tokens fail it. Used by handlers
120/// whose dispatch logic genuinely can't move to TOML (e.g. gh's
121/// sub × action matrix) but whose per-policy WordSets should live
122/// in TOML rather than as Rust `WordSet` constants.
123pub fn check_handler_policy(cmd_name: &str, key: &str, tokens: &[Token]) -> bool {
124    let Some(spec) = handler_spec(cmd_name) else { return false; };
125    let DispatchKind::Custom { handler_policies, .. } = &spec.kind else {
126        return false;
127    };
128    let Some(policy) = handler_policies.get(key) else { return false; };
129    dispatch::check_handler_policy_owned(tokens, policy)
130}
131
132fn handler_spec(cmd_name: &str) -> Option<&'static CommandSpec> {
133    CUSTOM_REGISTRY
134        .get(cmd_name)
135        .or_else(|| TOML_REGISTRY.get(cmd_name))
136}
137
138/// Returns true iff this invocation is tagged eval-safe — meaning its
139/// stdout is documented shell-init code that can safely be substituted
140/// inside `eval "$(...)"`.
141///
142/// The walker descends through `DispatchKind::Branching` AND
143/// `DispatchKind::Custom` matching subs token-by-token (handler-based
144/// commands such as `gh` can have tagged TOML-declared subs even though
145/// the handler does the actual dispatch). The leaf is the deepest matched
146/// node (where no further sub matches). `eval_safe` is checked only at
147/// the leaf — ancestor tags do NOT propagate. After confirming the leaf
148/// is tagged, every `-`-prefixed token in the remaining tail must appear
149/// in `eval_safe_flags`; positionals are unrestricted.
150///
151/// Tagged nodes are vetted manually per-command (see SAMPLE.toml). This
152/// function does not validate that `tokens` is syntactically allowed —
153/// callers must have already passed it through the regular dispatcher.
154pub fn is_eval_safe_invocation(tokens: &[Token]) -> bool {
155    if tokens.is_empty() {
156        return false;
157    }
158    let cmd = tokens[0].command_name();
159    let Some(spec) = CUSTOM_REGISTRY.get(cmd).or_else(|| TOML_REGISTRY.get(cmd)) else {
160        return false;
161    };
162    is_eval_safe_for_spec(spec, tokens)
163}
164
165/// Spec-local variant used by tests so they can build a `CommandSpec`
166/// via `load_toml` and exercise the walker without touching the global
167/// `TOML_REGISTRY`.
168pub(crate) fn is_eval_safe_for_spec(spec: &CommandSpec, tokens: &[Token]) -> bool {
169    if tokens.is_empty() {
170        return false;
171    }
172    walk_to_eval_safe_leaf(
173        &tokens[1..],
174        &spec.kind,
175        spec.eval_safe,
176        &spec.eval_safe_flags,
177        &spec.eval_safe_flag_values,
178        &spec.eval_safe_required_flags,
179    )
180}
181
182fn walk_to_eval_safe_leaf(
183    remaining: &[Token],
184    kind: &DispatchKind,
185    eval_safe: bool,
186    eval_safe_flags: &[String],
187    eval_safe_flag_values: &std::collections::HashMap<String, Vec<String>>,
188    eval_safe_required_flags: &[String],
189) -> bool {
190    let subs_opt = match kind {
191        DispatchKind::Branching { subs, .. } | DispatchKind::Custom { subs, .. } => Some(subs),
192        _ => None,
193    };
194    if let Some(subs) = subs_opt
195        && let Some(arg) = remaining.first()
196        && let Some(sub) = subs.iter().find(|s| s.name == arg.as_str())
197    {
198        return walk_to_eval_safe_leaf(
199            &remaining[1..],
200            &sub.kind,
201            sub.eval_safe,
202            &sub.eval_safe_flags,
203            &sub.eval_safe_flag_values,
204            &sub.eval_safe_required_flags,
205        );
206    }
207    if !eval_safe {
208        return false;
209    }
210    let mut i = 0;
211    let mut seen_required = false;
212    while i < remaining.len() {
213        let s = remaining[i].as_str();
214        if !s.starts_with('-') {
215            i += 1;
216            continue;
217        }
218        let (bare, eq_value) = match s.split_once('=') {
219            Some((k, v)) => (k, Some(v)),
220            None => (s, None),
221        };
222        if !eval_safe_flags.iter().any(|f| f == bare) {
223            return false;
224        }
225        if eval_safe_required_flags.iter().any(|f| f == bare) {
226            seen_required = true;
227        }
228        if let Some(allowed) = eval_safe_flag_values.get(bare) {
229            // Valued flag declared in eval_safe_flag_values. The value
230            // arrives either as `--flag=VALUE` (eq_value is Some) or as
231            // the next token (`--flag VALUE`); either way the walker
232            // consumes it because a flag in eval_safe_flag_values is
233            // structurally valued.
234            //
235            // `allowed` empty = explicit-unrestricted: contributor
236            // vetted that any bare-literal value preserves shell-init
237            // output. Non-empty = value must appear in the allowlist.
238            let value: &str = if let Some(v) = eq_value {
239                v
240            } else if let Some(next) = remaining.get(i + 1) {
241                let v = next.as_str();
242                i += 1;
243                v
244            } else {
245                return false;
246            };
247            // Empty value is denied even under the explicit-
248            // unrestricted (`= []`) posture: `--flag=` and an empty
249            // following token never represent a meaningful tool
250            // argument. The bare-literal alphabet check is per-char
251            // and vacuously passes empty strings, so the walker
252            // has to reject explicitly.
253            if value.is_empty() {
254                return false;
255            }
256            if !allowed.is_empty() && !allowed.iter().any(|av| av == value) {
257                return false;
258            }
259        }
260        i += 1;
261    }
262    if !eval_safe_required_flags.is_empty() && !seen_required {
263        return false;
264    }
265    true
266}
267
268pub fn toml_command_docs() -> Vec<crate::docs::CommandDoc> {
269    TOML_REGISTRY
270        .iter()
271        .filter(|(key, spec)| *key == &spec.name)
272        .map(|(_, spec)| spec.to_command_doc())
273        .collect()
274}
275
276#[cfg(test)]
277mod tests;