Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit c284f88

Browse files
committedDec 24, 2017
Auto merge of #46789 - Diggsey:command-env-capture, r=dtolnay
Capture `Command` environment at spawn Fixes #28975 This tracks a set of changes to the environment and then replays them at spawn time.
2 parents b9e4d34 + ccc91d7 commit c284f88

File tree

11 files changed

+321
-167
lines changed

11 files changed

+321
-167
lines changed
 

‎src/libstd/process.rs

Lines changed: 25 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -513,7 +513,7 @@ impl Command {
513513
pub fn env<K, V>(&mut self, key: K, val: V) -> &mut Command
514514
where K: AsRef<OsStr>, V: AsRef<OsStr>
515515
{
516-
self.inner.env(key.as_ref(), val.as_ref());
516+
self.inner.env_mut().set(key.as_ref(), val.as_ref());
517517
self
518518
}
519519

@@ -546,7 +546,7 @@ impl Command {
546546
where I: IntoIterator<Item=(K, V)>, K: AsRef<OsStr>, V: AsRef<OsStr>
547547
{
548548
for (ref key, ref val) in vars {
549-
self.inner.env(key.as_ref(), val.as_ref());
549+
self.inner.env_mut().set(key.as_ref(), val.as_ref());
550550
}
551551
self
552552
}
@@ -567,7 +567,7 @@ impl Command {
567567
/// ```
568568
#[stable(feature = "process", since = "1.0.0")]
569569
pub fn env_remove<K: AsRef<OsStr>>(&mut self, key: K) -> &mut Command {
570-
self.inner.env_remove(key.as_ref());
570+
self.inner.env_mut().remove(key.as_ref());
571571
self
572572
}
573573

@@ -587,7 +587,7 @@ impl Command {
587587
/// ```
588588
#[stable(feature = "process", since = "1.0.0")]
589589
pub fn env_clear(&mut self) -> &mut Command {
590-
self.inner.env_clear();
590+
self.inner.env_mut().clear();
591591
self
592592
}
593593

@@ -1715,6 +1715,27 @@ mod tests {
17151715
"didn't find RUN_TEST_NEW_ENV inside of:\n\n{}", output);
17161716
}
17171717

1718+
#[test]
1719+
fn test_capture_env_at_spawn() {
1720+
use env;
1721+
1722+
let mut cmd = env_cmd();
1723+
cmd.env("RUN_TEST_NEW_ENV1", "123");
1724+
1725+
// This variable will not be present if the environment has already
1726+
// been captured above.
1727+
env::set_var("RUN_TEST_NEW_ENV2", "456");
1728+
let result = cmd.output().unwrap();
1729+
env::remove_var("RUN_TEST_NEW_ENV2");
1730+
1731+
let output = String::from_utf8_lossy(&result.stdout).to_string();
1732+
1733+
assert!(output.contains("RUN_TEST_NEW_ENV1=123"),
1734+
"didn't find RUN_TEST_NEW_ENV1 inside of:\n\n{}", output);
1735+
assert!(output.contains("RUN_TEST_NEW_ENV2=456"),
1736+
"didn't find RUN_TEST_NEW_ENV2 inside of:\n\n{}", output);
1737+
}
1738+
17181739
// Regression tests for #30858.
17191740
#[test]
17201741
fn test_interior_nul_in_progname_is_error() {

‎src/libstd/sys/redox/process.rs

Lines changed: 7 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,7 @@
88
// option. This file may not be copied, modified, or distributed
99
// except according to those terms.
1010

11-
use collections::hash_map::HashMap;
12-
use env::{self, split_paths};
11+
use env::{split_paths};
1312
use ffi::OsStr;
1413
use os::unix::ffi::OsStrExt;
1514
use fmt;
@@ -19,6 +18,7 @@ use sys::fd::FileDesc;
1918
use sys::fs::{File, OpenOptions};
2019
use sys::pipe::{self, AnonPipe};
2120
use sys::{cvt, syscall};
21+
use sys_common::process::{CommandEnv, DefaultEnvKey};
2222

2323
////////////////////////////////////////////////////////////////////////////////
2424
// Command
@@ -44,7 +44,7 @@ pub struct Command {
4444
// other keys.
4545
program: String,
4646
args: Vec<String>,
47-
env: HashMap<String, String>,
47+
env: CommandEnv<DefaultEnvKey>,
4848

4949
cwd: Option<String>,
5050
uid: Option<u32>,
@@ -90,7 +90,7 @@ impl Command {
9090
Command {
9191
program: program.to_str().unwrap().to_owned(),
9292
args: Vec::new(),
93-
env: HashMap::new(),
93+
env: Default::default(),
9494
cwd: None,
9595
uid: None,
9696
gid: None,
@@ -106,16 +106,8 @@ impl Command {
106106
self.args.push(arg.to_str().unwrap().to_owned());
107107
}
108108

109-
pub fn env(&mut self, key: &OsStr, val: &OsStr) {
110-
self.env.insert(key.to_str().unwrap().to_owned(), val.to_str().unwrap().to_owned());
111-
}
112-
113-
pub fn env_remove(&mut self, key: &OsStr) {
114-
self.env.remove(key.to_str().unwrap());
115-
}
116-
117-
pub fn env_clear(&mut self) {
118-
self.env.clear();
109+
pub fn env_mut(&mut self) -> &mut CommandEnv<DefaultEnvKey> {
110+
&mut self.env
119111
}
120112

121113
pub fn cwd(&mut self, dir: &OsStr) {
@@ -309,9 +301,7 @@ impl Command {
309301
args.push([arg.as_ptr() as usize, arg.len()]);
310302
}
311303

312-
for (key, val) in self.env.iter() {
313-
env::set_var(key, val);
314-
}
304+
self.env.apply();
315305

316306
let program = if self.program.contains(':') || self.program.contains('/') {
317307
Some(PathBuf::from(&self.program))

‎src/libstd/sys/unix/process/process_common.rs

Lines changed: 60 additions & 83 deletions
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,6 @@
1010

1111
use os::unix::prelude::*;
1212

13-
use collections::hash_map::{HashMap, Entry};
14-
use env;
1513
use ffi::{OsString, OsStr, CString, CStr};
1614
use fmt;
1715
use io;
@@ -20,6 +18,8 @@ use ptr;
2018
use sys::fd::FileDesc;
2119
use sys::fs::{File, OpenOptions};
2220
use sys::pipe::{self, AnonPipe};
21+
use sys_common::process::{CommandEnv, DefaultEnvKey};
22+
use collections::BTreeMap;
2323

2424
////////////////////////////////////////////////////////////////////////////////
2525
// Command
@@ -45,9 +45,8 @@ pub struct Command {
4545
// other keys.
4646
program: CString,
4747
args: Vec<CString>,
48-
env: Option<HashMap<OsString, (usize, CString)>>,
4948
argv: Vec<*const c_char>,
50-
envp: Option<Vec<*const c_char>>,
49+
env: CommandEnv<DefaultEnvKey>,
5150

5251
cwd: Option<CString>,
5352
uid: Option<uid_t>,
@@ -96,8 +95,7 @@ impl Command {
9695
argv: vec![program.as_ptr(), ptr::null()],
9796
program,
9897
args: Vec::new(),
99-
env: None,
100-
envp: None,
98+
env: Default::default(),
10199
cwd: None,
102100
uid: None,
103101
gid: None,
@@ -121,68 +119,6 @@ impl Command {
121119
self.args.push(arg);
122120
}
123121

124-
fn init_env_map(&mut self) -> (&mut HashMap<OsString, (usize, CString)>,
125-
&mut Vec<*const c_char>) {
126-
if self.env.is_none() {
127-
let mut map = HashMap::new();
128-
let mut envp = Vec::new();
129-
for (k, v) in env::vars_os() {
130-
let s = pair_to_key(&k, &v, &mut self.saw_nul);
131-
envp.push(s.as_ptr());
132-
map.insert(k, (envp.len() - 1, s));
133-
}
134-
envp.push(ptr::null());
135-
self.env = Some(map);
136-
self.envp = Some(envp);
137-
}
138-
(self.env.as_mut().unwrap(), self.envp.as_mut().unwrap())
139-
}
140-
141-
pub fn env(&mut self, key: &OsStr, val: &OsStr) {
142-
let new_key = pair_to_key(key, val, &mut self.saw_nul);
143-
let (map, envp) = self.init_env_map();
144-
145-
// If `key` is already present then we just update `envp` in place
146-
// (and store the owned value), but if it's not there we override the
147-
// trailing NULL pointer, add a new NULL pointer, and store where we
148-
// were located.
149-
match map.entry(key.to_owned()) {
150-
Entry::Occupied(mut e) => {
151-
let (i, ref mut s) = *e.get_mut();
152-
envp[i] = new_key.as_ptr();
153-
*s = new_key;
154-
}
155-
Entry::Vacant(e) => {
156-
let len = envp.len();
157-
envp[len - 1] = new_key.as_ptr();
158-
envp.push(ptr::null());
159-
e.insert((len - 1, new_key));
160-
}
161-
}
162-
}
163-
164-
pub fn env_remove(&mut self, key: &OsStr) {
165-
let (map, envp) = self.init_env_map();
166-
167-
// If we actually ended up removing a key, then we need to update the
168-
// position of all keys that come after us in `envp` because they're all
169-
// one element sooner now.
170-
if let Some((i, _)) = map.remove(key) {
171-
envp.remove(i);
172-
173-
for (_, &mut (ref mut j, _)) in map.iter_mut() {
174-
if *j >= i {
175-
*j -= 1;
176-
}
177-
}
178-
}
179-
}
180-
181-
pub fn env_clear(&mut self) {
182-
self.env = Some(HashMap::new());
183-
self.envp = Some(vec![ptr::null()]);
184-
}
185-
186122
pub fn cwd(&mut self, dir: &OsStr) {
187123
self.cwd = Some(os2c(dir, &mut self.saw_nul));
188124
}
@@ -196,9 +132,6 @@ impl Command {
196132
pub fn saw_nul(&self) -> bool {
197133
self.saw_nul
198134
}
199-
pub fn get_envp(&self) -> &Option<Vec<*const c_char>> {
200-
&self.envp
201-
}
202135
pub fn get_argv(&self) -> &Vec<*const c_char> {
203136
&self.argv
204137
}
@@ -237,6 +170,15 @@ impl Command {
237170
self.stderr = Some(stderr);
238171
}
239172

173+
pub fn env_mut(&mut self) -> &mut CommandEnv<DefaultEnvKey> {
174+
&mut self.env
175+
}
176+
177+
pub fn capture_env(&mut self) -> Option<CStringArray> {
178+
let maybe_env = self.env.capture_if_changed();
179+
maybe_env.map(|env| construct_envp(env, &mut self.saw_nul))
180+
}
181+
240182
pub fn setup_io(&self, default: Stdio, needs_stdin: bool)
241183
-> io::Result<(StdioPipes, ChildPipes)> {
242184
let null = Stdio::Null;
@@ -268,6 +210,53 @@ fn os2c(s: &OsStr, saw_nul: &mut bool) -> CString {
268210
})
269211
}
270212

213+
// Helper type to manage ownership of the strings within a C-style array.
214+
pub struct CStringArray {
215+
items: Vec<CString>,
216+
ptrs: Vec<*const c_char>
217+
}
218+
219+
impl CStringArray {
220+
pub fn with_capacity(capacity: usize) -> Self {
221+
let mut result = CStringArray {
222+
items: Vec::with_capacity(capacity),
223+
ptrs: Vec::with_capacity(capacity+1)
224+
};
225+
result.ptrs.push(ptr::null());
226+
result
227+
}
228+
pub fn push(&mut self, item: CString) {
229+
let l = self.ptrs.len();
230+
self.ptrs[l-1] = item.as_ptr();
231+
self.ptrs.push(ptr::null());
232+
self.items.push(item);
233+
}
234+
pub fn as_ptr(&self) -> *const *const c_char {
235+
self.ptrs.as_ptr()
236+
}
237+
}
238+
239+
fn construct_envp(env: BTreeMap<DefaultEnvKey, OsString>, saw_nul: &mut bool) -> CStringArray {
240+
let mut result = CStringArray::with_capacity(env.len());
241+
for (k, v) in env {
242+
let mut k: OsString = k.into();
243+
244+
// Reserve additional space for '=' and null terminator
245+
k.reserve_exact(v.len() + 2);
246+
k.push("=");
247+
k.push(&v);
248+
249+
// Add the new entry into the array
250+
if let Ok(item) = CString::new(k.into_vec()) {
251+
result.push(item);
252+
} else {
253+
*saw_nul = true;
254+
}
255+
}
256+
257+
result
258+
}
259+
271260
impl Stdio {
272261
pub fn to_child_stdio(&self, readable: bool)
273262
-> io::Result<(ChildStdio, Option<AnonPipe>)> {
@@ -337,18 +326,6 @@ impl ChildStdio {
337326
}
338327
}
339328

340-
fn pair_to_key(key: &OsStr, value: &OsStr, saw_nul: &mut bool) -> CString {
341-
let (key, value) = (key.as_bytes(), value.as_bytes());
342-
let mut v = Vec::with_capacity(key.len() + value.len() + 1);
343-
v.extend(key);
344-
v.push(b'=');
345-
v.extend(value);
346-
CString::new(v).unwrap_or_else(|_e| {
347-
*saw_nul = true;
348-
CString::new("foo=bar").unwrap()
349-
})
350-
}
351-
352329
impl fmt::Debug for Command {
353330
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
354331
write!(f, "{:?}", self.program)?;

‎src/libstd/sys/unix/process/process_fuchsia.rs

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,14 +23,16 @@ use sys::process::process_common::*;
2323
impl Command {
2424
pub fn spawn(&mut self, default: Stdio, needs_stdin: bool)
2525
-> io::Result<(Process, StdioPipes)> {
26+
let envp = self.capture_env();
27+
2628
if self.saw_nul() {
2729
return Err(io::Error::new(io::ErrorKind::InvalidInput,
2830
"nul byte found in provided data"));
2931
}
3032

3133
let (ours, theirs) = self.setup_io(default, needs_stdin)?;
3234

33-
let process_handle = unsafe { self.do_exec(theirs)? };
35+
let process_handle = unsafe { self.do_exec(theirs, envp.as_ref())? };
3436

3537
Ok((Process { handle: Handle::new(process_handle) }, ours))
3638
}
@@ -50,13 +52,13 @@ impl Command {
5052
}
5153
}
5254

53-
unsafe fn do_exec(&mut self, stdio: ChildPipes)
55+
unsafe fn do_exec(&mut self, stdio: ChildPipes, maybe_envp: Option<&CStringArray>)
5456
-> io::Result<zx_handle_t> {
5557
use sys::process::zircon::*;
5658

5759
let job_handle = zx_job_default();
58-
let envp = match *self.get_envp() {
59-
Some(ref envp) => envp.as_ptr(),
60+
let envp = match maybe_envp {
61+
Some(envp) => envp.as_ptr(),
6062
None => ptr::null(),
6163
};
6264

‎src/libstd/sys/unix/process/process_unix.rs

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@ impl Command {
2626

2727
const CLOEXEC_MSG_FOOTER: &'static [u8] = b"NOEX";
2828

29+
let envp = self.capture_env();
30+
2931
if self.saw_nul() {
3032
return Err(io::Error::new(ErrorKind::InvalidInput,
3133
"nul byte found in provided data"));
@@ -38,7 +40,7 @@ impl Command {
3840
match cvt(libc::fork())? {
3941
0 => {
4042
drop(input);
41-
let err = self.do_exec(theirs);
43+
let err = self.do_exec(theirs, envp.as_ref());
4244
let errno = err.raw_os_error().unwrap_or(libc::EINVAL) as u32;
4345
let bytes = [
4446
(errno >> 24) as u8,
@@ -99,13 +101,15 @@ impl Command {
99101
}
100102

101103
pub fn exec(&mut self, default: Stdio) -> io::Error {
104+
let envp = self.capture_env();
105+
102106
if self.saw_nul() {
103107
return io::Error::new(ErrorKind::InvalidInput,
104108
"nul byte found in provided data")
105109
}
106110

107111
match self.setup_io(default, true) {
108-
Ok((_, theirs)) => unsafe { self.do_exec(theirs) },
112+
Ok((_, theirs)) => unsafe { self.do_exec(theirs, envp.as_ref()) },
109113
Err(e) => e,
110114
}
111115
}
@@ -140,7 +144,11 @@ impl Command {
140144
// allocation). Instead we just close it manually. This will never
141145
// have the drop glue anyway because this code never returns (the
142146
// child will either exec() or invoke libc::exit)
143-
unsafe fn do_exec(&mut self, stdio: ChildPipes) -> io::Error {
147+
unsafe fn do_exec(
148+
&mut self,
149+
stdio: ChildPipes,
150+
maybe_envp: Option<&CStringArray>
151+
) -> io::Error {
144152
use sys::{self, cvt_r};
145153

146154
macro_rules! t {
@@ -180,7 +188,7 @@ impl Command {
180188
if let Some(ref cwd) = *self.get_cwd() {
181189
t!(cvt(libc::chdir(cwd.as_ptr())));
182190
}
183-
if let Some(ref envp) = *self.get_envp() {
191+
if let Some(envp) = maybe_envp {
184192
*sys::os::environ() = envp.as_ptr();
185193
}
186194

‎src/libstd/sys/wasm/process.rs

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -14,12 +14,14 @@ use io;
1414
use sys::fs::File;
1515
use sys::pipe::AnonPipe;
1616
use sys::{unsupported, Void};
17+
use sys_common::process::{CommandEnv, DefaultEnvKey};
1718

1819
////////////////////////////////////////////////////////////////////////////////
1920
// Command
2021
////////////////////////////////////////////////////////////////////////////////
2122

2223
pub struct Command {
24+
env: CommandEnv<DefaultEnvKey>
2325
}
2426

2527
// passed back to std::process with the pipes connected to the child, if any
@@ -38,19 +40,16 @@ pub enum Stdio {
3840

3941
impl Command {
4042
pub fn new(_program: &OsStr) -> Command {
41-
Command {}
43+
Command {
44+
env: Default::default()
45+
}
4246
}
4347

4448
pub fn arg(&mut self, _arg: &OsStr) {
4549
}
4650

47-
pub fn env(&mut self, _key: &OsStr, _val: &OsStr) {
48-
}
49-
50-
pub fn env_remove(&mut self, _key: &OsStr) {
51-
}
52-
53-
pub fn env_clear(&mut self) {
51+
pub fn env_mut(&mut self) -> &mut CommandEnv<DefaultEnvKey> {
52+
&mut self.env
5453
}
5554

5655
pub fn cwd(&mut self, _dir: &OsStr) {

‎src/libstd/sys/windows/os_str.rs

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ use sys_common::wtf8::{Wtf8, Wtf8Buf};
1717
use mem;
1818
use rc::Rc;
1919
use sync::Arc;
20-
use sys_common::{AsInner, IntoInner};
20+
use sys_common::{AsInner, IntoInner, FromInner};
2121

2222
#[derive(Clone, Hash)]
2323
pub struct Buf {
@@ -30,6 +30,12 @@ impl IntoInner<Wtf8Buf> for Buf {
3030
}
3131
}
3232

33+
impl FromInner<Wtf8Buf> for Buf {
34+
fn from_inner(inner: Wtf8Buf) -> Self {
35+
Buf { inner }
36+
}
37+
}
38+
3339
impl AsInner<Wtf8> for Buf {
3440
fn as_inner(&self) -> &Wtf8 {
3541
&self.inner

‎src/libstd/sys/windows/process.rs

Lines changed: 52 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,10 @@
88
// option. This file may not be copied, modified, or distributed
99
// except according to those terms.
1010

11+
#![unstable(feature = "process_internals", issue = "0")]
12+
1113
use ascii::AsciiExt;
12-
use collections::HashMap;
13-
use collections;
14+
use collections::BTreeMap;
1415
use env::split_paths;
1516
use env;
1617
use ffi::{OsString, OsStr};
@@ -28,19 +29,42 @@ use sys::fs::{OpenOptions, File};
2829
use sys::handle::Handle;
2930
use sys::pipe::{self, AnonPipe};
3031
use sys::stdio;
31-
use sys::{self, cvt};
32-
use sys_common::{AsInner, FromInner};
32+
use sys::cvt;
33+
use sys_common::{AsInner, FromInner, IntoInner};
34+
use sys_common::process::{CommandEnv, EnvKey};
35+
use alloc::borrow::Borrow;
3336

3437
////////////////////////////////////////////////////////////////////////////////
3538
// Command
3639
////////////////////////////////////////////////////////////////////////////////
3740

38-
fn mk_key(s: &OsStr) -> OsString {
39-
FromInner::from_inner(sys::os_str::Buf {
40-
inner: s.as_inner().inner.to_ascii_uppercase()
41-
})
41+
#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)]
42+
#[doc(hidden)]
43+
pub struct WindowsEnvKey(OsString);
44+
45+
impl From<OsString> for WindowsEnvKey {
46+
fn from(k: OsString) -> Self {
47+
let mut buf = k.into_inner().into_inner();
48+
buf.make_ascii_uppercase();
49+
WindowsEnvKey(FromInner::from_inner(FromInner::from_inner(buf)))
50+
}
51+
}
52+
53+
impl From<WindowsEnvKey> for OsString {
54+
fn from(k: WindowsEnvKey) -> Self { k.0 }
4255
}
4356

57+
impl Borrow<OsStr> for WindowsEnvKey {
58+
fn borrow(&self) -> &OsStr { &self.0 }
59+
}
60+
61+
impl AsRef<OsStr> for WindowsEnvKey {
62+
fn as_ref(&self) -> &OsStr { &self.0 }
63+
}
64+
65+
impl EnvKey for WindowsEnvKey {}
66+
67+
4468
fn ensure_no_nuls<T: AsRef<OsStr>>(str: T) -> io::Result<T> {
4569
if str.as_ref().encode_wide().any(|b| b == 0) {
4670
Err(io::Error::new(ErrorKind::InvalidInput, "nul byte found in provided data"))
@@ -52,7 +76,7 @@ fn ensure_no_nuls<T: AsRef<OsStr>>(str: T) -> io::Result<T> {
5276
pub struct Command {
5377
program: OsString,
5478
args: Vec<OsString>,
55-
env: Option<HashMap<OsString, OsString>>,
79+
env: CommandEnv<WindowsEnvKey>,
5680
cwd: Option<OsString>,
5781
flags: u32,
5882
detach: bool, // not currently exposed in std::process
@@ -83,7 +107,7 @@ impl Command {
83107
Command {
84108
program: program.to_os_string(),
85109
args: Vec::new(),
86-
env: None,
110+
env: Default::default(),
87111
cwd: None,
88112
flags: 0,
89113
detach: false,
@@ -96,23 +120,8 @@ impl Command {
96120
pub fn arg(&mut self, arg: &OsStr) {
97121
self.args.push(arg.to_os_string())
98122
}
99-
fn init_env_map(&mut self){
100-
if self.env.is_none() {
101-
self.env = Some(env::vars_os().map(|(key, val)| {
102-
(mk_key(&key), val)
103-
}).collect());
104-
}
105-
}
106-
pub fn env(&mut self, key: &OsStr, val: &OsStr) {
107-
self.init_env_map();
108-
self.env.as_mut().unwrap().insert(mk_key(key), val.to_os_string());
109-
}
110-
pub fn env_remove(&mut self, key: &OsStr) {
111-
self.init_env_map();
112-
self.env.as_mut().unwrap().remove(&mk_key(key));
113-
}
114-
pub fn env_clear(&mut self) {
115-
self.env = Some(HashMap::new())
123+
pub fn env_mut(&mut self) -> &mut CommandEnv<WindowsEnvKey> {
124+
&mut self.env
116125
}
117126
pub fn cwd(&mut self, dir: &OsStr) {
118127
self.cwd = Some(dir.to_os_string())
@@ -132,13 +141,12 @@ impl Command {
132141

133142
pub fn spawn(&mut self, default: Stdio, needs_stdin: bool)
134143
-> io::Result<(Process, StdioPipes)> {
144+
let maybe_env = self.env.capture_if_changed();
135145
// To have the spawning semantics of unix/windows stay the same, we need
136146
// to read the *child's* PATH if one is provided. See #15149 for more
137147
// details.
138-
let program = self.env.as_ref().and_then(|env| {
139-
for (key, v) in env {
140-
if OsStr::new("PATH") != &**key { continue }
141-
148+
let program = maybe_env.as_ref().and_then(|env| {
149+
if let Some(v) = env.get(OsStr::new("PATH")) {
142150
// Split the value and test each path to see if the
143151
// program exists.
144152
for path in split_paths(&v) {
@@ -148,7 +156,6 @@ impl Command {
148156
return Some(path.into_os_string())
149157
}
150158
}
151-
break
152159
}
153160
None
154161
});
@@ -167,7 +174,7 @@ impl Command {
167174
flags |= c::DETACHED_PROCESS | c::CREATE_NEW_PROCESS_GROUP;
168175
}
169176

170-
let (envp, _data) = make_envp(self.env.as_ref())?;
177+
let (envp, _data) = make_envp(maybe_env)?;
171178
let (dirp, _data) = make_dirp(self.cwd.as_ref())?;
172179
let mut pi = zeroed_process_information();
173180

@@ -488,25 +495,24 @@ fn make_command_line(prog: &OsStr, args: &[OsString]) -> io::Result<Vec<u16>> {
488495
}
489496
}
490497

491-
fn make_envp(env: Option<&collections::HashMap<OsString, OsString>>)
498+
fn make_envp(maybe_env: Option<BTreeMap<WindowsEnvKey, OsString>>)
492499
-> io::Result<(*mut c_void, Vec<u16>)> {
493500
// On Windows we pass an "environment block" which is not a char**, but
494501
// rather a concatenation of null-terminated k=v\0 sequences, with a final
495502
// \0 to terminate.
496-
match env {
497-
Some(env) => {
498-
let mut blk = Vec::new();
499-
500-
for pair in env {
501-
blk.extend(ensure_no_nuls(pair.0)?.encode_wide());
502-
blk.push('=' as u16);
503-
blk.extend(ensure_no_nuls(pair.1)?.encode_wide());
504-
blk.push(0);
505-
}
503+
if let Some(env) = maybe_env {
504+
let mut blk = Vec::new();
505+
506+
for (k, v) in env {
507+
blk.extend(ensure_no_nuls(k.0)?.encode_wide());
508+
blk.push('=' as u16);
509+
blk.extend(ensure_no_nuls(v)?.encode_wide());
506510
blk.push(0);
507-
Ok((blk.as_mut_ptr() as *mut c_void, blk))
508511
}
509-
_ => Ok((ptr::null_mut(), Vec::new()))
512+
blk.push(0);
513+
Ok((blk.as_mut_ptr() as *mut c_void, blk))
514+
} else {
515+
Ok((ptr::null_mut(), Vec::new()))
510516
}
511517
}
512518

‎src/libstd/sys_common/mod.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ pub mod thread_local;
4444
pub mod util;
4545
pub mod wtf8;
4646
pub mod bytestring;
47+
pub mod process;
4748

4849
cfg_if! {
4950
if #[cfg(any(target_os = "redox", target_os = "l4re"))] {

‎src/libstd/sys_common/process.rs

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,124 @@
1+
// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
2+
// file at the top-level directory of this distribution and at
3+
// http://rust-lang.org/COPYRIGHT.
4+
//
5+
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
6+
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
7+
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
8+
// option. This file may not be copied, modified, or distributed
9+
// except according to those terms.
10+
11+
#![allow(dead_code)]
12+
#![unstable(feature = "process_internals", issue = "0")]
13+
14+
use ffi::{OsStr, OsString};
15+
use env;
16+
use collections::BTreeMap;
17+
use alloc::borrow::Borrow;
18+
19+
pub trait EnvKey:
20+
From<OsString> + Into<OsString> +
21+
Borrow<OsStr> + Borrow<Self> + AsRef<OsStr> +
22+
Ord + Clone {}
23+
24+
// Implement a case-sensitive environment variable key
25+
#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)]
26+
pub struct DefaultEnvKey(OsString);
27+
28+
impl From<OsString> for DefaultEnvKey {
29+
fn from(k: OsString) -> Self { DefaultEnvKey(k) }
30+
}
31+
32+
impl From<DefaultEnvKey> for OsString {
33+
fn from(k: DefaultEnvKey) -> Self { k.0 }
34+
}
35+
36+
impl Borrow<OsStr> for DefaultEnvKey {
37+
fn borrow(&self) -> &OsStr { &self.0 }
38+
}
39+
40+
impl AsRef<OsStr> for DefaultEnvKey {
41+
fn as_ref(&self) -> &OsStr { &self.0 }
42+
}
43+
44+
impl EnvKey for DefaultEnvKey {}
45+
46+
// Stores a set of changes to an environment
47+
#[derive(Clone, Debug)]
48+
pub struct CommandEnv<K> {
49+
clear: bool,
50+
vars: BTreeMap<K, Option<OsString>>
51+
}
52+
53+
impl<K: EnvKey> Default for CommandEnv<K> {
54+
fn default() -> Self {
55+
CommandEnv {
56+
clear: false,
57+
vars: Default::default()
58+
}
59+
}
60+
}
61+
62+
impl<K: EnvKey> CommandEnv<K> {
63+
// Capture the current environment with these changes applied
64+
pub fn capture(&self) -> BTreeMap<K, OsString> {
65+
let mut result = BTreeMap::<K, OsString>::new();
66+
if !self.clear {
67+
for (k, v) in env::vars_os() {
68+
result.insert(k.into(), v);
69+
}
70+
}
71+
for (k, maybe_v) in &self.vars {
72+
if let &Some(ref v) = maybe_v {
73+
result.insert(k.clone(), v.clone());
74+
} else {
75+
result.remove(k);
76+
}
77+
}
78+
result
79+
}
80+
81+
// Apply these changes directly to the current environment
82+
pub fn apply(&self) {
83+
if self.clear {
84+
for (k, _) in env::vars_os() {
85+
env::remove_var(k);
86+
}
87+
}
88+
for (key, maybe_val) in self.vars.iter() {
89+
if let &Some(ref val) = maybe_val {
90+
env::set_var(key, val);
91+
} else {
92+
env::remove_var(key);
93+
}
94+
}
95+
}
96+
97+
pub fn is_unchanged(&self) -> bool {
98+
!self.clear && self.vars.is_empty()
99+
}
100+
101+
pub fn capture_if_changed(&self) -> Option<BTreeMap<K, OsString>> {
102+
if self.is_unchanged() {
103+
None
104+
} else {
105+
Some(self.capture())
106+
}
107+
}
108+
109+
// The following functions build up changes
110+
pub fn set(&mut self, key: &OsStr, value: &OsStr) {
111+
self.vars.insert(key.to_owned().into(), Some(value.to_owned()));
112+
}
113+
pub fn remove(&mut self, key: &OsStr) {
114+
if self.clear {
115+
self.vars.remove(key);
116+
} else {
117+
self.vars.insert(key.to_owned().into(), None);
118+
}
119+
}
120+
pub fn clear(&mut self) {
121+
self.clear = true;
122+
self.vars.clear();
123+
}
124+
}

‎src/libstd/sys_common/wtf8.rs

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,12 @@ impl ops::Deref for Wtf8Buf {
134134
}
135135
}
136136

137+
impl ops::DerefMut for Wtf8Buf {
138+
fn deref_mut(&mut self) -> &mut Wtf8 {
139+
self.as_mut_slice()
140+
}
141+
}
142+
137143
/// Format the string with double quotes,
138144
/// and surrogates as `\u` followed by four hexadecimal digits.
139145
/// Example: `"a\u{D800}"` for a string with code points [U+0061, U+D800]
@@ -221,6 +227,11 @@ impl Wtf8Buf {
221227
unsafe { Wtf8::from_bytes_unchecked(&self.bytes) }
222228
}
223229

230+
#[inline]
231+
pub fn as_mut_slice(&mut self) -> &mut Wtf8 {
232+
unsafe { Wtf8::from_mut_bytes_unchecked(&mut self.bytes) }
233+
}
234+
224235
/// Reserves capacity for at least `additional` more bytes to be inserted
225236
/// in the given `Wtf8Buf`.
226237
/// The collection may reserve more space to avoid frequent reallocations.
@@ -486,6 +497,15 @@ impl Wtf8 {
486497
mem::transmute(value)
487498
}
488499

500+
/// Creates a mutable WTF-8 slice from a mutable WTF-8 byte slice.
501+
///
502+
/// Since the byte slice is not checked for valid WTF-8, this functions is
503+
/// marked unsafe.
504+
#[inline]
505+
unsafe fn from_mut_bytes_unchecked(value: &mut [u8]) -> &mut Wtf8 {
506+
mem::transmute(value)
507+
}
508+
489509
/// Returns the length, in WTF-8 bytes.
490510
#[inline]
491511
pub fn len(&self) -> usize {

0 commit comments

Comments
 (0)
Please sign in to comment.