Skip to content

Commit cd1c4b6

Browse files
miniksaminiksa
committed
IT WORKS!
1 parent 6cac096 commit cd1c4b6

File tree

4 files changed

+161
-37
lines changed

4 files changed

+161
-37
lines changed

‎src/inc/til/env.h

+107-36
Original file line numberDiff line numberDiff line change
@@ -50,10 +50,10 @@ namespace til // Terminal Implementation Library. Also: "Today I Learned"
5050
static constexpr std::wstring_view program_files_x86{ L"ProgramFiles(x86)" };
5151
static constexpr std::wstring_view program_files_arm64{ L"ProgramFiles(Arm)" };
5252
static constexpr std::wstring_view program_w6432{ L"ProgramW6432" };
53-
static constexpr std::wstring_view common_program_files{ L"ProgramW6432" };
54-
static constexpr std::wstring_view common_program_files_x86{ L"ProgramW6432" };
55-
static constexpr std::wstring_view common_program_files_arm64{ L"ProgramW6432" };
56-
static constexpr std::wstring_view common_program_w6432{ L"ProgramW6432" };
53+
static constexpr std::wstring_view common_program_files{ L"CommonProgramFiles" };
54+
static constexpr std::wstring_view common_program_files_x86{ L"CommonProgramFiles(x86)" };
55+
static constexpr std::wstring_view common_program_files_arm64{ L"CommonProgramFiles(Arm)" };
56+
static constexpr std::wstring_view common_program_w6432{ L"CommonProgramW6432" };
5757

