diff --git a/winml/test/image/imageTestHelper.cpp b/winml/test/image/imageTestHelper.cpp index 6d0c536fbe..d6578541db 100644 --- a/winml/test/image/imageTestHelper.cpp +++ b/winml/test/image/imageTestHelper.cpp @@ -41,7 +41,6 @@ namespace ImageTestHelper { wf::IMemoryBufferReference reference = spBitmapBuffer.CreateReference(); auto spByteAccess = reference.as<::Windows::Foundation::IMemoryBufferByteAccess>(); spByteAccess->GetBuffer(&pData, &size); - uint32_t height = softwareBitmap.PixelHeight(); uint32_t width = softwareBitmap.PixelWidth(); @@ -53,14 +52,15 @@ namespace ImageTestHelper { com_ptr itn = tf.as(); itn->GetBuffer(reinterpret_cast(&pCPUTensor), &uCapacity); if (BitmapPixelFormat::Bgra8 == GetPixelFormat(modelPixelFormat)) { - for (UINT32 i = 0; i < size; i += 4) { + // loop condition is i < size - 2 to avoid potential for extending past the memory buffer + for (UINT32 i = 0; i < size - 2; i += 4) { UINT32 pixelInd = i / 4; pCPUTensor[pixelInd] = (float)pData[i]; pCPUTensor[(height * width) + pixelInd] = (float)pData[i + 1]; pCPUTensor[(height * width * 2) + pixelInd] = (float)pData[i + 2]; } } else if (BitmapPixelFormat::Rgba8 == GetPixelFormat(modelPixelFormat)) { - for (UINT32 i = 0; i < size; i += 4) { + for (UINT32 i = 0; i < size - 2; i += 4) { UINT32 pixelInd = i / 4; pCPUTensor[pixelInd] = (float)pData[i + 2]; pCPUTensor[(height * width) + pixelInd] = (float)pData[i + 1]; @@ -99,14 +99,15 @@ namespace ImageTestHelper { uint32_t height = softwareBitmap.PixelHeight(); uint32_t width = softwareBitmap.PixelWidth(); if (BitmapPixelFormat::Bgra8 == GetPixelFormat(modelPixelFormat)) { - for (UINT32 i = 0; i < size; i += 4) { + // loop condition is i < size - 2 to avoid potential for extending past the memory buffer + for (UINT32 i = 0; i < size - 2; i += 4) { UINT32 pixelInd = i / 4; pCPUTensor[pixelInd] = (float)pData[i]; pCPUTensor[(height * width) + pixelInd] = (float)pData[i + 1]; pCPUTensor[(height * width * 2) + pixelInd] = (float)pData[i + 2]; } } else if (BitmapPixelFormat::Rgba8 == GetPixelFormat(modelPixelFormat)) { - for (UINT32 i = 0; i < size; i += 4) { + for (UINT32 i = 0; i < size - 2; i += 4) { UINT32 pixelInd = i / 4; pCPUTensor[pixelInd] = (float)pData[i + 2]; pCPUTensor[(height * width) + pixelInd] = (float)pData[i + 1]; diff --git a/winml/test/scenario/cppwinrt/scenariotestscppwinrt.cpp b/winml/test/scenario/cppwinrt/scenariotestscppwinrt.cpp index c55d411363..63d77a4ea2 100644 --- a/winml/test/scenario/cppwinrt/scenariotestscppwinrt.cpp +++ b/winml/test/scenario/cppwinrt/scenariotestscppwinrt.cpp @@ -816,7 +816,8 @@ static void Scenario22ImageBindingAsCPUTensor() { uint32_t height = softwareBitmap.PixelHeight(); uint32_t width = softwareBitmap.PixelWidth(); - for (UINT32 i = 0; i < size; i += 4) { + for (UINT32 i = 0; i < size - 2; i += 4) { + // loop condition is i < size - 2 to avoid potential for extending past the memory buffer UINT32 pixelInd = i / 4; pCPUTensor[pixelInd] = (float)pData[i]; pCPUTensor[(height * width) + pixelInd] = (float)pData[i + 1]; @@ -887,7 +888,8 @@ static void Scenario22ImageBindingAsGPUTensor() { uint32_t height = softwareBitmap.PixelHeight(); uint32_t width = softwareBitmap.PixelWidth(); - for (UINT32 i = 0; i < size; i += 4) { + for (UINT32 i = 0; i < size - 2; i += 4) { + // loop condition is i < size - 2 to avoid potential for extending past the memory buffer UINT32 pixelInd = i / 4; pCPUTensor[pixelInd] = (FLOAT)pData[i]; pCPUTensor[(height * width) + pixelInd] = (FLOAT)pData[i + 1]; @@ -1494,7 +1496,8 @@ static void BindMultipleCPUBuffersAsInputs(LearningModelDeviceKind kind) { green_byteaccess->GetBuffer(reinterpret_cast(&green_data), &frame_size); blue_byteaccess->GetBuffer(reinterpret_cast(&blue_data), &frame_size); - for (UINT32 i = 0; i < size; i += 4) { + for (UINT32 i = 0; i < size - 2 && i / 4 < frame_size; i += 4) { + // loop condition is i < size - 2 to avoid potential for extending past the memory buffer UINT32 pixelInd = i / 4; red_data[pixelInd] = (float)data[i]; green_data[pixelInd] = (float)data[i + 1];