WIP: Per directory env-variables (#1943)

* Add args in .nurc file to environment

* Working dummy version

* Add add_nurc to sync_env command

* Parse .nurc file

* Delete env vars after leaving directory

* Removing vals not working, strangely

* Refactoring, add comment

* Debugging

* Debug by logging to file

* Add and remove env var behavior appears correct

However, it does not use existing code that well.

* Move work to cli.rs

* Parse config directories

* I am in a state of distress

* Rename .nurc to .nu

* Some notes for me

* Refactoring

* Removing vars works, but not done in a very nice fashion

* Refactor env_vars_to_delete

* Refactor env_vars_to_add()

* Move directory environment code to separate file

* Refactor from_config

* Restore env values

* Working?

* Working?

* Update comments and change var name

* Formatting

* Remove vars after leaving dir

* Remove notes I made

* Rename config function

* Clippy

* Cleanup and handle errors

* cargo fmt

* Better error messages, remove last (?) unwrap

* FORMAT PLZ

* Rename whitelisted_directories to allowed_directories

* Add comment to clarify how overwritten values are restored.

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<Value>;
     fn path(&self) -> Option<Value>;
-
+    fn nu_env_dirs(&self) -> Option<Value>;
     fn reload(&self);
 }
 
@@ -13,6 +13,10 @@ impl Conf for Box<dyn Conf> {
         (**self).env()
     }
 
+    fn nu_env_dirs(&self) -> Option<Value> {
+        (**self).nu_env_dirs()
+    }
+
     fn path(&self) -> Option<Value> {
         (**self).path()
     }

crates/nu-cli/src/data/config/nuconfig.rs

@@ -20,6 +20,10 @@ impl Conf for NuConfig {
         self.path()
     }
 
+    fn nu_env_dirs(&self) -> Option<Value> {
+        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<Value> {
+        let vars = self.vars.lock();
+        if let Some(dirs) = vars.get("nu_env_dirs") {
+            return Some(dirs.clone());
+        }
+        None
+    }
+
     pub fn path(&self) -> Option<Value> {
         let vars = self.vars.lock();
 

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<Value> {
+        None
+    }
+
     fn path(&self) -> Option<Value> {
         self.config.path()
     }

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;

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<PathBuf>,
+
+    //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<PathBuf, Vec<String>>,
+
+    //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<PathBuf, Vec<(String, OsString)>>,
+}
+
+impl DirectorySpecificEnvironment {
+    pub fn new(allowed_directories: Option<Value>) -> 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<IndexMap<String, String>> {
+        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<IndexMap<String, String>> {
+        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::<toml::Value>()?;
+
+                    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<Vec<String>> {
+        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)
+    }
+}

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<Value>;
     fn path(&self) -> Option<Value>;
 
-    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<dyn Env> {
         (**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<dyn Env> {
 pub struct Environment {
     environment_vars: Option<Value>,
     path_vars: Option<Value>,
+    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<T: Conf>(&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(),

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