diff --git a/crates/nu-cli/src/data/config/conf.rs b/crates/nu-cli/src/data/config/conf.rs index b3504441e2..1f53f0c7a2 100644 --- a/crates/nu-cli/src/data/config/conf.rs +++ b/crates/nu-cli/src/data/config/conf.rs @@ -4,7 +4,7 @@ use std::fmt::Debug; pub trait Conf: Debug + Send { fn env(&self) -> Option; fn path(&self) -> Option; - + fn nu_env_dirs(&self) -> Option; fn reload(&self); } @@ -13,6 +13,10 @@ impl Conf for Box { (**self).env() } + fn nu_env_dirs(&self) -> Option { + (**self).nu_env_dirs() + } + fn path(&self) -> Option { (**self).path() } diff --git a/crates/nu-cli/src/data/config/nuconfig.rs b/crates/nu-cli/src/data/config/nuconfig.rs index 3444e1f7a9..54001026a2 100644 --- a/crates/nu-cli/src/data/config/nuconfig.rs +++ b/crates/nu-cli/src/data/config/nuconfig.rs @@ -20,6 +20,10 @@ impl Conf for NuConfig { self.path() } + fn nu_env_dirs(&self) -> Option { + self.nu_env_dirs() + } + fn reload(&self) { let mut vars = self.vars.lock(); @@ -52,6 +56,14 @@ impl NuConfig { None } + pub fn nu_env_dirs(&self) -> Option { + let vars = self.vars.lock(); + if let Some(dirs) = vars.get("nu_env_dirs") { + return Some(dirs.clone()); + } + None + } + pub fn path(&self) -> Option { let vars = self.vars.lock(); diff --git a/crates/nu-cli/src/data/config/tests.rs b/crates/nu-cli/src/data/config/tests.rs index ecbe0cb7f7..ea8712fe0f 100644 --- a/crates/nu-cli/src/data/config/tests.rs +++ b/crates/nu-cli/src/data/config/tests.rs @@ -16,6 +16,10 @@ impl Conf for FakeConfig { self.config.env() } + fn nu_env_dirs(&self) -> Option { + None + } + fn path(&self) -> Option { self.config.path() } diff --git a/crates/nu-cli/src/env.rs b/crates/nu-cli/src/env.rs index 3f1dd1c311..74ae277ab5 100644 --- a/crates/nu-cli/src/env.rs +++ b/crates/nu-cli/src/env.rs @@ -1,3 +1,4 @@ +pub(crate) mod directory_specific_environment; pub(crate) mod environment; pub(crate) mod environment_syncer; pub(crate) mod host; diff --git a/crates/nu-cli/src/env/directory_specific_environment.rs b/crates/nu-cli/src/env/directory_specific_environment.rs new file mode 100644 index 0000000000..ca201aed0f --- /dev/null +++ b/crates/nu-cli/src/env/directory_specific_environment.rs @@ -0,0 +1,195 @@ +use indexmap::IndexMap; +use nu_protocol::{Primitive, UntaggedValue, Value}; +use std::io::{Error, ErrorKind, Result}; +use std::{ffi::OsString, fmt::Debug, path::PathBuf}; + +#[derive(Debug, Default)] +pub struct DirectorySpecificEnvironment { + allowed_directories: Vec, + + //Directory -> Env key. If an environment var has been added from a .nu in a directory, we track it here so we can remove it when the user leaves the directory. + added_env_vars: IndexMap>, + + //Directory -> (env_key, value). If a .nu overwrites some existing environment variables, they are added here so that they can be restored later. + overwritten_env_values: IndexMap>, +} + +impl DirectorySpecificEnvironment { + pub fn new(allowed_directories: Option) -> DirectorySpecificEnvironment { + let mut allowed_directories = if let Some(Value { + value: UntaggedValue::Table(ref wrapped_directories), + tag: _, + }) = allowed_directories + { + wrapped_directories + .iter() + .filter_map(|dirval| { + if let Value { + value: UntaggedValue::Primitive(Primitive::String(ref dir)), + tag: _, + } = dirval + { + return Some(PathBuf::from(&dir)); + } + None + }) + .collect() + } else { + vec![] + }; + allowed_directories.sort(); + + DirectorySpecificEnvironment { + allowed_directories, + added_env_vars: IndexMap::new(), + overwritten_env_values: IndexMap::new(), + } + } + + //If we are no longer in a directory, we restore the values it overwrote. + pub fn overwritten_values_to_restore(&mut self) -> Result> { + let current_dir = std::env::current_dir()?; + + let mut keyvals_to_restore = IndexMap::new(); + let mut new_overwritten = IndexMap::new(); + + for (directory, keyvals) in &self.overwritten_env_values { + let mut working_dir = Some(current_dir.as_path()); + + let mut re_add_keyvals = true; + while let Some(wdir) = working_dir { + if wdir == directory.as_path() { + re_add_keyvals = false; + new_overwritten.insert(directory.clone(), keyvals.clone()); + break; + } else { + working_dir = working_dir //Keep going up in the directory structure with .parent() + .ok_or_else(|| { + Error::new(ErrorKind::NotFound, "Root directory has no parent") + })? + .parent(); + } + } + if re_add_keyvals { + for (k, v) in keyvals { + keyvals_to_restore.insert( + k.clone(), + v.to_str() + .ok_or_else(|| { + Error::new( + ErrorKind::Other, + format!("{:?} is not valid unicode", v), + ) + })? + .to_string(), + ); + } + } + } + + self.overwritten_env_values = new_overwritten; + Ok(keyvals_to_restore) + } + + pub fn env_vars_to_add(&mut self) -> Result> { + let current_dir = std::env::current_dir()?; + + let mut vars_to_add = IndexMap::new(); + for dir in &self.allowed_directories { + let mut working_dir = Some(current_dir.as_path()); + + //Start in the current directory, then traverse towards the root with working_dir to see if we are in a subdirectory of a valid directory. + while let Some(wdir) = working_dir { + if wdir == dir.as_path() { + let toml_doc = std::fs::read_to_string(wdir.join(".nu").as_path())? + .parse::()?; + + let vars_in_current_file = toml_doc + .get("env") + .ok_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + "No [env] section in .nu-file", + ) + })? + .as_table() + .ok_or_else(|| { + Error::new( + ErrorKind::InvalidData, + "Malformed [env] section in .nu-file", + ) + })?; + + let mut keys_in_current_nufile = vec![]; + for (k, v) in vars_in_current_file { + vars_to_add.insert( + k.clone(), + v.as_str() + .ok_or_else(|| { + Error::new( + ErrorKind::InvalidData, + format!("Could not read environment variable: {}", v), + ) + })? + .to_string(), + ); //This is used to add variables to the environment + keys_in_current_nufile.push(k.clone()); //this is used to keep track of which directory added which variables + } + + //If we are about to overwrite any environment variables, we save them first so they can be restored later. + self.overwritten_env_values.insert( + wdir.to_path_buf(), + keys_in_current_nufile + .iter() + .filter_map(|key| { + if let Some(val) = std::env::var_os(key) { + return Some((key.clone(), val)); + } + None + }) + .collect(), + ); + + self.added_env_vars + .insert(wdir.to_path_buf(), keys_in_current_nufile); + break; + } else { + working_dir = working_dir //Keep going up in the directory structure with .parent() + .ok_or_else(|| { + Error::new(ErrorKind::NotFound, "Root directory has no parent") + })? + .parent(); + } + } + } + + Ok(vars_to_add) + } + + //If the user has left directories which added env vars through .nu, we clear those vars + pub fn env_vars_to_delete(&mut self) -> Result> { + let current_dir = std::env::current_dir()?; + + //Gather up all environment variables that should be deleted. + //If we are not in a directory or one of its subdirectories, mark the env_vals it maps to for removal. + let vars_to_delete = self.added_env_vars.iter().fold( + Vec::new(), + |mut vars_to_delete, (directory, env_vars)| { + let mut working_dir = Some(current_dir.as_path()); + + while let Some(wdir) = working_dir { + if wdir == directory { + return vars_to_delete; + } else { + working_dir = working_dir.expect("Root directory has no parent").parent(); + } + } + //only delete vars from directories we are not in + vars_to_delete.extend(env_vars.clone()); + vars_to_delete + }, + ); + + Ok(vars_to_delete) + } +} diff --git a/crates/nu-cli/src/env/environment.rs b/crates/nu-cli/src/env/environment.rs index 9b586eded2..41e3270a61 100644 --- a/crates/nu-cli/src/env/environment.rs +++ b/crates/nu-cli/src/env/environment.rs @@ -1,4 +1,5 @@ use crate::data::config::Conf; +use crate::env::directory_specific_environment::*; use indexmap::{indexmap, IndexSet}; use nu_protocol::{UntaggedValue, Value}; use std::ffi::OsString; @@ -8,7 +9,7 @@ pub trait Env: Debug + Send { fn env(&self) -> Option; fn path(&self) -> Option; - fn add_env(&mut self, key: &str, value: &str); + fn add_env(&mut self, key: &str, value: &str, overwrite_existing: bool); fn add_path(&mut self, new_path: OsString); } @@ -21,8 +22,8 @@ impl Env for Box { (**self).path() } - fn add_env(&mut self, key: &str, value: &str) { - (**self).add_env(key, value); + fn add_env(&mut self, key: &str, value: &str, overwrite_existing: bool) { + (**self).add_env(key, value, overwrite_existing); } fn add_path(&mut self, new_path: OsString) { @@ -34,6 +35,7 @@ impl Env for Box { pub struct Environment { environment_vars: Option, path_vars: Option, + pub direnv: DirectorySpecificEnvironment, } impl Environment { @@ -41,6 +43,7 @@ impl Environment { Environment { environment_vars: None, path_vars: None, + direnv: DirectorySpecificEnvironment::new(None), } } @@ -51,9 +54,38 @@ impl Environment { Environment { environment_vars: env, path_vars: path, + direnv: DirectorySpecificEnvironment::new(configuration.nu_env_dirs()), } } + pub fn maintain_directory_environment(&mut self) -> std::io::Result<()> { + self.direnv.env_vars_to_delete()?.iter().for_each(|k| { + self.remove_env(&k); + }); + self.direnv.env_vars_to_add()?.iter().for_each(|(k, v)| { + self.add_env(&k, &v, true); + }); + + self.direnv + .overwritten_values_to_restore()? + .iter() + .for_each(|(k, v)| { + self.add_env(&k, &v, true); + }); + + Ok(()) + } + + fn remove_env(&mut self, key: &str) { + if let Some(Value { + value: UntaggedValue::Row(ref mut envs), + tag: _, + }) = self.environment_vars + { + envs.entries.remove(key); + }; + } + pub fn morph(&mut self, configuration: &T) { self.environment_vars = configuration.env(); self.path_vars = configuration.path(); @@ -77,7 +109,7 @@ impl Env for Environment { None } - fn add_env(&mut self, key: &str, value: &str) { + fn add_env(&mut self, key: &str, value: &str, overwrite_existing: bool) { let value = UntaggedValue::string(value); let new_envs = { @@ -88,7 +120,7 @@ impl Env for Environment { { let mut new_envs = envs.clone(); - if !new_envs.contains_key(key) { + if !new_envs.contains_key(key) || overwrite_existing { new_envs.insert_data_at_key(key, value.into_value(tag.clone())); } @@ -206,7 +238,7 @@ mod tests { let fake_config = FakeConfig::new(&file); let mut actual = Environment::from_config(&fake_config); - actual.add_env("USER", "NUNO"); + actual.add_env("USER", "NUNO", false); assert_eq!( actual.env(), @@ -239,7 +271,7 @@ mod tests { let fake_config = FakeConfig::new(&file); let mut actual = Environment::from_config(&fake_config); - actual.add_env("SHELL", "/usr/bin/sh"); + actual.add_env("SHELL", "/usr/bin/sh", false); assert_eq!( actual.env(), diff --git a/crates/nu-cli/src/env/environment_syncer.rs b/crates/nu-cli/src/env/environment_syncer.rs index 852ccfa75c..fbcd72e6a6 100644 --- a/crates/nu-cli/src/env/environment_syncer.rs +++ b/crates/nu-cli/src/env/environment_syncer.rs @@ -49,7 +49,9 @@ impl EnvironmentSyncer { if name != "path" && name != "PATH" { // account for new env vars present in the current session // that aren't loaded from config. - environment.add_env(&name, &value); + environment.add_env(&name, &value, false); + + environment.maintain_directory_environment().ok(); // clear the env var from the session // we are about to replace them