Skip to content

Commit

Permalink
Capture environment at spawn
Browse files Browse the repository at this point in the history
  • Loading branch information
Diggsey committed Dec 24, 2017
1 parent b058dc0 commit ccc91d7
Show file tree
Hide file tree
Showing 11 changed files with 321 additions and 167 deletions.
29 changes: 25 additions & 4 deletions src/libstd/process.rs
Original file line number Diff line number Diff line change
Expand Up @@ -513,7 +513,7 @@ impl Command {
pub fn env<K, V>(&mut self, key: K, val: V) -> &mut Command
where K: AsRef<OsStr>, V: AsRef<OsStr>
{
self.inner.env(key.as_ref(), val.as_ref());
self.inner.env_mut().set(key.as_ref(), val.as_ref());
self
}

Expand Down Expand Up @@ -546,7 +546,7 @@ impl Command {
where I: IntoIterator<Item=(K, V)>, K: AsRef<OsStr>, V: AsRef<OsStr>
{
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
}
Expand All @@ -567,7 +567,7 @@ impl Command {
/// ```
#[stable(feature = "process", since = "1.0.0")]
pub fn env_remove<K: AsRef<OsStr>>(&mut self, key: K) -> &mut Command {
self.inner.env_remove(key.as_ref());
self.inner.env_mut().remove(key.as_ref());
self
}

Expand All @@ -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
}

Expand Down Expand Up @@ -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() {
Expand Down
24 changes: 7 additions & 17 deletions src/libstd/sys/redox/process.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,7 @@
// 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 env::{split_paths};
use ffi::OsStr;
use os::unix::ffi::OsStrExt;
use fmt;
Expand All @@ -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
Expand All @@ -44,7 +44,7 @@ pub struct Command {
// other keys.
program: String,
args: Vec<String>,
env: HashMap<String, String>,
env: CommandEnv<DefaultEnvKey>,

cwd: Option<String>,
uid: Option<u32>,
Expand Down Expand Up @@ -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,
Expand All @@ -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<DefaultEnvKey> {
&mut self.env
}

pub fn cwd(&mut self, dir: &OsStr) {
Expand Down Expand Up @@ -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))
Expand Down
143 changes: 60 additions & 83 deletions src/libstd/sys/unix/process/process_common.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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
Expand All @@ -45,9 +45,8 @@ pub struct Command {
// other keys.
program: CString,
args: Vec<CString>,
env: Option<HashMap<OsString, (usize, CString)>>,
argv: Vec<*const c_char>,
envp: Option<Vec<*const c_char>>,
env: CommandEnv<DefaultEnvKey>,

cwd: Option<CString>,
uid: Option<uid_t>,
Expand Down Expand Up @@ -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,
Expand All @@ -121,68 +119,6 @@ impl Command {
self.args.push(arg);
}

fn init_env_map(&mut self) -> (&mut HashMap<OsString, (usize, CString)>,
&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));
}
Expand All @@ -196,9 +132,6 @@ impl Command {
pub fn saw_nul(&self) -> bool {
self.saw_nul
}
pub fn get_envp(&self) -> &Option<Vec<*const c_char>> {
&self.envp
}
pub fn get_argv(&self) -> &Vec<*const c_char> {
&self.argv
}
Expand Down Expand Up @@ -237,6 +170,15 @@ impl Command {
self.stderr = Some(stderr);
}

pub fn env_mut(&mut self) -> &mut CommandEnv<DefaultEnvKey> {
&mut self.env
}

pub fn capture_env(&mut self) -> Option<CStringArray> {
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;
Expand Down Expand Up @@ -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<CString>,
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<DefaultEnvKey, OsString>, 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<AnonPipe>)> {
Expand Down Expand Up @@ -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)?;
Expand Down
10 changes: 6 additions & 4 deletions src/libstd/sys/unix/process/process_fuchsia.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,14 +23,16 @@ 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"));
}

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))
}
Expand All @@ -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<zx_handle_t> {
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(),
};

Expand Down
Loading

0 comments on commit ccc91d7

Please sign in to comment.