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

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

include/atlex.h

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

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.
+    // Try with stack buffer first.
-    uiResult = ::MsiGetPropertyA(hInstall, szName, "", &dwSize);
+    uiResult = ::MsiGetPropertyA(hInstall, szName, szStackBuffer, &dwSize);
-    if (uiResult == ERROR_MORE_DATA) {
+    if (uiResult == NO_ERROR) {
-        // Prepare the buffer to read the string data into and read it.
+        // 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.
+    // Try with stack buffer first.
-    uiResult = ::MsiGetPropertyW(hInstall, szName, L"", &dwSize);
+    uiResult = ::MsiGetPropertyW(hInstall, szName, szStackBuffer, &dwSize);
-    if (uiResult == ERROR_MORE_DATA) {
+    if (uiResult == NO_ERROR) {
-        // Prepare the buffer to read the string data into and read it.
+        // 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.
+    // Try with stack buffer first.
-    uiResult = ::MsiRecordGetStringA(hRecord, iField, "", &dwSize);
+    uiResult = ::MsiRecordGetStringA(hRecord, iField, szStackBuffer, &dwSize);
-    if (uiResult == ERROR_MORE_DATA) {
+    if (uiResult == NO_ERROR) {
-        // Prepare the buffer to read the string data into and read it.
+        // 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.
+    // Try with stack buffer first.
-    uiResult = ::MsiRecordGetStringW(hRecord, iField, L"", &dwSize);
+    uiResult = ::MsiRecordGetStringW(hRecord, iField, szStackBuffer, &dwSize);
-    if (uiResult == ERROR_MORE_DATA) {
+    if (uiResult == NO_ERROR) {
-        // Prepare the buffer to read the string data into and read it.
+        // 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.
+    // Try with stack buffer first.
-    uiResult = ::MsiFormatRecordA(hInstall, hRecord, "", &dwSize);
+    uiResult = ::MsiFormatRecordA(hInstall, hRecord, szStackBuffer, &dwSize);
-    if (uiResult == ERROR_MORE_DATA) {
+    if (uiResult == NO_ERROR) {
-        // Prepare the buffer to format the string data into and read it.
+        // 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.
+    // Try with stack buffer first.
-    uiResult = ::MsiFormatRecordW(hInstall, hRecord, L"", &dwSize);
+    uiResult = ::MsiFormatRecordW(hInstall, hRecord, szStackBuffer, &dwSize);
-    if (uiResult == ERROR_MORE_DATA) {
+    if (uiResult == NO_ERROR) {
-        // Prepare the buffer to format the string data into and read it.
+        // 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.
+    // Try with stack buffer first.
-    uiResult = ::MsiGetTargetPathA(hInstall, szFolder, "", &dwSize);
+    uiResult = ::MsiGetTargetPathA(hInstall, szFolder, szStackBuffer, &dwSize);
-    if (uiResult == ERROR_MORE_DATA) {
+    if (uiResult == NO_ERROR) {
-        // Prepare the buffer to format the string data into and read it.
+        // 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.
+    // Try with stack buffer first.
-    uiResult = ::MsiGetTargetPathW(hInstall, szFolder, L"", &dwSize);
+    uiResult = ::MsiGetTargetPathW(hInstall, szFolder, szStackBuffer, &dwSize);
-    if (uiResult == ERROR_MORE_DATA) {
+    if (uiResult == NO_ERROR) {
-        // Prepare the buffer to format the string data into and read it.
+        // 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;

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.
+    // Try with stack buffer first.
-    if (!::GetUserNameExA(NameFormat, NULL, &ulSize)) {
+    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.
+    // Try with stack buffer first.
-    if (!::GetUserNameExW(NameFormat, NULL, &ulSize)) {
+    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;
     }
 }
 

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

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.
+    // Try with stack buffer first.
-    if ((lResult = ::RegQueryValueExA(hReg, pszName, NULL, &dwType, NULL, &dwSize)) == ERROR_SUCCESS) {
+    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));
+            LPSTR szBuffer = sValue.GetBuffer(dwSize / sizeof(CHAR));
-            if (!szTemp) return ERROR_OUTOFMEMORY;
+            if (!szBuffer) return ERROR_OUTOFMEMORY;
-            if ((lResult = ::RegQueryValueExA(hReg, pszName, NULL, NULL, (LPBYTE)szTemp, &dwSize)) == ERROR_SUCCESS) {
+            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.
+    // Try with stack buffer first.
-    if ((lResult = ::RegQueryValueExW(hReg, pszName, NULL, &dwType, NULL, &dwSize)) == ERROR_SUCCESS) {
+    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));
+            LPWSTR szBuffer = sValue.GetBuffer(dwSize / sizeof(WCHAR));
-            if (!szTemp) return ERROR_OUTOFMEMORY;
+            if (!szBuffer) return ERROR_OUTOFMEMORY;
-            if ((lResult = ::RegQueryValueExW(hReg, pszName, NULL, NULL, (LPBYTE)szTemp, &dwSize)) == ERROR_SUCCESS) {
+            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) {
+    // Try with stack buffer first.
-        if (!aData.SetCount(dwDataSize)) return ERROR_OUTOFMEMORY;
+    lResult = RegQueryValueExA(hKey, lpValueName, lpReserved, NULL, aStackBuffer, &dwSize);
-        if ((lResult = RegQueryValueExA(hKey, lpValueName, lpReserved, lpType, aData.GetData(), &dwDataSize)) != ERROR_SUCCESS)
+    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) {
+    // Try with stack buffer first.
-        if (!aData.SetCount(dwDataSize)) return ERROR_OUTOFMEMORY;
+    lResult = RegQueryValueExW(hKey, lpValueName, lpReserved, NULL, aStackBuffer, &dwSize);
-        if ((lResult = RegQueryValueExW(hKey, lpValueName, lpReserved, lpType, aData.GetData(), &dwDataSize)) != ERROR_SUCCESS)
+    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);
+    // Try with stack buffer first.
-    if (lResult == ERROR_MORE_DATA) {
+    lResult = RegLoadMUIStringA(hKey, pszValue, szStackBuffer, _countof(szStackBuffer), &dwSize, Flags, pszDirectory);
-        LPSTR szBuffer = sOut.GetBuffer(dwSize - 1);
+    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);
+        memcpy(szBuffer, szStackBuffer, dwSize);
-    } else if (lResult == ERROR_SUCCESS)
+        sOut.ReleaseBuffer(dwSize);
-        sOut.Empty();
+    } 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);
+    // Try with stack buffer first.
-    if (lResult == ERROR_MORE_DATA) {
+    lResult = RegLoadMUIStringW(hKey, pszValue, szStackBuffer, _countof(szStackBuffer), &dwSize, Flags, pszDirectory);
-        LPWSTR szBuffer = sOut.GetBuffer(dwSize - 1);
+    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);
+        wmemcpy(szBuffer, szStackBuffer, dwSize);
-    } else if (lResult == ERROR_SUCCESS)
+        sOut.ReleaseBuffer(dwSize);
-        sOut.Empty();
+    } 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;
 }