8 files changed, 130 insertions, 109 deletions

diff --git a/src/libs/dutil/WixToolset.DUtil/buffutil.cpp b/src/libs/dutil/WixToolset.DUtil/buffutil.cpp
index b6d58cc0..acde4dc9 100644
--- a/src/libs/dutil/WixToolset.DUtil/buffutil.cpp
+++ b/src/libs/dutil/WixToolset.DUtil/buffutil.cpp
@@ -508,10 +508,14 @@ static HRESULT EnsureBufferSize(
 {
     HRESULT hr = S_OK;
     SIZE_T cbTarget = ((cbSize / BUFFER_INCREMENT) + 1) * BUFFER_INCREMENT;
+    SIZE_T cbCurrent = 0;
 
     if (*ppbBuffer)
     {
-        if (MemSize(*ppbBuffer) < cbTarget)
+        hr = MemSizeChecked(*ppbBuffer, &cbCurrent);
+        BuffExitOnFailure(hr, "Failed to get current buffer size.");
+
+        if (cbCurrent < cbTarget)
         {
             LPVOID pv = MemReAlloc(*ppbBuffer, cbTarget, TRUE);
             BuffExitOnNull(pv, hr, E_OUTOFMEMORY, "Failed to reallocate buffer.");
diff --git a/src/libs/dutil/WixToolset.DUtil/inc/memutil.h b/src/libs/dutil/WixToolset.DUtil/inc/memutil.h
index b8269269..c4a3b7b8 100644
--- a/src/libs/dutil/WixToolset.DUtil/inc/memutil.h
+++ b/src/libs/dutil/WixToolset.DUtil/inc/memutil.h
@@ -80,6 +80,10 @@ HRESULT DAPI MemFree(
 SIZE_T DAPI MemSize(
     __in LPCVOID pv
     );
+HRESULT DAPI MemSizeChecked(
+    __in LPCVOID pv,
+    __out SIZE_T* pcb
+    );
 
 #ifdef __cplusplus
 }
diff --git a/src/libs/dutil/WixToolset.DUtil/inc/strutil.h b/src/libs/dutil/WixToolset.DUtil/inc/strutil.h
index 1cff9ab8..f2324a80 100644
--- a/src/libs/dutil/WixToolset.DUtil/inc/strutil.h
+++ b/src/libs/dutil/WixToolset.DUtil/inc/strutil.h
@@ -139,11 +139,15 @@ HRESULT DAPI StrAllocFromError(
 
 HRESULT DAPI StrMaxLength(
     __in LPCVOID p,
-    __out SIZE_T* pcbch
+    __out SIZE_T* pcch
+    );
+HRESULT DAPI StrMaxLengthAnsi(
+    __in LPCVOID p,
+    __out SIZE_T* pcch
     );
 HRESULT DAPI StrSize(
     __in LPCVOID p,
-    __out SIZE_T* pcbb
+    __out SIZE_T* pcb
     );
 
 HRESULT DAPI StrFree(
diff --git a/src/libs/dutil/WixToolset.DUtil/memutil.cpp b/src/libs/dutil/WixToolset.DUtil/memutil.cpp
index 977c189e..2ec04e5e 100644
--- a/src/libs/dutil/WixToolset.DUtil/memutil.cpp
+++ b/src/libs/dutil/WixToolset.DUtil/memutil.cpp
@@ -9,6 +9,7 @@
 #define MemExitWithLastError(x, s, ...) ExitWithLastErrorSource(DUTIL_SOURCE_MEMUTIL, x, s, __VA_ARGS__)
 #define MemExitOnFailure(x, s, ...) ExitOnFailureSource(DUTIL_SOURCE_MEMUTIL, x, s, __VA_ARGS__)
 #define MemExitOnRootFailure(x, s, ...) ExitOnRootFailureSource(DUTIL_SOURCE_MEMUTIL, x, s, __VA_ARGS__)
+#define MemExitWithRootFailure(x, e, s, ...) ExitWithRootFailureSource(DUTIL_SOURCE_MEMUTIL, x, e, s, __VA_ARGS__)
 #define MemExitOnFailureDebugTrace(x, s, ...) ExitOnFailureDebugTraceSource(DUTIL_SOURCE_MEMUTIL, x, s, __VA_ARGS__)
 #define MemExitOnNull(p, x, e, s, ...) ExitOnNullSource(DUTIL_SOURCE_MEMUTIL, p, x, e, s, __VA_ARGS__)
 #define MemExitOnNullWithLastError(p, x, s, ...) ExitOnNullWithLastErrorSource(DUTIL_SOURCE_MEMUTIL, p, x, s, __VA_ARGS__)
@@ -74,6 +75,7 @@ extern "C" HRESULT DAPI MemReAllocSecure(
     HRESULT hr = S_OK;
     DWORD dwFlags = HEAP_REALLOC_IN_PLACE_ONLY;
     LPVOID pvNew = NULL;
+    SIZE_T cb = 0;
 
     dwFlags |= fZero ? HEAP_ZERO_MEMORY : 0;
     pvNew = ::HeapReAlloc(::GetProcessHeap(), dwFlags, pv, cbSize);
@@ -82,18 +84,16 @@ extern "C" HRESULT DAPI MemReAllocSecure(
         pvNew = MemAlloc(cbSize, fZero);
         if (pvNew)
         {
-            const SIZE_T cbCurrent = MemSize(pv);
-            if (-1 == cbCurrent)
-            {
-                MemExitOnRootFailure(hr = E_INVALIDARG, "Failed to get memory size");
-            }
+            hr = MemSizeChecked(pv, &cb);
+            MemExitOnFailure(hr, "Failed to get current memory size.");
+
+            const SIZE_T cbCurrent = cb;
 
             // HeapReAlloc may allocate more memory than requested.
-            const SIZE_T cbNew = MemSize(pvNew);
-            if (-1 == cbNew)
-            {
-                MemExitOnRootFailure(hr = E_INVALIDARG, "Failed to get memory size");
-            }
+            hr = MemSizeChecked(pvNew, &cb);
+            MemExitOnFailure(hr, "Failed to get new memory size.");
+
+            const SIZE_T cbNew = cb;
 
             cbSize = cbNew;
             if (cbSize > cbCurrent)
@@ -149,7 +149,10 @@ extern "C" HRESULT DAPI MemReAllocArray(
 
     if (*ppvArray)
     {
-        SIZE_T cbCurrent = MemSize(*ppvArray);
+        SIZE_T cbCurrent = 0;
+        hr = MemSizeChecked(*ppvArray, &cbCurrent);
+        MemExitOnFailure(hr, "Failed to get current memory size.");
+
         if (cbCurrent < cbNew)
         {
             pvNew = MemReAlloc(*ppvArray, cbNew, TRUE);
@@ -192,7 +195,11 @@ extern "C" HRESULT DAPI MemEnsureArraySize(
     if (*ppvArray)
     {
         SIZE_T cbUsed = cArray * cbArrayType;
-        SIZE_T cbCurrent = MemSize(*ppvArray);
+        SIZE_T cbCurrent = 0;
+
+        hr = MemSizeChecked(*ppvArray, &cbCurrent);
+        MemExitOnFailure(hr, "Failed to get current memory size.");
+
         if (cbCurrent < cbUsed)
         {
             pvNew = MemReAlloc(*ppvArray, cbNew, TRUE);
@@ -355,3 +362,23 @@ extern "C" SIZE_T DAPI MemSize(
 //    AssertSz(vfMemInitialized, "MemInitialize() not called, this would normally crash");
     return ::HeapSize(::GetProcessHeap(), 0, pv);
 }
+
+
+extern "C" HRESULT DAPI MemSizeChecked(
+    __in LPCVOID pv,
+    __out SIZE_T* pcb
+    )
+{
+    HRESULT hr = S_OK;
+
+//    AssertSz(vfMemInitialized, "MemInitialize() not called, this would normally crash");
+    *pcb = MemSize(pv);
+
+    if (-1 == *pcb)
+    {
+        MemExitWithRootFailure(hr, E_INVALIDARG, "Failed to get memory size");
+    }
+
+LExit:
+    return hr;
+}
diff --git a/src/libs/dutil/WixToolset.DUtil/metautil.cpp b/src/libs/dutil/WixToolset.DUtil/metautil.cpp
index f313fc1c..c36aa96c 100644
--- a/src/libs/dutil/WixToolset.DUtil/metautil.cpp
+++ b/src/libs/dutil/WixToolset.DUtil/metautil.cpp
@@ -300,7 +300,10 @@ extern "C" HRESULT DAPI MetaGetValue(
     }
     else  // set the size of the data to the actual size of the memory
     {
-        SIZE_T cb = MemSize(pmr->pbMDData);
+        SIZE_T cb = 0;
+        hr = MemSizeChecked(pmr->pbMDData, &cb);
+        MetaExitOnFailure(hr, "failed to get metabase size");
+
         if (cb > DWORD_MAX)
         {
             MetaExitOnRootFailure(hr = E_INVALIDSTATE, "metabase data is too large: %Iu", cb);
diff --git a/src/libs/dutil/WixToolset.DUtil/strutil.cpp b/src/libs/dutil/WixToolset.DUtil/strutil.cpp
index e4fcc9c8..a483cf54 100644
--- a/src/libs/dutil/WixToolset.DUtil/strutil.cpp
+++ b/src/libs/dutil/WixToolset.DUtil/strutil.cpp
@@ -390,13 +390,8 @@ static HRESULT AllocStringHelper(
 
     if (*ppwz)
     {
-        cch = MemSize(*ppwz);  // get the count in bytes so we can check if it failed (returns -1)
-        if (-1 == cch)
-        {
-            hr = E_INVALIDARG;
-            StrExitOnFailure(hr, "failed to get size of destination string");
-        }
-        cch /= sizeof(WCHAR);  //convert the count in bytes to count in characters
+        hr = StrMaxLength(*ppwz, &cch);
+        StrExitOnFailure(hr, "failed to get size of destination string");
     }
 
     if (0 == cchSource && wzSource)
@@ -447,13 +442,8 @@ extern "C" HRESULT DAPI StrAnsiAllocString(
 
     if (*ppsz)
     {
-        cch = MemSize(*ppsz);  // get the count in bytes so we can check if it failed (returns -1)
-        if (-1 == cch)
-        {
-            hr = E_INVALIDARG;
-            StrExitOnFailure(hr, "failed to get size of destination string");
-        }
-        cch /= sizeof(CHAR);  //convert the count in bytes to count in characters
+        hr = StrMaxLengthAnsi(*ppsz, &cch);
+        StrExitOnFailure(hr, "failed to get size of destination string");
     }
 
     if (0 == cchSource)
@@ -527,13 +517,8 @@ extern "C" HRESULT DAPI StrAllocStringAnsi(
 
     if (*ppwz)
     {
-        cch = MemSize(*ppwz);  // get the count in bytes so we can check if it failed (returns -1)
-        if (-1 == cch)
-        {
-            hr = E_INVALIDARG;
-            StrExitOnFailure(hr, "failed to get size of destination string");
-        }
-        cch /= sizeof(WCHAR);  //convert the count in bytes to count in characters
+        hr = StrMaxLength(*ppwz, &cch);
+        StrExitOnFailure(hr, "failed to get size of destination string");
     }
 
     if (0 == cchSource)
@@ -605,13 +590,8 @@ HRESULT DAPI StrAnsiAllocStringAnsi(
 
     if (*ppsz)
     {
-        cch = MemSize(*ppsz);  // get the count in bytes so we can check if it failed (returns -1)
-        if (-1 == cch)
-        {
-            hr = E_INVALIDARG;
-            StrExitOnRootFailure(hr, "failed to get size of destination string");
-        }
-        cch /= sizeof(CHAR);  //convert the count in bytes to count in characters
+        hr = StrMaxLengthAnsi(*ppsz, &cch);
+        StrExitOnRootFailure(hr, "failed to get size of destination string");
     }
 
     if (0 == cchSource && szSource)
@@ -664,13 +644,8 @@ extern "C" HRESULT DAPI StrAllocPrefix(
 
     if (*ppwz)
     {
-        cch = MemSize(*ppwz);  // get the count in bytes so we can check if it failed (returns -1)
-        if (-1 == cch)
-        {
-            hr = E_INVALIDARG;
-            StrExitOnFailure(hr, "failed to get size of destination string");
-        }
-        cch /= sizeof(WCHAR);  //convert the count in bytes to count in characters
+        hr = StrMaxLength(*ppwz, &cch);
+        StrExitOnFailure(hr, "failed to get size of destination string");
 
         hr = ::StringCchLengthW(*ppwz, STRSAFE_MAX_CCH, reinterpret_cast<UINT_PTR*>(&cchLen));
         StrExitOnFailure(hr, "Failed to calculate length of string");
@@ -770,13 +745,8 @@ static HRESULT AllocConcatHelper(
 
     if (*ppwz)
     {
-        cch = MemSize(*ppwz);  // get the count in bytes so we can check if it failed (returns -1)
-        if (-1 == cch)
-        {
-            hr = E_INVALIDARG;
-            StrExitOnFailure(hr, "failed to get size of destination string");
-        }
-        cch /= sizeof(WCHAR);  //convert the count in bytes to count in characters
+        hr = StrMaxLength(*ppwz, &cch);
+        StrExitOnFailure(hr, "failed to get size of destination string");
 
         hr = ::StringCchLengthW(*ppwz, STRSAFE_MAX_CCH, reinterpret_cast<UINT_PTR*>(&cchLen));
         StrExitOnFailure(hr, "Failed to calculate length of string");
@@ -833,13 +803,8 @@ extern "C" HRESULT DAPI StrAnsiAllocConcat(
 
     if (*ppz)
     {
-        cch = MemSize(*ppz);  // get the count in bytes so we can check if it failed (returns -1)
-        if (-1 == cch)
-        {
-            hr = E_INVALIDARG;
-            StrExitOnFailure(hr, "failed to get size of destination string");
-        }
-        cch /= sizeof(CHAR);  // convert the count in bytes to count in characters
+        hr = StrMaxLengthAnsi(*ppz, &cch);
+        StrExitOnFailure(hr, "failed to get size of destination string");
 
 #pragma prefast(push)
 #pragma prefast(disable:25068)
@@ -1085,12 +1050,8 @@ static HRESULT AllocFormattedArgsHelper(
 
     if (*ppwz)
     {
-        cbOriginal = MemSize(*ppwz);  // get the count in bytes so we can check if it failed (returns -1)
-        if (-1 == cbOriginal)
-        {
-            hr = E_INVALIDARG;
-            StrExitOnRootFailure(hr, "failed to get size of destination string");
-        }
+        hr = StrSize(*ppwz, &cbOriginal);
+        StrExitOnFailure(hr, "failed to get size of destination string");
 
         cch = cbOriginal / sizeof(WCHAR);  //convert the count in bytes to count in characters
 
@@ -1161,19 +1122,14 @@ extern "C" HRESULT DAPI StrAnsiAllocFormattedArgs(
     Assert(ppsz && szFormat && *szFormat);
 
     HRESULT hr = S_OK;
-    SIZE_T cch = *ppsz ? MemSize(*ppsz) / sizeof(CHAR) : 0;
+    SIZE_T cch = 0;
     LPSTR pszOriginal = NULL;
     size_t cchOriginal = 0;
 
     if (*ppsz)
     {
-        cch = MemSize(*ppsz);  // get the count in bytes so we can check if it failed (returns -1)
-        if (-1 == cch)
-        {
-            hr = E_INVALIDARG;
-            StrExitOnRootFailure(hr, "failed to get size of destination string");
-        }
-        cch /= sizeof(CHAR);  //convert the count in bytes to count in characters
+        hr = StrMaxLengthAnsi(*ppsz, &cch);
+        StrExitOnFailure(hr, "failed to get size of destination string");
 
         hr = ::StringCchLengthA(*ppsz, STRSAFE_MAX_CCH, &cchOriginal);
         StrExitOnRootFailure(hr, "failed to get length of original string");
@@ -1280,11 +1236,8 @@ extern "C" HRESULT DAPI StrMaxLength(
 
     if (p)
     {
-        *pcch = MemSize(p);   // get size of entire buffer
-        if (-1 == *pcch)
-        {
-            ExitFunction1(hr = E_FAIL);
-        }
+        hr = StrSize(p, pcch);
+        StrExitOnFailure(hr, "Failed to get size of string buffer.");
 
         *pcch /= sizeof(WCHAR);   // reduce to count of characters
     }
@@ -1300,27 +1253,51 @@ LExit:
 
 
 /********************************************************************
-StrSize - returns count of bytes in dynamic string p
+StrMaxLengthAnsi - returns maximum number of characters that can be stored in dynamic string p
 
+NOTE:  assumes non-Unicode string
 ********************************************************************/
-extern "C" HRESULT DAPI StrSize(
+extern "C" HRESULT DAPI StrMaxLengthAnsi(
     __in LPCVOID p,
-    __out SIZE_T* pcb
+    __out SIZE_T* pcch
     )
 {
-    Assert(p && pcb);
+    Assert(pcch);
 
     HRESULT hr = S_OK;
 
-    *pcb = MemSize(p);
-    if (-1 == *pcb)
+    if (p)
+    {
+        hr = StrSize(p, pcch);
+        StrExitOnFailure(hr, "Failed to get size of string buffer.");
+
+        *pcch /= sizeof(CHAR);   // reduce to count of characters
+    }
+    else
     {
-        hr = E_FAIL;
+        *pcch = 0;
     }
+    Assert(S_OK == hr);
 
+LExit:
     return hr;
 }
 
+
+/********************************************************************
+StrSize - returns count of bytes in dynamic string p
+
+********************************************************************/
+extern "C" HRESULT DAPI StrSize(
+    __in LPCVOID p,
+    __out SIZE_T* pcb
+    )
+{
+    Assert(p && pcb);
+
+    return MemSizeChecked(p, pcb);
+}
+
 /********************************************************************
 StrFree - releases dynamic string memory allocated by any StrAlloc*() functions
 
@@ -2786,22 +2763,16 @@ extern "C" DAPI_(HRESULT) StrSecureZeroString(
     )
 {
     HRESULT hr = S_OK;
-    SIZE_T cch;
+    SIZE_T cb = 0;
 
     if (pwz)
     {
-        cch = MemSize(pwz);
-        if (-1 == cch)
-        {
-            hr = E_INVALIDARG;
-            StrExitOnFailure(hr, "Failed to get size of string");
-        }
-        else
-        {
-            SecureZeroMemory(pwz, cch);
-        }
+        hr = StrSize(pwz, &cb);
+        StrExitOnFailure(hr, "Failed to get size of string");
+
+        SecureZeroMemory(pwz, cb);
     }
-    
+
 LExit:
     return hr;
 }
diff --git a/src/libs/dutil/WixToolset.DUtil/thmutil.cpp b/src/libs/dutil/WixToolset.DUtil/thmutil.cpp
index d3d32176..52a52a1e 100644
--- a/src/libs/dutil/WixToolset.DUtil/thmutil.cpp
+++ b/src/libs/dutil/WixToolset.DUtil/thmutil.cpp
@@ -1189,6 +1189,7 @@ DAPI_(HRESULT) ThemeShowPageEx(
     BOOL fSaveEditboxes = FALSE;
     THEME_SAVEDVARIABLE* pSavedVariable = NULL;
     THEME_PAGE* pPage = ThemeGetPage(pTheme, dwPage);
+    SIZE_T cb = 0;
 
     if (pPage)
     {
@@ -1219,9 +1220,9 @@ DAPI_(HRESULT) ThemeShowPageEx(
             if (THEME_SHOW_PAGE_REASON_REFRESH != reason)
             {
                 pPage->cSavedVariables = 0;
-                if (pPage->rgSavedVariables)
+                if (pPage->rgSavedVariables && SUCCEEDED(MemSizeChecked(pPage->rgSavedVariables, &cb)))
                 {
-                    SecureZeroMemory(pPage->rgSavedVariables, MemSize(pPage->rgSavedVariables));
+                    SecureZeroMemory(pPage->rgSavedVariables, cb);
                 }
             }
 
@@ -1238,7 +1239,11 @@ DAPI_(HRESULT) ThemeShowPageEx(
                 hr = MemEnsureArraySize(reinterpret_cast<LPVOID*>(&pPage->rgSavedVariables), pPage->cControlIndices, sizeof(THEME_SAVEDVARIABLE), pPage->cControlIndices);
                 ThmExitOnFailure(hr, "Failed to allocate memory for saved variables.");
 
-                SecureZeroMemory(pPage->rgSavedVariables, MemSize(pPage->rgSavedVariables));
+                if (SUCCEEDED(MemSizeChecked(pPage->rgSavedVariables, &cb)))
+                {
+                    SecureZeroMemory(pPage->rgSavedVariables, cb);
+                }
+
                 pPage->cSavedVariables = pPage->cControlIndices;
 
                 // Save the variables in the loop below.
diff --git a/src/libs/dutil/test/DUtilUnitTest/MemUtilTest.cpp b/src/libs/dutil/test/DUtilUnitTest/MemUtilTest.cpp
index 09692bfb..520ed426 100644
--- a/src/libs/dutil/test/DUtilUnitTest/MemUtilTest.cpp
+++ b/src/libs/dutil/test/DUtilUnitTest/MemUtilTest.cpp
@@ -23,7 +23,7 @@ namespace DutilTests
         void MemUtilAppendTest()
         {
             HRESULT hr = S_OK;
-            DWORD dwSize;
+            SIZE_T cbSize = 0;
             ArrayValue *rgValues = NULL;
             DWORD cValues = 0;
 
@@ -65,8 +65,11 @@ namespace DutilTests
                 // and make sure it doesn't grow since we already have enough space
                 hr = MemEnsureArraySize(reinterpret_cast<LPVOID*>(&rgValues), cValues, sizeof(ArrayValue), 5);
                 NativeAssert::Succeeded(hr, "Failed to ensure array size matches what it should already be");
-                dwSize = MemSize(rgValues);
-                if (dwSize != 6 * sizeof(ArrayValue))
+
+                hr = MemSizeChecked(rgValues, &cbSize);
+                NativeAssert::Succeeded(hr, "Failed to get current array size");
+
+                if (cbSize != 6 * sizeof(ArrayValue))
                 {
                     hr = E_FAIL;
                     ExitOnFailure(hr, "MemEnsureArraySize is growing an array that is already big enough!");