Skip to content

Commit 08a6c4c

Browse files
andrewrkandrewrk
authored
Merge pull request #23272 from squeek502/getenvw-optim
Windows: Faster `getenvW` and a standalone environment variable test
2 parents 326f254 + 66dcebc commit 08a6c4c

File tree

5 files changed

+244
-26
lines changed

5 files changed

+244
-26
lines changed

‎lib/std/os/windows.zig

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4890,6 +4890,10 @@ pub const RTL_USER_PROCESS_PARAMETERS = extern struct {
48904890
DllPath: UNICODE_STRING,
48914891
ImagePathName: UNICODE_STRING,
48924892
CommandLine: UNICODE_STRING,
4893+
/// Points to a NUL-terminated sequence of NUL-terminated
4894+
/// WTF-16 LE encoded `name=value` sequences.
4895+
/// Example using string literal syntax:
4896+
/// `"NAME=value\x00foo=bar\x00\x00"`
48934897
Environment: [*:0]WCHAR,
48944898
dwX: ULONG,
48954899
dwY: ULONG,

‎lib/std/process.zig

Lines changed: 31 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -419,10 +419,10 @@ pub fn getEnvVarOwned(allocator: Allocator, key: []const u8) GetEnvVarOwnedError
419419
}
420420
}
421421

