[PATCH 2/2] shell32: add support for balloon icon copying

[Sergei Chernyadyev]

From: Sergei Chernyadyev

---
 dlls/shell32/systray.c      | 99 +++++++++++++++++++++++--------------
 programs/explorer/systray.c | 23 +++++++--
 2 files changed, 83 insertions(+), 39 deletions(-)

diff --git a/dlls/shell32/systray.c b/dlls/shell32/systray.c
index 04728e4bd29..c74fde4ed2f 100644
--- a/dlls/shell32/systray.c
+++ b/dlls/shell32/systray.c
@@ -134,6 +134,12 @@ BOOL WINAPI Shell_NotifyIconW(DWORD dwMessage, PNOTIFYICONDATAW nid)
     COPYDATASTRUCT cds;
     struct notify_data data_buffer;
     struct notify_data *data = &data_buffer;
+    ICONINFO iconinfo[2] = { { 0 }, { 0 } };
+    BITMAP bmMask[2];
+    BITMAP bmColour[2];
+    LONG cbMaskBits[2] = { 0, 0};
+    LONG cbColourBits[2] = { 0, 0 };
+    LONG iconCount = 0;
     BOOL ret;

TRACE("dwMessage = %ld, nid->cbSize=%ld\n", dwMessage, nid->cbSize);
@@ -168,60 +174,81 @@ BOOL WINAPI Shell_NotifyIconW(DWORD dwMessage, PNOTIFYICONDATAW nid)
      * icon handles */
     if (nid->uFlags & NIF_ICON)
     {
-        ICONINFO iconinfo;
-        BITMAP bmMask;
-        BITMAP bmColour;
-        LONG cbMaskBits;
-        LONG cbColourBits = 0;
-        char *buffer;
-
-        if (!GetIconInfo(nid->hIcon, &iconinfo))
+        if (!GetIconInfo(nid->hIcon, &iconinfo[iconCount]))
             goto noicon;

-        if (!GetObjectW(iconinfo.hbmMask, sizeof(bmMask), &bmMask) ||
-            (iconinfo.hbmColor && !GetObjectW(iconinfo.hbmColor, sizeof(bmColour), &bmColour)))
+        if (!GetObjectW(iconinfo[iconCount].hbmMask, sizeof(bmMask), &bmMask[iconCount]) ||
+            (iconinfo[iconCount].hbmColor && !GetObjectW(iconinfo[iconCount].hbmColor, sizeof(bmColour), &bmColour[iconCount])))
         {
-            DeleteObject(iconinfo.hbmMask);
-            if (iconinfo.hbmColor) DeleteObject(iconinfo.hbmColor);
+            DeleteObject(iconinfo[iconCount].hbmMask);
+            if (iconinfo[iconCount].hbmColor) DeleteObject(iconinfo[iconCount].hbmColor);
             goto noicon;
         }
+        cbMaskBits[iconCount] = (bmMask[iconCount].bmPlanes * bmMask[iconCount].bmWidth * bmMask[iconCount].bmHeight * bmMask[iconCount].bmBitsPixel + 15) / 16 * 2;
+        if (iconinfo[iconCount].hbmColor)
+            cbColourBits[iconCount] = (bmColour[iconCount].bmPlanes * bmColour[iconCount].bmWidth * bmColour[iconCount].bmHeight * bmColour[iconCount].bmBitsPixel + 15) / 16 * 2;
+        iconCount++;
+    }
+noicon:
+    if (nid->uFlags & NIF_INFO)
+    {
+        if (!GetIconInfo(nid->hBalloonIcon, &iconinfo[iconCount]))
+            goto noballoon;

-        cbMaskBits = (bmMask.bmPlanes * bmMask.bmWidth * bmMask.bmHeight * bmMask.bmBitsPixel + 15) / 16 * 2;
-        if (iconinfo.hbmColor)
-            cbColourBits = (bmColour.bmPlanes * bmColour.bmWidth * bmColour.bmHeight * bmColour.bmBitsPixel + 15) / 16 * 2;
-        cds.cbData = sizeof(*data) + sizeof(struct notify_data_icon) + cbMaskBits + cbColourBits;
+        if (!GetObjectW(iconinfo[iconCount].hbmMask, sizeof(bmMask), &bmMask[iconCount]) ||
+            (iconinfo[iconCount].hbmColor && !GetObjectW(iconinfo[iconCount].hbmColor, sizeof(bmColour), &bmColour[iconCount])))
+        {
+            DeleteObject(iconinfo[iconCount].hbmMask);
+            if (iconinfo[iconCount].hbmColor) DeleteObject(iconinfo[iconCount].hbmColor);
+            goto noballoon;
+        }
+        cbMaskBits[iconCount] = (bmMask[iconCount].bmPlanes * bmMask[iconCount].bmWidth * bmMask[iconCount].bmHeight * bmMask[iconCount].bmBitsPixel + 15) / 16 * 2;
+        if (iconinfo[iconCount].hbmColor)
+            cbColourBits[iconCount] = (bmColour[iconCount].bmPlanes * bmColour[iconCount].bmWidth * bmColour[iconCount].bmHeight * bmColour[iconCount].bmBitsPixel + 15) / 16 * 2;
+        iconCount++;
+    }
+noballoon:
+    if (iconCount > 0) {
+        char *buffer;
+        struct notify_data_icon *data_icon;
+        cds.cbData = sizeof(*data) + sizeof(struct notify_data_icon) * iconCount + cbMaskBits[0] + cbMaskBits[1] + cbColourBits[0] + cbColourBits[1];
         buffer = heap_alloc(cds.cbData);
         if (!buffer)
         {
-            DeleteObject(iconinfo.hbmMask);
-            if (iconinfo.hbmColor) DeleteObject(iconinfo.hbmColor);
+            for (unsigned int i = 0; i < iconCount; i++) {
+                if (iconinfo[i].hbmMask) DeleteObject(iconinfo[i].hbmMask);
+                if (iconinfo[i].hbmColor) DeleteObject(iconinfo[i].hbmColor);
+            }
             SetLastError(E_OUTOFMEMORY);
             return FALSE;
         }

data = (struct notify_data *)buffer;
+        data_icon = &data->icons[0];
         memset( data, 0, sizeof(*data) );
-        GetBitmapBits(iconinfo.hbmMask, cbMaskBits, &data->icons[0].buffer);
-        if (!iconinfo.hbmColor)
-        {
-            data->icons[0].width  = bmMask.bmWidth;
-            data->icons[0].height = bmMask.bmHeight / 2;
-            data->icons[0].planes = 1;
-            data->icons[0].bpp    = 1;
+        for (unsigned int i = 0; i < iconCount; i++) {
+            GetBitmapBits(iconinfo[i].hbmMask, cbMaskBits[i], &data_icon->buffer);
+            if (!iconinfo[i].hbmColor)
+            {
+                data_icon->width  = bmMask[i].bmWidth;
+                data_icon->height = bmMask[i].bmHeight / 2;
+                data_icon->planes = 1;
+                data_icon->bpp    = 1;
+            }
+            else
+            {
+                data_icon->width  = bmColour[i].bmWidth;
+                data_icon->height = bmColour[i].bmHeight;
+                data_icon->planes = bmColour[i].bmPlanes;
+                data_icon->bpp    = bmColour[i].bmBitsPixel;
+                GetBitmapBits(iconinfo[i].hbmColor, cbColourBits[i], &data_icon->buffer[cbMaskBits[i]]);
+                DeleteObject(iconinfo[i].hbmColor);
+            }
+            DeleteObject(iconinfo[i].hbmMask);
+            data_icon = (struct notify_data_icon*)&data_icon->buffer[cbMaskBits[i] + cbColourBits[i]];
         }
-        else
-        {
-            data->icons[0].width  = bmColour.bmWidth;
-            data->icons[0].height = bmColour.bmHeight;
-            data->icons[0].planes = bmColour.bmPlanes;
-            data->icons[0].bpp    = bmColour.bmBitsPixel;
-            GetBitmapBits(iconinfo.hbmColor, cbColourBits, &data->icons[0].buffer[cbMaskBits]);
-            DeleteObject(iconinfo.hbmColor);
-        }
-        DeleteObject(iconinfo.hbmMask);
     }

-noicon:
     data->hWnd   = HandleToLong( nid->hWnd );
     data->uID    = nid->uID;
     data->uFlags = nid->uFlags;
diff --git a/programs/explorer/systray.c b/programs/explorer/systray.c
index cabc1568b2f..1f80135364a 100644
--- a/programs/explorer/systray.c
+++ b/programs/explorer/systray.c
@@ -534,6 +534,7 @@ static BOOL handle_incoming(HWND hwndSource, COPYDATASTRUCT *cds)

if (cds->cbData < sizeof(*data)) return FALSE;
     data = cds->lpData;
+    data_icon = &data->icons[0];

nid.cbSize           = sizeof(nid);
     nid.hWnd             = LongToHandle( data->hWnd );
@@ -553,11 +554,10 @@ static BOOL handle_incoming(HWND hwndSource, COPYDATASTRUCT *cds)

/* FIXME: if statement only needed because we don't support interprocess
      * icon handles */
-    if ((nid.uFlags & NIF_ICON) && cds->cbData > (sizeof(*data) + sizeof(struct notify_data_icon)))
+    if ((nid.uFlags & NIF_ICON) && cds->cbData > ((char*)data_icon - (char*)data) + sizeof(struct notify_data_icon))
     {
         LONG cbMaskBits;
         LONG cbColourBits;
-        data_icon = &data->icons[0];

cbMaskBits = (data_icon->width * data_icon->height + 15) / 16 * 2;
         cbColourBits = (data_icon->planes * data_icon->width * data_icon->height * data_icon->bpp + 15) / 16 * 2;
@@ -569,9 +569,25 @@ static BOOL handle_incoming(HWND hwndSource, COPYDATASTRUCT *cds)
         }
         nid.hIcon = CreateIcon(NULL, data_icon->width, data_icon->height, data_icon->planes, data_icon->bpp,
                                &data_icon->buffer[0], &data_icon->buffer[cbMaskBits]);
-        data_icon = &data_icon->buffer[cbMaskBits + cbColourBits];
+        data_icon = (const struct notify_data_icon*)&data_icon->buffer[cbMaskBits + cbColourBits];
     }

