From da90b1da4a4ab5782ddb823146f546b374ae9ff7 Mon Sep 17 00:00:00 2001 From: Diggory Blake Date: Sun, 17 Dec 2017 15:21:47 +0000 Subject: [PATCH] Capture environment at spawn --- src/libstd/process.rs | 29 +++- src/libstd/sys/redox/process.rs | 22 +-- src/libstd/sys/unix/process/process_common.rs | 143 ++++++++---------- .../sys/unix/process/process_fuchsia.rs | 10 +- src/libstd/sys/unix/process/process_unix.rs | 16 +- src/libstd/sys/wasm/process.rs | 15 +- src/libstd/sys/windows/os_str.rs | 8 +- src/libstd/sys/windows/process.rs | 98 ++++++------ src/libstd/sys_common/mod.rs | 1 + src/libstd/sys_common/process.rs | 124 +++++++++++++++ src/libstd/sys_common/wtf8.rs | 20 +++ 11 files changed, 320 insertions(+), 166 deletions(-) create mode 100644 src/libstd/sys_common/process.rs diff --git a/src/libstd/process.rs b/src/libstd/process.rs index 2335695ae42d..1f6ddde00272 100644 --- a/src/libstd/process.rs +++ b/src/libstd/process.rs @@ -513,7 +513,7 @@ impl Command { pub fn env(&mut self, key: K, val: V) -> &mut Command where K: AsRef, V: AsRef { - self.inner.env(key.as_ref(), val.as_ref()); + self.inner.env_mut().set(key.as_ref(), val.as_ref()); self } @@ -546,7 +546,7 @@ impl Command { where I: IntoIterator, K: AsRef, V: AsRef { for (ref key, ref val) in vars { - self.inner.env(key.as_ref(), val.as_ref()); + self.inner.env_mut().set(key.as_ref(), val.as_ref()); } self } @@ -567,7 +567,7 @@ impl Command { /// ``` #[stable(feature = "process", since = "1.0.0")] pub fn env_remove>(&mut self, key: K) -> &mut Command { - self.inner.env_remove(key.as_ref()); + self.inner.env_mut().remove(key.as_ref()); self } @@ -587,7 +587,7 @@ impl Command { /// ``` #[stable(feature = "process", since = "1.0.0")] pub fn env_clear(&mut self) -> &mut Command { - self.inner.env_clear(); + self.inner.env_mut().clear(); self } @@ -1715,6 +1715,27 @@ mod tests { "didn't find RUN_TEST_NEW_ENV inside of:\n\n{}", output); } + #[test] + fn test_capture_env_at_spawn() { + use env; + + let mut cmd = env_cmd(); + cmd.env("RUN_TEST_NEW_ENV1", "123"); + + // This variable will not be present if the environment has already + // been captured above. + env::set_var("RUN_TEST_NEW_ENV2", "456"); + let result = cmd.output().unwrap(); + env::remove_var("RUN_TEST_NEW_ENV2"); + + let output = String::from_utf8_lossy(&result.stdout).to_string(); + + assert!(output.contains("RUN_TEST_NEW_ENV1=123"), + "didn't find RUN_TEST_NEW_ENV1 inside of:\n\n{}", output); + assert!(output.contains("RUN_TEST_NEW_ENV2=456"), + "didn't find RUN_TEST_NEW_ENV2 inside of:\n\n{}", output); + } + // Regression tests for #30858. #[test] fn test_interior_nul_in_progname_is_error() { diff --git a/src/libstd/sys/redox/process.rs b/src/libstd/sys/redox/process.rs index d87364b8121a..f0cb0f135879 100644 --- a/src/libstd/sys/redox/process.rs +++ b/src/libstd/sys/redox/process.rs @@ -8,7 +8,6 @@ // option. This file may not be copied, modified, or distributed // except according to those terms. -use collections::hash_map::HashMap; use env::{self, split_paths}; use ffi::OsStr; use os::unix::ffi::OsStrExt; @@ -19,6 +18,7 @@ use sys::fd::FileDesc; use sys::fs::{File, OpenOptions}; use sys::pipe::{self, AnonPipe}; use sys::{cvt, syscall}; +use sys_common::process::{CommandEnv, DefaultEnvKey}; //////////////////////////////////////////////////////////////////////////////// // Command @@ -44,7 +44,7 @@ pub struct Command { // other keys. program: String, args: Vec, - env: HashMap, + env: CommandEnv, cwd: Option, uid: Option, @@ -90,7 +90,7 @@ impl Command { Command { program: program.to_str().unwrap().to_owned(), args: Vec::new(), - env: HashMap::new(), + env: Default::default(), cwd: None, uid: None, gid: None, @@ -106,16 +106,8 @@ impl Command { self.args.push(arg.to_str().unwrap().to_owned()); } - pub fn env(&mut self, key: &OsStr, val: &OsStr) { - self.env.insert(key.to_str().unwrap().to_owned(), val.to_str().unwrap().to_owned()); - } - - pub fn env_remove(&mut self, key: &OsStr) { - self.env.remove(key.to_str().unwrap()); - } - - pub fn env_clear(&mut self) { - self.env.clear(); + pub fn env_mut(&mut self) -> &mut CommandEnv { + &mut self.env } pub fn cwd(&mut self, dir: &OsStr) { @@ -309,9 +301,7 @@ impl Command { args.push([arg.as_ptr() as usize, arg.len()]); } - for (key, val) in self.env.iter() { - env::set_var(key, val); - } + self.env.apply(); let program = if self.program.contains(':') || self.program.contains('/') { Some(PathBuf::from(&self.program)) diff --git a/src/libstd/sys/unix/process/process_common.rs b/src/libstd/sys/unix/process/process_common.rs index 383434b1cd87..c53bcdbf8e36 100644 --- a/src/libstd/sys/unix/process/process_common.rs +++ b/src/libstd/sys/unix/process/process_common.rs @@ -10,8 +10,6 @@ use os::unix::prelude::*; -use collections::hash_map::{HashMap, Entry}; -use env; use ffi::{OsString, OsStr, CString, CStr}; use fmt; use io; @@ -20,6 +18,8 @@ use ptr; use sys::fd::FileDesc; use sys::fs::{File, OpenOptions}; use sys::pipe::{self, AnonPipe}; +use sys_common::process::{CommandEnv, DefaultEnvKey}; +use collections::BTreeMap; //////////////////////////////////////////////////////////////////////////////// // Command @@ -45,9 +45,8 @@ pub struct Command { // other keys. program: CString, args: Vec, - env: Option>, argv: Vec<*const c_char>, - envp: Option>, + env: CommandEnv, cwd: Option, uid: Option, @@ -96,8 +95,7 @@ impl Command { argv: vec![program.as_ptr(), ptr::null()], program, args: Vec::new(), - env: None, - envp: None, + env: Default::default(), cwd: None, uid: None, gid: None, @@ -121,68 +119,6 @@ impl Command { self.args.push(arg); } - fn init_env_map(&mut self) -> (&mut HashMap, - &mut Vec<*const c_char>) { - if self.env.is_none() { - let mut map = HashMap::new(); - let mut envp = Vec::new(); - for (k, v) in env::vars_os() { - let s = pair_to_key(&k, &v, &mut self.saw_nul); - envp.push(s.as_ptr()); - map.insert(k, (envp.len() - 1, s)); - } - envp.push(ptr::null()); - self.env = Some(map); - self.envp = Some(envp); - } - (self.env.as_mut().unwrap(), self.envp.as_mut().unwrap()) - } - - pub fn env(&mut self, key: &OsStr, val: &OsStr) { - let new_key = pair_to_key(key, val, &mut self.saw_nul); - let (map, envp) = self.init_env_map(); - - // If `key` is already present then we just update `envp` in place - // (and store the owned value), but if it's not there we override the - // trailing NULL pointer, add a new NULL pointer, and store where we - // were located. - match map.entry(key.to_owned()) { - Entry::Occupied(mut e) => { - let (i, ref mut s) = *e.get_mut(); - envp[i] = new_key.as_ptr(); - *s = new_key; - } - Entry::Vacant(e) => { - let len = envp.len(); - envp[len - 1] = new_key.as_ptr(); - envp.push(ptr::null()); - e.insert((len - 1, new_key)); - } - } - } - - pub fn env_remove(&mut self, key: &OsStr) { - let (map, envp) = self.init_env_map(); - - // If we actually ended up removing a key, then we need to update the - // position of all keys that come after us in `envp` because they're all - // one element sooner now. - if let Some((i, _)) = map.remove(key) { - envp.remove(i); - - for (_, &mut (ref mut j, _)) in map.iter_mut() { - if *j >= i { - *j -= 1; - } - } - } - } - - pub fn env_clear(&mut self) { - self.env = Some(HashMap::new()); - self.envp = Some(vec![ptr::null()]); - } - pub fn cwd(&mut self, dir: &OsStr) { self.cwd = Some(os2c(dir, &mut self.saw_nul)); } @@ -196,9 +132,6 @@ impl Command { pub fn saw_nul(&self) -> bool { self.saw_nul } - pub fn get_envp(&self) -> &Option> { - &self.envp - } pub fn get_argv(&self) -> &Vec<*const c_char> { &self.argv } @@ -237,6 +170,15 @@ impl Command { self.stderr = Some(stderr); } + pub fn env_mut(&mut self) -> &mut CommandEnv { + &mut self.env + } + + pub fn capture_env(&mut self) -> Option { + let maybe_env = self.env.capture_if_changed(); + maybe_env.map(|env| construct_envp(env, &mut self.saw_nul)) + } + pub fn setup_io(&self, default: Stdio, needs_stdin: bool) -> io::Result<(StdioPipes, ChildPipes)> { let null = Stdio::Null; @@ -268,6 +210,53 @@ fn os2c(s: &OsStr, saw_nul: &mut bool) -> CString { }) } +// Helper type to manage ownership of the strings within a C-style array. +pub struct CStringArray { + items: Vec, + ptrs: Vec<*const c_char> +} + +impl CStringArray { + pub fn with_capacity(capacity: usize) -> Self { + let mut result = CStringArray { + items: Vec::with_capacity(capacity), + ptrs: Vec::with_capacity(capacity+1) + }; + result.ptrs.push(ptr::null()); + result + } + pub fn push(&mut self, item: CString) { + let l = self.ptrs.len(); + self.ptrs[l-1] = item.as_ptr(); + self.ptrs.push(ptr::null()); + self.items.push(item); + } + pub fn as_ptr(&self) -> *const *const c_char { + self.ptrs.as_ptr() + } +} + +fn construct_envp(env: BTreeMap, saw_nul: &mut bool) -> CStringArray { + let mut result = CStringArray::with_capacity(env.len()); + for (k, v) in env { + let mut k: OsString = k.into(); + + // Reserve additional space for '=' and null terminator + k.reserve_exact(v.len() + 2); + k.push("="); + k.push(&v); + + // Add the new entry into the array + if let Ok(item) = CString::new(k.into_vec()) { + result.push(item); + } else { + *saw_nul = true; + } + } + + result +} + impl Stdio { pub fn to_child_stdio(&self, readable: bool) -> io::Result<(ChildStdio, Option)> { @@ -337,18 +326,6 @@ impl ChildStdio { } } -fn pair_to_key(key: &OsStr, value: &OsStr, saw_nul: &mut bool) -> CString { - let (key, value) = (key.as_bytes(), value.as_bytes()); - let mut v = Vec::with_capacity(key.len() + value.len() + 1); - v.extend(key); - v.push(b'='); - v.extend(value); - CString::new(v).unwrap_or_else(|_e| { - *saw_nul = true; - CString::new("foo=bar").unwrap() - }) -} - impl fmt::Debug for Command { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { write!(f, "{:?}", self.program)?; diff --git a/src/libstd/sys/unix/process/process_fuchsia.rs b/src/libstd/sys/unix/process/process_fuchsia.rs index a7a67ed36e82..06c0540fec0d 100644 --- a/src/libstd/sys/unix/process/process_fuchsia.rs +++ b/src/libstd/sys/unix/process/process_fuchsia.rs @@ -23,6 +23,8 @@ use sys::process::process_common::*; impl Command { pub fn spawn(&mut self, default: Stdio, needs_stdin: bool) -> io::Result<(Process, StdioPipes)> { + let envp = self.capture_env(); + if self.saw_nul() { return Err(io::Error::new(io::ErrorKind::InvalidInput, "nul byte found in provided data")); @@ -30,7 +32,7 @@ impl Command { let (ours, theirs) = self.setup_io(default, needs_stdin)?; - let process_handle = unsafe { self.do_exec(theirs)? }; + let process_handle = unsafe { self.do_exec(theirs, envp.as_ref())? }; Ok((Process { handle: Handle::new(process_handle) }, ours)) } @@ -50,13 +52,13 @@ impl Command { } } - unsafe fn do_exec(&mut self, stdio: ChildPipes) + unsafe fn do_exec(&mut self, stdio: ChildPipes, maybe_envp: Option<&CStringArray>) -> io::Result { use sys::process::zircon::*; let job_handle = zx_job_default(); - let envp = match *self.get_envp() { - Some(ref envp) => envp.as_ptr(), + let envp = match maybe_envp { + Some(envp) => envp.as_ptr(), None => ptr::null(), }; diff --git a/src/libstd/sys/unix/process/process_unix.rs b/src/libstd/sys/unix/process/process_unix.rs index 743c458d580c..189280a4ba9a 100644 --- a/src/libstd/sys/unix/process/process_unix.rs +++ b/src/libstd/sys/unix/process/process_unix.rs @@ -26,6 +26,8 @@ impl Command { const CLOEXEC_MSG_FOOTER: &'static [u8] = b"NOEX"; + let envp = self.capture_env(); + if self.saw_nul() { return Err(io::Error::new(ErrorKind::InvalidInput, "nul byte found in provided data")); @@ -38,7 +40,7 @@ impl Command { match cvt(libc::fork())? { 0 => { drop(input); - let err = self.do_exec(theirs); + let err = self.do_exec(theirs, envp.as_ref()); let errno = err.raw_os_error().unwrap_or(libc::EINVAL) as u32; let bytes = [ (errno >> 24) as u8, @@ -99,13 +101,15 @@ impl Command { } pub fn exec(&mut self, default: Stdio) -> io::Error { + let envp = self.capture_env(); + if self.saw_nul() { return io::Error::new(ErrorKind::InvalidInput, "nul byte found in provided data") } match self.setup_io(default, true) { - Ok((_, theirs)) => unsafe { self.do_exec(theirs) }, + Ok((_, theirs)) => unsafe { self.do_exec(theirs, envp.as_ref()) }, Err(e) => e, } } @@ -140,7 +144,11 @@ impl Command { // allocation). Instead we just close it manually. This will never // have the drop glue anyway because this code never returns (the // child will either exec() or invoke libc::exit) - unsafe fn do_exec(&mut self, stdio: ChildPipes) -> io::Error { + unsafe fn do_exec( + &mut self, + stdio: ChildPipes, + maybe_envp: Option<&CStringArray> + ) -> io::Error { use sys::{self, cvt_r}; macro_rules! t { @@ -180,7 +188,7 @@ impl Command { if let Some(ref cwd) = *self.get_cwd() { t!(cvt(libc::chdir(cwd.as_ptr()))); } - if let Some(ref envp) = *self.get_envp() { + if let Some(envp) = maybe_envp { *sys::os::environ() = envp.as_ptr(); } diff --git a/src/libstd/sys/wasm/process.rs b/src/libstd/sys/wasm/process.rs index 4febe8a14638..08ae1b253d84 100644 --- a/src/libstd/sys/wasm/process.rs +++ b/src/libstd/sys/wasm/process.rs @@ -14,12 +14,14 @@ use io; use sys::fs::File; use sys::pipe::AnonPipe; use sys::{unsupported, Void}; +use sys_common::process::{CommandEnv, DefaultEnvKey}; //////////////////////////////////////////////////////////////////////////////// // Command //////////////////////////////////////////////////////////////////////////////// pub struct Command { + env: CommandEnv } // passed back to std::process with the pipes connected to the child, if any @@ -38,19 +40,16 @@ pub enum Stdio { impl Command { pub fn new(_program: &OsStr) -> Command { - Command {} + Command { + env: Default::default() + } } pub fn arg(&mut self, _arg: &OsStr) { } - pub fn env(&mut self, _key: &OsStr, _val: &OsStr) { - } - - pub fn env_remove(&mut self, _key: &OsStr) { - } - - pub fn env_clear(&mut self) { + pub fn env_mut(&mut self) -> CommandEnv { + &mut self.env } pub fn cwd(&mut self, _dir: &OsStr) { diff --git a/src/libstd/sys/windows/os_str.rs b/src/libstd/sys/windows/os_str.rs index b8d2f7bc53ce..414c9c5418e4 100644 --- a/src/libstd/sys/windows/os_str.rs +++ b/src/libstd/sys/windows/os_str.rs @@ -17,7 +17,7 @@ use sys_common::wtf8::{Wtf8, Wtf8Buf}; use mem; use rc::Rc; use sync::Arc; -use sys_common::{AsInner, IntoInner}; +use sys_common::{AsInner, IntoInner, FromInner}; #[derive(Clone, Hash)] pub struct Buf { @@ -30,6 +30,12 @@ impl IntoInner for Buf { } } +impl FromInner for Buf { + fn from_inner(inner: Wtf8Buf) -> Self { + Buf { inner } + } +} + impl AsInner for Buf { fn as_inner(&self) -> &Wtf8 { &self.inner diff --git a/src/libstd/sys/windows/process.rs b/src/libstd/sys/windows/process.rs index 631d69b05e11..c93179869a65 100644 --- a/src/libstd/sys/windows/process.rs +++ b/src/libstd/sys/windows/process.rs @@ -8,9 +8,10 @@ // option. This file may not be copied, modified, or distributed // except according to those terms. +#![unstable(feature = "process_internals", issue = "0")] + use ascii::AsciiExt; -use collections::HashMap; -use collections; +use collections::BTreeMap; use env::split_paths; use env; use ffi::{OsString, OsStr}; @@ -28,19 +29,42 @@ use sys::fs::{OpenOptions, File}; use sys::handle::Handle; use sys::pipe::{self, AnonPipe}; use sys::stdio; -use sys::{self, cvt}; -use sys_common::{AsInner, FromInner}; +use sys::cvt; +use sys_common::{AsInner, FromInner, IntoInner}; +use sys_common::process::{CommandEnv, EnvKey}; +use alloc::borrow::Borrow; //////////////////////////////////////////////////////////////////////////////// // Command //////////////////////////////////////////////////////////////////////////////// -fn mk_key(s: &OsStr) -> OsString { - FromInner::from_inner(sys::os_str::Buf { - inner: s.as_inner().inner.to_ascii_uppercase() - }) +#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)] +#[doc(hidden)] +pub struct WindowsEnvKey(OsString); + +impl From for WindowsEnvKey { + fn from(k: OsString) -> Self { + let mut buf = k.into_inner().into_inner(); + buf.make_ascii_uppercase(); + WindowsEnvKey(FromInner::from_inner(FromInner::from_inner(buf))) + } +} + +impl From for OsString { + fn from(k: WindowsEnvKey) -> Self { k.0 } } +impl Borrow for WindowsEnvKey { + fn borrow(&self) -> &OsStr { &self.0 } +} + +impl AsRef for WindowsEnvKey { + fn as_ref(&self) -> &OsStr { &self.0 } +} + +impl EnvKey for WindowsEnvKey {} + + fn ensure_no_nuls>(str: T) -> io::Result { if str.as_ref().encode_wide().any(|b| b == 0) { Err(io::Error::new(ErrorKind::InvalidInput, "nul byte found in provided data")) @@ -52,7 +76,7 @@ fn ensure_no_nuls>(str: T) -> io::Result { pub struct Command { program: OsString, args: Vec, - env: Option>, + env: CommandEnv, cwd: Option, flags: u32, detach: bool, // not currently exposed in std::process @@ -83,7 +107,7 @@ impl Command { Command { program: program.to_os_string(), args: Vec::new(), - env: None, + env: Default::default(), cwd: None, flags: 0, detach: false, @@ -96,23 +120,8 @@ impl Command { pub fn arg(&mut self, arg: &OsStr) { self.args.push(arg.to_os_string()) } - fn init_env_map(&mut self){ - if self.env.is_none() { - self.env = Some(env::vars_os().map(|(key, val)| { - (mk_key(&key), val) - }).collect()); - } - } - pub fn env(&mut self, key: &OsStr, val: &OsStr) { - self.init_env_map(); - self.env.as_mut().unwrap().insert(mk_key(key), val.to_os_string()); - } - pub fn env_remove(&mut self, key: &OsStr) { - self.init_env_map(); - self.env.as_mut().unwrap().remove(&mk_key(key)); - } - pub fn env_clear(&mut self) { - self.env = Some(HashMap::new()) + pub fn env_mut(&mut self) -> &mut CommandEnv { + &mut self.env } pub fn cwd(&mut self, dir: &OsStr) { self.cwd = Some(dir.to_os_string()) @@ -132,13 +141,12 @@ impl Command { pub fn spawn(&mut self, default: Stdio, needs_stdin: bool) -> io::Result<(Process, StdioPipes)> { + let maybe_env = self.env.capture_if_changed(); // To have the spawning semantics of unix/windows stay the same, we need // to read the *child's* PATH if one is provided. See #15149 for more // details. - let program = self.env.as_ref().and_then(|env| { - for (key, v) in env { - if OsStr::new("PATH") != &**key { continue } - + let program = maybe_env.as_ref().and_then(|env| { + if let Some(v) = env.get(OsStr::new("PATH")) { // Split the value and test each path to see if the // program exists. for path in split_paths(&v) { @@ -148,7 +156,6 @@ impl Command { return Some(path.into_os_string()) } } - break } None }); @@ -167,7 +174,7 @@ impl Command { flags |= c::DETACHED_PROCESS | c::CREATE_NEW_PROCESS_GROUP; } - let (envp, _data) = make_envp(self.env.as_ref())?; + let (envp, _data) = make_envp(maybe_env)?; let (dirp, _data) = make_dirp(self.cwd.as_ref())?; let mut pi = zeroed_process_information(); @@ -488,25 +495,24 @@ fn make_command_line(prog: &OsStr, args: &[OsString]) -> io::Result> { } } -fn make_envp(env: Option<&collections::HashMap>) +fn make_envp(maybe_env: Option>) -> io::Result<(*mut c_void, Vec)> { // On Windows we pass an "environment block" which is not a char**, but // rather a concatenation of null-terminated k=v\0 sequences, with a final // \0 to terminate. - match env { - Some(env) => { - let mut blk = Vec::new(); - - for pair in env { - blk.extend(ensure_no_nuls(pair.0)?.encode_wide()); - blk.push('=' as u16); - blk.extend(ensure_no_nuls(pair.1)?.encode_wide()); - blk.push(0); - } + if let Some(env) = maybe_env { + let mut blk = Vec::new(); + + for (k, v) in env { + blk.extend(ensure_no_nuls(k.0)?.encode_wide()); + blk.push('=' as u16); + blk.extend(ensure_no_nuls(v)?.encode_wide()); blk.push(0); - Ok((blk.as_mut_ptr() as *mut c_void, blk)) } - _ => Ok((ptr::null_mut(), Vec::new())) + blk.push(0); + Ok((blk.as_mut_ptr() as *mut c_void, blk)) + } else { + Ok((ptr::null_mut(), Vec::new())) } } diff --git a/src/libstd/sys_common/mod.rs b/src/libstd/sys_common/mod.rs index 5c4d7b52754c..b16299a1d636 100644 --- a/src/libstd/sys_common/mod.rs +++ b/src/libstd/sys_common/mod.rs @@ -44,6 +44,7 @@ pub mod thread_local; pub mod util; pub mod wtf8; pub mod bytestring; +pub mod process; cfg_if! { if #[cfg(any(target_os = "redox", target_os = "l4re"))] { diff --git a/src/libstd/sys_common/process.rs b/src/libstd/sys_common/process.rs new file mode 100644 index 000000000000..fd1a5fdb4109 --- /dev/null +++ b/src/libstd/sys_common/process.rs @@ -0,0 +1,124 @@ +// Copyright 2014 The Rust Project Developers. See the COPYRIGHT +// file at the top-level directory of this distribution and at +// http://rust-lang.org/COPYRIGHT. +// +// Licensed under the Apache License, Version 2.0 or the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. + +#![allow(dead_code)] +#![unstable(feature = "process_internals", issue = "0")] + +use ffi::{OsStr, OsString}; +use env; +use collections::BTreeMap; +use alloc::borrow::Borrow; + +pub trait EnvKey: + From + Into + + Borrow + Borrow + AsRef + + Ord + Clone {} + +// Implement a case-sensitive environment variable key +#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)] +pub struct DefaultEnvKey(OsString); + +impl From for DefaultEnvKey { + fn from(k: OsString) -> Self { DefaultEnvKey(k) } +} + +impl From for OsString { + fn from(k: DefaultEnvKey) -> Self { k.0 } +} + +impl Borrow for DefaultEnvKey { + fn borrow(&self) -> &OsStr { &self.0 } +} + +impl AsRef for DefaultEnvKey { + fn as_ref(&self) -> &OsStr { &self.0 } +} + +impl EnvKey for DefaultEnvKey {} + +// Stores a set of changes to an environment +#[derive(Clone, Debug)] +pub struct CommandEnv { + clear: bool, + vars: BTreeMap> +} + +impl Default for CommandEnv { + fn default() -> Self { + CommandEnv { + clear: false, + vars: Default::default() + } + } +} + +impl CommandEnv { + // Capture the current environment with these changes applied + pub fn capture(&self) -> BTreeMap { + let mut result = BTreeMap::::new(); + if !self.clear { + for (k, v) in env::vars_os() { + result.insert(k.into(), v); + } + } + for (k, maybe_v) in &self.vars { + if let &Some(ref v) = maybe_v { + result.insert(k.clone(), v.clone()); + } else { + result.remove(k); + } + } + result + } + + // Apply these changes directly to the current environment + pub fn apply(&self) { + if self.clear { + for (k, _) in env::vars_os() { + env::remove_var(k); + } + } + for (key, maybe_val) in self.vars.iter() { + if let &Some(ref val) = maybe_val { + env::set_var(key, val); + } else { + env::remove_var(key); + } + } + } + + pub fn is_unchanged(&self) -> bool { + !self.clear && self.vars.is_empty() + } + + pub fn capture_if_changed(&self) -> Option> { + if self.is_unchanged() { + None + } else { + Some(self.capture()) + } + } + + // The following functions build up changes + pub fn set(&mut self, key: &OsStr, value: &OsStr) { + self.vars.insert(key.to_owned().into(), Some(value.to_owned())); + } + pub fn remove(&mut self, key: &OsStr) { + if self.clear { + self.vars.remove(key); + } else { + self.vars.insert(key.to_owned().into(), None); + } + } + pub fn clear(&mut self) { + self.clear = true; + self.vars.clear(); + } +} diff --git a/src/libstd/sys_common/wtf8.rs b/src/libstd/sys_common/wtf8.rs index b2fc559bb37f..46d554d6411b 100644 --- a/src/libstd/sys_common/wtf8.rs +++ b/src/libstd/sys_common/wtf8.rs @@ -134,6 +134,12 @@ impl ops::Deref for Wtf8Buf { } } +impl ops::DerefMut for Wtf8Buf { + fn deref_mut(&mut self) -> &mut Wtf8 { + self.as_mut_slice() + } +} + /// Format the string with double quotes, /// and surrogates as `\u` followed by four hexadecimal digits. /// Example: `"a\u{D800}"` for a string with code points [U+0061, U+D800] @@ -221,6 +227,11 @@ impl Wtf8Buf { unsafe { Wtf8::from_bytes_unchecked(&self.bytes) } } + #[inline] + pub fn as_mut_slice(&mut self) -> &mut Wtf8 { + unsafe { Wtf8::from_mut_bytes_unchecked(&mut self.bytes) } + } + /// Reserves capacity for at least `additional` more bytes to be inserted /// in the given `Wtf8Buf`. /// The collection may reserve more space to avoid frequent reallocations. @@ -486,6 +497,15 @@ impl Wtf8 { mem::transmute(value) } + /// Creates a mutable WTF-8 slice from a mutable WTF-8 byte slice. + /// + /// Since the byte slice is not checked for valid WTF-8, this functions is + /// marked unsafe. + #[inline] + unsafe fn from_mut_bytes_unchecked(value: &mut [u8]) -> &mut Wtf8 { + mem::transmute(value) + } + /// Returns the length, in WTF-8 bytes. #[inline] pub fn len(&self) -> usize {