422-
/// On Windows, `key` must be valid UTF-8.
422+
/// On Windows, `key` must be valid WTF-8.
423423
pub fn hasEnvVarConstant(comptime key: []const u8) bool {
424424
if (native_os == .windows) {
425-
const key_w = comptime unicode.utf8ToUtf16LeStringLiteral(key);
425+
const key_w = comptime unicode.wtf8ToWtf16LeStringLiteral(key);
426426
return getenvW(key_w) != null;
427427
} else if (native_os == .wasi and !builtin.link_libc) {
428428
@compileError("hasEnvVarConstant is not supported for WASI without libc");
@@ -431,10 +431,10 @@ pub fn hasEnvVarConstant(comptime key: []const u8) bool {
431431
}
432432
}
433433

434-
/// On Windows, `key` must be valid UTF-8.
434+
/// On Windows, `key` must be valid WTF-8.
435435
pub fn hasNonEmptyEnvVarConstant(comptime key: []const u8) bool {
436436
if (native_os == .windows) {
437-
const key_w = comptime unicode.utf8ToUtf16LeStringLiteral(key);
437+
const key_w = comptime unicode.wtf8ToWtf16LeStringLiteral(key);
438438
const value = getenvW(key_w) orelse return false;
439439
return value.len != 0;
440440
} else if (native_os == .wasi and !builtin.link_libc) {
@@ -451,10 +451,10 @@ pub const ParseEnvVarIntError = std.fmt.ParseIntError || error{EnvironmentVariab
451451
///
452452
/// Since the key is comptime-known, no allocation is needed.
453453
///
454-
/// On Windows, `key` must be valid UTF-8.
454+
/// On Windows, `key` must be valid WTF-8.
455455
pub fn parseEnvVarInt(comptime key: []const u8, comptime I: type, base: u8) ParseEnvVarIntError!I {
456456
if (native_os == .windows) {
457-
const key_w = comptime std.unicode.utf8ToUtf16LeStringLiteral(key);
457+
const key_w = comptime std.unicode.wtf8ToWtf16LeStringLiteral(key);
458458
const text = getenvW(key_w) orelse return error.EnvironmentVariableNotFound;
459459
return std.fmt.parseIntWithGenericCharacter(I, u16, text, base);
460460
} else if (native_os == .wasi and !builtin.link_libc) {
@@ -527,31 +527,33 @@ pub fn getenvW(key: [*:0]const u16) ?[:0]const u16 {
527527
@compileError("Windows-only");
528528
}
529529
const key_slice = mem.sliceTo(key, 0);
530+
// '=' anywhere but the start makes this an invalid environment variable name
531+
if (key_slice.len > 0 and std.mem.indexOfScalar(u16, key_slice[1..], '=') != null) {
532+
return null;
533+
}
530534
const ptr = windows.peb().ProcessParameters.Environment;
531535
var i: usize = 0;
532536
while (ptr[i] != 0) {
533-
const key_start = i;
537+
const key_value = mem.sliceTo(ptr[i..], 0);
534538

535539
// There are some special environment variables that start with =,
536540
// so we need a special case to not treat = as a key/value separator
537541
// if it's the first character.
538542
// https://devblogs.microsoft.com/oldnewthing/20100506-00/?p=14133
539-
if (ptr[key_start] == '=') i += 1;
540-
541-
while (ptr[i] != 0 and ptr[i] != '=') : (i += 1) {}
542-
const this_key = ptr[key_start..i];
543-
544-
if (ptr[i] == '=') i += 1;
545-
546-
const value_start = i;
547-
while (ptr[i] != 0) : (i += 1) {}
548-
const this_value = ptr[value_start..i :0];
543+
const equal_search_start: usize = if (key_value[0] == '=') 1 else 0;
544+
const equal_index = std.mem.indexOfScalarPos(u16, key_value, equal_search_start, '=') orelse {
545+
// This is enforced by CreateProcess.
546+
// If violated, CreateProcess will fail with INVALID_PARAMETER.
547+
unreachable; // must contain a =
548+
};
549549

550+
const this_key = key_value[0..equal_index];
550551
if (windows.eqlIgnoreCaseWTF16(key_slice, this_key)) {
551-
return this_value;
552+
return key_value[equal_index + 1 ..];
552553
}
553554

554-
i += 1; // skip over null byte
555+
// skip past the NUL terminator
556+
i += key_value.len + 1;
555557
}
556558
return null;
557559
}
@@ -2037,7 +2039,8 @@ test createNullDelimitedEnvMap {
20372039
pub fn createWindowsEnvBlock(allocator: mem.Allocator, env_map: *const EnvMap) ![]u16 {
20382040
// count bytes needed
20392041
const max_chars_needed = x: {
2040-
var max_chars_needed: usize = 4; // 4 for the final 4 null bytes
2042+
// Only need 2 trailing NUL code units for an empty environment
2043+
var max_chars_needed: usize = if (env_map.count() == 0) 2 else 1;
20412044
var it = env_map.iterator();
20422045
while (it.next()) |pair| {
20432046
// +1 for '='
@@ -2061,12 +2064,14 @@ pub fn createWindowsEnvBlock(allocator: mem.Allocator, env_map: *const EnvMap) !
20612064
}
20622065
result[i] = 0;
20632066
i += 1;
2064-
result[i] = 0;
2065-
i += 1;
2066-
result[i] = 0;
2067-
i += 1;
2068-
result[i] = 0;
2069-
i += 1;
2067+
// An empty environment is a special case that requires a redundant
2068+
// NUL terminator. CreateProcess will read the second code unit even
2069+
// though theoretically the first should be enough to recognize that the
2070+
// environment is empty (see https://nullprogram.com/blog/2023/08/23/)
2071+
if (env_map.count() == 0) {
2072+
result[i] = 0;
2073+
i += 1;
2074+
}
20702075
return try allocator.realloc(result, i);
20712076
}
20722077

‎test/standalone/build.zig.zon

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,9 @@
9696
.empty_env = .{
9797
.path = "empty_env",
9898
},
99+
.env_vars = .{
100+
.path = "env_vars",
101+
},
99102
.issue_11595 = .{
100103
.path = "issue_11595",
101104
},

‎test/standalone/env_vars/build.zig

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
const std = @import("std");
2+
const builtin = @import("builtin");
3+
4+
pub fn build(b: *std.Build) void {
5+
const test_step = b.step("test", "Test it");
6+
b.default_step = test_step;
7+
8+
const optimize: std.builtin.OptimizeMode = .Debug;
9+
10+
const main = b.addExecutable(.{
11+
.name = "main",
12+
.root_module = b.createModule(.{
13+
.root_source_file = b.path("main.zig"),
14+
.target = b.graph.host,
15+
.optimize = optimize,
16+
}),
17+
});
18+
19+
const run = b.addRunArtifact(main);
20+
run.clearEnvironment();
21+
run.setEnvironmentVariable("FOO", "123");
22+
run.setEnvironmentVariable("EQUALS", "ABC=123");
23+
run.setEnvironmentVariable("NO_VALUE", "");
24+
run.setEnvironmentVariable("КИРиллИЦА", "non-ascii አማርኛ \u{10FFFF}");
25+
if (b.graph.host.result.os.tag == .windows) {
26+
run.setEnvironmentVariable("=Hidden", "hi");
27+
// \xed\xa0\x80 is a WTF-8 encoded unpaired surrogate code point
28+
run.setEnvironmentVariable("INVALID_UTF16_\xed\xa0\x80", "\xed\xa0\x80");
29+
}
30+
run.disable_zig_progress = true;
31+
32+
test_step.dependOn(&run.step);
33+
}

‎test/standalone/env_vars/main.zig

Lines changed: 173 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,173 @@
1+
const std = @import("std");
2+
const builtin = @import("builtin");
3+
4+
// Note: the environment variables under test are set by the build.zig
5+
pub fn main() !void {
6+
var gpa: std.heap.GeneralPurposeAllocator(.{}) = .init;
7+
defer _ = gpa.deinit();
8+
const allocator = gpa.allocator();
9+
10+
var arena_state = std.heap.ArenaAllocator.init(allocator);
11+
defer arena_state.deinit();
12+
const arena = arena_state.allocator();
13+
14+
// hasNonEmptyEnvVar
15+
{
16+
try std.testing.expect(try std.process.hasNonEmptyEnvVar(allocator, "FOO"));
17+
try std.testing.expect(!(try std.process.hasNonEmptyEnvVar(allocator, "FOO=")));
18+
try std.testing.expect(!(try std.process.hasNonEmptyEnvVar(allocator, "FO")));
19+
try std.testing.expect(!(try std.process.hasNonEmptyEnvVar(allocator, "FOOO")));
20+
if (builtin.os.tag == .windows) {
21+
try std.testing.expect(try std.process.hasNonEmptyEnvVar(allocator, "foo"));
22+
}
23+
try std.testing.expect(try std.process.hasNonEmptyEnvVar(allocator, "EQUALS"));
24+
try std.testing.expect(!(try std.process.hasNonEmptyEnvVar(allocator, "EQUALS=ABC")));
25+
try std.testing.expect(try std.process.hasNonEmptyEnvVar(allocator, "КИРиллИЦА"));
26+
if (builtin.os.tag == .windows) {
27+
try std.testing.expect(try std.process.hasNonEmptyEnvVar(allocator, "кирИЛЛица"));
28+
}
29+
try std.testing.expect(!(try std.process.hasNonEmptyEnvVar(allocator, "NO_VALUE")));
30+
try std.testing.expect(!(try std.process.hasNonEmptyEnvVar(allocator, "NOT_SET")));
31+
if (builtin.os.tag == .windows) {
32+
try std.testing.expect(try std.process.hasNonEmptyEnvVar(allocator, "=HIDDEN"));
33+
try std.testing.expect(try std.process.hasNonEmptyEnvVar(allocator, "INVALID_UTF16_\xed\xa0\x80"));
34+
}
35+
}
36+
37+
// hasNonEmptyEnvVarContstant
38+
{
39+
try std.testing.expect(std.process.hasNonEmptyEnvVarConstant("FOO"));
40+
try std.testing.expect(!std.process.hasNonEmptyEnvVarConstant("FOO="));
41+
try std.testing.expect(!std.process.hasNonEmptyEnvVarConstant("FO"));
42+
try std.testing.expect(!std.process.hasNonEmptyEnvVarConstant("FOOO"));
43+
if (builtin.os.tag == .windows) {
44+
try std.testing.expect(std.process.hasNonEmptyEnvVarConstant("foo"));
45+
}
46+
try std.testing.expect(std.process.hasNonEmptyEnvVarConstant("EQUALS"));
47+
try std.testing.expect(!std.process.hasNonEmptyEnvVarConstant("EQUALS=ABC"));
48+
try std.testing.expect(std.process.hasNonEmptyEnvVarConstant("КИРиллИЦА"));
49+
if (builtin.os.tag == .windows) {
50+
try std.testing.expect(std.process.hasNonEmptyEnvVarConstant("кирИЛЛица"));
51+
}
52+
try std.testing.expect(!(std.process.hasNonEmptyEnvVarConstant("NO_VALUE")));
53+
try std.testing.expect(!(std.process.hasNonEmptyEnvVarConstant("NOT_SET")));
54+
if (builtin.os.tag == .windows) {
55+
try std.testing.expect(std.process.hasNonEmptyEnvVarConstant("=HIDDEN"));
56+
try std.testing.expect(std.process.hasNonEmptyEnvVarConstant("INVALID_UTF16_\xed\xa0\x80"));
57+
}
58+
}
59+
60+
// hasEnvVar
61+
{
62+
try std.testing.expect(try std.process.hasEnvVar(allocator, "FOO"));
63+
try std.testing.expect(!(try std.process.hasEnvVar(allocator, "FOO=")));
64+
try std.testing.expect(!(try std.process.hasEnvVar(allocator, "FO")));
65+
try std.testing.expect(!(try std.process.hasEnvVar(allocator, "FOOO")));
66+
if (builtin.os.tag == .windows) {
67+
try std.testing.expect(try std.process.hasEnvVar(allocator, "foo"));
68+
}
69+
try std.testing.expect(try std.process.hasEnvVar(allocator, "EQUALS"));
70+
try std.testing.expect(!(try std.process.hasEnvVar(allocator, "EQUALS=ABC")));
71+
try std.testing.expect(try std.process.hasEnvVar(allocator, "КИРиллИЦА"));
72+
if (builtin.os.tag == .windows) {
73+
try std.testing.expect(try std.process.hasEnvVar(allocator, "кирИЛЛица"));
74+
}
75+
try std.testing.expect(try std.process.hasEnvVar(allocator, "NO_VALUE"));
76+
try std.testing.expect(!(try std.process.hasEnvVar(allocator, "NOT_SET")));
77+
if (builtin.os.tag == .windows) {
78+
try std.testing.expect(try std.process.hasEnvVar(allocator, "=HIDDEN"));
79+
try std.testing.expect(try std.process.hasEnvVar(allocator, "INVALID_UTF16_\xed\xa0\x80"));
80+
}
81+
}
82+
83+
// hasEnvVarConstant
84+
{
85+
try std.testing.expect(std.process.hasEnvVarConstant("FOO"));
86+
try std.testing.expect(!std.process.hasEnvVarConstant("FOO="));
87+
try std.testing.expect(!std.process.hasEnvVarConstant("FO"));
88+
try std.testing.expect(!std.process.hasEnvVarConstant("FOOO"));
89+
if (builtin.os.tag == .windows) {
90+
try std.testing.expect(std.process.hasEnvVarConstant("foo"));
91+
}
92+
try std.testing.expect(std.process.hasEnvVarConstant("EQUALS"));
93+
try std.testing.expect(!std.process.hasEnvVarConstant("EQUALS=ABC"));
94+
try std.testing.expect(std.process.hasEnvVarConstant("КИРиллИЦА"));
95+
if (builtin.os.tag == .windows) {
96+
try std.testing.expect(std.process.hasEnvVarConstant("кирИЛЛица"));
97+
}
98+
try std.testing.expect(std.process.hasEnvVarConstant("NO_VALUE"));
99+
try std.testing.expect(!(std.process.hasEnvVarConstant("NOT_SET")));
100+
if (builtin.os.tag == .windows) {
101+
try std.testing.expect(std.process.hasEnvVarConstant("=HIDDEN"));
102+
try std.testing.expect(std.process.hasEnvVarConstant("INVALID_UTF16_\xed\xa0\x80"));
103+
}
104+
}
105+
106+
// getEnvVarOwned
107+
{
108+
try std.testing.expectEqualSlices(u8, "123", try std.process.getEnvVarOwned(arena, "FOO"));
109+
try std.testing.expectError(error.EnvironmentVariableNotFound, std.process.getEnvVarOwned(arena, "FOO="));
110+
try std.testing.expectError(error.EnvironmentVariableNotFound, std.process.getEnvVarOwned(arena, "FO"));
111+
try std.testing.expectError(error.EnvironmentVariableNotFound, std.process.getEnvVarOwned(arena, "FOOO"));
112+
if (builtin.os.tag == .windows) {
113+
try std.testing.expectEqualSlices(u8, "123", try std.process.getEnvVarOwned(arena, "foo"));
114+
}
115+
try std.testing.expectEqualSlices(u8, "ABC=123", try std.process.getEnvVarOwned(arena, "EQUALS"));
116+
try std.testing.expectError(error.EnvironmentVariableNotFound, std.process.getEnvVarOwned(arena, "EQUALS=ABC"));
117+
try std.testing.expectEqualSlices(u8, "non-ascii አማርኛ \u{10FFFF}", try std.process.getEnvVarOwned(arena, "КИРиллИЦА"));
118+
if (builtin.os.tag == .windows) {
119+
try std.testing.expectEqualSlices(u8, "non-ascii አማርኛ \u{10FFFF}", try std.process.getEnvVarOwned(arena, "кирИЛЛица"));
120+
}
121+
try std.testing.expectEqualSlices(u8, "", try std.process.getEnvVarOwned(arena, "NO_VALUE"));
122+
try std.testing.expectError(error.EnvironmentVariableNotFound, std.process.getEnvVarOwned(arena, "NOT_SET"));
123+
if (builtin.os.tag == .windows) {
124+
try std.testing.expectEqualSlices(u8, "hi", try std.process.getEnvVarOwned(arena, "=HIDDEN"));
125+
try std.testing.expectEqualSlices(u8, "\xed\xa0\x80", try std.process.getEnvVarOwned(arena, "INVALID_UTF16_\xed\xa0\x80"));
126+
}
127+
}
128+
129+
// parseEnvVarInt
130+
{
131+
try std.testing.expectEqual(123, try std.process.parseEnvVarInt("FOO", u32, 10));
132+
try std.testing.expectError(error.EnvironmentVariableNotFound, std.process.parseEnvVarInt("FO", u32, 10));
133+
try std.testing.expectError(error.EnvironmentVariableNotFound, std.process.parseEnvVarInt("FOOO", u32, 10));
134+
try std.testing.expectEqual(0x123, try std.process.parseEnvVarInt("FOO", u32, 16));
135+
if (builtin.os.tag == .windows) {
136+
try std.testing.expectEqual(123, try std.process.parseEnvVarInt("foo", u32, 10));
137+
}
138+
try std.testing.expectError(error.InvalidCharacter, std.process.parseEnvVarInt("EQUALS", u32, 10));
139+
try std.testing.expectError(error.EnvironmentVariableNotFound, std.process.parseEnvVarInt("EQUALS=ABC", u32, 10));
140+
try std.testing.expectError(error.InvalidCharacter, std.process.parseEnvVarInt("КИРиллИЦА", u32, 10));
141+
try std.testing.expectError(error.InvalidCharacter, std.process.parseEnvVarInt("NO_VALUE", u32, 10));
142+
try std.testing.expectError(error.EnvironmentVariableNotFound, std.process.parseEnvVarInt("NOT_SET", u32, 10));
143+
if (builtin.os.tag == .windows) {
144+
try std.testing.expectError(error.InvalidCharacter, std.process.parseEnvVarInt("=HIDDEN", u32, 10));
145+
try std.testing.expectError(error.InvalidCharacter, std.process.parseEnvVarInt("INVALID_UTF16_\xed\xa0\x80", u32, 10));
146+
}
147+
}
148+
149+
// EnvMap
150+
{
151+
var env_map = try std.process.getEnvMap(allocator);
152+
defer env_map.deinit();
153+
154+
try std.testing.expectEqualSlices(u8, "123", env_map.get("FOO").?);
155+
try std.testing.expectEqual(null, env_map.get("FO"));
156+
try std.testing.expectEqual(null, env_map.get("FOOO"));
157+
if (builtin.os.tag == .windows) {
158+
try std.testing.expectEqualSlices(u8, "123", env_map.get("foo").?);
159+
}
160+
try std.testing.expectEqualSlices(u8, "ABC=123", env_map.get("EQUALS").?);
161+
try std.testing.expectEqual(null, env_map.get("EQUALS=ABC"));
162+
try std.testing.expectEqualSlices(u8, "non-ascii አማርኛ \u{10FFFF}", env_map.get("КИРиллИЦА").?);
163+
if (builtin.os.tag == .windows) {
164+
try std.testing.expectEqualSlices(u8, "non-ascii አማርኛ \u{10FFFF}", env_map.get("кирИЛЛица").?);
165+
}
166+
try std.testing.expectEqualSlices(u8, "", env_map.get("NO_VALUE").?);
167+
try std.testing.expectEqual(null, env_map.get("NOT_SET"));
168+
if (builtin.os.tag == .windows) {
169+
try std.testing.expectEqualSlices(u8, "hi", env_map.get("=HIDDEN").?);
170+
try std.testing.expectEqualSlices(u8, "\xed\xa0\x80", env_map.get("INVALID_UTF16_\xed\xa0\x80").?);
171+
}
172+
}
173+
}

0 commit comments

Comments
 (0)