+    if ((nid.uFlags & NIF_INFO) && cds->cbData > ((char*)data_icon - (char*)data) + sizeof(struct notify_data_icon)) {
+        /* Balloon icon */
+        LONG cbMaskBits;
+        LONG cbColourBits;
+
+        cbMaskBits = (data_icon->width * data_icon->height + 15) / 16 * 2;
+        cbColourBits = (data_icon->planes * data_icon->width * data_icon->height * data_icon->bpp + 15) / 16 * 2;
+
+        if (cds->cbData < sizeof(*data) + sizeof(*data_icon) + cbMaskBits + cbColourBits)
+        {
+            WINE_ERR("buffer underflow\n");
+            return FALSE;
+        }
+        nid.hBalloonIcon = CreateIcon(NULL, data_icon->width, data_icon->height, data_icon->planes, data_icon->bpp,
+                               &data_icon->buffer[0], &data_icon->buffer[cbMaskBits]);
+    }
     /* try forwarding to the display driver first */
     if (cds->dwData == NIM_ADD || !(icon = get_icon( nid.hWnd, nid.uID )))
     {
@@ -607,6 +623,7 @@ static BOOL handle_incoming(HWND hwndSource, COPYDATASTRUCT *cds)
     }

if (nid.hIcon) DestroyIcon( nid.hIcon );
+    if (nid.hBalloonIcon) DestroyIcon( nid.hBalloonIcon );
     sync_taskbar_buttons();
     return ret;
 }