@joebigelow / wix-1 / commits / 0f993110

Add MemSizeChecked.

Sean Hall committed May 26, 2022 at 17:32 UTC 0f9931107ecf9e1f6714e6fd2cabc76d2ddb1153

8 files changed +130 -109

src/libs/dutil/WixToolset.DUtil/buffutil.cpp

+5 -1

@@ -508,10 +508,14 @@ static HRESULT EnsureBufferSize(

508 {

509 HRESULT hr = S_OK;

510 SIZE_T cbTarget = ((cbSize / BUFFER_INCREMENT) + 1) * BUFFER_INCREMENT;

511 + SIZE_T cbCurrent = 0;

512

513 if (*ppbBuffer)

514 {

514 - if (MemSize(*ppbBuffer) < cbTarget)

515 + hr = MemSizeChecked(*ppbBuffer, &cbCurrent);

516 + BuffExitOnFailure(hr, "Failed to get current buffer size.");

517 +

518 + if (cbCurrent < cbTarget)

519 {

520 LPVOID pv = MemReAlloc(*ppbBuffer, cbTarget, TRUE);

521 BuffExitOnNull(pv, hr, E_OUTOFMEMORY, "Failed to reallocate buffer.");

src/libs/dutil/WixToolset.DUtil/inc/memutil.h

@@ -80,6 +80,10 @@ HRESULT DAPI MemFree(

80 SIZE_T DAPI MemSize(

81 __in LPCVOID pv

82 );

83 +HRESULT DAPI MemSizeChecked(

84 + __in LPCVOID pv,

85 + __out SIZE_T* pcb

86 + );

88 #ifdef __cplusplus

89 }

src/libs/dutil/WixToolset.DUtil/inc/strutil.h

+6 -2

@@ -139,11 +139,15 @@ HRESULT DAPI StrAllocFromError(

139

140 HRESULT DAPI StrMaxLength(

141 __in LPCVOID p,

142 - __out SIZE_T* pcbch

142 + __out SIZE_T* pcch

143 + );

144 +HRESULT DAPI StrMaxLengthAnsi(

145 + __in LPCVOID p,

146 + __out SIZE_T* pcch

147 );

148 HRESULT DAPI StrSize(

149 __in LPCVOID p,

146 - __out SIZE_T* pcbb

150 + __out SIZE_T* pcb

151 );

152

153 HRESULT DAPI StrFree(

src/libs/dutil/WixToolset.DUtil/memutil.cpp

+39 -12

@@ -9,6 +9,7 @@

9 #define MemExitWithLastError(x, s, ...) ExitWithLastErrorSource(DUTIL_SOURCE_MEMUTIL, x, s, __VA_ARGS__)

10 #define MemExitOnFailure(x, s, ...) ExitOnFailureSource(DUTIL_SOURCE_MEMUTIL, x, s, __VA_ARGS__)

11 #define MemExitOnRootFailure(x, s, ...) ExitOnRootFailureSource(DUTIL_SOURCE_MEMUTIL, x, s, __VA_ARGS__)

12 +#define MemExitWithRootFailure(x, e, s, ...) ExitWithRootFailureSource(DUTIL_SOURCE_MEMUTIL, x, e, s, __VA_ARGS__)

13 #define MemExitOnFailureDebugTrace(x, s, ...) ExitOnFailureDebugTraceSource(DUTIL_SOURCE_MEMUTIL, x, s, __VA_ARGS__)

14 #define MemExitOnNull(p, x, e, s, ...) ExitOnNullSource(DUTIL_SOURCE_MEMUTIL, p, x, e, s, __VA_ARGS__)

15 #define MemExitOnNullWithLastError(p, x, s, ...) ExitOnNullWithLastErrorSource(DUTIL_SOURCE_MEMUTIL, p, x, s, __VA_ARGS__)

@@ -74,6 +75,7 @@ extern "C" HRESULT DAPI MemReAllocSecure(

75 HRESULT hr = S_OK;

76 DWORD dwFlags = HEAP_REALLOC_IN_PLACE_ONLY;

77 LPVOID pvNew = NULL;

78 + SIZE_T cb = 0;

80 dwFlags |= fZero ? HEAP_ZERO_MEMORY : 0;

81 pvNew = ::HeapReAlloc(::GetProcessHeap(), dwFlags, pv, cbSize);

@@ -82,18 +84,16 @@ extern "C" HRESULT DAPI MemReAllocSecure(

84 pvNew = MemAlloc(cbSize, fZero);

85 if (pvNew)

86 {

85 - const SIZE_T cbCurrent = MemSize(pv);

86 - if (-1 == cbCurrent)

87 - {

88 - MemExitOnRootFailure(hr = E_INVALIDARG, "Failed to get memory size");

89 - }

87 + hr = MemSizeChecked(pv, &cb);

88 + MemExitOnFailure(hr, "Failed to get current memory size.");

89 +

90 + const SIZE_T cbCurrent = cb;

92 // HeapReAlloc may allocate more memory than requested.

92 - const SIZE_T cbNew = MemSize(pvNew);

93 - if (-1 == cbNew)

94 - {

95 - MemExitOnRootFailure(hr = E_INVALIDARG, "Failed to get memory size");

96 - }

93 + hr = MemSizeChecked(pvNew, &cb);

94 + MemExitOnFailure(hr, "Failed to get new memory size.");

95 +

96 + const SIZE_T cbNew = cb;

98 cbSize = cbNew;

99 if (cbSize > cbCurrent)

@@ -149,7 +149,10 @@ extern "C" HRESULT DAPI MemReAllocArray(

149

150 if (*ppvArray)

151 {

152 - SIZE_T cbCurrent = MemSize(*ppvArray);

152 + SIZE_T cbCurrent = 0;

153 + hr = MemSizeChecked(*ppvArray, &cbCurrent);

154 + MemExitOnFailure(hr, "Failed to get current memory size.");

155 +

156 if (cbCurrent < cbNew)

157 {

158 pvNew = MemReAlloc(*ppvArray, cbNew, TRUE);

@@ -192,7 +195,11 @@ extern "C" HRESULT DAPI MemEnsureArraySize(

195 if (*ppvArray)

196 {

197 SIZE_T cbUsed = cArray * cbArrayType;

195 - SIZE_T cbCurrent = MemSize(*ppvArray);

198 + SIZE_T cbCurrent = 0;

199 +

200 + hr = MemSizeChecked(*ppvArray, &cbCurrent);

201 + MemExitOnFailure(hr, "Failed to get current memory size.");

202 +

203 if (cbCurrent < cbUsed)

204 {

205 pvNew = MemReAlloc(*ppvArray, cbNew, TRUE);

@@ -355,3 +362,23 @@ extern "C" SIZE_T DAPI MemSize(

362 // AssertSz(vfMemInitialized, "MemInitialize() not called, this would normally crash");

363 return ::HeapSize(::GetProcessHeap(), 0, pv);

364 }

365 +

366 +

367 +extern "C" HRESULT DAPI MemSizeChecked(

368 + __in LPCVOID pv,

369 + __out SIZE_T* pcb

370 + )

371 +{

372 + HRESULT hr = S_OK;

373 +

374 +// AssertSz(vfMemInitialized, "MemInitialize() not called, this would normally crash");

375 + *pcb = MemSize(pv);

376 +

377 + if (-1 == *pcb)

378 + {

379 + MemExitWithRootFailure(hr, E_INVALIDARG, "Failed to get memory size");

380 + }

381 +

382 +LExit:

383 + return hr;

384 +}

src/libs/dutil/WixToolset.DUtil/metautil.cpp

+4 -1

@@ -300,7 +300,10 @@ extern "C" HRESULT DAPI MetaGetValue(

300 }

301 else // set the size of the data to the actual size of the memory

302 {

303 - SIZE_T cb = MemSize(pmr->pbMDData);

303 + SIZE_T cb = 0;

304 + hr = MemSizeChecked(pmr->pbMDData, &cb);

305 + MetaExitOnFailure(hr, "failed to get metabase size");

306 +

307 if (cb > DWORD_MAX)

308 {

309 MetaExitOnRootFailure(hr = E_INVALIDSTATE, "metabase data is too large: %Iu", cb);

src/libs/dutil/WixToolset.DUtil/strutil.cpp

+58 -87

@@ -390,13 +390,8 @@ static HRESULT AllocStringHelper(

390

391 if (*ppwz)

392 {

393 - cch = MemSize(*ppwz); // get the count in bytes so we can check if it failed (returns -1)

394 - if (-1 == cch)

395 - {

396 - hr = E_INVALIDARG;

397 - StrExitOnFailure(hr, "failed to get size of destination string");

398 - }

399 - cch /= sizeof(WCHAR); //convert the count in bytes to count in characters

393 + hr = StrMaxLength(*ppwz, &cch);

394 + StrExitOnFailure(hr, "failed to get size of destination string");

395 }

396

397 if (0 == cchSource && wzSource)

@@ -447,13 +442,8 @@ extern "C" HRESULT DAPI StrAnsiAllocString(

442

443 if (*ppsz)

444 {

450 - cch = MemSize(*ppsz); // get the count in bytes so we can check if it failed (returns -1)

451 - if (-1 == cch)

452 - {

453 - hr = E_INVALIDARG;

454 - StrExitOnFailure(hr, "failed to get size of destination string");

455 - }

456 - cch /= sizeof(CHAR); //convert the count in bytes to count in characters

445 + hr = StrMaxLengthAnsi(*ppsz, &cch);

446 + StrExitOnFailure(hr, "failed to get size of destination string");

447 }

448

449 if (0 == cchSource)

@@ -527,13 +517,8 @@ extern "C" HRESULT DAPI StrAllocStringAnsi(

517

518 if (*ppwz)

519 {

530 - cch = MemSize(*ppwz); // get the count in bytes so we can check if it failed (returns -1)

531 - if (-1 == cch)

532 - {

533 - hr = E_INVALIDARG;

534 - StrExitOnFailure(hr, "failed to get size of destination string");

535 - }

536 - cch /= sizeof(WCHAR); //convert the count in bytes to count in characters

520 + hr = StrMaxLength(*ppwz, &cch);

521 + StrExitOnFailure(hr, "failed to get size of destination string");

522 }

523

524 if (0 == cchSource)

@@ -605,13 +590,8 @@ HRESULT DAPI StrAnsiAllocStringAnsi(

590

591 if (*ppsz)

592 {

608 - cch = MemSize(*ppsz); // get the count in bytes so we can check if it failed (returns -1)

609 - if (-1 == cch)

610 - {

611 - hr = E_INVALIDARG;

612 - StrExitOnRootFailure(hr, "failed to get size of destination string");

613 - }

614 - cch /= sizeof(CHAR); //convert the count in bytes to count in characters

593 + hr = StrMaxLengthAnsi(*ppsz, &cch);

594 + StrExitOnRootFailure(hr, "failed to get size of destination string");

595 }

596

597 if (0 == cchSource && szSource)

@@ -664,13 +644,8 @@ extern "C" HRESULT DAPI StrAllocPrefix(

644

645 if (*ppwz)

646 {

667 - cch = MemSize(*ppwz); // get the count in bytes so we can check if it failed (returns -1)

668 - if (-1 == cch)

669 - {

670 - hr = E_INVALIDARG;

671 - StrExitOnFailure(hr, "failed to get size of destination string");

672 - }

673 - cch /= sizeof(WCHAR); //convert the count in bytes to count in characters

647 + hr = StrMaxLength(*ppwz, &cch);

648 + StrExitOnFailure(hr, "failed to get size of destination string");

649

650 hr = ::StringCchLengthW(*ppwz, STRSAFE_MAX_CCH, reinterpret_cast<UINT_PTR*>(&cchLen));

651 StrExitOnFailure(hr, "Failed to calculate length of string");

@@ -770,13 +745,8 @@ static HRESULT AllocConcatHelper(

745

746 if (*ppwz)

747 {

773 - cch = MemSize(*ppwz); // get the count in bytes so we can check if it failed (returns -1)

774 - if (-1 == cch)

775 - {

776 - hr = E_INVALIDARG;

777 - StrExitOnFailure(hr, "failed to get size of destination string");

778 - }

779 - cch /= sizeof(WCHAR); //convert the count in bytes to count in characters

748 + hr = StrMaxLength(*ppwz, &cch);

749 + StrExitOnFailure(hr, "failed to get size of destination string");

750

751 hr = ::StringCchLengthW(*ppwz, STRSAFE_MAX_CCH, reinterpret_cast<UINT_PTR*>(&cchLen));

752 StrExitOnFailure(hr, "Failed to calculate length of string");

@@ -833,13 +803,8 @@ extern "C" HRESULT DAPI StrAnsiAllocConcat(

803

804 if (*ppz)

805 {

836 - cch = MemSize(*ppz); // get the count in bytes so we can check if it failed (returns -1)

837 - if (-1 == cch)

838 - {

839 - hr = E_INVALIDARG;

840 - StrExitOnFailure(hr, "failed to get size of destination string");

841 - }

842 - cch /= sizeof(CHAR); // convert the count in bytes to count in characters

806 + hr = StrMaxLengthAnsi(*ppz, &cch);

807 + StrExitOnFailure(hr, "failed to get size of destination string");

808

809 #pragma prefast(push)

810 #pragma prefast(disable:25068)

@@ -1085,12 +1050,8 @@ static HRESULT AllocFormattedArgsHelper(

1050

1051 if (*ppwz)

1052 {

1088 - cbOriginal = MemSize(*ppwz); // get the count in bytes so we can check if it failed (returns -1)

1089 - if (-1 == cbOriginal)

1090 - {

1091 - hr = E_INVALIDARG;

1092 - StrExitOnRootFailure(hr, "failed to get size of destination string");

1093 - }

1053 + hr = StrSize(*ppwz, &cbOriginal);

1054 + StrExitOnFailure(hr, "failed to get size of destination string");

1055

1056 cch = cbOriginal / sizeof(WCHAR); //convert the count in bytes to count in characters

1057

@@ -1161,19 +1122,14 @@ extern "C" HRESULT DAPI StrAnsiAllocFormattedArgs(

1122 Assert(ppsz && szFormat && *szFormat);

1123

1124 HRESULT hr = S_OK;

1164 - SIZE_T cch = *ppsz ? MemSize(*ppsz) / sizeof(CHAR) : 0;

1125 + SIZE_T cch = 0;

1126 LPSTR pszOriginal = NULL;

1127 size_t cchOriginal = 0;

1128

1129 if (*ppsz)

1130 {

1170 - cch = MemSize(*ppsz); // get the count in bytes so we can check if it failed (returns -1)

1171 - if (-1 == cch)

1172 - {

1173 - hr = E_INVALIDARG;

1174 - StrExitOnRootFailure(hr, "failed to get size of destination string");

1175 - }

1176 - cch /= sizeof(CHAR); //convert the count in bytes to count in characters

1131 + hr = StrMaxLengthAnsi(*ppsz, &cch);

1132 + StrExitOnFailure(hr, "failed to get size of destination string");

1133

1134 hr = ::StringCchLengthA(*ppsz, STRSAFE_MAX_CCH, &cchOriginal);

1135 StrExitOnRootFailure(hr, "failed to get length of original string");

@@ -1280,11 +1236,8 @@ extern "C" HRESULT DAPI StrMaxLength(

1236

1237 if (p)

1238 {

1283 - *pcch = MemSize(p); // get size of entire buffer

1284 - if (-1 == *pcch)

1285 - {

1286 - ExitFunction1(hr = E_FAIL);

1287 - }

1239 + hr = StrSize(p, pcch);

1240 + StrExitOnFailure(hr, "Failed to get size of string buffer.");

1241

1242 *pcch /= sizeof(WCHAR); // reduce to count of characters

1243 }

@@ -1300,27 +1253,51 @@ LExit:

1253

1254

1255 /********************************************************************

1303 -StrSize - returns count of bytes in dynamic string p

1256 +StrMaxLengthAnsi - returns maximum number of characters that can be stored in dynamic string p

1257

1258 +NOTE: assumes non-Unicode string

1259 ********************************************************************/

1306 -extern "C" HRESULT DAPI StrSize(

1260 +extern "C" HRESULT DAPI StrMaxLengthAnsi(

1261 __in LPCVOID p,

1308 - __out SIZE_T* pcb

1262 + __out SIZE_T* pcch

1263 )

1264 {

1311 - Assert(p && pcb);

1265 + Assert(pcch);

1266

1267 HRESULT hr = S_OK;

1268

1315 - *pcb = MemSize(p);

1316 - if (-1 == *pcb)

1269 + if (p)

1270 + {

1271 + hr = StrSize(p, pcch);

1272 + StrExitOnFailure(hr, "Failed to get size of string buffer.");

1273 +

1274 + *pcch /= sizeof(CHAR); // reduce to count of characters

1275 + }

1276 + else

1277 {

1318 - hr = E_FAIL;

1278 + *pcch = 0;

1279 }

1280 + Assert(S_OK == hr);

1281

1282 +LExit:

1283 return hr;

1284 }

1285

1286 +

1287 +/********************************************************************

1288 +StrSize - returns count of bytes in dynamic string p

1289 +

1290 +********************************************************************/

1291 +extern "C" HRESULT DAPI StrSize(

1292 + __in LPCVOID p,

1293 + __out SIZE_T* pcb

1294 + )

1295 +{

1296 + Assert(p && pcb);

1297 +

1298 + return MemSizeChecked(p, pcb);

1299 +}

1300 +

1301 /********************************************************************

1302 StrFree - releases dynamic string memory allocated by any StrAlloc*() functions

1303

@@ -2786,22 +2763,16 @@ extern "C" DAPI_(HRESULT) StrSecureZeroString(

2763 )

2764 {

2765 HRESULT hr = S_OK;

2789 - SIZE_T cch;

2766 + SIZE_T cb = 0;

2767

2768 if (pwz)

2769 {

2793 - cch = MemSize(pwz);

2794 - if (-1 == cch)

2795 - {

2796 - hr = E_INVALIDARG;

2797 - StrExitOnFailure(hr, "Failed to get size of string");

2798 - }

2799 - else

2800 - {

2801 - SecureZeroMemory(pwz, cch);

2802 - }

2770 + hr = StrSize(pwz, &cb);

2771 + StrExitOnFailure(hr, "Failed to get size of string");

2772 +

2773 + SecureZeroMemory(pwz, cb);

2774 }

2804 -

2775 +

2776 LExit:

2777 return hr;

2778 }

src/libs/dutil/WixToolset.DUtil/thmutil.cpp

+8 -3

@@ -1189,6 +1189,7 @@ DAPI_(HRESULT) ThemeShowPageEx(

1189 BOOL fSaveEditboxes = FALSE;

1190 THEME_SAVEDVARIABLE* pSavedVariable = NULL;

1191 THEME_PAGE* pPage = ThemeGetPage(pTheme, dwPage);

1192 + SIZE_T cb = 0;

1193

1194 if (pPage)

1195 {

@@ -1219,9 +1220,9 @@ DAPI_(HRESULT) ThemeShowPageEx(

1220 if (THEME_SHOW_PAGE_REASON_REFRESH != reason)

1221 {

1222 pPage->cSavedVariables = 0;

1222 - if (pPage->rgSavedVariables)

1223 + if (pPage->rgSavedVariables && SUCCEEDED(MemSizeChecked(pPage->rgSavedVariables, &cb)))

1224 {

1224 - SecureZeroMemory(pPage->rgSavedVariables, MemSize(pPage->rgSavedVariables));

1225 + SecureZeroMemory(pPage->rgSavedVariables, cb);

1226 }

1227 }

1228

@@ -1238,7 +1239,11 @@ DAPI_(HRESULT) ThemeShowPageEx(

1239 hr = MemEnsureArraySize(reinterpret_cast<LPVOID*>(&pPage->rgSavedVariables), pPage->cControlIndices, sizeof(THEME_SAVEDVARIABLE), pPage->cControlIndices);

1240 ThmExitOnFailure(hr, "Failed to allocate memory for saved variables.");

1241

1241 - SecureZeroMemory(pPage->rgSavedVariables, MemSize(pPage->rgSavedVariables));

1242 + if (SUCCEEDED(MemSizeChecked(pPage->rgSavedVariables, &cb)))

1243 + {

1244 + SecureZeroMemory(pPage->rgSavedVariables, cb);

1245 + }

1246 +

1247 pPage->cSavedVariables = pPage->cControlIndices;

1248

1249 // Save the variables in the loop below.

src/libs/dutil/test/DUtilUnitTest/MemUtilTest.cpp

+6 -3

@@ -23,7 +23,7 @@ namespace DutilTests

23 void MemUtilAppendTest()

24 {

25 HRESULT hr = S_OK;

26 - DWORD dwSize;

26 + SIZE_T cbSize = 0;

27 ArrayValue *rgValues = NULL;

28 DWORD cValues = 0;

@@ -65,8 +65,11 @@ namespace DutilTests

65 // and make sure it doesn't grow since we already have enough space

66 hr = MemEnsureArraySize(reinterpret_cast<LPVOID*>(&rgValues), cValues, sizeof(ArrayValue), 5);

67 NativeAssert::Succeeded(hr, "Failed to ensure array size matches what it should already be");

68 - dwSize = MemSize(rgValues);

69 - if (dwSize != 6 * sizeof(ArrayValue))

68 +

69 + hr = MemSizeChecked(rgValues, &cbSize);

70 + NativeAssert::Succeeded(hr, "Failed to get current array size");

71 +

72 + if (cbSize != 6 * sizeof(ArrayValue))

73 {

74 hr = E_FAIL;

75 ExitOnFailure(hr, "MemEnsureArraySize is growing an array that is already big enough!");