lean_ctx/core/
command_vars.rs1use std::collections::HashMap;
13
14use super::shell_allowlist::Separator;
15
16#[derive(Debug, Clone, Default)]
18pub struct KnownVars {
19 literals: HashMap<String, String>,
20 tmpdir_rebound: bool,
25}
26
27impl KnownVars {
28 fn set(&mut self, name: &str, value: Option<&str>) {
29 if let Some(value) = value {
30 self.literals.insert(name.to_string(), value.to_string());
31 } else {
32 self.literals.remove(name);
33 if name == "TMPDIR" {
34 self.tmpdir_rebound = true;
35 }
36 }
37 }
38
39 fn forget_all(&mut self) {
40 self.literals.clear();
41 self.tmpdir_rebound = true;
42 }
43
44 fn get(&self, name: &str, env_tmpdir: impl FnOnce() -> Option<String>) -> Option<String> {
47 if let Some(value) = self.literals.get(name) {
48 return Some(value.clone());
49 }
50 if name == "TMPDIR" && !self.tmpdir_rebound {
51 return env_tmpdir().filter(|v| !v.is_empty());
52 }
53 None
54 }
55}
56
57pub fn literal_assignments_before_each_segment(
68 command: &str,
69 segment_count: usize,
70) -> Vec<KnownVars> {
71 let separated = super::shell_allowlist::segments_with_separators(command);
72 if separated.len() != segment_count {
73 let nothing = KnownVars {
76 tmpdir_rebound: true,
77 ..KnownVars::default()
78 };
79 return vec![nothing; segment_count];
80 }
81 literal_assignments_before_each(&separated)
82}
83
84pub(crate) fn literal_assignments_before_each(
86 separated: &[(String, Option<Separator>)],
87) -> Vec<KnownVars> {
88 let mut known = KnownVars::default();
89 let mut prev: Option<Separator> = None;
90 let mut before_each = Vec::with_capacity(separated.len());
91 for (segment, sep) in separated {
92 let sep = *sep;
93 before_each.push(known.clone());
94 let tokens = super::shell_allowlist::shell_tokenize(segment);
95 let words: &[String] = match tokens.first().map(String::as_str) {
96 Some("export") => &tokens[1..],
97 _ => &tokens,
98 };
99 let assignments: Option<Vec<(&str, &str)>> = words
100 .iter()
101 .map(|word| {
102 let (name, value) = word.split_once('=')?;
103 let valid_name = name
104 .chars()
105 .next()
106 .is_some_and(|c| c.is_ascii_alphabetic() || c == '_')
107 && name.chars().all(|c| c.is_ascii_alphanumeric() || c == '_');
108 valid_name.then_some((name, value))
109 })
110 .collect();
111 match assignments {
112 Some(pairs) if !pairs.is_empty() => {
113 let in_current_shell = matches!(prev, None | Some(Separator::Sequence))
114 && matches!(sep, None | Some(Separator::Sequence | Separator::And));
115 for (name, value) in pairs {
116 let literal = !value.contains(['$', '`']) && !value.starts_with('~');
117 known.set(name, (in_current_shell && literal).then_some(value));
118 }
119 }
120 _ => {
121 let rebinds = matches!(
122 tokens.first().map(String::as_str),
123 Some(
124 "read"
125 | "for"
126 | "select"
127 | "eval"
128 | "source"
129 | "."
130 | "unset"
131 | "declare"
132 | "typeset"
133 | "local"
134 | "readonly"
135 | "export"
136 | "mapfile"
137 | "readarray"
138 | "getopts"
139 | "printf"
140 )
141 );
142 if rebinds {
143 known.forget_all();
144 }
145 }
146 }
147 prev = sep;
148 }
149 before_each
150}
151
152pub fn expand_leading_variable(path: &str, known: &KnownVars) -> Option<String> {
156 expand_with(path, known, || std::env::var("TMPDIR").ok())
157}
158
159fn expand_with(
160 path: &str,
161 known: &KnownVars,
162 env_tmpdir: impl FnOnce() -> Option<String>,
163) -> Option<String> {
164 let rest = path.strip_prefix('$')?;
165 let (name, tail) = if let Some(braced) = rest.strip_prefix('{') {
166 let (name, tail) = braced.split_once('}')?;
167 (name, tail)
168 } else {
169 let end = rest
170 .find(|c: char| !(c.is_ascii_alphanumeric() || c == '_'))
171 .unwrap_or(rest.len());
172 rest.split_at(end)
173 };
174 if name.is_empty() || !(tail.is_empty() || tail.starts_with('/')) {
175 return None;
176 }
177 Some(format!("{}{tail}", known.get(name, env_tmpdir)?))
178}
179
180#[cfg(test)]
181mod tests {
182 use super::*;
183
184 fn before_last(command: &str) -> KnownVars {
185 let separated = crate::core::shell_allowlist::segments_with_separators(command);
186 literal_assignments_before_each(&separated)
187 .pop()
188 .expect("at least one segment")
189 }
190
191 #[test]
192 fn a_literal_assignment_in_the_current_shell_is_known() {
193 let known = before_last("D=/tmp/x && ls");
194 assert_eq!(
195 expand_leading_variable("$D/f", &known).as_deref(),
196 Some("/tmp/x/f")
197 );
198 assert_eq!(
199 expand_leading_variable("${D}/f", &known).as_deref(),
200 Some("/tmp/x/f")
201 );
202 assert_eq!(expand_leading_variable("$Dx/f", &known), None);
203 }
204
205 fn expand_tmpdir(command: &str) -> Option<String> {
206 expand_with("$TMPDIR/a", &before_last(command), || {
207 Some("/var/folders/xy/T".to_string())
208 })
209 }
210
211 #[test]
213 fn tmpdir_is_known_only_while_the_command_has_not_rebound_it() {
214 assert_eq!(expand_tmpdir("ls").as_deref(), Some("/var/folders/xy/T/a"));
215 assert_eq!(
216 expand_tmpdir("TMPDIR=/tmp/own; ls").as_deref(),
217 Some("/tmp/own/a"),
218 "a literal rebinding is followed"
219 );
220 for command in [
221 "TMPDIR=$PWD; ls",
222 "export TMPDIR=$(pwd); ls",
223 "true && TMPDIR=/tmp/maybe; ls",
224 "read TMPDIR; ls",
225 "unset TMPDIR; ls",
226 ] {
227 assert_eq!(expand_tmpdir(command), None, "{command}");
228 }
229 }
230}