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