| ... | ... | @@ -52,341 +52,128 @@ test "getCwdAlloc" { |
| 52 | 52 | testing.allocator.free(cwd); |
| 53 | 53 | } |
| 54 | 54 | |
| 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). |
| 62 | | pub 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, |
| 55 | pub const EnvMap = struct { |
| 56 | hash_map: HashMap, |
| 57 | |
| 58 | const HashMap = std.HashMap( |
| 59 | []const u8, |
| 60 | []const u8, |
| 61 | EnvNameHashContext, |
| 62 | std.hash_map.default_max_load_percentage, |
| 63 | ); |
| 64 | |
| 65 | pub const EnvNameHashContext = struct { |
| 66 | pub fn hash(self: @This(), s: []const u8) u64 { |
| 67 | _ = self; |
| 68 | if (builtin.os.tag == .windows) { |
| 69 | const h = std.hash.Wyhash.init(0); |
| 70 | // TODO: improve this, instead of iterating over ascii, |
| 71 | // iterate over with unicode |
| 72 | for (s) |c| { |
| 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(); |
| 77 | | |
| 78 | | /// Deinitialize with `deinit`. |
| 79 | | pub fn init(allocator: Allocator) Self { |
| 80 | | return .{ |
| 81 | | .allocator = allocator, |
| 82 | | .uppercased_map = std.StringHashMapUnmanaged(EnvValue){}, |
| 83 | | }; |
| 91 | /// Create a EnvMap backed by a specific allocator. |
| 92 | /// That allocator will be used for both backing allocations |
| 93 | /// and string deduplication. |
| 94 | pub fn init(allocator: Allocator) EnvMap { |
| 95 | return EnvMap{ .hash_map = HashMap.init(allocator) }; |
| 84 | 96 | } |
| 85 | 97 | |
| 86 | | pub fn deinit(self: *Self) void { |
| 87 | | var it = self.uppercased_map.iterator(); |
| 98 | /// Free the backing storage of the map, as well as all |
| 99 | /// of the stored keys and values. |
| 100 | pub fn deinit(self: *EnvMap) void { |
| 101 | var it = self.hash_map.iterator(); |
| 88 | 102 | 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); |
| 103 | self.free(entry.key_ptr.*); |
| 104 | self.free(entry.value_ptr.*); |
| 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. |
| 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 | | } |
| 107 | self.hash_map.deinit(); |
| 103 | 108 | } |
| 104 | 109 | |
| 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 |
| 110 | /// Same as `put` but the key and value become owned by the EnvMap rather |
| 111 | /// than being copied. |
| 112 | /// If `putMove` fails, the ownership of key and value does not transfer. |
| 113 | pub fn putMove(self: *EnvMap, key: []u8, value: []u8) !void { |
| 114 | const get_or_put = try self.hash_map.getOrPut(key); |
| 115 | if (get_or_put.found_existing) { |
| 116 | self.free(get_or_put.key_ptr.*); |
| 117 | self.free(get_or_put.value_ptr.*); |
| 118 | get_or_put.key_ptr.* = key; |
| 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 |
| 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 | | const bytes_written = try std.unicode.utf8ToUtf16LeWriter(name_utf16_buf.writer(), name); |
| 145 | | var name_utf16 = name_utf16_buf.items[0..bytes_written]; |
| 146 | | |
| 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, |
| 123 | /// `key` and `value` are copied into the EnvMap. |
| 124 | pub fn put(self: *EnvMap, key: []const u8, value: []const u8) !void { |
| 125 | const value_copy = try self.copy(value); |
| 126 | errdefer self.free(value_copy); |
| 127 | const get_or_put = try self.hash_map.getOrPut(key); |
| 128 | if (get_or_put.found_existing) { |
| 129 | self.free(get_or_put.value_ptr.*); |
| 130 | } else { |
| 131 | get_or_put.key_ptr.* = self.copy(key) catch |err| { |
| 132 | _ = self.hash_map.remove(key); |
| 133 | return err; |
| 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 | | }; |
| 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; |
| 136 | get_or_put.value_ptr.* = value_copy; |
| 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; |
| 316 | | |
| 317 | | const Self = @This(); |
| 318 | | |
| 319 | | /// Deinitialize with `deinit`. |
| 320 | | pub fn init(allocator: Allocator) Self { |
| 321 | | return Self{ .storage = StorageType.init(allocator) }; |
| 139 | /// Find the address of the value associated with a key. |
| 140 | /// The returned pointer is invalidated if the map resizes. |
| 141 | pub fn getPtr(self: EnvMap, key: []const u8) ?*[]const u8 { |
| 142 | return self.hash_map.getPtr(key); |
| 322 | 143 | } |
| 323 | 144 | |
| 324 | | pub fn deinit(self: *Self) void { |
| 325 | | self.storage.deinit(); |
| 145 | /// Return the map's copy of the value associated with |
| 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 { |
| 329 | | switch (builtin.os.tag) { |
| 330 | | .windows => { |
| 331 | | if (self.storage.get(name)) |entry| { |
| 332 | | return entry.value; |
| 333 | | } else { |
| 334 | | return null; |
| 335 | | } |
| 336 | | }, |
| 337 | | else => return self.storage.get(name), |
| 338 | | } |
| 152 | /// Removes the item from the map and frees its value. |
| 153 | /// This invalidates the value returned by get() for this key. |
| 154 | pub fn remove(self: *EnvMap, key: []const u8) void { |
| 155 | const kv = self.hash_map.fetchRemove(key) orelse return; |
| 156 | self.free(kv.key); |
| 157 | self.free(kv.value); |
| 339 | 158 | } |
| 340 | 159 | |
| 341 | | pub fn count(self: Self) Size { |
| 342 | | return self.storage.count(); |
| 160 | /// Returns the number of KV pairs stored in the map. |
| 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 { |
| 346 | | return .{ .storage_iterator = self.storage.iterator() }; |
| 165 | /// Returns an iterator over entries in the map. |
| 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 { |
| 350 | | switch (builtin.os.tag) { |
| 351 | | .windows => { |
| 352 | | try self.storage.putUtf8(name, value); |
| 353 | | try self.storage.reallocUppercaseBuf(); |
| 354 | | }, |
| 355 | | else => return self.storage.put(name, value), |
| 356 | | } |
| 170 | fn free(self: EnvMap, value: []const u8) void { |
| 171 | self.hash_map.allocator.free(value); |
| 357 | 172 | } |
| 358 | 173 | |
| 359 | | pub fn remove(self: *Self, name: []const u8) void { |
| 360 | | _ = self.storage.remove(name); |
| 174 | fn copy(self: EnvMap, value: []const u8) ![]u8 { |
| 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 | 179 | test "EnvMap" { |
Jonathan Marler <johnnymarler@gmail.com> 2022-02-04 12:08:38-07:00