Change list of allowed dirs to indexmap

This commit is contained in:
Sam Hedin 2020-06-09 07:00:22 +02:00
parent 5a85a3448e
commit a40ed5fef2

View File

@ -1,11 +1,11 @@
use indexmap::IndexMap; use indexmap::{IndexMap, IndexSet};
use nu_protocol::{Primitive, UntaggedValue, Value}; use nu_protocol::{Primitive, UntaggedValue, Value};
use std::io::{Error, ErrorKind, Result}; use std::io::{Error, ErrorKind, Result};
use std::{ffi::OsString, fmt::Debug, path::PathBuf}; use std::{ffi::OsString, fmt::Debug, path::PathBuf};
#[derive(Debug, Default)] #[derive(Debug, Default)]
pub struct DirectorySpecificEnvironment { pub struct DirectorySpecificEnvironment {
allowed_directories: Vec<PathBuf>, allowed_directories: IndexSet<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. //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>>, added_env_vars: IndexMap<PathBuf, Vec<String>>,
@ -38,9 +38,13 @@ impl DirectorySpecificEnvironment {
vec![] vec![]
}; };
allowed_directories.sort(); allowed_directories.sort();
let mut allowed = IndexSet::new();
for d in allowed_directories {
allowed.insert(d);
}
DirectorySpecificEnvironment { DirectorySpecificEnvironment {
allowed_directories, allowed_directories: allowed,
added_env_vars: IndexMap::new(), added_env_vars: IndexMap::new(),
overwritten_env_values: IndexMap::new(), overwritten_env_values: IndexMap::new(),
} }
@ -93,16 +97,16 @@ impl DirectorySpecificEnvironment {
pub fn env_vars_to_add(&mut self) -> Result<IndexMap<String, String>> { pub fn env_vars_to_add(&mut self) -> Result<IndexMap<String, String>> {
let current_dir = std::env::current_dir()?; let current_dir = std::env::current_dir()?;
let mut vars_to_add = IndexMap::new(); let mut vars_to_add = IndexMap::new();
for dir in &self.allowed_directories {
let mut working_dir = Some(current_dir.as_path()); 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. //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 { while let Some(wdir) = working_dir {
if wdir == dir.as_path() { if self.allowed_directories.contains(wdir) {
let toml_doc = std::fs::read_to_string(wdir.join(".nu").as_path())? let toml_doc = match std::fs::read_to_string(wdir.join(".nu").as_path()) {
.parse::<toml::Value>()?; Ok(doc) => doc.parse::<toml::Value>()?,
Err(_) => return Ok(vars_to_add),
};
let vars_in_current_file = toml_doc let vars_in_current_file = toml_doc
.get("env") .get("env")
@ -122,17 +126,19 @@ impl DirectorySpecificEnvironment {
let mut keys_in_current_nufile = vec![]; let mut keys_in_current_nufile = vec![];
for (k, v) in vars_in_current_file { for (k, v) in vars_in_current_file {
if !vars_to_add.contains_key(k) {
vars_to_add.insert( vars_to_add.insert(
k.clone(), k.clone(),
v.as_str() v.as_str()
.ok_or_else(|| { .ok_or_else(|| {
Error::new( Error::new(
ErrorKind::InvalidData, ErrorKind::InvalidData,
format!("Could not read environment variable: {}", v), format!("Could not read environment variable: {}\n", v),
) )
})? })?
.to_string(), .to_string(),
); //This is used to add variables to the environment ); //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 keys_in_current_nufile.push(k.clone()); //this is used to keep track of which directory added which variables
} }
@ -152,16 +158,14 @@ impl DirectorySpecificEnvironment {
self.added_env_vars self.added_env_vars
.insert(wdir.to_path_buf(), keys_in_current_nufile); .insert(wdir.to_path_buf(), keys_in_current_nufile);
break; }
} else { working_dir =
working_dir = working_dir //Keep going up in the directory structure with .parent() working_dir //Keep going up in the directory structure with .parent()
.ok_or_else(|| { .ok_or_else(|| {
Error::new(ErrorKind::NotFound, "Root directory has no parent") Error::new(ErrorKind::NotFound, "Root directory has no parent")
})? })?
.parent(); .parent();
} }
}
}
Ok(vars_to_add) Ok(vars_to_add)
} }