authorgravatar for squeek502@hotmail.comRyan Liptak <squeek502@hotmail.com> 2022-01-16 20:11:08-08:00
committergravatar for johnnymarler@gmail.comJonathan Marler <johnnymarler@gmail.com> 2022-05-11 18:40:53-06:00
log15d5988e692c182892a118115fd7025048e06c29
tree8181463a2a8983373a0b8b53f82caae785f0b5ee
parentd383b940c2e9c9d0f2e8ef7607b38b7a74021b47

Add `process.EnvMap`, a platform-independent environment variable map

EnvMap provides the same API as the previously used BufMap (besides `putMove` and `getPtr`), so usage sites of `getEnvMap` can usually remain unchanged. For non-Windows, EnvMap is a wrapper around BufMap. On Windows, it uses a new EnvMapWindows to handle some Windows-specific behavior: - Lookups use Unicode-aware case insensitivity (but `get` cannot return an error because EnvMapWindows has an internal buffer to use for lookup conversions) - Canonical names are returned when iterating the EnvMap Fixes #10561, closes #4603

4 files changed, 421 insertions(+), 13 deletions(-)

lib/std/buf_map.zig+1-1
......@@ -9,7 +9,7 @@ const testing = std.testing;
99pub const BufMap = struct {
1010 hash_map: BufMapHashMap,
1111
12 const BufMapHashMap = StringHashMap([]const u8);
12 pub const BufMapHashMap = StringHashMap([]const u8);
1313
1414 /// Create a BufMap backed by a specific allocator.
1515 /// That allocator will be used for both backing allocations
lib/std/os/windows/ntdll.zig+6
......@@ -229,6 +229,12 @@ pub extern "ntdll" fn RtlEqualUnicodeString(
229229 CaseInSensitive: BOOLEAN,
230230) callconv(WINAPI) BOOLEAN;
231231
232pub extern "NtDll" fn RtlUpcaseUnicodeString(
233 DestinationString: *UNICODE_STRING,
234 SourceString: *const UNICODE_STRING,
235 AllocateDestinationString: BOOLEAN,
236) callconv(WINAPI) NTSTATUS;
237
232238pub extern "ntdll" fn NtLockFile(
233239 FileHandle: HANDLE,
234240 Event: ?HANDLE,
lib/std/process.zig+391-12
......@@ -2,7 +2,6 @@ const std = @import("std.zig");
22const builtin = @import("builtin");
33const os = std.os;
44const fs = std.fs;
5const BufMap = std.BufMap;
65const mem = std.mem;
76const math = std.math;
87const Allocator = mem.Allocator;
......@@ -53,9 +52,385 @@ test "getCwdAlloc" {
5352 testing.allocator.free(cwd);
5453}
5554
56/// Caller owns resulting `BufMap`.
57pub fn getEnvMap(allocator: Allocator) !BufMap {
58 var result = BufMap.init(allocator);
55/// EnvMap for Windows that handles Unicode-aware case insensitivity for lookups, while also
56/// providing the canonical environment variable names when iterating.
57///
58/// Allows for zero-allocation lookups (even though it needs to do UTF-8 -> UTF-16 -> uppercase
59/// conversions) by allocating a buffer large enough to fit the largest environment variable
60/// name, and using that when doing lookups (i.e. anything that overflows the buffer can be treated
61/// as the environment variable not being found).
62pub const EnvMapWindows = struct {
63 allocator: Allocator,
64 /// Keys are UTF-16le stored as []const u8
65 uppercased_map: std.StringHashMapUnmanaged(EnvValue),
66 /// Buffer for converting to uppercased UTF-16 on key lookups
67 /// Must call `reallocUppercaseBuf` before doing any lookups after a `put` call.
68 uppercase_buf_utf16: []u16 = &[_]u16{},
69 max_name_utf16_length: usize = 0,
70
71 pub const EnvValue = struct {
72 value: []const u8,
73 canonical_name: []const u8,
74 };
75
76 const Self = @This();
77
78 /// Deinitialize with `deinit`.
79 pub fn init(allocator: Allocator) Self {
80 return .{
81 .allocator = allocator,
82 .uppercased_map = std.StringHashMapUnmanaged(EnvValue){},
83 };
84 }
85
86 pub fn deinit(self: *Self) void {
87 var it = self.uppercased_map.iterator();
88 while (it.next()) |entry| {
89 self.allocator.free(entry.key_ptr.*);
90 self.allocator.free(entry.value_ptr.value);
91 self.allocator.free(entry.value_ptr.canonical_name);
92 }
93 self.uppercased_map.deinit(self.allocator);
94 self.allocator.free(self.uppercase_buf_utf16);
95 }
96
97 /// Increases the size of the uppercase buffer if the maximum name size has increased.
98 /// Must be called before any `get` calls after any number of `put` calls.
99 pub fn reallocUppercaseBuf(self: *Self) !void {
100 if (self.max_name_utf16_length > self.uppercase_buf_utf16.len) {
101 self.uppercase_buf_utf16 = try self.allocator.realloc(self.uppercase_buf_utf16, self.max_name_utf16_length);
102 }
103 }
104
105 /// Converts `src` to uppercase using `RtlUpcaseUnicodeString` and puts the result in `dest`.
106 /// Returns the length of the converted UTF-16 string. `dest.len` must be >= `src.len`.
107 ///
108 /// Note: As of now, RtlUpcaseUnicodeString does not seem to handle codepoints above 0x10000
109 /// (i.e. those that require a surrogate pair), so this function will always return a length
110 /// equal to `src.len`. However, if RtlUpcaseUnicodeString is updated to handle codepoints above
111 /// 0x10000, this property would still hold unless there are lowercase <-> uppercase conversions
112 /// that cross over the boundary between codepoints >= 0x10000 and < 0x10000.
113 /// TODO: Is it feasible that Unicode lowercase <-> uppercase conversions could cross that boundary?
114 fn uppercaseName(dest: []u16, src: []const u16) u16 {
115 assert(dest.len >= src.len);
116
117 const dest_bytes = @intCast(u16, dest.len * 2);
118 var dest_string = os.windows.UNICODE_STRING{
119 .Length = dest_bytes,
120 .MaximumLength = dest_bytes,
121 .Buffer = @intToPtr([*]u16, @ptrToInt(dest.ptr)),
122 };
123 const src_bytes = @intCast(u16, src.len * 2);
124 const src_string = os.windows.UNICODE_STRING{
125 .Length = src_bytes,
126 .MaximumLength = src_bytes,
127 .Buffer = @intToPtr([*]u16, @ptrToInt(src.ptr)),
128 };
129 const rc = os.windows.ntdll.RtlUpcaseUnicodeString(&dest_string, &src_string, os.windows.FALSE);
130 switch (rc) {
131 .SUCCESS => return dest_string.Length / 2,
132 else => unreachable, // we are not allocating, so no errors should be possible
133 }
134 }
135
136 /// Note: Does not realloc the uppercase buf to allow for calling put for many variables and
137 /// only allocating the uppercase buf afterwards.
138 pub fn putUtf8(self: *Self, name: []const u8, value: []const u8) !void {
139 const uppercased_len = len: {
140 const name_uppercased_utf16 = uppercased: {
141 var name_utf16_buf = try std.ArrayListAligned(u8, @alignOf(u16)).initCapacity(self.allocator, name.len);
142 errdefer name_utf16_buf.deinit();
143
144 var uppercased_len = try std.unicode.utf8ToUtf16LeWriter(name_utf16_buf.writer(), name);
145 assert(uppercased_len == name_utf16_buf.items.len);
146
147 break :uppercased name_utf16_buf.toOwnedSlice();
148 };
149 errdefer self.allocator.free(name_uppercased_utf16);
150
151 const name_canonical = try self.allocator.dupe(u8, name);
152 errdefer self.allocator.free(name_canonical);
153
154 const value_dupe = try self.allocator.dupe(u8, value);
155 errdefer self.allocator.free(value_dupe);
156
157 const get_or_put = try self.uppercased_map.getOrPut(self.allocator, name_uppercased_utf16);
158 if (get_or_put.found_existing) {
159 // note: this is only safe from UAF because the errdefer that frees this value above
160 // no longer has a possibility of being triggered after this point
161 self.allocator.free(name_uppercased_utf16);
162 self.allocator.free(get_or_put.value_ptr.value);
163 self.allocator.free(get_or_put.value_ptr.canonical_name);
164 } else {
165 get_or_put.key_ptr.* = name_uppercased_utf16;
166 }
167 get_or_put.value_ptr.value = value_dupe;
168 get_or_put.value_ptr.canonical_name = name_canonical;
169
170 break :len name_uppercased_utf16.len;
171 };
172
173 // The buffer for case conversion for key lookups will need to be as big as the largest
174 // key stored in the hash map.
175 self.max_name_utf16_length = @maximum(self.max_name_utf16_length, uppercased_len);
176 }
177
178 /// Asserts that the name does not already exist in the map.
179 /// Note: Does not realloc the uppercase buf to allow for calling put for many variables and
180 /// only allocating the uppercase buf afterwards.
181 pub fn putUtf16NoClobber(self: *Self, name_utf16: []const u16, value_utf16: []const u16) !void {
182 const uppercased_len = len: {
183 const name_canonical = try std.unicode.utf16leToUtf8Alloc(self.allocator, name_utf16);
184 errdefer self.allocator.free(name_canonical);
185
186 const value = try std.unicode.utf16leToUtf8Alloc(self.allocator, value_utf16);
187 errdefer self.allocator.free(value);
188
189 const name_uppercased_utf16 = try self.allocator.alloc(u16, name_utf16.len);
190 errdefer self.allocator.free(name_uppercased_utf16);
191
192 const uppercased_len = uppercaseName(name_uppercased_utf16, name_utf16);
193 assert(uppercased_len == name_uppercased_utf16.len);
194
195 try self.uppercased_map.putNoClobber(self.allocator, std.mem.sliceAsBytes(name_uppercased_utf16), EnvValue{
196 .value = value,
197 .canonical_name = name_canonical,
198 });
199 break :len name_uppercased_utf16.len;
200 };
201
202 // The buffer for case conversion for key lookups will need to be as big as the largest
203 // key stored in the hash map.
204 self.max_name_utf16_length = @maximum(self.max_name_utf16_length, uppercased_len);
205 }
206
207 /// Attempts to convert a UTF-8 name into a uppercased UTF-16le name for a lookup. If the
208 /// name cannot be converted, this function will return `null`.
209 fn utf8ToUppercasedUtf16(self: Self, name: []const u8) ?[]u16 {
210 const name_utf16: []u16 = to_utf16: {
211 var utf16_buf_stream = std.io.fixedBufferStream(std.mem.sliceAsBytes(self.uppercase_buf_utf16));
212 _ = std.unicode.utf8ToUtf16LeWriter(utf16_buf_stream.writer(), name) catch |err| switch (err) {
213 // If the buffer isn't large enough, we can treat that as 'env var not found', as we
214 // know anything too large for the buffer can't be found in the map.
215 error.NoSpaceLeft => return null,
216 // Anything with invalid UTF-8 will also not be found in the map, so treat that as
217 // 'env var not found' too
218 error.InvalidUtf8 => return null,
219 };
220 break :to_utf16 std.mem.bytesAsSlice(u16, utf16_buf_stream.getWritten());
221 };
222
223 // uppercase in place
224 const uppercased_len = uppercaseName(name_utf16, name_utf16);
225 assert(uppercased_len == name_utf16.len);
226
227 return name_utf16;
228 }
229
230 /// Returns true if an entry was found and deleted, false otherwise.
231 pub fn remove(self: *Self, name: []const u8) bool {
232 const name_utf16 = self.utf8ToUppercasedUtf16(name) orelse return false;
233 const kv = self.uppercased_map.fetchRemove(std.mem.sliceAsBytes(name_utf16)) orelse return false;
234 self.allocator.free(kv.key);
235 self.allocator.free(kv.value.value);
236 self.allocator.free(kv.value.canonical_name);
237 return true;
238 }
239
240 pub fn get(self: Self, name: []const u8) ?EnvValue {
241 const name_utf16 = self.utf8ToUppercasedUtf16(name) orelse return null;
242 return self.uppercased_map.get(std.mem.sliceAsBytes(name_utf16));
243 }
244
245 pub fn count(self: Self) EnvMap.Size {
246 return self.uppercased_map.count();
247 }
248
249 pub fn iterator(self: *const Self) Iterator {
250 return .{
251 .env_map = self,
252 .uppercased_map_iterator = self.uppercased_map.iterator(),
253 };
254 }
255
256 pub const Iterator = struct {
257 env_map: *const Self,
258 uppercased_map_iterator: std.StringHashMapUnmanaged(EnvValue).Iterator,
259
260 pub fn next(it: *Iterator) ?EnvMap.Entry {
261 if (it.uppercased_map_iterator.next()) |uppercased_entry| {
262 return EnvMap.Entry{
263 .name = uppercased_entry.value_ptr.canonical_name,
264 .value = uppercased_entry.value_ptr.value,
265 };
266 } else {
267 return null;
268 }
269 }
270 };
271};
272
273test "EnvMapWindows" {
274 if (builtin.os.tag != .windows) return error.SkipZigTest;
275
276 var env_map = EnvMapWindows.init(testing.allocator);
277 defer env_map.deinit();
278
279 // both put methods
280 try env_map.putUtf16NoClobber(std.unicode.utf8ToUtf16LeStringLiteral("Path"), std.unicode.utf8ToUtf16LeStringLiteral("something"));
281 try env_map.putUtf8("КИРИЛЛИЦА", "something else");
282 try env_map.reallocUppercaseBuf();
283
284 try testing.expectEqual(@as(EnvMap.Size, 2), env_map.count());
285
286 // unicode-aware case-insensitive lookups
287 try testing.expectEqualStrings("something", env_map.get("PATH").?.value);
288 try testing.expectEqualStrings("something else", env_map.get("кириллица").?.value);
289 try testing.expect(env_map.get("missing") == null);
290
291 // canonical names when iterating
292 var it = env_map.iterator();
293 var count: EnvMap.Size = 0;
294 while (it.next()) |entry| {
295 const is_an_expected_name = std.mem.eql(u8, "Path", entry.name) or std.mem.eql(u8, "КИРИЛЛИЦА", entry.name);
296 try testing.expect(is_an_expected_name);
297 count += 1;
298 }
299 try testing.expectEqual(@as(EnvMap.Size, 2), count);
300}
301
302pub const EnvMap = struct {
303 storage: StorageType,
304
305 pub const StorageType = switch (builtin.os.tag) {
306 .windows => EnvMapWindows,
307 else => std.BufMap,
308 };
309
310 /// Matches what BufMap uses for its internal HashMap Size
311 pub const Size = u32;
312
313 const Self = @This();
314
315 /// Deinitialize with `deinit`.
316 pub fn init(allocator: Allocator) Self {
317 return Self{ .storage = StorageType.init(allocator) };
318 }
319
320 pub fn deinit(self: *Self) void {
321 self.storage.deinit();
322 }
323
324 pub fn get(self: Self, name: []const u8) ?[]const u8 {
325 switch (builtin.os.tag) {
326 .windows => {
327 if (self.storage.get(name)) |entry| {
328 return entry.value;
329 } else {
330 return null;
331 }
332 },
333 else => return self.storage.get(name),
334 }
335 }
336
337 pub fn count(self: Self) Size {
338 return self.storage.count();
339 }
340
341 pub fn iterator(self: *const Self) Iterator {
342 return .{ .storage_iterator = self.storage.iterator() };
343 }
344
345 pub fn put(self: *Self, name: []const u8, value: []const u8) !void {
346 switch (builtin.os.tag) {
347 .windows => {
348 try self.storage.putUtf8(name, value);
349 try self.storage.reallocUppercaseBuf();
350 },
351 else => return self.storage.put(name, value),
352 }
353 }
354
355 pub fn remove(self: *Self, name: []const u8) void {
356 _ = self.storage.remove(name);
357 }
358
359 pub const Entry = struct {
360 name: []const u8,
361 value: []const u8,
362 };
363
364 pub const Iterator = struct {
365 storage_iterator: switch (builtin.os.tag) {
366 .windows => EnvMapWindows.Iterator,
367 else => std.BufMap.BufMapHashMap.Iterator,
368 },
369
370 pub fn next(it: *Iterator) ?Entry {
371 switch (builtin.os.tag) {
372 .windows => return it.storage_iterator.next(),
373 else => {
374 if (it.storage_iterator.next()) |entry| {
375 return Entry{
376 .name = entry.key_ptr.*,
377 .value = entry.value_ptr.*,
378 };
379 } else {
380 return null;
381 }
382 },
383 }
384 }
385 };
386};
387
388test "EnvMap" {
389 var env = EnvMap.init(testing.allocator);
390 defer env.deinit();
391
392 try env.put("SOMETHING_NEW", "hello");
393 try testing.expectEqualStrings("hello", env.get("SOMETHING_NEW").?);
394 try testing.expectEqual(@as(EnvMap.Size, 1), env.count());
395
396 // overwrite
397 try env.put("SOMETHING_NEW", "something");
398 try testing.expectEqualStrings("something", env.get("SOMETHING_NEW").?);
399 try testing.expectEqual(@as(EnvMap.Size, 1), env.count());
400
401 // a new longer name to test the Windows-specific conversion buffer
402 try env.put("SOMETHING_NEW_AND_LONGER", "1");
403 try testing.expectEqualStrings("1", env.get("SOMETHING_NEW_AND_LONGER").?);
404 try testing.expectEqual(@as(EnvMap.Size, 2), env.count());
405
406 // case insensitivity on Windows only
407 if (builtin.os.tag == .windows) {
408 try testing.expectEqualStrings("1", env.get("something_New_aNd_LONGER").?);
409 } else {
410 try testing.expect(null == env.get("something_New_aNd_LONGER"));
411 }
412
413 var it = env.iterator();
414 var count: EnvMap.Size = 0;
415 while (it.next()) |entry| {
416 const is_an_expected_name = std.mem.eql(u8, "SOMETHING_NEW", entry.name) or std.mem.eql(u8, "SOMETHING_NEW_AND_LONGER", entry.name);
417 try testing.expect(is_an_expected_name);
418 count += 1;
419 }
420 try testing.expectEqual(@as(EnvMap.Size, 2), count);
421
422 env.remove("SOMETHING_NEW");
423 try testing.expect(env.get("SOMETHING_NEW") == null);
424
425 try testing.expectEqual(@as(EnvMap.Size, 1), env.count());
426}
427
428/// Returns a snapshot of the environment variables of the current process.
429/// Any modifications to the resulting EnvMap will not be not reflected in the environment, and
430/// likewise, any future modifications to the environment will not be reflected in the EnvMap.
431/// Caller owns resulting `EnvMap` and should call its `deinit` fn when done.
432pub fn getEnvMap(allocator: Allocator) !EnvMap {
433 var result = EnvMap.init(allocator);
59434 errdefer result.deinit();
60435
61436 if (builtin.os.tag == .windows) {
......@@ -65,23 +440,27 @@ pub fn getEnvMap(allocator: Allocator) !BufMap {
65440 while (ptr[i] != 0) {
66441 const key_start = i;
67442
443 // There are some special environment variables that start with =,
444 // so we need a special case to not treat = as a key/value separator
445 // if it's the first character.
446 // https://devblogs.microsoft.com/oldnewthing/20100506-00/?p=14133
447 if (ptr[key_start] == '=') i += 1;
448
68449 while (ptr[i] != 0 and ptr[i] != '=') : (i += 1) {}
69450 const key_w = ptr[key_start..i];
70 const key = try std.unicode.utf16leToUtf8Alloc(allocator, key_w);
71 errdefer allocator.free(key);
72451
73452 if (ptr[i] == '=') i += 1;
74453
75454 const value_start = i;
76455 while (ptr[i] != 0) : (i += 1) {}
77456 const value_w = ptr[value_start..i];
78 const value = try std.unicode.utf16leToUtf8Alloc(allocator, value_w);
79 errdefer allocator.free(value);
80457
81 i += 1; // skip over null byte
458 try result.storage.putUtf16NoClobber(key_w, value_w);
82459
83 try result.putMove(key, value);
460 i += 1; // skip over null byte
84461 }
462
463 try result.storage.reallocUppercaseBuf();
85464 return result;
86465 } else if (builtin.os.tag == .wasi and !builtin.link_libc) {
87466 var environ_count: usize = undefined;
......@@ -140,8 +519,8 @@ pub fn getEnvMap(allocator: Allocator) !BufMap {
140519 }
141520}
142521
143test "os.getEnvMap" {
144 var env = try getEnvMap(std.testing.allocator);
522test "getEnvMap" {
523 var env = try getEnvMap(testing.allocator);
145524 defer env.deinit();
146525}
147526
lib/std/unicode.zig+23
......@@ -710,6 +710,29 @@ pub fn utf8ToUtf16Le(utf16le: []u16, utf8: []const u8) !usize {
710710 return dest_i;
711711}
712712
713pub fn utf8ToUtf16LeWriter(writer: anytype, utf8: []const u8) !usize {
714 var src_i: usize = 0;
715 var bytes_written: usize = 0;
716 while (src_i < utf8.len) {
717 const n = utf8ByteSequenceLength(utf8[src_i]) catch return error.InvalidUtf8;
718 const next_src_i = src_i + n;
719 const codepoint = utf8Decode(utf8[src_i..next_src_i]) catch return error.InvalidUtf8;
720 if (codepoint < 0x10000) {
721 const short = @intCast(u16, codepoint);
722 try writer.writeIntLittle(u16, short);
723 bytes_written += 2;
724 } else {
725 const high = @intCast(u16, (codepoint - 0x10000) >> 10) + 0xD800;
726 const low = @intCast(u16, codepoint & 0x3FF) + 0xDC00;
727 try writer.writeIntLittle(u16, high);
728 try writer.writeIntLittle(u16, low);
729 bytes_written += 4;
730 }
731 src_i = next_src_i;
732 }
733 return bytes_written;
734}
735
713736test "utf8ToUtf16Le" {
714737 var utf16le: [2]u16 = [_]u16{0} ** 2;
715738 {