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