| ... | @@ -52,341 +52,128 @@ test "getCwdAlloc" { | ... | @@ -52,341 +52,128 @@ test "getCwdAlloc" { |
| 52 | testing.allocator.free(cwd); | 52 | testing.allocator.free(cwd); |
| 53 | } | 53 | } |
| 54 | | 54 | |
| 55 | /// EnvMap for Windows that handles Unicode-aware case insensitivity for lookups, while also | 55 | pub const EnvMap = struct { |
| 56 | /// providing the canonical environment variable names when iterating. | 56 | hash_map: HashMap, |
| 57 | /// | 57 | |
| 58 | /// Allows for zero-allocation lookups (even though it needs to do UTF-8 -> UTF-16 -> uppercase | 58 | const HashMap = std.HashMap( |
| 59 | /// conversions) by allocating a buffer large enough to fit the largest environment variable | 59 | []const u8, |
| 60 | /// name, and using that when doing lookups (i.e. anything that overflows the buffer can be treated | 60 | []const u8, |
| 61 | /// as the environment variable not being found). | 61 | EnvNameHashContext, |
| 62 | pub const EnvMapWindows = struct { | 62 | std.hash_map.default_max_load_percentage, |
| 63 | allocator: Allocator, | 63 | ); |
| 64 | /// Keys are UTF-16le stored as []const u8 | 64 | |
| 65 | uppercased_map: std.StringHashMapUnmanaged(EnvValue), | 65 | pub const EnvNameHashContext = struct { |
| 66 | /// Buffer for converting to uppercased UTF-16 on key lookups | 66 | pub fn hash(self: @This(), s: []const u8) u64 { |
| 67 | /// Must call `reallocUppercaseBuf` before doing any lookups after a `put` call. | 67 | _ = self; |
| 68 | uppercase_buf_utf16: []u16 = &[_]u16{}, | 68 | if (builtin.os.tag == .windows) { |
| 69 | max_name_utf16_length: usize = 0, | 69 | const h = std.hash.Wyhash.init(0); |
| 70 | | 70 | // TODO: improve this, instead of iterating over ascii, |
| 71 | pub const EnvValue = struct { | 71 | // iterate over with unicode |
| 72 | value: []const u8, | 72 | for (s) |c| { |
| 73 | canonical_name: []const u8, | 73 | var s_upper = [_]u8 { std.ascii.toLower(c) }; |
| | 74 | h.update(s_upper); |
| | 75 | } |
| | 76 | return h.final(); |
| | 77 | } |
| | 78 | return std.hash_map.hashString(s); |
| | 79 | } |
| | 80 | pub fn eql(self: @This(), a: []const u8, b: []const u8) bool { |
| | 81 | _ = self; |
| | 82 | if (builtin.os.tag == .windows) { |
| | 83 | // TODO: improve this, instead of comparing ascii |
| | 84 | // compare with unicode |
| | 85 | return std.ascii.eqlIgnoreCase(a, b); |
| | 86 | } |
| | 87 | return std.hash_map.eqlString(a, b); |
| | 88 | } |
| 74 | }; | 89 | }; |
| 75 | | 90 | |
| 76 | const Self = @This(); | 91 | /// Create a EnvMap backed by a specific allocator. |
| 77 | | 92 | /// That allocator will be used for both backing allocations |
| 78 | /// Deinitialize with `deinit`. | 93 | /// and string deduplication. |
| 79 | pub fn init(allocator: Allocator) Self { | 94 | pub fn init(allocator: Allocator) EnvMap { |
| 80 | return .{ | 95 | return EnvMap{ .hash_map = HashMap.init(allocator) }; |
| 81 | .allocator = allocator, | | |
| 82 | .uppercased_map = std.StringHashMapUnmanaged(EnvValue){}, | | |
| 83 | }; | | |
| 84 | } | 96 | } |
| 85 | | 97 | |
| 86 | pub fn deinit(self: *Self) void { | 98 | /// Free the backing storage of the map, as well as all |
| 87 | var it = self.uppercased_map.iterator(); | 99 | /// of the stored keys and values. |
| | 100 | pub fn deinit(self: *EnvMap) void { |
| | 101 | var it = self.hash_map.iterator(); |
| 88 | while (it.next()) |entry| { | 102 | while (it.next()) |entry| { |
| 89 | self.allocator.free(entry.key_ptr.*); | 103 | self.free(entry.key_ptr.*); |
| 90 | self.allocator.free(entry.value_ptr.value); | 104 | self.free(entry.value_ptr.*); |
| 91 | self.allocator.free(entry.value_ptr.canonical_name); | | |
| 92 | } | 105 | } |
| 93 | self.uppercased_map.deinit(self.allocator); | | |
| 94 | self.allocator.free(self.uppercase_buf_utf16); | | |
| 95 | } | | |
| 96 | | 106 | |
| 97 | /// Increases the size of the uppercase buffer if the maximum name size has increased. | 107 | self.hash_map.deinit(); |
| 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 | } | 108 | } |
| 104 | | 109 | |
| 105 | /// Converts `src` to uppercase using `RtlUpcaseUnicodeString` and puts the result in `dest`. | 110 | /// Same as `put` but the key and value become owned by the EnvMap rather |
| 106 | /// Returns the length of the converted UTF-16 string. `dest.len` must be >= `src.len`. | 111 | /// than being copied. |
| 107 | /// | 112 | /// If `putMove` fails, the ownership of key and value does not transfer. |
| 108 | /// Note: As of now, RtlUpcaseUnicodeString does not seem to handle codepoints above 0x10000 | 113 | pub fn putMove(self: *EnvMap, key: []u8, value: []u8) !void { |
| 109 | /// (i.e. those that require a surrogate pair), so this function will always return a length | 114 | const get_or_put = try self.hash_map.getOrPut(key); |
| 110 | /// equal to `src.len`. However, if RtlUpcaseUnicodeString is updated to handle codepoints above | 115 | if (get_or_put.found_existing) { |
| 111 | /// 0x10000, this property would still hold unless there are lowercase <-> uppercase conversions | 116 | self.free(get_or_put.key_ptr.*); |
| 112 | /// that cross over the boundary between codepoints >= 0x10000 and < 0x10000. | 117 | self.free(get_or_put.value_ptr.*); |
| 113 | /// TODO: Is it feasible that Unicode lowercase <-> uppercase conversions could cross that boundary? | 118 | get_or_put.key_ptr.* = key; |
| 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 | } | 119 | } |
| | 120 | get_or_put.value_ptr.* = value; |
| 134 | } | 121 | } |
| 135 | | 122 | |
| 136 | /// Note: Does not realloc the uppercase buf to allow for calling put for many variables and | 123 | /// `key` and `value` are copied into the EnvMap. |
| 137 | /// only allocating the uppercase buf afterwards. | 124 | pub fn put(self: *EnvMap, key: []const u8, value: []const u8) !void { |
| 138 | pub fn putUtf8(self: *Self, name: []const u8, value: []const u8) !void { | 125 | const value_copy = try self.copy(value); |
| 139 | const uppercased_len = len: { | 126 | errdefer self.free(value_copy); |
| 140 | const name_uppercased_utf16 = uppercased: { | 127 | const get_or_put = try self.hash_map.getOrPut(key); |
| 141 | var name_utf16_buf = try std.ArrayListAligned(u8, @alignOf(u16)).initCapacity(self.allocator, name.len); | 128 | if (get_or_put.found_existing) { |
| 142 | errdefer name_utf16_buf.deinit(); | 129 | self.free(get_or_put.value_ptr.*); |
| 143 | | 130 | } else { |
| 144 | const bytes_written = try std.unicode.utf8ToUtf16LeWriter(name_utf16_buf.writer(), name); | 131 | get_or_put.key_ptr.* = self.copy(key) catch |err| { |
| 145 | var name_utf16 = name_utf16_buf.items[0..bytes_written]; | 132 | _ = self.hash_map.remove(key); |
| 146 | | 133 | return err; |
| 147 | // uppercase in place | | |
| 148 | var name_uppercased_utf16 = std.mem.bytesAsSlice(u16, name_utf16); | | |
| 149 | const uppercased_len = uppercaseName(name_uppercased_utf16, name_uppercased_utf16); | | |
| 150 | assert(uppercased_len == name_uppercased_utf16.len); | | |
| 151 | | | |
| 152 | break :uppercased name_utf16_buf.toOwnedSlice(); | | |
| 153 | }; | | |
| 154 | errdefer self.allocator.free(name_uppercased_utf16); | | |
| 155 | | | |
| 156 | const name_canonical = try self.allocator.dupe(u8, name); | | |
| 157 | errdefer self.allocator.free(name_canonical); | | |
| 158 | | | |
| 159 | const value_dupe = try self.allocator.dupe(u8, value); | | |
| 160 | errdefer self.allocator.free(value_dupe); | | |
| 161 | | | |
| 162 | const get_or_put = try self.uppercased_map.getOrPut(self.allocator, name_uppercased_utf16); | | |
| 163 | if (get_or_put.found_existing) { | | |
| 164 | // note: this is only safe from UAF because the errdefer that frees this value above | | |
| 165 | // no longer has a possibility of being triggered after this point | | |
| 166 | self.allocator.free(name_uppercased_utf16); | | |
| 167 | self.allocator.free(get_or_put.value_ptr.value); | | |
| 168 | self.allocator.free(get_or_put.value_ptr.canonical_name); | | |
| 169 | } else { | | |
| 170 | get_or_put.key_ptr.* = name_uppercased_utf16; | | |
| 171 | } | | |
| 172 | get_or_put.value_ptr.value = value_dupe; | | |
| 173 | get_or_put.value_ptr.canonical_name = name_canonical; | | |
| 174 | | | |
| 175 | break :len name_uppercased_utf16.len; | | |
| 176 | }; | | |
| 177 | | | |
| 178 | // The buffer for case conversion for key lookups will need to be as big as the largest | | |
| 179 | // key stored in the hash map. | | |
| 180 | self.max_name_utf16_length = @maximum(self.max_name_utf16_length, uppercased_len); | | |
| 181 | } | | |
| 182 | | | |
| 183 | /// Asserts that the name does not already exist in the map. | | |
| 184 | /// Note: Does not realloc the uppercase buf to allow for calling put for many variables and | | |
| 185 | /// only allocating the uppercase buf afterwards. | | |
| 186 | pub fn putUtf16NoClobber(self: *Self, name_utf16: []const u16, value_utf16: []const u16) !void { | | |
| 187 | const uppercased_len = len: { | | |
| 188 | const name_canonical = try std.unicode.utf16leToUtf8Alloc(self.allocator, name_utf16); | | |
| 189 | errdefer self.allocator.free(name_canonical); | | |
| 190 | | | |
| 191 | const value = try std.unicode.utf16leToUtf8Alloc(self.allocator, value_utf16); | | |
| 192 | errdefer self.allocator.free(value); | | |
| 193 | | | |
| 194 | const name_uppercased_utf16 = try self.allocator.alloc(u16, name_utf16.len); | | |
| 195 | errdefer self.allocator.free(name_uppercased_utf16); | | |
| 196 | | | |
| 197 | const uppercased_len = uppercaseName(name_uppercased_utf16, name_utf16); | | |
| 198 | assert(uppercased_len == name_uppercased_utf16.len); | | |
| 199 | | | |
| 200 | try self.uppercased_map.putNoClobber(self.allocator, std.mem.sliceAsBytes(name_uppercased_utf16), EnvValue{ | | |
| 201 | .value = value, | | |
| 202 | .canonical_name = name_canonical, | | |
| 203 | }); | | |
| 204 | break :len name_uppercased_utf16.len; | | |
| 205 | }; | | |
| 206 | | | |
| 207 | // The buffer for case conversion for key lookups will need to be as big as the largest | | |
| 208 | // key stored in the hash map. | | |
| 209 | self.max_name_utf16_length = @maximum(self.max_name_utf16_length, uppercased_len); | | |
| 210 | } | | |
| 211 | | | |
| 212 | /// Attempts to convert a UTF-8 name into a uppercased UTF-16le name for a lookup. If the | | |
| 213 | /// name cannot be converted, this function will return `null`. | | |
| 214 | fn utf8ToUppercasedUtf16(self: Self, name: []const u8) ?[]u16 { | | |
| 215 | const name_utf16: []u16 = to_utf16: { | | |
| 216 | var utf16_buf_stream = std.io.fixedBufferStream(std.mem.sliceAsBytes(self.uppercase_buf_utf16)); | | |
| 217 | _ = std.unicode.utf8ToUtf16LeWriter(utf16_buf_stream.writer(), name) catch |err| switch (err) { | | |
| 218 | // If the buffer isn't large enough, we can treat that as 'env var not found', as we | | |
| 219 | // know anything too large for the buffer can't be found in the map. | | |
| 220 | error.NoSpaceLeft => return null, | | |
| 221 | // Anything with invalid UTF-8 will also not be found in the map, so treat that as | | |
| 222 | // 'env var not found' too | | |
| 223 | error.InvalidUtf8 => return null, | | |
| 224 | }; | 134 | }; |
| 225 | break :to_utf16 std.mem.bytesAsSlice(u16, utf16_buf_stream.getWritten()); | | |
| 226 | }; | | |
| 227 | | | |
| 228 | // uppercase in place | | |
| 229 | const uppercased_len = uppercaseName(name_utf16, name_utf16); | | |
| 230 | assert(uppercased_len == name_utf16.len); | | |
| 231 | | | |
| 232 | return name_utf16; | | |
| 233 | } | | |
| 234 | | | |
| 235 | /// Returns true if an entry was found and deleted, false otherwise. | | |
| 236 | pub fn remove(self: *Self, name: []const u8) bool { | | |
| 237 | const name_utf16 = self.utf8ToUppercasedUtf16(name) orelse return false; | | |
| 238 | const kv = self.uppercased_map.fetchRemove(std.mem.sliceAsBytes(name_utf16)) orelse return false; | | |
| 239 | self.allocator.free(kv.key); | | |
| 240 | self.allocator.free(kv.value.value); | | |
| 241 | self.allocator.free(kv.value.canonical_name); | | |
| 242 | return true; | | |
| 243 | } | | |
| 244 | | | |
| 245 | pub fn get(self: Self, name: []const u8) ?EnvValue { | | |
| 246 | const name_utf16 = self.utf8ToUppercasedUtf16(name) orelse return null; | | |
| 247 | return self.uppercased_map.get(std.mem.sliceAsBytes(name_utf16)); | | |
| 248 | } | | |
| 249 | | | |
| 250 | pub fn count(self: Self) EnvMap.Size { | | |
| 251 | return self.uppercased_map.count(); | | |
| 252 | } | | |
| 253 | | | |
| 254 | pub fn iterator(self: *const Self) Iterator { | | |
| 255 | return .{ | | |
| 256 | .env_map = self, | | |
| 257 | .uppercased_map_iterator = self.uppercased_map.iterator(), | | |
| 258 | }; | | |
| 259 | } | | |
| 260 | | | |
| 261 | pub const Iterator = struct { | | |
| 262 | env_map: *const Self, | | |
| 263 | uppercased_map_iterator: std.StringHashMapUnmanaged(EnvValue).Iterator, | | |
| 264 | | | |
| 265 | pub fn next(it: *Iterator) ?EnvMap.Entry { | | |
| 266 | if (it.uppercased_map_iterator.next()) |uppercased_entry| { | | |
| 267 | return EnvMap.Entry{ | | |
| 268 | .name = uppercased_entry.value_ptr.canonical_name, | | |
| 269 | .value = uppercased_entry.value_ptr.value, | | |
| 270 | }; | | |
| 271 | } else { | | |
| 272 | return null; | | |
| 273 | } | | |
| 274 | } | 135 | } |
| 275 | }; | 136 | get_or_put.value_ptr.* = value_copy; |
| 276 | }; | | |
| 277 | | | |
| 278 | test "EnvMapWindows" { | | |
| 279 | if (builtin.os.tag != .windows) return error.SkipZigTest; | | |
| 280 | | | |
| 281 | var env_map = EnvMapWindows.init(testing.allocator); | | |
| 282 | defer env_map.deinit(); | | |
| 283 | | | |
| 284 | // both put methods | | |
| 285 | try env_map.putUtf16NoClobber(std.unicode.utf8ToUtf16LeStringLiteral("Path"), std.unicode.utf8ToUtf16LeStringLiteral("something")); | | |
| 286 | try env_map.putUtf8("КИРиллИЦА", "something else"); | | |
| 287 | try env_map.reallocUppercaseBuf(); | | |
| 288 | | | |
| 289 | try testing.expectEqual(@as(EnvMap.Size, 2), env_map.count()); | | |
| 290 | | | |
| 291 | // unicode-aware case-insensitive lookups | | |
| 292 | try testing.expectEqualStrings("something", env_map.get("PATH").?.value); | | |
| 293 | try testing.expectEqualStrings("something else", env_map.get("кириллица").?.value); | | |
| 294 | try testing.expect(env_map.get("missing") == null); | | |
| 295 | | | |
| 296 | // canonical names when iterating | | |
| 297 | var it = env_map.iterator(); | | |
| 298 | var count: EnvMap.Size = 0; | | |
| 299 | while (it.next()) |entry| { | | |
| 300 | const is_an_expected_name = std.mem.eql(u8, "Path", entry.name) or std.mem.eql(u8, "КИРиллИЦА", entry.name); | | |
| 301 | try testing.expect(is_an_expected_name); | | |
| 302 | count += 1; | | |
| 303 | } | 137 | } |
| 304 | try testing.expectEqual(@as(EnvMap.Size, 2), count); | | |
| 305 | } | | |
| 306 | | | |
| 307 | pub const EnvMap = struct { | | |
| 308 | storage: StorageType, | | |
| 309 | | | |
| 310 | pub const StorageType = switch (builtin.os.tag) { | | |
| 311 | .windows => EnvMapWindows, | | |
| 312 | else => std.BufMap, | | |
| 313 | }; | | |
| 314 | | 138 | |
| 315 | pub const Size = std.BufMap.BufMapHashMap.Size; | 139 | /// Find the address of the value associated with a key. |
| 316 | | 140 | /// The returned pointer is invalidated if the map resizes. |
| 317 | const Self = @This(); | 141 | pub fn getPtr(self: EnvMap, key: []const u8) ?*[]const u8 { |
| 318 | | 142 | return self.hash_map.getPtr(key); |
| 319 | /// Deinitialize with `deinit`. | | |
| 320 | pub fn init(allocator: Allocator) Self { | | |
| 321 | return Self{ .storage = StorageType.init(allocator) }; | | |
| 322 | } | 143 | } |
| 323 | | 144 | |
| 324 | pub fn deinit(self: *Self) void { | 145 | /// Return the map's copy of the value associated with |
| 325 | self.storage.deinit(); | 146 | /// a key. The returned string is invalidated if this |
| | 147 | /// key is removed from the map. |
| | 148 | pub fn get(self: EnvMap, key: []const u8) ?[]const u8 { |
| | 149 | return self.hash_map.get(key); |
| 326 | } | 150 | } |
| 327 | | 151 | |
| 328 | pub fn get(self: Self, name: []const u8) ?[]const u8 { | 152 | /// Removes the item from the map and frees its value. |
| 329 | switch (builtin.os.tag) { | 153 | /// This invalidates the value returned by get() for this key. |
| 330 | .windows => { | 154 | pub fn remove(self: *EnvMap, key: []const u8) void { |
| 331 | if (self.storage.get(name)) |entry| { | 155 | const kv = self.hash_map.fetchRemove(key) orelse return; |
| 332 | return entry.value; | 156 | self.free(kv.key); |
| 333 | } else { | 157 | self.free(kv.value); |
| 334 | return null; | | |
| 335 | } | | |
| 336 | }, | | |
| 337 | else => return self.storage.get(name), | | |
| 338 | } | | |
| 339 | } | 158 | } |
| 340 | | 159 | |
| 341 | pub fn count(self: Self) Size { | 160 | /// Returns the number of KV pairs stored in the map. |
| 342 | return self.storage.count(); | 161 | pub fn count(self: EnvMap) HashMap.Size { |
| | 162 | return self.hash_map.count(); |
| 343 | } | 163 | } |
| 344 | | 164 | |
| 345 | pub fn iterator(self: *const Self) Iterator { | 165 | /// Returns an iterator over entries in the map. |
| 346 | return .{ .storage_iterator = self.storage.iterator() }; | 166 | pub fn iterator(self: *const EnvMap) HashMap.Iterator { |
| | 167 | return self.hash_map.iterator(); |
| 347 | } | 168 | } |
| 348 | | 169 | |
| 349 | pub fn put(self: *Self, name: []const u8, value: []const u8) !void { | 170 | fn free(self: EnvMap, value: []const u8) void { |
| 350 | switch (builtin.os.tag) { | 171 | self.hash_map.allocator.free(value); |
| 351 | .windows => { | | |
| 352 | try self.storage.putUtf8(name, value); | | |
| 353 | try self.storage.reallocUppercaseBuf(); | | |
| 354 | }, | | |
| 355 | else => return self.storage.put(name, value), | | |
| 356 | } | | |
| 357 | } | 172 | } |
| 358 | | 173 | |
| 359 | pub fn remove(self: *Self, name: []const u8) void { | 174 | fn copy(self: EnvMap, value: []const u8) ![]u8 { |
| 360 | _ = self.storage.remove(name); | 175 | return self.hash_map.allocator.dupe(u8, value); |
| 361 | } | 176 | } |
| 362 | | | |
| 363 | pub const Entry = struct { | | |
| 364 | name: []const u8, | | |
| 365 | value: []const u8, | | |
| 366 | }; | | |
| 367 | | | |
| 368 | pub const Iterator = struct { | | |
| 369 | storage_iterator: switch (builtin.os.tag) { | | |
| 370 | .windows => EnvMapWindows.Iterator, | | |
| 371 | else => std.BufMap.BufMapHashMap.Iterator, | | |
| 372 | }, | | |
| 373 | | | |
| 374 | pub fn next(it: *Iterator) ?Entry { | | |
| 375 | switch (builtin.os.tag) { | | |
| 376 | .windows => return it.storage_iterator.next(), | | |
| 377 | else => { | | |
| 378 | if (it.storage_iterator.next()) |entry| { | | |
| 379 | return Entry{ | | |
| 380 | .name = entry.key_ptr.*, | | |
| 381 | .value = entry.value_ptr.*, | | |
| 382 | }; | | |
| 383 | } else { | | |
| 384 | return null; | | |
| 385 | } | | |
| 386 | }, | | |
| 387 | } | | |
| 388 | } | | |
| 389 | }; | | |
| 390 | }; | 177 | }; |
| 391 | | 178 | |
| 392 | test "EnvMap" { | 179 | test "EnvMap" { |
Jonathan Marler <johnnymarler@gmail.com> 2022-02-04 12:08:38-07:00