env_hooks/
lib.rs

1pub mod shells;
2pub mod state;
3
4use std::{
5    collections::{BTreeMap, HashSet},
6    env, fs, num,
7    path::PathBuf,
8    process::ExitStatus,
9};
10
11use bstr::{B, BString, ByteSlice};
12use duct::cmd;
13use indexmap::IndexSet;
14use once_cell::sync::Lazy;
15use shell_quote::Bash;
16
17pub type EnvVars = BTreeMap<String, String>;
18pub type EnvVarsState = BTreeMap<String, Option<String>>;
19
20pub fn get_old_env_vars_to_be_updated(old_env_vars: EnvVars, new_env_vars: &EnvVars) -> EnvVars {
21    old_env_vars
22        .into_iter()
23        .fold(BTreeMap::new(), |mut acc, (key, value)| {
24            if new_env_vars.contains_key(&key) && new_env_vars.get(&key) != Some(&value) {
25                acc.insert(key, value);
26            }
27            acc
28        })
29}
30
31// TODO: Create new type patern for `EnvVars` and `EnvVarsState`?
32// impl From<EnvVars> for EnvVarsState {
33//   fn from(value: EnvVars) -> Self {
34//     value
35//       .into_iter()
36//       .map(|(key, value)| (key, Some(value)))
37//       .collect()
38//   }
39// }
40pub fn env_vars_state_from_env_vars(env_vars: EnvVars) -> EnvVarsState {
41    env_vars
42        .into_iter()
43        .map(|(key, value)| (key, Some(value)))
44        .collect()
45}
46
47pub fn get_env_vars_reset(
48    mut old_env_vars_that_were_updated: EnvVars,
49    new_env_vars: HashSet<String>,
50    env_state_var_key: String,
51) -> EnvVarsState {
52    let mut env_vars_state = new_env_vars
53        .into_iter()
54        .fold(EnvVarsState::new(), |mut acc, key| {
55            let value = old_env_vars_that_were_updated.remove(&key);
56            acc.insert(key, value);
57            acc
58        });
59    env_vars_state.insert(env_state_var_key, None);
60    env_vars_state
61}
62
63pub fn get_env_vars_from_current_process() -> EnvVars {
64    env::vars().collect::<EnvVars>()
65}
66
67pub enum BashSource {
68    File(PathBuf),
69    Script(BString),
70}
71
72impl AsRef<BashSource> for BashSource {
73    fn as_ref(&self) -> &BashSource {
74        self
75    }
76}
77
78impl BashSource {
79    fn to_command_string(&self) -> BString {
80        match &self {
81            Self::File(path) => bstr::join(" ", [B("source"), &Bash::quote_vec(path)]).into(),
82            Self::Script(script) => bstr::join(" ", [B("eval"), &Bash::quote_vec(script)]).into(),
83        }
84    }
85}
86
87pub(crate) trait SimplifiedExitOk {
88    fn simplified_exit_ok(&self) -> anyhow::Result<()>;
89}
90
91impl SimplifiedExitOk for ExitStatus {
92    /// Simplified implementation of <https://github.com/rust-lang/rust/issues/84908>
93    // TODO: Remove this and use `exit_ok` when it's stabilized.
94    fn simplified_exit_ok(&self) -> anyhow::Result<()> {
95        match num::NonZero::try_from(self.code().unwrap_or(-1)) {
96            Ok(_) => Err(anyhow::format_err!(
97                "process exited unsuccessfully: {}",
98                &self
99            )),
100            Err(_) => Ok(()),
101        }
102    }
103}
104
105pub fn get_env_vars_from_bash(
106    source: impl AsRef<BashSource>,
107    env_vars: Option<EnvVars>,
108) -> anyhow::Result<EnvVars> {
109    let bash_env_vars_file = tempfile::NamedTempFile::new()?;
110
111    let command_string = bstr::join(
112        " ",
113        [
114            &source.as_ref().to_command_string(),
115            B("&& env -0 >"),
116            &Bash::quote_vec(bash_env_vars_file.path()),
117        ],
118    );
119    let handle = cmd!("bash", "-c", command_string.to_os_str()?)
120        .full_env(env_vars.unwrap_or_default())
121        .stdout_to_stderr()
122        .start()?;
123    let output = handle.wait()?;
124    output
125        .status
126        .simplified_exit_ok()
127        .map_err(|e| anyhow::format_err!("Bash command to retrieve env vars failed:\n{e}"))?;
128
129    let bash_env_vars_string = fs::read_to_string(bash_env_vars_file.path())?;
130
131    let bash_env_vars = bash_env_vars_string
132        .split('\0')
133        .filter_map(|env_var| env_var.split_once('='))
134        .map(|(key, value)| (String::from(key), String::from(value)))
135        .collect::<EnvVars>();
136
137    Ok(bash_env_vars)
138}
139
140pub fn merge_delimited_env_var(
141    env_var: &str,
142    split_delimiter: char,
143    join_delimiter: char,
144    old_env_vars: &BTreeMap<String, String>,
145    new_env_vars: &mut BTreeMap<String, String>,
146) {
147    if let (Some(old_value), Some(new_value)) =
148        (old_env_vars.get(env_var), new_env_vars.get_mut(env_var))
149    {
150        *new_value = merge_delimited_values(split_delimiter, join_delimiter, old_value, new_value);
151    }
152}
153
154pub fn merge_delimited_values(
155    split_delimiter: char,
156    join_delimiter: char,
157    old_value: &str,
158    new_value: &str,
159) -> String {
160    new_value
161        .split(split_delimiter)
162        .chain(old_value.split(split_delimiter))
163        .collect::<IndexSet<_>>()
164        .into_iter()
165        .collect::<Vec<_>>()
166        .join(&join_delimiter.to_string())
167}
168
169const IGNORED_ENV_VAR_PREFIXES: &[&str] = &["__fish", "BASH_FUNC_"];
170
171static IGNORED_ENV_VAR_KEYS: Lazy<HashSet<&str>> = Lazy::new(|| {
172    HashSet::from([
173        // direnv env config
174        "DIRENV_CONFIG",
175        "DIRENV_BASH",
176        // should only be available inside of the .envrc or .env
177        "DIRENV_IN_ENVRC",
178        "COMP_WORDBREAKS", // Avoids segfaults in bash
179        "PS1",             // PS1 should not be exported, fixes problem in bash
180        // variables that should change freely
181        "OLDPWD",
182        "PWD",
183        "SHELL",
184        "SHELLOPTS",
185        "SHLVL",
186        "_",
187    ])
188});
189
190pub fn ignored_env_var_key(env_var_key: &str) -> bool {
191    for ignored_env_var_prefix in IGNORED_ENV_VAR_PREFIXES {
192        if env_var_key.starts_with(ignored_env_var_prefix) {
193            return true;
194        }
195    }
196    IGNORED_ENV_VAR_KEYS.contains(env_var_key)
197}
198
199pub fn remove_ignored_env_vars(env_vars: &mut EnvVars) {
200    let env_var_keys = env_vars.keys().cloned().collect::<Vec<_>>();
201    env_var_keys.into_iter().for_each(|env_var_key| {
202        if ignored_env_var_key(&env_var_key) {
203            env_vars.remove(&env_var_key);
204        }
205    });
206}
env_hooks/lib.rs

env_hooks/
lib.rs
