Amebis/atlex
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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.

Browse Source
This commit is contained in:
Simon Rozman
2015-11-12 19:35:26 +01:00
parent e20ab835ed
commit 7fef07d0fd
6 changed files with 255 additions and 126 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
include/atlcrypt.h
View File

@@ -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);

17
include/atlex.h
View File

@@ -23,6 +23,23 @@
#include <atldef.h>
#include <atlstr.h>
#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
{

176
include/atlmsi.h
View File

@@ -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;

46
include/atlsec.h
View File

@@ -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;
}
}

4
include/atlshlwapi.h
View File

@@ -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);

134
include/atlwin.h
View File

@@ -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<BYTE> &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<BYTE> &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;
}