From 0f9931107ecf9e1f6714e6fd2cabc76d2ddb1153 Mon Sep 17 00:00:00 2001 From: Sean Hall Date: Thu, 26 May 2022 17:32:45 -0500 Subject: Add MemSizeChecked. --- src/libs/dutil/WixToolset.DUtil/buffutil.cpp | 6 +- src/libs/dutil/WixToolset.DUtil/inc/memutil.h | 4 + src/libs/dutil/WixToolset.DUtil/inc/strutil.h | 8 +- src/libs/dutil/WixToolset.DUtil/memutil.cpp | 51 ++++++-- src/libs/dutil/WixToolset.DUtil/metautil.cpp | 5 +- src/libs/dutil/WixToolset.DUtil/strutil.cpp | 145 +++++++++------------- src/libs/dutil/WixToolset.DUtil/thmutil.cpp | 11 +- src/libs/dutil/test/DUtilUnitTest/MemUtilTest.cpp | 9 +- 8 files changed, 130 insertions(+), 109 deletions(-) (limited to 'src') 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(&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(&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(&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(&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!"); -- cgit v1.2.3-55-g6feb