From 7fef07d0fd78b83a73c87ba6901b93ce73a8a807 Mon Sep 17 00:00:00 2001 From: Simon Rozman Date: Thu, 12 Nov 2015 19:35:26 +0100 Subject: [PATCH] Selected system functions are called with fixed size stack buffers first to eliminate the need to retry again when the stack buffer proved sufficient now. --- include/atlcrypt.h | 4 +- include/atlex.h | 17 +++++ include/atlmsi.h | 176 +++++++++++++++++++++++++------------------ include/atlsec.h | 46 +++++++---- include/atlshlwapi.h | 4 +- include/atlwin.h | 134 +++++++++++++++++++++++--------- 6 files changed, 255 insertions(+), 126 deletions(-) diff --git a/include/atlcrypt.h b/include/atlcrypt.h index 674a171..de94b5d 100644 --- a/include/atlcrypt.h +++ b/include/atlcrypt.h @@ -40,7 +40,7 @@ inline DWORD CertGetNameStringA(_In_ PCCERT_CONTEXT pCertContext, _In_ DWORD dwT // Query the final string length first. DWORD dwSize = ::CertGetNameStringA(pCertContext, dwType, dwFlags, pvTypePara, NULL, 0); - // Prepare the buffer to format the string data into and read it. + // Allocate buffer on heap to format the string data into and read it. LPSTR szBuffer = sNameString.GetBuffer(dwSize); if (!szBuffer) return ERROR_OUTOFMEMORY; dwSize = ::CertGetNameStringA(pCertContext, dwType, dwFlags, pvTypePara, szBuffer, dwSize); @@ -59,7 +59,7 @@ inline DWORD CertGetNameStringW(_In_ PCCERT_CONTEXT pCertContext, _In_ DWORD dwT // Query the final string length first. DWORD dwSize = ::CertGetNameStringW(pCertContext, dwType, dwFlags, pvTypePara, NULL, 0); - // Prepare the buffer to format the string data into and read it. + // Allocate buffer on heap to format the string data into and read it. LPWSTR szBuffer = sNameString.GetBuffer(dwSize); if (!szBuffer) return ERROR_OUTOFMEMORY; dwSize = ::CertGetNameStringW(pCertContext, dwType, dwFlags, pvTypePara, szBuffer, dwSize); diff --git a/include/atlex.h b/include/atlex.h index 4d6ce61..94da856 100644 --- a/include/atlex.h +++ b/include/atlex.h @@ -23,6 +23,23 @@ #include #include +#ifndef ATL_STACK_BUFFER_BYTES +/// +/// Size of the stack buffer in bytes used for initial system function call +/// +/// Some system functions with variable length output data fail for +/// insufficient buffer sizes, and return an exact buffer length required. +/// The function helpers use a fixed size stack buffer first. If the stack +/// buffer really prooved sufficient, the helper allocates the exact length +/// output on heap and copies the data without calling the system function +/// again. Otherwise it allocates the exact length output on heap and retries. +/// +/// \note +/// Decrease this value in case of stack overflow. +/// +#define ATL_STACK_BUFFER_BYTES 1024 +#endif + namespace ATL { diff --git a/include/atlmsi.h b/include/atlmsi.h index c5cc901..2e70182 100644 --- a/include/atlmsi.h +++ b/include/atlmsi.h @@ -36,22 +36,26 @@ /// inline UINT MsiGetPropertyA(_In_ MSIHANDLE hInstall, _In_ LPCSTR szName, _Out_ ATL::CAtlStringA &sValue) { - DWORD dwSize = 0; + CHAR szStackBuffer[ATL_STACK_BUFFER_BYTES/sizeof(CHAR)]; + DWORD dwSize = _countof(szStackBuffer); UINT uiResult; - // Query the actual string length first. - uiResult = ::MsiGetPropertyA(hInstall, szName, "", &dwSize); - if (uiResult == ERROR_MORE_DATA) { - // Prepare the buffer to read the string data into and read it. + // Try with stack buffer first. + uiResult = ::MsiGetPropertyA(hInstall, szName, szStackBuffer, &dwSize); + if (uiResult == NO_ERROR) { + // Allocate buffer on heap, copy from stack buffer, and zero terminate. + LPSTR szBuffer = sValue.GetBuffer(dwSize + 1); + if (!szBuffer) return ERROR_OUTOFMEMORY; + memcpy(szBuffer, szStackBuffer, dwSize); szBuffer[dwSize] = 0; + sValue.ReleaseBuffer(dwSize + 1); + return NO_ERROR; + } else if (uiResult == ERROR_MORE_DATA) { + // Allocate buffer on heap to read the string data into and read it. LPSTR szBuffer = sValue.GetBuffer(dwSize++); if (!szBuffer) return ERROR_OUTOFMEMORY; uiResult = ::MsiGetPropertyA(hInstall, szName, szBuffer, &dwSize); sValue.ReleaseBuffer(uiResult == NO_ERROR ? dwSize : 0); return uiResult; - } else if (uiResult == NO_ERROR) { - // The string in database is empty. - sValue.Empty(); - return NO_ERROR; } else { // Return error code. return uiResult; @@ -66,22 +70,26 @@ inline UINT MsiGetPropertyA(_In_ MSIHANDLE hInstall, _In_ LPCSTR szName, _Out_ A /// inline UINT MsiGetPropertyW(_In_ MSIHANDLE hInstall, _In_ LPCWSTR szName, _Out_ ATL::CAtlStringW &sValue) { - DWORD dwSize = 0; + WCHAR szStackBuffer[ATL_STACK_BUFFER_BYTES/sizeof(WCHAR)]; + DWORD dwSize = _countof(szStackBuffer); UINT uiResult; - // Query the actual string length first. - uiResult = ::MsiGetPropertyW(hInstall, szName, L"", &dwSize); - if (uiResult == ERROR_MORE_DATA) { - // Prepare the buffer to read the string data into and read it. + // Try with stack buffer first. + uiResult = ::MsiGetPropertyW(hInstall, szName, szStackBuffer, &dwSize); + if (uiResult == NO_ERROR) { + // Allocate buffer on heap, copy from stack buffer, and zero terminate. + LPWSTR szBuffer = sValue.GetBuffer(dwSize + 1); + if (!szBuffer) return ERROR_OUTOFMEMORY; + wmemcpy(szBuffer, szStackBuffer, dwSize); szBuffer[dwSize] = 0; + sValue.ReleaseBuffer(dwSize + 1); + return NO_ERROR; + } else if (uiResult == ERROR_MORE_DATA) { + // Allocate buffer on heap to read the string data into and read it. LPWSTR szBuffer = sValue.GetBuffer(dwSize++); if (!szBuffer) return ERROR_OUTOFMEMORY; uiResult = ::MsiGetPropertyW(hInstall, szName, szBuffer, &dwSize); sValue.ReleaseBuffer(uiResult == NO_ERROR ? dwSize : 0); return uiResult; - } else if (uiResult == NO_ERROR) { - // The string in database is empty. - sValue.Empty(); - return NO_ERROR; } else { // Return error code. return uiResult; @@ -96,22 +104,26 @@ inline UINT MsiGetPropertyW(_In_ MSIHANDLE hInstall, _In_ LPCWSTR szName, _Out_ /// inline UINT MsiRecordGetStringA(_In_ MSIHANDLE hRecord, _In_ unsigned int iField, _Out_ ATL::CAtlStringA &sValue) { - DWORD dwSize = 0; + CHAR szStackBuffer[ATL_STACK_BUFFER_BYTES/sizeof(CHAR)]; + DWORD dwSize = _countof(szStackBuffer); UINT uiResult; - // Query the actual string length first. - uiResult = ::MsiRecordGetStringA(hRecord, iField, "", &dwSize); - if (uiResult == ERROR_MORE_DATA) { - // Prepare the buffer to read the string data into and read it. + // Try with stack buffer first. + uiResult = ::MsiRecordGetStringA(hRecord, iField, szStackBuffer, &dwSize); + if (uiResult == NO_ERROR) { + // Allocate buffer on heap, copy from stack buffer, and zero terminate. + LPSTR szBuffer = sValue.GetBuffer(dwSize + 1); + if (!szBuffer) return ERROR_OUTOFMEMORY; + memcpy(szBuffer, szStackBuffer, dwSize); szBuffer[dwSize] = 0; + sValue.ReleaseBuffer(dwSize + 1); + return NO_ERROR; + } else if (uiResult == ERROR_MORE_DATA) { + // Allocate buffer on heap to read the string data into and read it. LPSTR szBuffer = sValue.GetBuffer(dwSize++); if (!szBuffer) return ERROR_OUTOFMEMORY; uiResult = ::MsiRecordGetStringA(hRecord, iField, szBuffer, &dwSize); sValue.ReleaseBuffer(uiResult == NO_ERROR ? dwSize : 0); return uiResult; - } else if (uiResult == NO_ERROR) { - // The string in database is empty. - sValue.Empty(); - return NO_ERROR; } else { // Return error code. return uiResult; @@ -126,22 +138,26 @@ inline UINT MsiRecordGetStringA(_In_ MSIHANDLE hRecord, _In_ unsigned int iField /// inline UINT MsiRecordGetStringW(_In_ MSIHANDLE hRecord, _In_ unsigned int iField, _Out_ ATL::CAtlStringW &sValue) { - DWORD dwSize = 0; + WCHAR szStackBuffer[ATL_STACK_BUFFER_BYTES/sizeof(WCHAR)]; + DWORD dwSize = _countof(szStackBuffer); UINT uiResult; - // Query the actual string length first. - uiResult = ::MsiRecordGetStringW(hRecord, iField, L"", &dwSize); - if (uiResult == ERROR_MORE_DATA) { - // Prepare the buffer to read the string data into and read it. + // Try with stack buffer first. + uiResult = ::MsiRecordGetStringW(hRecord, iField, szStackBuffer, &dwSize); + if (uiResult == NO_ERROR) { + // Allocate buffer on heap, copy from stack buffer, and zero terminate. + LPWSTR szBuffer = sValue.GetBuffer(dwSize + 1); + if (!szBuffer) return ERROR_OUTOFMEMORY; + wmemcpy(szBuffer, szStackBuffer, dwSize); szBuffer[dwSize] = 0; + sValue.ReleaseBuffer(dwSize + 1); + return NO_ERROR; + } else if (uiResult == ERROR_MORE_DATA) { + // Allocate buffer on heap to read the string data into and read it. LPWSTR szBuffer = sValue.GetBuffer(dwSize++); if (!szBuffer) return ERROR_OUTOFMEMORY; uiResult = ::MsiRecordGetStringW(hRecord, iField, szBuffer, &dwSize); sValue.ReleaseBuffer(uiResult == NO_ERROR ? dwSize : 0); return uiResult; - } else if (uiResult == NO_ERROR) { - // The string in database is empty. - sValue.Empty(); - return NO_ERROR; } else { // Return error code. return uiResult; @@ -156,22 +172,26 @@ inline UINT MsiRecordGetStringW(_In_ MSIHANDLE hRecord, _In_ unsigned int iField /// inline UINT MsiFormatRecordA(MSIHANDLE hInstall, MSIHANDLE hRecord, ATL::CAtlStringA &sValue) { - DWORD dwSize = 0; + CHAR szStackBuffer[ATL_STACK_BUFFER_BYTES/sizeof(CHAR)]; + DWORD dwSize = _countof(szStackBuffer); UINT uiResult; - // Query the final string length first. - uiResult = ::MsiFormatRecordA(hInstall, hRecord, "", &dwSize); - if (uiResult == ERROR_MORE_DATA) { - // Prepare the buffer to format the string data into and read it. + // Try with stack buffer first. + uiResult = ::MsiFormatRecordA(hInstall, hRecord, szStackBuffer, &dwSize); + if (uiResult == NO_ERROR) { + // Allocate buffer on heap, copy from stack buffer, and zero terminate. + LPSTR szBuffer = sValue.GetBuffer(dwSize + 1); + if (!szBuffer) return ERROR_OUTOFMEMORY; + memcpy(szBuffer, szStackBuffer, dwSize); szBuffer[dwSize] = 0; + sValue.ReleaseBuffer(dwSize + 1); + return NO_ERROR; + } else if (uiResult == ERROR_MORE_DATA) { + // Allocate buffer on heap to format the string data into and read it. LPSTR szBuffer = sValue.GetBuffer(dwSize++); if (!szBuffer) return ERROR_OUTOFMEMORY; uiResult = ::MsiFormatRecordA(hInstall, hRecord, szBuffer, &dwSize); sValue.ReleaseBuffer(uiResult == NO_ERROR ? dwSize : 0); return uiResult; - } else if (uiResult == NO_ERROR) { - // The result is empty. - sValue.Empty(); - return NO_ERROR; } else { // Return error code. return uiResult; @@ -186,22 +206,26 @@ inline UINT MsiFormatRecordA(MSIHANDLE hInstall, MSIHANDLE hRecord, ATL::CAtlStr /// inline UINT MsiFormatRecordW(MSIHANDLE hInstall, MSIHANDLE hRecord, ATL::CAtlStringW &sValue) { - DWORD dwSize = 0; + WCHAR szStackBuffer[ATL_STACK_BUFFER_BYTES/sizeof(WCHAR)]; + DWORD dwSize = _countof(szStackBuffer); UINT uiResult; - // Query the final string length first. - uiResult = ::MsiFormatRecordW(hInstall, hRecord, L"", &dwSize); - if (uiResult == ERROR_MORE_DATA) { - // Prepare the buffer to format the string data into and read it. + // Try with stack buffer first. + uiResult = ::MsiFormatRecordW(hInstall, hRecord, szStackBuffer, &dwSize); + if (uiResult == NO_ERROR) { + // Allocate buffer on heap, copy from stack buffer, and zero terminate. + LPWSTR szBuffer = sValue.GetBuffer(dwSize + 1); + if (!szBuffer) return ERROR_OUTOFMEMORY; + wmemcpy(szBuffer, szStackBuffer, dwSize); szBuffer[dwSize] = 0; + sValue.ReleaseBuffer(dwSize + 1); + return NO_ERROR; + } else if (uiResult == ERROR_MORE_DATA) { + // Allocate buffer on heap to format the string data into and read it. LPWSTR szBuffer = sValue.GetBuffer(dwSize++); if (!szBuffer) return ERROR_OUTOFMEMORY; uiResult = ::MsiFormatRecordW(hInstall, hRecord, szBuffer, &dwSize); sValue.ReleaseBuffer(uiResult == NO_ERROR ? dwSize : 0); return uiResult; - } else if (uiResult == NO_ERROR) { - // The result is empty. - sValue.Empty(); - return NO_ERROR; } else { // Return error code. return uiResult; @@ -238,22 +262,26 @@ inline UINT MsiRecordReadStream(_In_ MSIHANDLE hRecord, _In_ unsigned int iField /// inline UINT MsiGetTargetPathA(_In_ MSIHANDLE hInstall, _In_ LPCSTR szFolder, _Out_ ATL::CAtlStringA &sValue) { - DWORD dwSize = 0; + CHAR szStackBuffer[ATL_STACK_BUFFER_BYTES/sizeof(CHAR)]; + DWORD dwSize = _countof(szStackBuffer); UINT uiResult; - // Query the final string length first. - uiResult = ::MsiGetTargetPathA(hInstall, szFolder, "", &dwSize); - if (uiResult == ERROR_MORE_DATA) { - // Prepare the buffer to format the string data into and read it. + // Try with stack buffer first. + uiResult = ::MsiGetTargetPathA(hInstall, szFolder, szStackBuffer, &dwSize); + if (uiResult == NO_ERROR) { + // Allocate buffer on heap, copy from stack buffer, and zero terminate. + LPSTR szBuffer = sValue.GetBuffer(dwSize + 1); + if (!szBuffer) return ERROR_OUTOFMEMORY; + memcpy(szBuffer, szStackBuffer, dwSize); szBuffer[dwSize] = 0; + sValue.ReleaseBuffer(dwSize + 1); + return NO_ERROR; + } else if (uiResult == ERROR_MORE_DATA) { + // Allocate buffer on heap to format the string data into and read it. LPSTR szBuffer = sValue.GetBuffer(dwSize++); if (!szBuffer) return ERROR_OUTOFMEMORY; uiResult = ::MsiGetTargetPathA(hInstall, szFolder, szBuffer, &dwSize); sValue.ReleaseBuffer(uiResult == NO_ERROR ? dwSize : 0); return uiResult; - } else if (uiResult == NO_ERROR) { - // The result is empty. - sValue.Empty(); - return NO_ERROR; } else { // Return error code. return uiResult; @@ -268,22 +296,26 @@ inline UINT MsiGetTargetPathA(_In_ MSIHANDLE hInstall, _In_ LPCSTR szFolder, _Ou /// inline UINT MsiGetTargetPathW(_In_ MSIHANDLE hInstall, _In_ LPCWSTR szFolder, _Out_ ATL::CAtlStringW &sValue) { - DWORD dwSize = 0; + WCHAR szStackBuffer[ATL_STACK_BUFFER_BYTES/sizeof(WCHAR)]; + DWORD dwSize = _countof(szStackBuffer); UINT uiResult; - // Query the final string length first. - uiResult = ::MsiGetTargetPathW(hInstall, szFolder, L"", &dwSize); - if (uiResult == ERROR_MORE_DATA) { - // Prepare the buffer to format the string data into and read it. + // Try with stack buffer first. + uiResult = ::MsiGetTargetPathW(hInstall, szFolder, szStackBuffer, &dwSize); + if (uiResult == NO_ERROR) { + // Allocate buffer on heap, copy from stack buffer, and zero terminate. + LPWSTR szBuffer = sValue.GetBuffer(dwSize + 1); + if (!szBuffer) return ERROR_OUTOFMEMORY; + wmemcpy(szBuffer, szStackBuffer, dwSize); szBuffer[dwSize] = 0; + sValue.ReleaseBuffer(dwSize + 1); + return NO_ERROR; + } else if (uiResult == ERROR_MORE_DATA) { + // Allocate buffer on heap to format the string data into and read it. LPWSTR szBuffer = sValue.GetBuffer(dwSize++); if (!szBuffer) return ERROR_OUTOFMEMORY; uiResult = ::MsiGetTargetPathW(hInstall, szFolder, szBuffer, &dwSize); sValue.ReleaseBuffer(uiResult == NO_ERROR ? dwSize : 0); return uiResult; - } else if (uiResult == NO_ERROR) { - // The result is empty. - sValue.Empty(); - return NO_ERROR; } else { // Return error code. return uiResult; diff --git a/include/atlsec.h b/include/atlsec.h index 09bf6be..5d20ff2 100644 --- a/include/atlsec.h +++ b/include/atlsec.h @@ -32,12 +32,23 @@ /// BOOLEAN GetUserNameExA(_In_ EXTENDED_NAME_FORMAT NameFormat, _Out_ ATL::CAtlStringA &sName) { - ULONG ulSize = 0; + CHAR szStackBuffer[ATL_STACK_BUFFER_BYTES/sizeof(CHAR)]; + ULONG ulSize = _countof(szStackBuffer); - // Query the final string length first. - if (!::GetUserNameExA(NameFormat, NULL, &ulSize)) { + // Try with stack buffer first. + if (::GetUserNameExA(NameFormat, szStackBuffer, &ulSize)) { + // Allocate buffer on heap, copy from stack buffer, and zero terminate. + LPSTR szBuffer = sName.GetBuffer(ulSize + 1); + if (!szBuffer) { + SetLastError(ERROR_OUTOFMEMORY); + return FALSE; + } + memcpy(szBuffer, szStackBuffer, ulSize); szBuffer[ulSize] = 0; + sName.ReleaseBuffer(ulSize + 1); + return TRUE; + } else { if (::GetLastError() == ERROR_MORE_DATA) { - // Prepare the buffer and retry. + // Allocate buffer on heap and retry. LPSTR szBuffer = sName.GetBuffer(ulSize - 1); if (!szBuffer) { SetLastError(ERROR_OUTOFMEMORY); @@ -54,10 +65,6 @@ BOOLEAN GetUserNameExA(_In_ EXTENDED_NAME_FORMAT NameFormat, _Out_ ATL::CAtlStri // Return error. return FALSE; } - } else { - // The result is empty. - sName.Empty(); - return NO_ERROR; } } @@ -69,12 +76,23 @@ BOOLEAN GetUserNameExA(_In_ EXTENDED_NAME_FORMAT NameFormat, _Out_ ATL::CAtlStri /// BOOLEAN GetUserNameExW(_In_ EXTENDED_NAME_FORMAT NameFormat, _Out_ ATL::CAtlStringW &sName) { - ULONG ulSize = 0; + WCHAR szStackBuffer[ATL_STACK_BUFFER_BYTES/sizeof(WCHAR)]; + ULONG ulSize = _countof(szStackBuffer); - // Query the final string length first. - if (!::GetUserNameExW(NameFormat, NULL, &ulSize)) { + // Try with stack buffer first. + if (::GetUserNameExW(NameFormat, szStackBuffer, &ulSize)) { + // Allocate buffer on heap, copy from stack buffer, and zero terminate. + LPWSTR szBuffer = sName.GetBuffer(ulSize + 1); + if (!szBuffer) { + SetLastError(ERROR_OUTOFMEMORY); + return FALSE; + } + wmemcpy(szBuffer, szStackBuffer, ulSize); szBuffer[ulSize] = 0; + sName.ReleaseBuffer(ulSize + 1); + return TRUE; + } else { if (::GetLastError() == ERROR_MORE_DATA) { - // Prepare the buffer and retry. + // Allocate buffer on heap and retry. LPWSTR szBuffer = sName.GetBuffer(ulSize - 1); if (!szBuffer) { SetLastError(ERROR_OUTOFMEMORY); @@ -91,10 +109,6 @@ BOOLEAN GetUserNameExW(_In_ EXTENDED_NAME_FORMAT NameFormat, _Out_ ATL::CAtlStri // Return error. return FALSE; } - } else { - // The result is empty. - sName.Empty(); - return NO_ERROR; } } diff --git a/include/atlshlwapi.h b/include/atlshlwapi.h index 3524129..b1e4105 100644 --- a/include/atlshlwapi.h +++ b/include/atlshlwapi.h @@ -35,7 +35,7 @@ /// inline BOOL PathCanonicalizeA(__out ATL::CAtlStringA &sValue, __in LPCSTR pszPath) { - // Prepare the buffer data and read into it. + // Allocate buffer on heap and read into it. LPSTR szBuffer = sValue.GetBuffer(MAX_PATH); if (!szBuffer) { ::SetLastError(ERROR_OUTOFMEMORY); @@ -55,7 +55,7 @@ inline BOOL PathCanonicalizeA(__out ATL::CAtlStringA &sValue, __in LPCSTR pszPat /// inline BOOL PathCanonicalizeW(__out ATL::CAtlStringW &sValue, __in LPCWSTR pszPath) { - // Prepare the buffer data and read into it. + // Allocate buffer on heap and read into it. LPWSTR szBuffer = sValue.GetBuffer(MAX_PATH); if (!szBuffer) { ::SetLastError(ERROR_OUTOFMEMORY); diff --git a/include/atlwin.h b/include/atlwin.h index 463bfd8..9240251 100644 --- a/include/atlwin.h +++ b/include/atlwin.h @@ -120,7 +120,7 @@ inline int GetWindowTextA(_In_ HWND hWnd, _Out_ ATL::CAtlStringA &sValue) // Query the final string length first. iResult = ::GetWindowTextLengthA(hWnd); if (iResult > 0) { - // Prepare the buffer and read the string data into it. + // Allocate buffer on heap and read the string data into it. LPSTR szBuffer = sValue.GetBuffer(iResult++); if (!szBuffer) { SetLastError(ERROR_OUTOFMEMORY); @@ -149,7 +149,7 @@ inline int GetWindowTextW(_In_ HWND hWnd, _Out_ ATL::CAtlStringW &sValue) // Query the final string length first. iResult = ::GetWindowTextLengthW(hWnd); if (iResult > 0) { - // Prepare the buffer and read the string data into it. + // Allocate buffer on heap and read the string data into it. LPWSTR szBuffer = sValue.GetBuffer(iResult++); if (!szBuffer) { SetLastError(ERROR_OUTOFMEMORY); @@ -331,15 +331,32 @@ inline VOID GuidToString(_In_ LPCGUID lpGuid, _Out_ ATL::CAtlStringW &str) inline LSTATUS RegQueryStringValue(_In_ HKEY hReg, _In_z_ LPCSTR pszName, _Out_ ATL::CAtlStringA &sValue) { LSTATUS lResult; - DWORD dwSize = 0, dwType; + BYTE aStackBuffer[ATL_STACK_BUFFER_BYTES]; + DWORD dwSize = sizeof(aStackBuffer), dwType; - // Determine the type and size first. - if ((lResult = ::RegQueryValueExA(hReg, pszName, NULL, &dwType, NULL, &dwSize)) == ERROR_SUCCESS) { + // Try with stack buffer first. + lResult = ::RegQueryValueExA(hReg, pszName, NULL, &dwType, aStackBuffer, &dwSize); + if (lResult == ERROR_SUCCESS) { + if (dwType == REG_SZ || dwType == REG_MULTI_SZ) { + // The value is REG_SZ or REG_MULTI_SZ. Allocate buffer on heap, copy from stack buffer. + LPSTR szBuffer = sValue.GetBuffer(dwSize / sizeof(CHAR)); + if (!szBuffer) return ERROR_OUTOFMEMORY; + memcpy(szBuffer, aStackBuffer, dwSize); + sValue.ReleaseBuffer(); + } else if (dwType == REG_EXPAND_SZ) { + // The value is REG_EXPAND_SZ. Expand it from stack buffer. + if (::ExpandEnvironmentStringsA((const CHAR*)aStackBuffer, sValue) == 0) + lResult = ::GetLastError(); + } else { + // The value is not a string type. + lResult = ERROR_INVALID_DATA; + } + } else if (lResult == ERROR_MORE_DATA) { if (dwType == REG_SZ || dwType == REG_MULTI_SZ) { // The value is REG_SZ or REG_MULTI_SZ. Read it now. - LPSTR szTemp = sValue.GetBuffer(dwSize / sizeof(CHAR)); - if (!szTemp) return ERROR_OUTOFMEMORY; - if ((lResult = ::RegQueryValueExA(hReg, pszName, NULL, NULL, (LPBYTE)szTemp, &dwSize)) == ERROR_SUCCESS) { + LPSTR szBuffer = sValue.GetBuffer(dwSize / sizeof(CHAR)); + if (!szBuffer) return ERROR_OUTOFMEMORY; + if ((lResult = ::RegQueryValueExA(hReg, pszName, NULL, NULL, (LPBYTE)szBuffer, &dwSize)) == ERROR_SUCCESS) { sValue.ReleaseBuffer(); } else { // Reading of the value failed. @@ -380,15 +397,32 @@ inline LSTATUS RegQueryStringValue(_In_ HKEY hReg, _In_z_ LPCSTR pszName, _Out_ inline LSTATUS RegQueryStringValue(_In_ HKEY hReg, _In_z_ LPCWSTR pszName, _Out_ ATL::CAtlStringW &sValue) { LSTATUS lResult; - DWORD dwSize = 0, dwType; + BYTE aStackBuffer[ATL_STACK_BUFFER_BYTES]; + DWORD dwSize = sizeof(aStackBuffer), dwType; - // Determine the type and size first. - if ((lResult = ::RegQueryValueExW(hReg, pszName, NULL, &dwType, NULL, &dwSize)) == ERROR_SUCCESS) { + // Try with stack buffer first. + lResult = ::RegQueryValueExW(hReg, pszName, NULL, &dwType, aStackBuffer, &dwSize); + if (lResult == ERROR_SUCCESS) { + if (dwType == REG_SZ || dwType == REG_MULTI_SZ) { + // The value is REG_SZ or REG_MULTI_SZ. Allocate buffer on heap, copy from stack buffer. + LPWSTR szBuffer = sValue.GetBuffer(dwSize / sizeof(WCHAR)); + if (!szBuffer) return ERROR_OUTOFMEMORY; + memcpy(szBuffer, aStackBuffer, dwSize); + sValue.ReleaseBuffer(); + } else if (dwType == REG_EXPAND_SZ) { + // The value is REG_EXPAND_SZ. Expand it from stack buffer. + if (::ExpandEnvironmentStringsW((const WCHAR*)aStackBuffer, sValue) == 0) + lResult = ::GetLastError(); + } else { + // The value is not a string type. + lResult = ERROR_INVALID_DATA; + } + } else if (lResult == ERROR_MORE_DATA) { if (dwType == REG_SZ || dwType == REG_MULTI_SZ) { // The value is REG_SZ or REG_MULTI_SZ. Read it now. - LPWSTR szTemp = sValue.GetBuffer(dwSize / sizeof(WCHAR)); - if (!szTemp) return ERROR_OUTOFMEMORY; - if ((lResult = ::RegQueryValueExW(hReg, pszName, NULL, NULL, (LPBYTE)szTemp, &dwSize)) == ERROR_SUCCESS) { + LPWSTR szBuffer = sValue.GetBuffer(dwSize / sizeof(WCHAR)); + if (!szBuffer) return ERROR_OUTOFMEMORY; + if ((lResult = ::RegQueryValueExW(hReg, pszName, NULL, NULL, (LPBYTE)szBuffer, &dwSize)) == ERROR_SUCCESS) { sValue.ReleaseBuffer(); } else { // Reading of the value failed. @@ -418,11 +452,19 @@ inline LSTATUS RegQueryStringValue(_In_ HKEY hReg, _In_z_ LPCWSTR pszName, _Out_ inline LSTATUS RegQueryValueExA(_In_ HKEY hKey, _In_opt_ LPCSTR lpValueName, __reserved LPDWORD lpReserved, _Out_opt_ LPDWORD lpType, _Out_ ATL::CAtlArray &aData) { LSTATUS lResult; - DWORD dwDataSize; + BYTE aStackBuffer[ATL_STACK_BUFFER_BYTES]; + DWORD dwSize = sizeof(aStackBuffer); - if ((lResult = RegQueryValueExA(hKey, lpValueName, lpReserved, NULL, NULL, &dwDataSize)) == ERROR_SUCCESS) { - if (!aData.SetCount(dwDataSize)) return ERROR_OUTOFMEMORY; - if ((lResult = RegQueryValueExA(hKey, lpValueName, lpReserved, lpType, aData.GetData(), &dwDataSize)) != ERROR_SUCCESS) + // Try with stack buffer first. + lResult = RegQueryValueExA(hKey, lpValueName, lpReserved, NULL, aStackBuffer, &dwSize); + if (lResult == ERROR_SUCCESS) { + // Allocate buffer on heap, copy from stack buffer. + if (!aData.SetCount(dwSize)) return ERROR_OUTOFMEMORY; + memcpy(aData.GetData(), aStackBuffer, dwSize); + } else if (lResult == ERROR_MORE_DATA) { + // Allocate buffer on heap and retry. + if (!aData.SetCount(dwSize)) return ERROR_OUTOFMEMORY; + if ((lResult = RegQueryValueExA(hKey, lpValueName, lpReserved, lpType, aData.GetData(), &dwSize)) != ERROR_SUCCESS) aData.SetCount(0); } @@ -438,11 +480,19 @@ inline LSTATUS RegQueryValueExA(_In_ HKEY hKey, _In_opt_ LPCSTR lpValueName, __r inline LSTATUS RegQueryValueExW(_In_ HKEY hKey, _In_opt_ LPCWSTR lpValueName, __reserved LPDWORD lpReserved, _Out_opt_ LPDWORD lpType, _Out_ ATL::CAtlArray &aData) { LSTATUS lResult; - DWORD dwDataSize; + BYTE aStackBuffer[ATL_STACK_BUFFER_BYTES]; + DWORD dwSize = sizeof(aStackBuffer); - if ((lResult = RegQueryValueExW(hKey, lpValueName, lpReserved, NULL, NULL, &dwDataSize)) == ERROR_SUCCESS) { - if (!aData.SetCount(dwDataSize)) return ERROR_OUTOFMEMORY; - if ((lResult = RegQueryValueExW(hKey, lpValueName, lpReserved, lpType, aData.GetData(), &dwDataSize)) != ERROR_SUCCESS) + // Try with stack buffer first. + lResult = RegQueryValueExW(hKey, lpValueName, lpReserved, NULL, aStackBuffer, &dwSize); + if (lResult == ERROR_SUCCESS) { + // Allocate buffer on heap, copy from stack buffer. + if (!aData.SetCount(dwSize)) return ERROR_OUTOFMEMORY; + memcpy(aData.GetData(), aStackBuffer, dwSize); + } else if (lResult == ERROR_MORE_DATA) { + // Allocate buffer on heap and retry. + if (!aData.SetCount(dwSize)) return ERROR_OUTOFMEMORY; + if ((lResult = RegQueryValueExW(hKey, lpValueName, lpReserved, lpType, aData.GetData(), &dwSize)) != ERROR_SUCCESS) aData.SetCount(0); } @@ -460,17 +510,25 @@ inline LSTATUS RegQueryValueExW(_In_ HKEY hKey, _In_opt_ LPCWSTR lpValueName, __ inline LSTATUS RegLoadMUIStringA(_In_ HKEY hKey, _In_opt_ LPCSTR pszValue, _Out_ ATL::CAtlStringA &sOut, _In_ DWORD Flags, _In_opt_ LPCSTR pszDirectory) { LSTATUS lResult; + CHAR szStackBuffer[ATL_STACK_BUFFER_BYTES/sizeof(CHAR)]; DWORD dwSize; Flags &= ~REG_MUI_STRING_TRUNCATE; - lResult = RegLoadMUIStringA(hKey, pszValue, NULL, 0, &dwSize, Flags, pszDirectory); - if (lResult == ERROR_MORE_DATA) { - LPSTR szBuffer = sOut.GetBuffer(dwSize - 1); + // Try with stack buffer first. + lResult = RegLoadMUIStringA(hKey, pszValue, szStackBuffer, _countof(szStackBuffer), &dwSize, Flags, pszDirectory); + if (lResult == ERROR_SUCCESS) { + // Allocate buffer on heap, copy from stack buffer. + LPSTR szBuffer = sOut.GetBuffer(dwSize); if (!szBuffer) return ERROR_OUTOFMEMORY; - sOut.ReleaseBuffer((lResult = RegLoadMUIStringA(hKey, pszValue, szBuffer, dwSize, &dwSize, Flags, pszDirectory)) == ERROR_SUCCESS ? dwSize - 1 : 0); - } else if (lResult == ERROR_SUCCESS) - sOut.Empty(); + memcpy(szBuffer, szStackBuffer, dwSize); + sOut.ReleaseBuffer(dwSize); + } else if (lResult == ERROR_MORE_DATA) { + // Allocate buffer on heap and retry. + LPSTR szBuffer = sOut.GetBuffer(dwSize); + if (!szBuffer) return ERROR_OUTOFMEMORY; + sOut.ReleaseBuffer((lResult = RegLoadMUIStringA(hKey, pszValue, szBuffer, dwSize, &dwSize, Flags, pszDirectory)) == ERROR_SUCCESS ? dwSize : 0); + } return lResult; } @@ -484,17 +542,25 @@ inline LSTATUS RegLoadMUIStringA(_In_ HKEY hKey, _In_opt_ LPCSTR pszValue, _Out_ inline LSTATUS RegLoadMUIStringW(_In_ HKEY hKey, _In_opt_ LPCWSTR pszValue, _Out_ ATL::CAtlStringW &sOut, _In_ DWORD Flags, _In_opt_ LPCWSTR pszDirectory) { LSTATUS lResult; + WCHAR szStackBuffer[ATL_STACK_BUFFER_BYTES/sizeof(WCHAR)]; DWORD dwSize; Flags &= ~REG_MUI_STRING_TRUNCATE; - lResult = RegLoadMUIStringW(hKey, pszValue, NULL, 0, &dwSize, Flags, pszDirectory); - if (lResult == ERROR_MORE_DATA) { - LPWSTR szBuffer = sOut.GetBuffer(dwSize - 1); + // Try with stack buffer first. + lResult = RegLoadMUIStringW(hKey, pszValue, szStackBuffer, _countof(szStackBuffer), &dwSize, Flags, pszDirectory); + if (lResult == ERROR_SUCCESS) { + // Allocate buffer on heap, copy from stack buffer. + LPWSTR szBuffer = sOut.GetBuffer(dwSize); if (!szBuffer) return ERROR_OUTOFMEMORY; - sOut.ReleaseBuffer((lResult = RegLoadMUIStringW(hKey, pszValue, szBuffer, dwSize, &dwSize, Flags, pszDirectory)) == ERROR_SUCCESS ? dwSize - 1 : 0); - } else if (lResult == ERROR_SUCCESS) - sOut.Empty(); + wmemcpy(szBuffer, szStackBuffer, dwSize); + sOut.ReleaseBuffer(dwSize); + } else if (lResult == ERROR_MORE_DATA) { + // Allocate buffer on heap and retry. + LPWSTR szBuffer = sOut.GetBuffer(dwSize); + if (!szBuffer) return ERROR_OUTOFMEMORY; + sOut.ReleaseBuffer((lResult = RegLoadMUIStringW(hKey, pszValue, szBuffer, dwSize, &dwSize, Flags, pszDirectory)) == ERROR_SUCCESS ? dwSize : 0); + } return lResult; }