[v2 PATCH] wincodecs/icoformat: Improve input data validation on decoder initialization.

[Nikolay Sivov]

Signed-off-by: Nikolay Sivov <nsivov(a)codeweavers.com>
---
 dlls/windowscodecs/icoformat.c       |  34 ++++-
 dlls/windowscodecs/tests/icoformat.c | 208 ++++++++++++++++++---------
 2 files changed, 174 insertions(+), 68 deletions(-)

diff --git a/dlls/windowscodecs/icoformat.c b/dlls/windowscodecs/icoformat.c
index 5e38ee0d0f..dc08fd5576 100644
--- a/dlls/windowscodecs/icoformat.c
+++ b/dlls/windowscodecs/icoformat.c
@@ -511,6 +511,9 @@ static HRESULT WINAPI IcoDecoder_Initialize(IWICBitmapDecoder *iface, IStream *p
     LARGE_INTEGER seek;
     HRESULT hr;
     ULONG bytesread;
+    STATSTG statstg;
+    unsigned int i;
+
     TRACE("(%p,%p,%x)\n", iface, pIStream, cacheOptions);
 
     EnterCriticalSection(&This->lock);
@@ -527,14 +530,41 @@ static HRESULT WINAPI IcoDecoder_Initialize(IWICBitmapDecoder *iface, IStream *p
 
     hr = IStream_Read(pIStream, &This->header, sizeof(ICONHEADER), &bytesread);
     if (FAILED(hr)) goto end;
-    if (bytesread != sizeof(ICONHEADER) ||
-        This->header.idReserved != 0 ||
+
+    if (bytesread != sizeof(ICONHEADER))
+    {
+        hr = WINCODEC_ERR_STREAMREAD;
+        goto end;
+    }
+
+    if (This->header.idReserved != 0 ||
         This->header.idType != 1)
     {
         hr = E_FAIL;
         goto end;
     }
 
+    hr = IStream_Stat(pIStream, &statstg, STATFLAG_NONAME);
+    if (FAILED(hr))
+    {
+        WARN("Stat() failed, hr %#x.\n", hr);
+        goto end;
+    }
+
+    for (i = 0; i < This->header.idCount; i++)
+    {
+        ICONDIRENTRY direntry;
+
+        hr = IStream_Read(pIStream, &direntry, sizeof(direntry), &bytesread);
+        if (FAILED(hr)) goto end;
+
+        if (bytesread != sizeof(direntry) || (direntry.dwDIBSize + direntry.dwDIBOffset > statstg.cbSize.QuadPart))
+        {
+            hr = WINCODEC_ERR_BADIMAGE;
+            goto end;
+        }
+    }
+
     This->initialized = TRUE;
     This->stream = pIStream;
     IStream_AddRef(pIStream);
diff --git a/dlls/windowscodecs/tests/icoformat.c b/dlls/windowscodecs/tests/icoformat.c
index c53739dbd6..3d1de70419 100644
--- a/dlls/windowscodecs/tests/icoformat.c
+++ b/dlls/windowscodecs/tests/icoformat.c
@@ -26,72 +26,115 @@
 #include "wincodec.h"
 #include "wine/test.h"
 
-static unsigned char testico_bad_icondirentry_size[] = {
-    /* ICONDIR */
-    0, 0, /* reserved */
-    1, 0, /* type */
-    1, 0, /* count */
+#include "pshpack1.h"
+
+struct ICONHEADER
+{
+    WORD idReserved;
+    WORD idType;
+    WORD idCount;
+};
+
+struct ICONDIRENTRY
+{
+    BYTE bWidth;
+    BYTE bHeight;
+    BYTE bColorCount;
+    BYTE bReserved;
+    WORD wPlanes;
+    WORD wBitCount;
+    DWORD dwDIBSize;
+    DWORD dwDIBOffset;
+};
+
+struct test_ico
+{
+    struct ICONHEADER header;
+    struct ICONDIRENTRY direntry;
+    BITMAPINFOHEADER bmi;
+    unsigned char data[512];
+};
+
+static const struct test_ico ico_1 =
+{
+    /* ICONHEADER */
+    {
+      0, /* reserved */
+      1, /* type */
+      1, /* count */
+    },
     /* ICONDIRENTRY */
-    2, /* width */
-    2, /* height */
-    2, /* colorCount */
-    0, /* reserved */
-    1,0, /* planes */
-    8,0, /* bitCount */
-    (40+2*4+16*16+16*4) & 0xFF,((40+2*4+16*16+16*4) >> 8) & 0xFF,0,0, /* bytesInRes */
-    22,0,0,0, /* imageOffset */
+    {
+      16, /* width */
+      16, /* height */
+      2, /* color count */
+      0, /* reserved */
+      1, /* planes */
+      8, /* bitcount*/
+      40 + 2*4 + 16 * 16 + 16 * 4, /* data size */
+      22 /* data offset */
+    },
     /* BITMAPINFOHEADER */
-    40,0,0,0, /* header size */
-    16,0,0,0, /* width */
-    2*16,0,0,0, /* height (XOR+AND rows) */
-    1,0, /* planes */
-    8,0, /* bit count */
-    0,0,0,0, /* compression */
-    0,0,0,0, /* sizeImage */
-    0,0,0,0, /* x pels per meter */
-    0,0,0,0, /* y pels per meter */
-    2,0,0,0, /* clrUsed */
-    0,0,0,0, /* clrImportant */
-    /* palette */
-    0,0,0,0,
-    0xFF,0xFF,0xFF,0,
-    /* XOR mask */
-    0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
-    0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
-    0,1,0,0,0,0,0,0,0,0,0,0,0,1,0,0,
-    0,1,0,0,0,0,0,0,0,0,0,0,0,1,0,0,
-    0,1,0,0,0,0,0,0,0,0,0,0,0,1,0,0,
-    0,0,1,0,0,0,0,0,0,0,0,0,1,0,0,0,
-    0,0,1,0,0,0,0,0,0,0,0,0,1,0,0,0,
-    0,0,1,0,0,0,0,0,0,0,0,0,1,0,0,0,
-    0,0,0,1,0,0,0,0,0,0,0,1,0,0,0,0,
-    0,0,0,1,0,0,0,0,0,0,0,1,0,0,0,0,
-    0,0,0,1,0,0,0,1,0,0,0,1,0,0,0,0,
-    0,0,0,0,1,0,1,0,1,0,1,0,0,0,0,0,
-    0,0,0,0,0,1,0,0,0,1,0,0,0,0,0,0,
-    0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
-    0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
-    0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
-    /* AND mask */
-    0,0,0,0,
-    0,0,0,0,
-    0,0,0,0,
-    0,0,0,0,
-    0,0,0,0,
-    0,0,0,0,
-    0,0,0,0,
-    0,0,0,0,
-    0,0,0,0,
-    0,0,0,0,
-    0,0,0,0,
-    0,0,0,0,
-    0,0,0,0,
-    0,0,0,0,
-    0,0,0,0,
-    0,0,0,0
+    {
+      sizeof(BITMAPINFOHEADER), /* header size */
+      16, /* width */
+      2*16, /* height (XOR+AND rows) */
+      1, /* planes */
+      8, /* bit count */
+      0, /* compression */
+      0, /* sizeImage */
+      0, /* x pels per meter */
+      0, /* y pels per meter */
+      2, /* clrUsed */
+      0, /* clrImportant */
+    },
+    {
+      /* palette */
+      0,0,0,0,
+      0xFF,0xFF,0xFF,0,
+      /* XOR mask */
+      0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
+      0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
+      0,1,0,0,0,0,0,0,0,0,0,0,0,1,0,0,
+      0,1,0,0,0,0,0,0,0,0,0,0,0,1,0,0,
+      0,1,0,0,0,0,0,0,0,0,0,0,0,1,0,0,
+      0,0,1,0,0,0,0,0,0,0,0,0,1,0,0,0,
+      0,0,1,0,0,0,0,0,0,0,0,0,1,0,0,0,
+      0,0,1,0,0,0,0,0,0,0,0,0,1,0,0,0,
+      0,0,0,1,0,0,0,0,0,0,0,1,0,0,0,0,
+      0,0,0,1,0,0,0,0,0,0,0,1,0,0,0,0,
+      0,0,0,1,0,0,0,1,0,0,0,1,0,0,0,0,
+      0,0,0,0,1,0,1,0,1,0,1,0,0,0,0,0,
+      0,0,0,0,0,1,0,0,0,1,0,0,0,0,0,0,
+      0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
+      0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
+      0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
+      /* AND mask */
+      0,0,0,0,
+      0,0,0,0,
+      0,0,0,0,
+      0,0,0,0,
+      0,0,0,0,
+      0,0,0,0,
+      0,0,0,0,
+      0,0,0,0,
+      0,0,0,0,
+      0,0,0,0,
+      0,0,0,0,
+      0,0,0,0,
+      0,0,0,0,
+      0,0,0,0,
+      0,0,0,0,
+      0,0,0,0
+    }
 };
 
-static void test_bad_icondirentry_size(void)
+#include "poppack.h"
+
+#define test_ico_data(a, b, c, d, e) test_ico_data_(a, b, c, d, e, 0, __LINE__)
+#define test_ico_data_todo(a, b, c, d, e) test_ico_data_(a, b, c, d, e, 1, __LINE__)
+static void test_ico_data_(void *data, DWORD data_size, UINT expected_width, UINT expected_height,
+    HRESULT init_hr, int todo, unsigned int line)
 {
     IWICBitmapDecoder *decoder;
     IWICImagingFactory *factory;
@@ -108,8 +151,7 @@ static void test_bad_icondirentry_size(void)
     ok(hr == S_OK, "CreateStream failed, hr=%x\n", hr);
     if (SUCCEEDED(hr))
     {
-        hr = IWICStream_InitializeFromMemory(icostream, testico_bad_icondirentry_size,
-            sizeof(testico_bad_icondirentry_size));
+        hr = IWICStream_InitializeFromMemory(icostream, data, data_size);
         ok(hr == S_OK, "InitializeFromMemory failed, hr=%x\n", hr);
 
         if (SUCCEEDED(hr))
@@ -123,7 +165,8 @@ static void test_bad_icondirentry_size(void)
         {
             hr = IWICBitmapDecoder_Initialize(decoder, (IStream*)icostream,
                 WICDecodeMetadataCacheOnDemand);
-            ok(hr == S_OK, "Initialize failed, hr=%x\n", hr);
+        todo_wine_if(todo)
+            ok_(__FILE__, line)(hr == init_hr, "Initialize failed, hr=%x\n", hr);
 
             if (SUCCEEDED(hr))
             {
@@ -163,11 +206,44 @@ static void test_bad_icondirentry_size(void)
     IWICImagingFactory_Release(factory);
 }
 
+static void test_decoder(void)
+{
+    struct test_ico ico;
+
+    /* Icon size specified in ICONDIRENTRY does not match bitmap header. */
+    ico = ico_1;
+    ico.direntry.bWidth = 2;
+    ico.direntry.bHeight = 2;
+    test_ico_data(&ico, sizeof(ico), 16, 16, S_OK);
+
+    /* Invalid DIRENTRY data size/offset. */
+    ico = ico_1;
+    ico.direntry.dwDIBOffset = sizeof(ico);
+    test_ico_data(&ico, sizeof(ico), 0, 0, WINCODEC_ERR_BADIMAGE);
+
+    ico = ico_1;
+    ico.direntry.dwDIBSize = sizeof(ico);
+    test_ico_data(&ico, sizeof(ico), 0, 0, WINCODEC_ERR_BADIMAGE);
+
+    /* Header fields validation. */
+    ico = ico_1;
+    ico.header.idReserved = 1;
+    test_ico_data_todo(&ico, sizeof(ico), 16, 16, S_OK);
+    ico.header.idReserved = 0;
+    ico.header.idType = 100;
+    test_ico_data_todo(&ico, sizeof(ico), 16, 16, S_OK);
+
+    /* Premature end of data. */
+    ico = ico_1;
+    test_ico_data(&ico, sizeof(ico.header) - 1, 0, 0, WINCODEC_ERR_STREAMREAD);
+    test_ico_data(&ico, sizeof(ico.header) + sizeof(ico.direntry) - 1, 0, 0, WINCODEC_ERR_BADIMAGE);
+}
+
 START_TEST(icoformat)
 {
     CoInitializeEx(NULL, COINIT_APARTMENTTHREADED);
 
-    test_bad_icondirentry_size();
+    test_decoder();
 
     CoUninitialize();
 }
-- 
2.19.1