1use anyhow::{Context as _, Result};
2use collections::HashMap;
3use std::borrow::Cow;
4use std::ffi::OsStr;
5use std::io::Read;
6use std::path::{Path, PathBuf};
7use std::process::Command;
8use tempfile::NamedTempFile;
9
10/// Capture all environment variables from the login shell.
11pub fn capture(change_dir: Option<impl AsRef<Path>>) -> Result<HashMap<String, String>> {
12 let shell_path = std::env::var("SHELL").map(PathBuf::from)?;
13 let shell_name = shell_path.file_name().and_then(OsStr::to_str);
14
15 let mut command_string = String::new();
16
17 // What we're doing here is to spawn a shell and then `cd` into
18 // the project directory to get the env in there as if the user
19 // `cd`'d into it. We do that because tools like direnv, asdf, ...
20 // hook into `cd` and only set up the env after that.
21 if let Some(dir) = change_dir {
22 let dir_str = dir.as_ref().to_string_lossy();
23 command_string.push_str(&format!("cd '{dir_str}';"));
24 }
25
26 // In certain shells we need to execute additional_command in order to
27 // trigger the behavior of direnv, etc.
28 command_string.push_str(match shell_name {
29 Some("fish") => "emit fish_prompt;",
30 _ => "",
31 });
32
33 let mut env_output_file = NamedTempFile::new()?;
34 command_string.push_str(&format!(
35 "sh -c 'export -p' > '{}';",
36 env_output_file.path().to_string_lossy(),
37 ));
38
39 let mut command = Command::new(&shell_path);
40
41 // For csh/tcsh, the login shell option is set by passing `-` as
42 // the 0th argument instead of using `-l`.
43 if let Some("tcsh" | "csh") = shell_name {
44 #[cfg(unix)]
45 std::os::unix::process::CommandExt::arg0(&mut command, "-");
46 } else {
47 command.arg("-l");
48 }
49
50 command.args(["-i", "-c", &command_string]);
51
52 let process_output = super::set_pre_exec_to_start_new_session(&mut command).output()?;
53 anyhow::ensure!(
54 process_output.status.success(),
55 "login shell exited with {}. stdout: {:?}, stderr: {:?}",
56 process_output.status,
57 String::from_utf8_lossy(&process_output.stdout),
58 String::from_utf8_lossy(&process_output.stderr),
59 );
60
61 let mut env_output = String::new();
62 env_output_file.read_to_string(&mut env_output)?;
63
64 parse(&env_output)
65 .filter_map(|entry| match entry {
66 Ok((name, value)) => Some(Ok((name.into(), value?.into()))),
67 Err(err) => Some(Err(err)),
68 })
69 .collect::<Result<HashMap<String, String>>>()
70}
71
72/// Parse the result of calling `sh -c 'export -p'`.
73///
74/// https://www.man7.org/linux/man-pages/man1/export.1p.html
75fn parse(mut input: &str) -> impl Iterator<Item = Result<(Cow<'_, str>, Option<Cow<'_, str>>)>> {
76 std::iter::from_fn(move || {
77 if input.is_empty() {
78 return None;
79 }
80 match parse_declaration(input) {
81 Ok((entry, rest)) => {
82 input = rest;
83 Some(Ok(entry))
84 }
85 Err(err) => Some(Err(err)),
86 }
87 })
88}
89
90fn parse_declaration(input: &str) -> Result<((Cow<'_, str>, Option<Cow<'_, str>>), &str)> {
91 let rest = input
92 .strip_prefix("export ")
93 .context("expected 'export ' prefix")?;
94
95 if let Some((name, rest)) = parse_name_and_terminator(rest, '\n') {
96 Ok(((name, None), rest))
97 } else {
98 let (name, rest) = parse_name_and_terminator(rest, '=').context("invalid name")?;
99 let (value, rest) = parse_literal_and_terminator(rest, '\n').context("invalid value")?;
100 Ok(((name, Some(value)), rest))
101 }
102}
103
104fn parse_name_and_terminator(input: &str, terminator: char) -> Option<(Cow<'_, str>, &str)> {
105 let (name, rest) = parse_literal_and_terminator(input, terminator)?;
106 (!name.is_empty() && !name.contains('=')).then_some((name, rest))
107}
108
109fn parse_literal_and_terminator(input: &str, terminator: char) -> Option<(Cow<'_, str>, &str)> {
110 if let Some((literal, rest)) = parse_literal_single_quoted(input) {
111 let rest = rest.strip_prefix(terminator)?;
112 Some((Cow::Borrowed(literal), rest))
113 } else if let Some((literal, rest)) = parse_literal_double_quoted(input) {
114 let rest = rest.strip_prefix(terminator)?;
115 Some((Cow::Owned(literal), rest))
116 } else {
117 let (literal, rest) = input.split_once(terminator)?;
118 (!literal.contains(|c: char| c.is_ascii_whitespace()))
119 .then_some((Cow::Borrowed(literal), rest))
120 }
121}
122
123/// https://www.gnu.org/software/bash/manual/html_node/Single-Quotes.html
124fn parse_literal_single_quoted(input: &str) -> Option<(&str, &str)> {
125 input.strip_prefix('\'')?.split_once('\'')
126}
127
128/// https://www.gnu.org/software/bash/manual/html_node/Double-Quotes.html
129fn parse_literal_double_quoted(input: &str) -> Option<(String, &str)> {
130 let rest = input.strip_prefix('"')?;
131
132 let mut char_indices = rest.char_indices();
133 let mut escaping = false;
134 let (literal, rest) = loop {
135 let (index, char) = char_indices.next()?;
136 if char == '"' && !escaping {
137 break (&rest[..index], &rest[index + 1..]);
138 } else {
139 escaping = !escaping && char == '\\';
140 }
141 };
142
143 let literal = literal
144 .replace("\\$", "$")
145 .replace("\\`", "`")
146 .replace("\\\"", "\"")
147 .replace("\\\n", "")
148 .replace("\\\\", "\\");
149
150 Some((literal, rest))
151}
152
153#[cfg(test)]
154mod tests {
155 use super::*;
156
157 #[test]
158 fn test_parse() {
159 let input = indoc::indoc! {r#"
160 export foo
161 export 'foo'
162 export "foo"
163 export foo=
164 export 'foo'=
165 export "foo"=
166 export foo=bar
167 export foo='bar'
168 export foo="bar"
169 export foo='b
170 a
171 z'
172 export foo="b
173 a
174 z"
175 export foo='b\
176 a\
177 z'
178 export foo="b\
179 a\
180 z"
181 export foo='\`Hello\`
182 \"wo\
183 rld\"\n!\\
184 !'
185 export foo="\`Hello\`
186 \"wo\
187 rld\"\n!\\
188 !"
189 "#};
190
191 let expected_values = [
192 None,
193 None,
194 None,
195 Some(""),
196 Some(""),
197 Some(""),
198 Some("bar"),
199 Some("bar"),
200 Some("bar"),
201 Some("b\na\nz"),
202 Some("b\na\nz"),
203 Some("b\\\na\\\nz"),
204 Some("baz"),
205 Some(indoc::indoc! {r#"
206 \`Hello\`
207 \"wo\
208 rld\"\n!\\
209 !"#}),
210 Some(indoc::indoc! {r#"
211 `Hello`
212 "world"\n!\!"#}),
213 ];
214 let expected = expected_values
215 .into_iter()
216 .map(|value| ("foo".into(), value.map(Into::into)))
217 .collect::<Vec<_>>();
218
219 let actual = parse(input).collect::<Result<Vec<_>>>().unwrap();
220 assert_eq!(expected, actual);
221 }
222
223 #[test]
224 fn test_parse_declaration() {
225 let ((name, value), rest) = parse_declaration("export foo\nrest").unwrap();
226 assert_eq!(name, "foo");
227 assert_eq!(value, None);
228 assert_eq!(rest, "rest");
229
230 let ((name, value), rest) = parse_declaration("export foo=bar\nrest").unwrap();
231 assert_eq!(name, "foo");
232 assert_eq!(value.as_deref(), Some("bar"));
233 assert_eq!(rest, "rest");
234 }
235
236 #[test]
237 fn test_parse_literal_single_quoted() {
238 let input = indoc::indoc! {r#"
239 '\`Hello\`
240 \"wo\
241 rld\"\n!\\
242 !'
243 rest"#};
244
245 let expected = indoc::indoc! {r#"
246 \`Hello\`
247 \"wo\
248 rld\"\n!\\
249 !"#};
250
251 let (actual, rest) = parse_literal_single_quoted(input).unwrap();
252 assert_eq!(expected, actual);
253 assert_eq!(rest, "\nrest");
254 }
255
256 #[test]
257 fn test_parse_literal_double_quoted() {
258 let input = indoc::indoc! {r#"
259 "\`Hello\`
260 \"wo\
261 rld\"\n!\\
262 !"
263 rest"#};
264
265 let expected = indoc::indoc! {r#"
266 `Hello`
267 "world"\n!\!"#};
268
269 let (actual, rest) = parse_literal_double_quoted(input).unwrap();
270 assert_eq!(expected, actual);
271 assert_eq!(rest, "\nrest");
272 }
273}