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