@@ -3,15 +3,15 @@ use crate::io::{self, Error, ErrorKind};
3
3
use crate :: mem;
4
4
use crate :: num:: NonZeroI32 ;
5
5
use crate :: sys;
6
- use crate :: sys:: cvt;
7
6
use crate :: sys:: process:: process_common:: * ;
7
+ use crate :: sys:: { cvt, memchr} ;
8
8
use core:: ffi:: NonZero_c_int ;
9
9
10
10
#[ cfg( target_os = "linux" ) ]
11
11
use crate :: os:: linux:: process:: PidFd ;
12
12
#[ cfg( target_os = "linux" ) ]
13
13
use crate :: os:: unix:: io:: AsRawFd ;
14
-
14
+ use crate :: sys :: os :: { env_read_lock , environ } ;
15
15
#[ cfg( any(
16
16
target_os = "macos" ,
17
17
target_os = "watchos" ,
@@ -29,6 +29,8 @@ use libc::RTP_ID as pid_t;
29
29
#[ cfg( not( target_os = "vxworks" ) ) ]
30
30
use libc:: { c_int, pid_t} ;
31
31
32
+ use crate :: collections:: HashSet ;
33
+ use crate :: ffi:: { CStr , CString } ;
32
34
#[ cfg( not( any(
33
35
target_os = "vxworks" ,
34
36
target_os = "l4re" ,
@@ -68,6 +70,90 @@ cfg_if::cfg_if! {
68
70
// Command
69
71
////////////////////////////////////////////////////////////////////////////////
70
72
73
+ #[ cfg( target_os = "linux" ) ]
74
+ fn count_env_vars ( ) -> usize {
75
+ let mut count = 0 ;
76
+ unsafe {
77
+ let _guard = env_read_lock ( ) ;
78
+ let mut environ = * environ ( ) ;
79
+ while !( * environ) . is_null ( ) {
80
+ environ = environ. add ( 1 ) ;
81
+ count += 1 ;
82
+ }
83
+ }
84
+ count
85
+ }
86
+
87
+ /// Super-duper optimized version of capturing environment variables, that tries to avoid
88
+ /// unnecessary allocations and sorting.
89
+ #[ cfg( target_os = "linux" ) ]
90
+ fn capture_envp ( cmd : & mut Command ) -> CStringArray {
91
+ use crate :: os:: unix:: ffi:: OsStrExt ;
92
+
93
+ // Count the upper bound of environment variables (vars from the environ + vars coming from the
94
+ // command).
95
+ let env_count_upper_bound = count_env_vars ( ) + cmd. env . vars . len ( ) ;
96
+
97
+ let mut env_array = CStringArray :: with_capacity ( env_count_upper_bound) ;
98
+
99
+ // Remember which vars were already set by the user.
100
+ // If the user value is Some, we will add the variable to `env_array` and modify `visited`.
101
+ // If the user value is None, we will only modify `visited`.
102
+ // In either case, a variable with the same name from `environ` will not be added to `env_array`.
103
+ let mut visited: HashSet < & [ u8 ] > = HashSet :: with_capacity ( cmd. env . vars . len ( ) ) ;
104
+
105
+ // First, add user defined variables to `env_array`, and mark the visited ones.
106
+ for ( key, maybe_value) in cmd. env . vars . iter ( ) {
107
+ if let Some ( value) = maybe_value {
108
+ // One extra byte for '=', and one extra byte for the NULL terminator.
109
+ let mut env_var: Vec < u8 > =
110
+ Vec :: with_capacity ( key. as_bytes ( ) . len ( ) + value. as_bytes ( ) . len ( ) + 2 ) ;
111
+ env_var. extend_from_slice ( key. as_bytes ( ) ) ;
112
+ env_var. push ( b'=' ) ;
113
+ env_var. extend_from_slice ( value. as_bytes ( ) ) ;
114
+
115
+ if let Ok ( item) = CString :: new ( env_var) {
116
+ env_array. push ( item) ;
117
+ } else {
118
+ cmd. saw_nul = true ;
119
+ return env_array;
120
+ }
121
+ }
122
+ visited. insert ( key. as_bytes ( ) ) ;
123
+ }
124
+
125
+ // Then, if we're not clearing the original environment, go through it, and add each variable
126
+ // to env_array if we haven't seen it yet.
127
+ if !cmd. env . clear {
128
+ unsafe {
129
+ let _guard = env_read_lock ( ) ;
130
+ let mut environ = * environ ( ) ;
131
+ if !environ. is_null ( ) {
132
+ while !( * environ) . is_null ( ) {
133
+ let c_str = CStr :: from_ptr ( * environ) ;
134
+ let key_value = c_str. to_bytes ( ) ;
135
+ if !key_value. is_empty ( ) {
136
+ if let Some ( pos) = memchr:: memchr ( b'=' , & key_value[ 1 ..] ) . map ( |p| p + 1 ) {
137
+ let key = & key_value[ ..pos] ;
138
+ if !visited. contains ( & key) {
139
+ env_array. push ( CString :: from ( c_str) ) ;
140
+ }
141
+ }
142
+ }
143
+ environ = environ. add ( 1 ) ;
144
+ }
145
+ }
146
+ }
147
+ }
148
+
149
+ env_array
150
+ }
151
+
152
+ #[ cfg( target_os = "linux" ) ]
153
+ pub fn capture_env_linux ( cmd : & mut Command ) -> Option < CStringArray > {
154
+ if cmd. env . is_unchanged ( ) { None } else { Some ( capture_envp ( cmd) ) }
155
+ }
156
+
71
157
impl Command {
72
158
pub fn spawn (
73
159
& mut self ,
@@ -76,6 +162,9 @@ impl Command {
76
162
) -> io:: Result < ( Process , StdioPipes ) > {
77
163
const CLOEXEC_MSG_FOOTER : [ u8 ; 4 ] = * b"NOEX" ;
78
164
165
+ #[ cfg( target_os = "linux" ) ]
166
+ let envp = capture_env_linux ( self ) ;
167
+ #[ cfg( not( target_os = "linux" ) ) ]
79
168
let envp = self . capture_env ( ) ;
80
169
81
170
if self . saw_nul ( ) {
0 commit comments