From 584213c5ffeca09b3fe24bd5e92f73fd057ac642 Mon Sep 17 00:00:00 2001 From: Sean Hall Date: Fri, 3 Jun 2022 17:49:15 -0500 Subject: Add RegReadUnexpandedString to get an unexpanded REG_EXPAND_SZ value. --- src/libs/dutil/WixToolset.DUtil/regutil.cpp | 420 ++++++++++++++++++---------- 1 file changed, 265 insertions(+), 155 deletions(-) (limited to 'src/libs/dutil/WixToolset.DUtil/regutil.cpp') diff --git a/src/libs/dutil/WixToolset.DUtil/regutil.cpp b/src/libs/dutil/WixToolset.DUtil/regutil.cpp index 64224d42..584966ed 100644 --- a/src/libs/dutil/WixToolset.DUtil/regutil.cpp +++ b/src/libs/dutil/WixToolset.DUtil/regutil.cpp @@ -9,6 +9,7 @@ #define RegExitWithLastError(x, s, ...) ExitWithLastErrorSource(DUTIL_SOURCE_REGUTIL, x, s, __VA_ARGS__) #define RegExitOnFailure(x, s, ...) ExitOnFailureSource(DUTIL_SOURCE_REGUTIL, x, s, __VA_ARGS__) #define RegExitOnRootFailure(x, s, ...) ExitOnRootFailureSource(DUTIL_SOURCE_REGUTIL, x, s, __VA_ARGS__) +#define RegExitWithRootFailure(x, e, s, ...) ExitWithRootFailureSource(DUTIL_SOURCE_REGUTIL, x, e, s, __VA_ARGS__) #define RegExitOnFailureDebugTrace(x, s, ...) ExitOnFailureDebugTraceSource(DUTIL_SOURCE_REGUTIL, x, s, __VA_ARGS__) #define RegExitOnNull(p, x, e, s, ...) ExitOnNullSource(DUTIL_SOURCE_REGUTIL, p, x, e, s, __VA_ARGS__) #define RegExitOnNullWithLastError(p, x, s, ...) ExitOnNullWithLastErrorSource(DUTIL_SOURCE_REGUTIL, p, x, s, __VA_ARGS__) @@ -28,10 +29,19 @@ static PFN_REGQUERYINFOKEYW vpfnRegQueryInfoKeyW = ::RegQueryInfoKeyW; static PFN_REGQUERYVALUEEXW vpfnRegQueryValueExW = ::RegQueryValueExW; static PFN_REGSETVALUEEXW vpfnRegSetValueExW = ::RegSetValueExW; static PFN_REGDELETEVALUEW vpfnRegDeleteValueW = ::RegDeleteValueW; +static PFN_REGGETVALUEW vpfnRegGetValueW = NULL; +static PFN_REGGETVALUEW vpfnRegGetValueWFromLibrary = NULL; static HMODULE vhAdvApi32Dll = NULL; static BOOL vfRegInitialized = FALSE; +static HRESULT GetRegValue( + __in HKEY hk, + __in_z_opt LPCWSTR wzName, + __deref_out_bcount_opt(*pcbBuffer) BYTE* pbBuffer, + __inout SIZE_T* pcbBuffer, + __out DWORD* pdwType + ); static HRESULT WriteStringToRegistry( __in HKEY hk, __in_z_opt LPCWSTR wzName, @@ -46,14 +56,22 @@ DAPI_(HRESULT) RegInitialize() hr = LoadSystemLibrary(L"AdvApi32.dll", &vhAdvApi32Dll); RegExitOnFailure(hr, "Failed to load AdvApi32.dll"); - // ignore failures - if this doesn't exist, we'll fall back to RegDeleteKeyW + // Ignore failures - if this doesn't exist, we'll fall back to RegDeleteKeyW. vpfnRegDeleteKeyExWFromLibrary = reinterpret_cast(::GetProcAddress(vhAdvApi32Dll, "RegDeleteKeyExW")); - if (NULL == vpfnRegDeleteKeyExW) + // Ignore failures - if this doesn't exist, we'll fall back to RegQueryValueExW. + vpfnRegGetValueWFromLibrary = reinterpret_cast(::GetProcAddress(vhAdvApi32Dll, "RegGetValueW")); + + if (!vpfnRegDeleteKeyExW) { vpfnRegDeleteKeyExW = vpfnRegDeleteKeyExWFromLibrary; } + if (!vpfnRegGetValueW) + { + vpfnRegGetValueW = vpfnRegGetValueWFromLibrary; + } + vfRegInitialized = TRUE; LExit: @@ -68,7 +86,9 @@ DAPI_(void) RegUninitialize() ::FreeLibrary(vhAdvApi32Dll); vhAdvApi32Dll = NULL; vpfnRegDeleteKeyExWFromLibrary = NULL; + vpfnRegGetValueWFromLibrary = NULL; vpfnRegDeleteKeyExW = NULL; + vpfnRegGetValueW = NULL; } vfRegInitialized = FALSE; @@ -84,7 +104,8 @@ DAPI_(void) RegFunctionOverride( __in_opt PFN_REGQUERYINFOKEYW pfnRegQueryInfoKeyW, __in_opt PFN_REGQUERYVALUEEXW pfnRegQueryValueExW, __in_opt PFN_REGSETVALUEEXW pfnRegSetValueExW, - __in_opt PFN_REGDELETEVALUEW pfnRegDeleteValueW + __in_opt PFN_REGDELETEVALUEW pfnRegDeleteValueW, + __in_opt PFN_REGGETVALUEW pfnRegGetValueW ) { vpfnRegCreateKeyExW = pfnRegCreateKeyExW ? pfnRegCreateKeyExW : ::RegCreateKeyExW; @@ -96,6 +117,14 @@ DAPI_(void) RegFunctionOverride( vpfnRegQueryValueExW = pfnRegQueryValueExW ? pfnRegQueryValueExW : ::RegQueryValueExW; vpfnRegSetValueExW = pfnRegSetValueExW ? pfnRegSetValueExW : ::RegSetValueExW; vpfnRegDeleteValueW = pfnRegDeleteValueW ? pfnRegDeleteValueW : ::RegDeleteValueW; + vpfnRegGetValueW = pfnRegGetValueW ? pfnRegGetValueW : vpfnRegGetValueWFromLibrary; +} + + +DAPI_(void) RegFunctionForceFallback() +{ + vpfnRegDeleteKeyExW = NULL; + vpfnRegGetValueW = NULL; } @@ -358,54 +387,104 @@ LExit: return hr; } -DAPI_(HRESULT) RegReadBinary( +DAPI_(HRESULT) RegReadValue( __in HKEY hk, __in_z_opt LPCWSTR wzName, + __in BOOL fExpand, __deref_out_bcount_opt(*pcbBuffer) BYTE** ppbBuffer, - __out SIZE_T *pcbBuffer - ) + __inout SIZE_T* pcbBuffer, + __out DWORD* pdwType + ) { HRESULT hr = S_OK; - LPBYTE pbBuffer = NULL; - DWORD er = ERROR_SUCCESS; - DWORD cb = 0; - DWORD dwType = 0; - - er = vpfnRegQueryValueExW(hk, wzName, NULL, &dwType, NULL, &cb); - RegExitOnWin32Error(er, hr, "Failed to get size of registry value."); + DWORD dwAttempts = 0; + LPWSTR sczExpand = NULL; - // Zero-length binary values can exist - if (0 < cb) + hr = GetRegValue(hk, wzName, *ppbBuffer, pcbBuffer, pdwType); + if (E_FILENOTFOUND == hr) { - pbBuffer = static_cast(MemAlloc(cb, FALSE)); - RegExitOnNull(pbBuffer, hr, E_OUTOFMEMORY, "Failed to allocate buffer for binary registry value."); + ExitFunction(); + } + else if (E_MOREDATA != hr) + { + RegExitOnFailure(hr, "Failed to get size of raw registry value."); - er = vpfnRegQueryValueExW(hk, wzName, NULL, &dwType, pbBuffer, &cb); - if (E_FILENOTFOUND == HRESULT_FROM_WIN32(er)) + // Zero-length raw values can exist + if (!*ppbBuffer && 0 < *pcbBuffer) { - ExitFunction1(hr = E_FILENOTFOUND); + hr = E_MOREDATA; } - RegExitOnWin32Error(er, hr, "Failed to read registry value."); } - if (REG_BINARY == dwType) + while (E_MOREDATA == hr && dwAttempts < 10) { - *ppbBuffer = pbBuffer; - pbBuffer = NULL; - *pcbBuffer = cb; + ++dwAttempts; + + if (*ppbBuffer) + { + *ppbBuffer = static_cast(MemReAlloc(*ppbBuffer, *pcbBuffer, FALSE)); + } + else + { + *ppbBuffer = static_cast(MemAlloc(*pcbBuffer, FALSE)); + } + RegExitOnNull(*ppbBuffer, hr, E_OUTOFMEMORY, "Failed to allocate buffer for raw registry value."); + + hr = GetRegValue(hk, wzName, *ppbBuffer, pcbBuffer, pdwType); + if (E_FILENOTFOUND == hr) + { + ExitFunction(); + } + else if (E_MOREDATA != hr) + { + RegExitOnFailure(hr, "Failed to read raw registry value."); + } } - else + + if (fExpand && SUCCEEDED(hr) && REG_EXPAND_SZ == *pdwType) { - hr = HRESULT_FROM_WIN32(ERROR_INVALID_DATATYPE); - RegExitOnRootFailure(hr, "Error reading binary registry value due to unexpected data type: %u", dwType); + LPWSTR sczValue = reinterpret_cast(*ppbBuffer); + hr = PathExpand(&sczExpand, sczValue, PATH_EXPAND_ENVIRONMENT); + RegExitOnFailure(hr, "Failed to expand registry value: %ls", sczValue); + + *ppbBuffer = reinterpret_cast(sczExpand); + *pcbBuffer = (lstrlenW(sczExpand) + 1) * sizeof(WCHAR); + sczExpand = NULL; + ReleaseMem(sczValue); } LExit: - ReleaseMem(pbBuffer); + ReleaseStr(sczExpand); return hr; } +DAPI_(HRESULT) RegReadBinary( + __in HKEY hk, + __in_z_opt LPCWSTR wzName, + __deref_out_bcount_opt(*pcbBuffer) BYTE** ppbBuffer, + __out SIZE_T* pcbBuffer + ) +{ + HRESULT hr = S_OK; + DWORD dwType = 0; + + hr = RegReadValue(hk, wzName, FALSE, ppbBuffer, pcbBuffer, &dwType); + if (E_FILENOTFOUND == hr) + { + ExitFunction(); + } + RegExitOnFailure(hr, "Failed to read binary registry value."); + + if (REG_BINARY != dwType) + { + RegExitWithRootFailure(hr, HRESULT_FROM_WIN32(ERROR_INVALID_DATATYPE), "Error reading binary registry value due to unexpected data type: %u", dwType); + } + +LExit: + return hr; +} + DAPI_(HRESULT) RegReadString( __in HKEY hk, @@ -414,65 +493,64 @@ DAPI_(HRESULT) RegReadString( ) { HRESULT hr = S_OK; - DWORD er = ERROR_SUCCESS; SIZE_T cbValue = 0; - DWORD cch = 0; - DWORD cb = 0; DWORD dwType = 0; - LPWSTR sczExpand = NULL; if (psczValue && *psczValue) { - hr = StrMaxLength(*psczValue, &cbValue); - RegExitOnFailure(hr, "Failed to determine length of string."); - - cch = (DWORD)min(DWORD_MAX, cbValue); + hr = MemSizeChecked(*psczValue, &cbValue); + RegExitOnFailure(hr, "Failed to get size of input buffer."); } - if (2 > cch) + hr = RegReadValue(hk, wzName, TRUE, reinterpret_cast(psczValue), &cbValue, &dwType); + if (E_FILENOTFOUND == hr) { - cch = 2; - - hr = StrAlloc(psczValue, cch); - RegExitOnFailure(hr, "Failed to allocate string to minimum size."); + ExitFunction(); } + RegExitOnFailure(hr, "Failed to read string registry value."); - cb = sizeof(WCHAR) * (cch - 1); // subtract one to ensure there will be a space at the end of the string for the null terminator. - er = vpfnRegQueryValueExW(hk, wzName, NULL, &dwType, reinterpret_cast(*psczValue), &cb); - if (ERROR_MORE_DATA == er) + if (REG_EXPAND_SZ != dwType && REG_SZ != dwType) { - cch = cb / sizeof(WCHAR) + 1; // add one to ensure there will be space at the end for the null terminator - hr = StrAlloc(psczValue, cch); - RegExitOnFailure(hr, "Failed to allocate string bigger for registry value."); - - er = vpfnRegQueryValueExW(hk, wzName, NULL, &dwType, reinterpret_cast(*psczValue), &cb); + RegExitWithRootFailure(hr, HRESULT_FROM_WIN32(ERROR_INVALID_DATATYPE), "Error reading string registry value due to unexpected data type: %u", dwType); } - if (E_FILENOTFOUND == HRESULT_FROM_WIN32(er)) + +LExit: + return hr; +} + + +DAPI_(HRESULT) RegReadUnexpandedString( + __in HKEY hk, + __in_z_opt LPCWSTR wzName, + __inout BOOL* pfNeedsExpansion, + __deref_out_z LPWSTR* psczValue + ) +{ + HRESULT hr = S_OK; + SIZE_T cbValue = 0; + DWORD dwType = 0; + LPWSTR sczExpand = NULL; + + if (psczValue && *psczValue) { - ExitFunction1(hr = E_FILENOTFOUND); + hr = MemSizeChecked(*psczValue, &cbValue); + RegExitOnFailure(hr, "Failed to get size of input buffer."); } - RegExitOnWin32Error(er, hr, "Failed to read registry key."); - if (REG_SZ == dwType || REG_EXPAND_SZ == dwType) + hr = RegReadValue(hk, wzName, FALSE, reinterpret_cast(psczValue), &cbValue, &dwType); + if (E_FILENOTFOUND == hr) { - // Always ensure the registry value is null terminated. - (*psczValue)[cch - 1] = L'\0'; - - if (REG_EXPAND_SZ == dwType) - { - hr = StrAllocString(&sczExpand, *psczValue, 0); - RegExitOnFailure(hr, "Failed to copy registry value to expand."); - - hr = PathExpand(psczValue, sczExpand, PATH_EXPAND_ENVIRONMENT); - RegExitOnFailure(hr, "Failed to expand registry value: %ls", *psczValue); - } + ExitFunction(); } - else + RegExitOnFailure(hr, "Failed to read expand string registry value."); + + if (REG_EXPAND_SZ != dwType && REG_SZ != dwType) { - hr = HRESULT_FROM_WIN32(ERROR_INVALID_DATATYPE); - RegExitOnRootFailure(hr, "Error reading string registry value due to unexpected data type: %u", dwType); + RegExitWithRootFailure(hr, HRESULT_FROM_WIN32(ERROR_INVALID_DATATYPE), "Error reading expand string registry value due to unexpected data type: %u", dwType); } + *pfNeedsExpansion = REG_EXPAND_SZ == dwType; + LExit: ReleaseStr(sczExpand); @@ -488,67 +566,57 @@ DAPI_(HRESULT) RegReadStringArray( ) { HRESULT hr = S_OK; - DWORD er = ERROR_SUCCESS; DWORD dwNullCharacters = 0; DWORD dwType = 0; - DWORD cb = 0; - DWORD cch = 0; + SIZE_T cb = 0; + SIZE_T cch = 0; LPCWSTR wzSource = NULL; LPWSTR sczValue = NULL; - er = vpfnRegQueryValueExW(hk, wzName, NULL, &dwType, reinterpret_cast(sczValue), &cb); - if (0 < cb) - { - cch = cb / sizeof(WCHAR); - hr = StrAlloc(&sczValue, cch); - RegExitOnFailure(hr, "Failed to allocate string for registry value."); - - er = vpfnRegQueryValueExW(hk, wzName, NULL, &dwType, reinterpret_cast(sczValue), &cb); - } - if (E_FILENOTFOUND == HRESULT_FROM_WIN32(er)) + hr = RegReadValue(hk, wzName, FALSE, reinterpret_cast(&sczValue), &cb, &dwType); + if (E_FILENOTFOUND == hr) { - ExitFunction1(hr = E_FILENOTFOUND); + ExitFunction(); } - RegExitOnWin32Error(er, hr, "Failed to read registry key."); + RegExitOnFailure(hr, "Failed to read string array registry value."); - if (cb / sizeof(WCHAR) != cch) + if (REG_MULTI_SZ != dwType) { - hr = E_UNEXPECTED; - RegExitOnFailure(hr, "The size of registry value %ls unexpected changed between 2 reads", wzName); + RegExitWithRootFailure(hr, HRESULT_FROM_WIN32(ERROR_INVALID_DATATYPE), "Tried to read string array, but registry value %ls is of an incorrect type", wzName); } - if (REG_MULTI_SZ != dwType) + cch = cb / sizeof(WCHAR); + + for (DWORD i = 0; i < cch; ++i) { - hr = HRESULT_FROM_WIN32(ERROR_INVALID_DATATYPE); - RegExitOnRootFailure(hr, "Tried to read string array, but registry value %ls is of an incorrect type", wzName); + if (L'\0' == sczValue[i]) + { + ++dwNullCharacters; + } } // Value exists, but is empty, so no strings to return. - if (2 > cch) + if (!cb || 1 == dwNullCharacters && 1 == cch || 2 == dwNullCharacters && 2 == cch) { *prgsczStrings = NULL; *pcStrings = 0; ExitFunction1(hr = S_OK); } - // The docs specifically say if the value was written without double-null-termination, it'll get read back without it too. - if (L'\0' != sczValue[cch-1] || L'\0' != sczValue[cch-2]) + if (sczValue[cch - 1] != L'\0') { - hr = E_INVALIDARG; - RegExitOnFailure(hr, "Tried to read string array, but registry value %ls is invalid (isn't double-null-terminated)", wzName); + // Count the terminating null character that RegReadValue added past the end. + ++dwNullCharacters; + Assert(!sczValue[cch]); } - - cch = cb / sizeof(WCHAR); - for (DWORD i = 0; i < cch; ++i) + else if (cch > 1 && sczValue[cch - 2] == L'\0') { - if (L'\0' == sczValue[i]) - { - ++dwNullCharacters; - } + // Don't count the extra 1 at the end of the properly double-null terminated string. + --dwNullCharacters; } - // There's one string for every null character encountered (except the extra 1 at the end of the string) - *pcStrings = dwNullCharacters - 1; + // There's one string for every null character encountered. + *pcStrings = dwNullCharacters; hr = MemEnsureArraySize(reinterpret_cast(prgsczStrings), *pcStrings, sizeof(LPWSTR), 0); RegExitOnFailure(hr, "Failed to resize array while reading REG_MULTI_SZ value"); @@ -557,11 +625,13 @@ DAPI_(HRESULT) RegReadStringArray( wzSource = sczValue; for (DWORD i = 0; i < *pcStrings; ++i) { - hr = StrAllocString(&(*prgsczStrings)[i], wzSource, 0); + cch = lstrlenW(wzSource); + + hr = StrAllocString(&(*prgsczStrings)[i], wzSource, cch); RegExitOnFailure(hr, "Failed to allocate copy of string"); // Skip past this string - wzSource += lstrlenW(wzSource) + 1; + wzSource += cch + 1; } #pragma prefast(pop) @@ -579,38 +649,36 @@ DAPI_(HRESULT) RegReadVersion( ) { HRESULT hr = S_OK; - DWORD er = ERROR_SUCCESS; DWORD dwType = 0; - DWORD cb = 0; - LPWSTR sczVersion = NULL; + SIZE_T cb = 0; + LPWSTR sczValue = NULL; - cb = sizeof(DWORD64); - er = vpfnRegQueryValueExW(hk, wzName, NULL, &dwType, reinterpret_cast(*pdw64Version), &cb); - if (E_FILENOTFOUND == HRESULT_FROM_WIN32(er)) + hr = RegReadValue(hk, wzName, TRUE, reinterpret_cast(&sczValue), &cb, &dwType); + if (E_FILENOTFOUND == hr) { - ExitFunction1(hr = E_FILENOTFOUND); + ExitFunction(); } + RegExitOnFailure(hr, "Failed to read version registry value."); if (REG_SZ == dwType || REG_EXPAND_SZ == dwType) { - hr = RegReadString(hk, wzName, &sczVersion); - RegExitOnFailure(hr, "Failed to read registry version as string."); - - hr = FileVersionFromStringEx(sczVersion, 0, pdw64Version); + hr = FileVersionFromStringEx(sczValue, 0, pdw64Version); RegExitOnFailure(hr, "Failed to convert registry string to version."); } else if (REG_QWORD == dwType) { - RegExitOnWin32Error(er, hr, "Failed to read registry key."); + if (memcpy_s(pdw64Version, sizeof(DWORD64), sczValue, cb)) + { + RegExitWithRootFailure(hr, E_INVALIDARG, "Failed to copy QWORD version value."); + } } else // unexpected data type { - hr = HRESULT_FROM_WIN32(ERROR_INVALID_DATATYPE); - RegExitOnRootFailure(hr, "Error reading version registry value due to unexpected data type: %u", dwType); + RegExitWithRootFailure(hr, HRESULT_FROM_WIN32(ERROR_INVALID_DATATYPE), "Error reading version registry value due to unexpected data type: %u", dwType); } LExit: - ReleaseStr(sczVersion); + ReleaseStr(sczValue); return hr; } @@ -623,37 +691,38 @@ DAPI_(HRESULT) RegReadWixVersion( ) { HRESULT hr = S_OK; - DWORD er = ERROR_SUCCESS; DWORD dwType = 0; - DWORD cb = 0; + SIZE_T cb = 0; DWORD64 dw64Version = 0; - LPWSTR sczVersion = NULL; + LPWSTR sczValue = NULL; VERUTIL_VERSION* pVersion = NULL; - cb = sizeof(DWORD64); - er = vpfnRegQueryValueExW(hk, wzName, NULL, &dwType, reinterpret_cast(&dw64Version), &cb); - if (E_FILENOTFOUND == HRESULT_FROM_WIN32(er)) + hr = RegReadValue(hk, wzName, TRUE, reinterpret_cast(&sczValue), &cb, &dwType); + if (E_FILENOTFOUND == hr) { - ExitFunction1(hr = E_FILENOTFOUND); + ExitFunction(); } + RegExitOnFailure(hr, "Failed to read wix version registry value."); if (REG_SZ == dwType || REG_EXPAND_SZ == dwType) { - hr = RegReadString(hk, wzName, &sczVersion); - RegExitOnFailure(hr, "Failed to read registry version as string."); - - hr = VerParseVersion(sczVersion, 0, FALSE, &pVersion); - RegExitOnFailure(hr, "Failed to convert registry string to version."); + hr = VerParseVersion(sczValue, 0, FALSE, &pVersion); + RegExitOnFailure(hr, "Failed to convert registry string to wix version."); } else if (REG_QWORD == dwType) { + if (memcpy_s(&dw64Version, sizeof(DWORD64), sczValue, cb)) + { + RegExitWithRootFailure(hr, E_INVALIDARG, "Failed to copy QWORD wix version value."); + } + hr = VerVersionFromQword(dw64Version, &pVersion); - RegExitOnFailure(hr, "Failed to convert registry string to version."); + RegExitOnFailure(hr, "Failed to convert registry string to wix version."); } else // unexpected data type { hr = HRESULT_FROM_WIN32(ERROR_INVALID_DATATYPE); - RegExitOnRootFailure(hr, "Error reading version registry value due to unexpected data type: %u", dwType); + RegExitOnRootFailure(hr, "Error reading wix version registry value due to unexpected data type: %u", dwType); } *ppVersion = pVersion; @@ -661,7 +730,7 @@ DAPI_(HRESULT) RegReadWixVersion( LExit: ReleaseVerutilVersion(pVersion); - ReleaseStr(sczVersion); + ReleaseStr(sczValue); return hr; } @@ -673,20 +742,18 @@ DAPI_(HRESULT) RegReadNone( ) { HRESULT hr = S_OK; - DWORD er = ERROR_SUCCESS; DWORD dwType = 0; - er = vpfnRegQueryValueExW(hk, wzName, NULL, &dwType, NULL, NULL); - if (E_FILENOTFOUND == HRESULT_FROM_WIN32(er)) + hr = RegGetType(hk, wzName, &dwType); + if (E_FILENOTFOUND == hr) { - ExitFunction1(hr = E_FILENOTFOUND); + ExitFunction(); } - RegExitOnWin32Error(er, hr, "Failed to query registry key value."); + RegExitOnFailure(hr, "Error reading none registry value type."); if (REG_NONE != dwType) { - hr = HRESULT_FROM_WIN32(ERROR_INVALID_DATATYPE); - RegExitOnRootFailure(hr, "Error reading none registry value due to unexpected data type: %u", dwType); + RegExitWithRootFailure(hr, HRESULT_FROM_WIN32(ERROR_INVALID_DATATYPE), "Error reading none registry value due to unexpected data type: %u", dwType); } LExit: @@ -829,12 +896,9 @@ DAPI_(HRESULT) RegWriteStringArray( DWORD dwTotalStringSize = 0; DWORD cbTotalStringSize = 0; DWORD dwTemp = 0; + DWORD dwRemainingStringSize = 0; - if (0 == cValues) - { - wzWriteValue = L"\0"; - } - else + if (cValues) { // Add space for the null terminator dwTotalStringSize = 1; @@ -850,21 +914,22 @@ DAPI_(HRESULT) RegWriteStringArray( RegExitOnFailure(hr, "Failed to allocate space for string while writing REG_MULTI_SZ"); wzCopyDestination = sczWriteValue; - dwTemp = dwTotalStringSize; + dwRemainingStringSize = dwTotalStringSize; for (DWORD i = 0; i < cValues; ++i) { - hr = ::StringCchCopyW(wzCopyDestination, dwTotalStringSize, rgwzValues[i]); + hr = ::StringCchCopyW(wzCopyDestination, dwRemainingStringSize, rgwzValues[i]); RegExitOnFailure(hr, "failed to copy string: %ls", rgwzValues[i]); - dwTemp -= lstrlenW(rgwzValues[i]) + 1; - wzCopyDestination += lstrlenW(rgwzValues[i]) + 1; + dwTemp = lstrlenW(rgwzValues[i]) + 1; + dwRemainingStringSize -= dwTemp; + wzCopyDestination += dwTemp; } wzWriteValue = sczWriteValue; - } - hr = ::DWordMult(dwTotalStringSize, sizeof(WCHAR), &cbTotalStringSize); - RegExitOnFailure(hr, "Failed to get total string size in bytes"); + hr = ::DWordMult(dwTotalStringSize, sizeof(WCHAR), &cbTotalStringSize); + RegExitOnFailure(hr, "Failed to get total string size in bytes"); + } er = vpfnRegSetValueExW(hk, wzName, 0, REG_MULTI_SZ, reinterpret_cast(wzWriteValue), cbTotalStringSize); RegExitOnWin32Error(er, hr, "Failed to set registry value to array of strings (first string of which is): %ls", wzWriteValue); @@ -1000,6 +1065,48 @@ DAPI_(REGSAM) RegTranslateKeyBitness( } } +static HRESULT GetRegValue( + __in HKEY hk, + __in_z_opt LPCWSTR wzName, + __deref_out_bcount_opt(*pcbBuffer) BYTE* pbBuffer, + __inout SIZE_T* pcbBuffer, + __out DWORD* pdwType + ) +{ + HRESULT hr = S_OK; + DWORD cb = (DWORD)min(DWORD_MAX, *pcbBuffer); + + if (vpfnRegGetValueW) + { + hr = HRESULT_FROM_WIN32(vpfnRegGetValueW(hk, NULL, wzName, RRF_RT_ANY | RRF_NOEXPAND, pdwType, pbBuffer, &cb)); + } + else + { + hr = HRESULT_FROM_WIN32(vpfnRegQueryValueExW(hk, wzName, NULL, pdwType, pbBuffer, &cb)); + + if (REG_SZ == *pdwType || REG_EXPAND_SZ == *pdwType || REG_MULTI_SZ == *pdwType) + { + if (E_MOREDATA == hr || S_OK == hr && (cb + sizeof(WCHAR)) > *pcbBuffer) + { + // Make sure there's room for a null terminator at the end. + HRESULT hrAdd = ::DWordAdd(cb, sizeof(WCHAR), &cb); + + hr = FAILED(hrAdd) ? hrAdd : (pbBuffer ? E_MOREDATA : S_OK); + } + else if (S_OK == hr && pbBuffer) + { + // Always ensure the registry value is null terminated. + WCHAR* pch = reinterpret_cast(pbBuffer + cb); + *pch = L'\0'; + } + } + } + + *pcbBuffer = cb; + + return hr; +} + static HRESULT WriteStringToRegistry( __in HKEY hk, __in_z_opt LPCWSTR wzName, @@ -1013,9 +1120,12 @@ static HRESULT WriteStringToRegistry( if (wzValue) { - hr = ::StringCbLengthW(wzValue, STRSAFE_MAX_CCH * sizeof(TCHAR), &cbValue); + hr = ::StringCbLengthW(wzValue, DWORD_MAX, &cbValue); RegExitOnFailure(hr, "Failed to determine length of registry value: %ls", wzName); + // Need to include the null terminator. + cbValue += sizeof(WCHAR); + er = vpfnRegSetValueExW(hk, wzName, 0, dwType, reinterpret_cast(wzValue), static_cast(cbValue)); RegExitOnWin32Error(er, hr, "Failed to set registry value: %ls", wzName); } -- cgit v1.2.3-55-g6feb