5858
const std::map<std::wstring, std::wstring_view> program_files_map{
5959
{ L"ProgramFilesDir", program_files },
@@ -100,7 +100,7 @@ namespace til // Terminal Implementation Library. Also: "Today I Learned"
100100

101101
::SetLastError(ERROR_SUCCESS);
102102

103-
DWORD length = valueLength;
103+
DWORD length = static_cast<DWORD>(valueLength);
104104

105105
auto result = ::GetComputerNameW(value, &length);
106106
*valueLengthNeededWithNul = length;
@@ -117,7 +117,7 @@ namespace til // Terminal Implementation Library. Also: "Today I Learned"
117117
template<typename string_type, size_t initialBufferLength = MAX_COMPUTERNAME_LENGTH + 1>
118118
HRESULT TryGetComputerNameW(string_type& result) WI_NOEXCEPT
119119
{
120-
const auto hr = wiltmp::TryGetComputerNameW<string_type, initialBufferLength>(result);
120+
const auto hr = wiltmp::GetComputerNameW<string_type, initialBufferLength>(result);
121121
RETURN_HR_IF(hr, FAILED(hr) && (hr != HRESULT_FROM_WIN32(ERROR_ENVVAR_NOT_FOUND)));
122122
return S_OK;
123123
}
@@ -144,7 +144,7 @@ namespace til // Terminal Implementation Library. Also: "Today I Learned"
144144

145145
/** Looks up a registry value from 'key' and fails if it is not found. */
146146
template<typename string_type, size_t initialBufferLength = 256>
147-
inline HRESULT RegQueryValueExW(HANDLE key, PCWSTR valueName, string_type& result) WI_NOEXCEPT
147+
inline HRESULT RegQueryValueExW(HKEY key, PCWSTR valueName, string_type& result) WI_NOEXCEPT
148148
{
149149
return wil::AdaptFixedSizeToAllocatedResult<string_type, initialBufferLength>(result,
150150
[&](_Out_writes_(valueLength) PWSTR value, size_t valueLength, _Out_ size_t* valueLengthNeededWithNul) -> HRESULT {
@@ -153,14 +153,14 @@ namespace til // Terminal Implementation Library. Also: "Today I Learned"
153153
// length will receive the number of bytes including trailing null byte. Convert to a number of wchar_t's.
154154
// AdaptFixedSizeToAllocatedResult will then resize buffer to valueLengthNeededWithNull.
155155
// We're rounding up to prevent infinite loops if the data isn't a REG_SZ and length isn't divisible by 2.
156-
*valueLengthNeededWithNull = (length + sizeof(wchar_t) - 1) / sizeof(wchar_t);
156+
*valueLengthNeededWithNul = (length + sizeof(wchar_t) - 1) / sizeof(wchar_t);
157157
return status == ERROR_MORE_DATA ? S_OK : HRESULT_FROM_WIN32(status);
158158
});
159159
}
160160

161161
/** Looks up a registry value from 'key' and returns null if it is not found. */
162162
template<typename string_type, size_t initialBufferLength = 256>
163-
HRESULT TryRegQueryValueExW(HANDLE key, PCWSTR valueName, string_type& result) WI_NOEXCEPT
163+
HRESULT TryRegQueryValueExW(HKEY key, PCWSTR valueName, string_type& result) WI_NOEXCEPT
164164
{
165165
const auto hr = wiltmp::TryRegQueryValueExW<string_type, initialBufferLength>(key, valueName, result);
166166
RETURN_HR_IF(hr, FAILED(hr) && (hr != HRESULT_FROM_WIN32(ERROR_ENVVAR_NOT_FOUND)));
@@ -170,7 +170,7 @@ namespace til // Terminal Implementation Library. Also: "Today I Learned"
170170
#ifdef WIL_ENABLE_EXCEPTIONS
171171
/** Looks up a registry value from 'key' and fails if it is not found. */
172172
template<typename string_type = wil::unique_cotaskmem_string, size_t initialBufferLength = 256>
173-
string_type RegQueryValueExW(HANDLE key, PCWSTR valueName)
173+
string_type RegQueryValueExW(HKEY key, PCWSTR valueName)
174174
{
175175
string_type result;
176176
THROW_IF_FAILED((wiltmp::RegQueryValueExW<string_type, initialBufferLength>(key, valueName, result)));
@@ -179,7 +179,7 @@ namespace til // Terminal Implementation Library. Also: "Today I Learned"
179179

180180
/** Looks up a registry value from 'key' and returns null if it is not found. */
181181
template<typename string_type = wil::unique_cotaskmem_string, size_t initialBufferLength = 256>
182-
string_type TryRegQueryValueExW(HANDLE key, PCWSTR valueName)
182+
string_type TryRegQueryValueExW(HKEY key, PCWSTR valueName)
183183
{
184184
string_type result;
185185
THROW_IF_FAILED((wiltmp::TryRegQueryValueExW<string_type, initialBufferLength>(key, valueName, result)));
@@ -192,12 +192,11 @@ namespace til // Terminal Implementation Library. Also: "Today I Learned"
192192
template<typename string_type, size_t stackBufferLength = 256>
193193
HRESULT GetShortPathNameW(PCWSTR file, string_type& path)
194194
{
195-
wil::assign_null_to_opt_param(filePart);
196195
const auto hr = wil::AdaptFixedSizeToAllocatedResult<string_type, stackBufferLength>(path,
197196
[&](_Out_writes_(valueLength) PWSTR value, size_t valueLength, _Out_ size_t* valueLengthNeededWithNull) -> HRESULT {
198197
// Note that GetShortPathNameW() is not limited to MAX_PATH
199198
// but it does take a fixed size buffer.
200-
*valueLengthNeededWithNull = ::GetShortPathNameW(file, static_cast<DWORD>(valueLength), value);
199+
*valueLengthNeededWithNull = ::GetShortPathNameW(file, value, static_cast<DWORD>(valueLength));
201200
RETURN_LAST_ERROR_IF(*valueLengthNeededWithNull == 0);
202201
WI_ASSERT((*value != L'\0') == (*valueLengthNeededWithNull < valueLength));
203202
if (*valueLengthNeededWithNull < valueLength)
@@ -235,15 +234,15 @@ namespace til // Terminal Implementation Library. Also: "Today I Learned"
235234
{
236235
if (auto value = wil::TryGetEnvironmentVariableW(variable.c_str()))
237236
{
238-
insert_or_assign(variable, value.get());
237+
save_to_map(variable, value.get());
239238
}
240239
}
241240

242241
void get_computer_name()
243242
{
244243
if (auto value = til::details::wiltmp::TryGetComputerNameW())
245244
{
246-
insert_or_assign(std::wstring{ til::details::vars::computer_name }, value.get());
245+
save_to_map(std::wstring{ til::details::vars::computer_name }, value.get());
247246
}
248247
}
249248

@@ -256,17 +255,17 @@ namespace til // Terminal Implementation Library. Also: "Today I Learned"
256255
DWORD accountNameSize = 0, userDomainSize = 0;
257256
SID_NAME_USE sidNameUse;
258257
SetLastError(ERROR_SUCCESS);
259-
if (LookupAccountSidLocalW(user.get()->User.Sid, nullptr, &accountNameSize, nullptr, &userDomainSize, &sidNameUse) || GetLastError() == ERROR_INSUFFICIENT_BUFFER)
258+
if (LookupAccountSidW(nullptr, user.get()->User.Sid, nullptr, &accountNameSize, nullptr, &userDomainSize, &sidNameUse) || GetLastError() == ERROR_INSUFFICIENT_BUFFER)
260259
{
261260
std::wstring accountName, userDomain;
262261
accountName.resize(accountNameSize);
263262
userDomain.resize(userDomainSize);
264263

265264
SetLastError(ERROR_SUCCESS);
266-
if (LookupAccountSidLocalW(user.get()->User.Sid, accountName.data(), &accountNameSize, userDomain.data(), &userDomainSize, &sidNameUse))
265+
if (LookupAccountSidW(nullptr, user.get()->User.Sid, accountName.data(), &accountNameSize, userDomain.data(), &userDomainSize, &sidNameUse))
267266
{
268-
insert_or_assign(std::wstring{ til::details::vars::user_name }, accountName);
269-
insert_or_assign(std::wstring{ til::details::vars::user_domain }, userDomain);
267+
save_to_map(std::wstring{ til::details::vars::user_name }, accountName);
268+
save_to_map(std::wstring{ til::details::vars::user_domain }, userDomain);
270269
}
271270
}
272271
}
@@ -277,10 +276,10 @@ namespace til // Terminal Implementation Library. Also: "Today I Learned"
277276
wil::unique_hkey key;
278277
if (RegOpenKeyExW(HKEY_LOCAL_MACHINE, til::details::vars::reg::program_files_root.data(), 0, KEY_READ, &key) == ERROR_SUCCESS)
279278
{
280-
for (auto& pair : til::details::vars::program_files_map)
279+
for (auto& [keyName, varName] : til::details::vars::program_files_map)
281280
{
282-
auto value = til::details::wiltmp::RegQueryValueExW<std::wstring, 256>(key.get(), pair.first.c_str());
283-
set_user_environment_var(std::wstring{ pair.second }, value);
281+
auto value = til::details::wiltmp::RegQueryValueExW<std::wstring, 256>(key.get(), keyName.c_str());
282+
set_user_environment_var(std::wstring{ varName }, value);
284283
}
285284
}
286285
}
@@ -291,23 +290,29 @@ namespace til // Terminal Implementation Library. Also: "Today I Learned"
291290
if (RegOpenKeyExW(rootKey, subkey.data(), 0, KEY_READ, &key) == ERROR_SUCCESS)
292291
{
293292
DWORD maxValueNameSize = 0, maxValueDataSize = 0;
294-
if (RegQueryInfoKeyW(key.get(), nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, &maxValueNameSize, &maxValueDataSize, nullptr, nullptr))
293+
if (RegQueryInfoKeyW(key.get(), nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, &maxValueNameSize, &maxValueDataSize, nullptr, nullptr) == ERROR_SUCCESS)
295294
{
295+
maxValueNameSize++; // the query does not include the terminating null, but needs that space to enum
296296
std::wstring valueName;
297297
std::basic_string<BYTE> valueData;
298298
valueName.resize(maxValueNameSize);
299299
valueData.resize(maxValueDataSize);
300300

301-
DWORD valueNameSize = maxValueNameSize;
302-
DWORD valueDataSize = maxValueDataSize;
303-
304-
DWORD index = 0;
305-
DWORD type = 0;
306-
307301
for (DWORD pass = 0; pass < 2; ++pass)
308302
{
309-
while (!RegEnumValueW(key.get(), index, valueName.data(), &valueNameSize, nullptr, &type, valueData.data(), &valueDataSize))
303+
DWORD valueNameSize = maxValueNameSize;
304+
DWORD valueDataSize = maxValueDataSize;
305+
306+
DWORD index = 0;
307+
DWORD type = 0;
308+
309+
while (true)
310310
{
311+
DWORD result = RegEnumValueW(key.get(), index, valueName.data(), &valueNameSize, nullptr, &type, valueData.data(), &valueDataSize);
312+
if (result != ERROR_SUCCESS)
313+
{
314+
break;
315+
}
311316
valueName.resize(valueNameSize);
312317
valueData.resize(valueDataSize);
313318

@@ -320,7 +325,6 @@ namespace til // Terminal Implementation Library. Also: "Today I Learned"
320325
reinterpret_cast<wchar_t*>(valueData.data()), valueData.size() / sizeof(wchar_t)
321326
};
322327
}
323-
// TODO: This one is supposed to be using some of the vars we just expanded from the first one.
324328
else if (pass == 1 && (type == REG_EXPAND_SZ) && valueDataSize >= sizeof(wchar_t))
325329
{
326330
data = {
@@ -329,9 +333,20 @@ namespace til // Terminal Implementation Library. Also: "Today I Learned"
329333
data = expand_environment_strings(data.data());
330334
}
331335

336+
// Because Registry data may or may not be null terminated... check if we've managed
337+
// to store an extra null in the wstring by telling it to create itself from pointer and size.
338+
// If we did, pull it off.
339+
if (!data.empty())
340+
{
341+
if (data.back() == L'\0')
342+
{
343+
data = data.substr(0, data.size() - 1);
344+
}
345+
}
346+
332347
if (!data.empty())
333348
{
334-
if (is_path_var(data))
349+
if (is_path_var(valueName))
335350
{
336351
concat_var(valueName, std::wstring{ data });
337352
}
@@ -348,6 +363,7 @@ namespace til // Terminal Implementation Library. Also: "Today I Learned"
348363
valueDataSize = maxValueDataSize;
349364
index++;
350365
}
366+
index = 0;
351367
}
352368
}
353369
}
@@ -363,25 +379,35 @@ namespace til // Terminal Implementation Library. Also: "Today I Learned"
363379
{
364380
value = expand_environment_strings(value);
365381
value = check_for_temp(value);
366-
insert_or_assign(var, value);
382+
save_to_map(var, value);
367383
}
368384

369385
void concat_var(std::wstring var, std::wstring value)
370386
{
371387
// I wanted contains() but this isn't C++20... yet.
372388
if (find(var) != end())
373389
{
374-
auto existing = at(var);
390+
std::wstring existing = at(var);
391+
// If it doesn't already trail with a ;, add one.
392+
// Otherwise, just take advantage of the one it has.
375393
if (existing.back() != L';')
376394
{
377395
existing.append(L";");
378396
}
379397
existing.append(value);
380-
insert_or_assign(var, existing);
398+
save_to_map(var, existing);
381399
}
382400
else
383401
{
384-
insert(var, value);
402+
save_to_map(var, value);
403+
}
404+
}
405+
406+
void save_to_map(std::wstring var, std::wstring value)
407+
{
408+
if (!var.empty() && !value.empty())
409+
{
410+
insert_or_assign(var, value);
385411
}
386412
}
387413

@@ -408,9 +434,40 @@ namespace til // Terminal Implementation Library. Also: "Today I Learned"
408434
return !_wcsicmp(input.data(), path.data()) || !_wcsicmp(input.data(), libpath.data()) || !_wcsicmp(input.data(), os2libpath.data());
409435
}
410436

437+
void parse(wchar_t* block)
438+
{
439+
for (wchar_t const* lastCh{ block }; *lastCh != '\0'; ++lastCh)
440+
{
441+
// Copy current entry into temporary map.
442+
const size_t cchEntry{ ::wcslen(lastCh) };
443+
const std::wstring_view entry{ lastCh, cchEntry };
444+
445+
// Every entry is of the form "name=value\0".
446+
const auto pos = entry.find_first_of(L"=", 0, 1);
447+
THROW_HR_IF(E_UNEXPECTED, pos == std::wstring::npos);
448+
449+
std::wstring name{ entry.substr(0, pos) }; // portion before '='
450+
std::wstring value{ entry.substr(pos + 1) }; // portion after '='
451+
452+
// Don't replace entries that already exist.
453+
try_emplace(std::move(name), std::move(value));
454+
lastCh += cchEntry;
455+
}
456+
}
457+
411458
public:
459+
env()
460+
{
461+
}
462+
463+
env(wchar_t* block)
464+
{
465+
parse(block);
466+
}
467+
412468
void regenerate()
413469
{
470+
// Generally replicates the behavior of shell32!RegenerateUserEnvironment
414471
get(std::wstring{ til::details::vars::system_root });
415472
get(std::wstring{ til::details::vars::system_drive });
416473
get(std::wstring{ til::details::vars::all_users_profile });
@@ -432,5 +489,19 @@ namespace til // Terminal Implementation Library. Also: "Today I Learned"
432489
get_vars_from_registry(HKEY_CURRENT_USER, til::details::vars::reg::user_volatile_env_var_root);
433490
get_vars_from_registry(HKEY_CURRENT_USER, fmt::format(til::details::vars::reg::user_volatile_session_env_var_root_pattern, NtCurrentTeb()->ProcessEnvironmentBlock->SessionId));
434491
}
492+
493+
std::wstring to_string()
494+
{
495+
std::wstring result;
496+
for (const auto& [name, value] : *this)
497+
{
498+
result += name;
499+
result += L"=";
500+
result += value;
501+
result.append(L"\0", 1); // Override string's natural propensity to stop at \0
502+
}
503+
result.append(L"\0", 1);
504+
return result;
505+
}
435506
};
436507
};

0 commit comments

Comments
 (0)