env_hooks/
lib.rs

1pub mod shells;
2pub mod state;
3
4use std::{
5    collections::HashSet,
6    env, fs, num,
7    ops::{Deref, DerefMut},
8    path::PathBuf,
9    process::ExitStatus,
10};
11
12use bstr::{B, BString, ByteSlice};
13use duct::cmd;
14use indexmap::{IndexMap, IndexSet, map::IntoIter};
15use once_cell::sync::Lazy;
16use serde::{Deserialize, Serialize};
17use shell_quote::Bash;
18
19type EnvVarsInner = IndexMap<String, String>;
20
21#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
22pub struct EnvVars(EnvVarsInner);
23
24impl EnvVars {
25    pub fn new() -> Self {
26        Self::default()
27    }
28}
29
30impl Deref for EnvVars {
31    type Target = EnvVarsInner;
32    fn deref(&self) -> &Self::Target {
33        &self.0
34    }
35}
36
37impl DerefMut for EnvVars {
38    fn deref_mut(&mut self) -> &mut Self::Target {
39        &mut self.0
40    }
41}
42
43impl IntoIterator for EnvVars {
44    type Item = (String, String);
45    type IntoIter = EnvVarIntoIter<String, String>;
46
47    fn into_iter(self) -> Self::IntoIter {
48        EnvVarIntoIter(self.0.into_iter())
49    }
50}
51
52impl FromIterator<(String, String)> for EnvVars {
53    fn from_iter<I: IntoIterator<Item = (String, String)>>(iter: I) -> Self {
54        EnvVars(EnvVarsInner::from_iter(iter))
55    }
56}
57
58impl From<EnvVars> for EnvVarsState {
59    fn from(value: EnvVars) -> Self {
60        EnvVarsState(value.0.into_iter().map(|(k, v)| (k, Some(v))).collect())
61    }
62}
63
64type EnvVarsStateInner = IndexMap<String, Option<String>>;
65
66#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
67pub struct EnvVarsState(EnvVarsStateInner);
68
69impl EnvVarsState {
70    pub fn new() -> Self {
71        Self::default()
72    }
73}
74
75impl Deref for EnvVarsState {
76    type Target = EnvVarsStateInner;
77    fn deref(&self) -> &Self::Target {
78        &self.0
79    }
80}
81
82impl DerefMut for EnvVarsState {
83    fn deref_mut(&mut self) -> &mut Self::Target {
84        &mut self.0
85    }
86}
87
88impl IntoIterator for EnvVarsState {
89    type Item = (String, Option<String>);
90    type IntoIter = EnvVarIntoIter<String, Option<String>>;
91
92    fn into_iter(self) -> Self::IntoIter {
93        EnvVarIntoIter(self.0.into_iter())
94    }
95}
96
97impl FromIterator<(String, Option<String>)> for EnvVarsState {
98    fn from_iter<I: IntoIterator<Item = (String, Option<String>)>>(iter: I) -> Self {
99        EnvVarsState(EnvVarsStateInner::from_iter(iter))
100    }
101}
102
103type EnvVarInner<K, V> = IntoIter<K, V>;
104
105pub struct EnvVarIntoIter<K, V>(EnvVarInner<K, V>);
106
107impl<K, V> Deref for EnvVarIntoIter<K, V> {
108    type Target = EnvVarInner<K, V>;
109    fn deref(&self) -> &Self::Target {
110        &self.0
111    }
112}
113
114impl<K, V> DerefMut for EnvVarIntoIter<K, V> {
115    fn deref_mut(&mut self) -> &mut Self::Target {
116        &mut self.0
117    }
118}
119
120impl<K, V> Iterator for EnvVarIntoIter<K, V> {
121    type Item = (K, V);
122    fn next(&mut self) -> Option<Self::Item> {
123        self.0.next()
124    }
125}
126
127pub fn get_old_env_vars_to_be_updated(old_env_vars: EnvVars, new_env_vars: &EnvVars) -> EnvVars {
128    old_env_vars
129        .into_iter()
130        .fold(EnvVars::new(), |mut acc, (key, value)| {
131            if new_env_vars.contains_key(&key) && new_env_vars.get(&key) != Some(&value) {
132                acc.insert(key, value);
133            }
134            acc
135        })
136}
137
138pub fn get_env_vars_reset(
139    mut old_env_vars_that_were_updated: EnvVars,
140    new_env_vars: HashSet<String>,
141    env_state_var_key: String,
142) -> EnvVarsState {
143    let mut env_vars_state = new_env_vars
144        .into_iter()
145        .fold(EnvVarsState::new(), |mut acc, key| {
146            let value = old_env_vars_that_were_updated.shift_remove(&key);
147            acc.insert(key, value);
148            acc
149        });
150    env_vars_state.insert(env_state_var_key, None);
151    env_vars_state
152}
153
154pub fn get_env_vars_from_current_process() -> EnvVars {
155    EnvVars(env::vars().collect::<EnvVarsInner>())
156}
157
158pub enum BashSource {
159    File(PathBuf),
160    Script(BString),
161}
162
163impl AsRef<BashSource> for BashSource {
164    fn as_ref(&self) -> &BashSource {
165        self
166    }
167}
168
169impl BashSource {
170    fn to_command_string(&self) -> BString {
171        match &self {
172            Self::File(path) => bstr::join(" ", [B("source"), &Bash::quote_vec(path)]).into(),
173            Self::Script(script) => bstr::join(" ", [B("eval"), &Bash::quote_vec(script)]).into(),
174        }
175    }
176}
177
178pub(crate) trait SimplifiedExitOk {
179    fn simplified_exit_ok(&self) -> anyhow::Result<()>;
180}
181
182impl SimplifiedExitOk for ExitStatus {
183    /// Simplified implementation of <https://github.com/rust-lang/rust/issues/84908>
184    // TODO: Remove this and use `exit_ok` when it's stabilized.
185    fn simplified_exit_ok(&self) -> anyhow::Result<()> {
186        match num::NonZero::try_from(self.code().unwrap_or(-1)) {
187            Ok(_) => Err(anyhow::format_err!(
188                "process exited unsuccessfully: {}",
189                &self
190            )),
191            Err(_) => Ok(()),
192        }
193    }
194}
195
196pub fn get_env_vars_from_bash(
197    source: impl AsRef<BashSource>,
198    env_vars: Option<EnvVars>,
199) -> anyhow::Result<EnvVars> {
200    let bash_env_vars_file = tempfile::NamedTempFile::new()?;
201
202    let command_string = bstr::join(
203        " ",
204        [
205            &source.as_ref().to_command_string(),
206            B("&& env -0 >"),
207            &Bash::quote_vec(bash_env_vars_file.path()),
208        ],
209    );
210    let handle = cmd!("bash", "-c", command_string.to_os_str()?)
211        .full_env(env_vars.unwrap_or_default())
212        .stdout_to_stderr()
213        .start()?;
214    let output = handle.wait()?;
215    output
216        .status
217        .simplified_exit_ok()
218        .map_err(|e| anyhow::format_err!("Bash command to retrieve env vars failed:\n{e}"))?;
219
220    let bash_env_vars_string = fs::read_to_string(bash_env_vars_file.path())?;
221
222    let bash_env_vars = EnvVars(
223        bash_env_vars_string
224            .split('\0')
225            .filter_map(|env_var| env_var.split_once('='))
226            .map(|(key, value)| (String::from(key), String::from(value)))
227            .collect::<EnvVarsInner>(),
228    );
229
230    Ok(bash_env_vars)
231}
232
233pub fn merge_delimited_env_var(
234    env_var: &str,
235    split_delimiter: char,
236    join_delimiter: char,
237    old_env_vars: &EnvVars,
238    new_env_vars: &mut EnvVars,
239) {
240    if let (Some(old_value), Some(new_value)) =
241        (old_env_vars.get(env_var), new_env_vars.get_mut(env_var))
242    {
243        *new_value = merge_delimited_values(split_delimiter, join_delimiter, old_value, new_value);
244    }
245}
246
247pub fn merge_delimited_values(
248    split_delimiter: char,
249    join_delimiter: char,
250    old_value: &str,
251    new_value: &str,
252) -> String {
253    new_value
254        .split(split_delimiter)
255        .chain(old_value.split(split_delimiter))
256        .collect::<IndexSet<_>>()
257        .into_iter()
258        .collect::<Vec<_>>()
259        .join(&join_delimiter.to_string())
260}
261
262const IGNORED_ENV_VAR_PREFIXES: &[&str] = &["__fish", "BASH_FUNC_"];
263
264static IGNORED_ENV_VAR_KEYS: Lazy<HashSet<&str>> = Lazy::new(|| {
265    HashSet::from([
266        // direnv env config
267        "DIRENV_CONFIG",
268        "DIRENV_BASH",
269        // should only be available inside of the .envrc or .env
270        "DIRENV_IN_ENVRC",
271        "COMP_WORDBREAKS", // Avoids segfaults in bash
272        "PS1",             // PS1 should not be exported, fixes problem in bash
273        // variables that should change freely
274        "OLDPWD",
275        "PWD",
276        "SHELL",
277        "SHELLOPTS",
278        "SHLVL",
279        "_",
280        "__CF_USER_TEXT_ENCODING",
281    ])
282});
283
284pub fn ignored_env_var_key(env_var_key: &str) -> bool {
285    for ignored_env_var_prefix in IGNORED_ENV_VAR_PREFIXES {
286        if env_var_key.starts_with(ignored_env_var_prefix) {
287            return true;
288        }
289    }
290    IGNORED_ENV_VAR_KEYS.contains(env_var_key)
291}
292
293pub fn remove_ignored_env_vars(env_vars: &mut EnvVars) {
294    let env_var_keys = env_vars.keys().cloned().collect::<Vec<_>>();
295    env_var_keys.into_iter().for_each(|env_var_key| {
296        if ignored_env_var_key(&env_var_key) {
297            env_vars.shift_remove(&env_var_key);
298        }
299    });
300}
env_hooks/lib.rs

env_hooks/
lib.rs
