Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Capture Command environment at spawn #46789

Merged
merged 1 commit into from
Dec 24, 2017
Merged

Capture Command environment at spawn #46789

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 25 additions & 4 deletions src/libstd/process.rs
View file
Open in desktop
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
View file
Open in desktop
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
View file
Open in desktop
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
View file
Open in desktop
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